| import os |
| from typing import BinaryIO, Optional, Union |
|
|
| import fsspec |
| import pyarrow.parquet as pq |
|
|
| from .. import Dataset, Features, NamedSplit, config |
| from ..arrow_writer import get_writer_batch_size |
| from ..formatting import query_table |
| from ..packaged_modules import _PACKAGED_DATASETS_MODULES |
| from ..packaged_modules.parquet.parquet import Parquet |
| from ..utils import tqdm as hf_tqdm |
| from ..utils.typing import NestedDataStructureLike, PathLike |
| from .abc import AbstractDatasetReader |
|
|
|
|
| class ParquetDatasetReader(AbstractDatasetReader): |
| def __init__( |
| self, |
| path_or_paths: NestedDataStructureLike[PathLike], |
| split: Optional[NamedSplit] = None, |
| features: Optional[Features] = None, |
| cache_dir: str = None, |
| keep_in_memory: bool = False, |
| streaming: bool = False, |
| num_proc: Optional[int] = None, |
| **kwargs, |
| ): |
| super().__init__( |
| path_or_paths, |
| split=split, |
| features=features, |
| cache_dir=cache_dir, |
| keep_in_memory=keep_in_memory, |
| streaming=streaming, |
| num_proc=num_proc, |
| **kwargs, |
| ) |
| path_or_paths = path_or_paths if isinstance(path_or_paths, dict) else {self.split: path_or_paths} |
| hash = _PACKAGED_DATASETS_MODULES["parquet"][1] |
| self.builder = Parquet( |
| cache_dir=cache_dir, |
| data_files=path_or_paths, |
| features=features, |
| hash=hash, |
| **kwargs, |
| ) |
|
|
| def read(self): |
| |
| if self.streaming: |
| dataset = self.builder.as_streaming_dataset(split=self.split) |
| |
| else: |
| download_config = None |
| download_mode = None |
| verification_mode = None |
| base_path = None |
|
|
| self.builder.download_and_prepare( |
| download_config=download_config, |
| download_mode=download_mode, |
| verification_mode=verification_mode, |
| base_path=base_path, |
| num_proc=self.num_proc, |
| ) |
| dataset = self.builder.as_dataset( |
| split=self.split, verification_mode=verification_mode, in_memory=self.keep_in_memory |
| ) |
| return dataset |
|
|
|
|
| class ParquetDatasetWriter: |
| def __init__( |
| self, |
| dataset: Dataset, |
| path_or_buf: Union[PathLike, BinaryIO], |
| batch_size: Optional[int] = None, |
| storage_options: Optional[dict] = None, |
| **parquet_writer_kwargs, |
| ): |
| self.dataset = dataset |
| self.path_or_buf = path_or_buf |
| self.batch_size = batch_size or get_writer_batch_size(dataset.features) |
| self.storage_options = storage_options or {} |
| self.parquet_writer_kwargs = parquet_writer_kwargs |
|
|
| def write(self) -> int: |
| batch_size = self.batch_size if self.batch_size else config.DEFAULT_MAX_BATCH_SIZE |
|
|
| if isinstance(self.path_or_buf, (str, bytes, os.PathLike)): |
| with fsspec.open(self.path_or_buf, "wb", **(self.storage_options or {})) as buffer: |
| written = self._write(file_obj=buffer, batch_size=batch_size, **self.parquet_writer_kwargs) |
| else: |
| written = self._write(file_obj=self.path_or_buf, batch_size=batch_size, **self.parquet_writer_kwargs) |
| return written |
|
|
| def _write(self, file_obj: BinaryIO, batch_size: int, **parquet_writer_kwargs) -> int: |
| """Writes the pyarrow table as Parquet to a binary file handle. |
| |
| Caller is responsible for opening and closing the handle. |
| """ |
| written = 0 |
| _ = parquet_writer_kwargs.pop("path_or_buf", None) |
| schema = self.dataset.features.arrow_schema |
|
|
| writer = pq.ParquetWriter(file_obj, schema=schema, **parquet_writer_kwargs) |
|
|
| for offset in hf_tqdm( |
| range(0, len(self.dataset), batch_size), |
| unit="ba", |
| desc="Creating parquet from Arrow format", |
| ): |
| batch = query_table( |
| table=self.dataset._data, |
| key=slice(offset, offset + batch_size), |
| indices=self.dataset._indices, |
| ) |
| writer.write_table(batch) |
| written += batch.nbytes |
| writer.close() |
| return written |
|
|
