shard parquet in download_and_prepare

Author: lhoestq

src/datasets/builder.py

@@ -30,6 +30,7 @@
 from typing import Dict, Mapping, Optional, Tuple, Union
 
 import fsspec
+from tqdm.contrib.concurrent import thread_map
 
 from . import config, utils
 from .arrow_dataset import Dataset
@@ -62,6 +63,7 @@
 from .utils.info_utils import get_size_checksum_dict, verify_checksums, verify_splits
 from .utils.py_utils import (
     classproperty,
+    convert_file_size_to_int,
     has_sufficient_disk_space,
     map_nested,
     memoize,
@@ -575,6 +577,14 @@ def get_imported_module_dir(cls):
         """Return the path of the module of this class or subclass."""
         return os.path.dirname(inspect.getfile(inspect.getmodule(cls)))
 
+    def _rename(self, src: str, dst: str):
+        is_local = not is_remote_filesystem(self._fs)
+        if is_local:
+            # LocalFileSystem.mv does copy + rm, it is more efficient to simply rename a local directory
+            os.rename(self._fs._strip_protocol(src), self._fs._strip_protocol(dst))
+        else:
+            self._fs.mv(src, dst, recursive=True)
+
     def download_and_prepare(
         self,
         download_config: Optional[DownloadConfig] = None,
@@ -672,11 +682,7 @@ def incomplete_dir(dirname):
                     yield tmp_dir
                     if self._fs.isdir(dirname):
                         self._fs.rm(dirname, recursive=True)
-                    if is_local:
-                        # LocalFileSystem.mv does copy + rm, it is more efficient to simply rename a local directory
-                        os.rename(self._fs._strip_protocol(tmp_dir), self._fs._strip_protocol(dirname))
-                    else:
-                        self._fs.mv(tmp_dir, dirname, recursive=True)
+                    self._rename(tmp_dir, dirname)
                 finally:
                     if self._fs.exists(tmp_dir):
                         self._fs.rm(tmp_dir, recursive=True)
@@ -1224,51 +1230,90 @@ def _generate_examples(self, **kwargs):
         """
         raise NotImplementedError()
 
-    def _prepare_split(self, split_generator, check_duplicate_keys, file_format=None):
+    def _prepare_split(self, split_generator, check_duplicate_keys, file_format=None, max_shard_size=None):
         is_local = not is_remote_filesystem(self._fs)
         path_join = os.path.join if is_local else posixpath.join
+        file_format = file_format or "arrow"
+
+        if max_shard_size is not None:
+            max_shard_size = convert_file_size_to_int(max_shard_size)
+            if file_format == "arrow":
+                raise NotImplementedError(
+                    "Writing sharded arrow files is not supported. Please don't use max_shard_size or use parquet."
+                )
 
         if self.info.splits is not None:
             split_info = self.info.splits[split_generator.name]
         else:
             split_info = split_generator.split_info
 
-        file_format = file_format or "arrow"
-        suffix = "-00000-of-00001" if file_format == "parquet" else ""
+        suffix = "-SSSSS-of-NNNNN" if file_format == "parquet" else ""
         fname = f"{self.name}-{split_generator.name}{suffix}.{file_format}"
         fpath = path_join(self._cache_dir, fname)
 
         generator = self._generate_examples(**split_generator.gen_kwargs)
 
         writer_class = ParquetWriter if file_format == "parquet" else ArrowWriter
-        with writer_class(
+
+        shard_id = 0
+        writer = writer_class(
             features=self.info.features,
-            path=fpath,
+            path=fpath.replace("SSSSS", f"{shard_id:05d}"),
             writer_batch_size=self._writer_batch_size,
             hash_salt=split_info.name,
             check_duplicates=check_duplicate_keys,
             storage_options=self._fs.storage_options,
-        ) as writer:
-            try:
-                for key, record in logging.tqdm(
-                    generator,
-                    unit=" examples",
-                    total=split_info.num_examples,
-                    leave=False,
-                    disable=not logging.is_progress_bar_enabled(),
-                    desc=f"Generating {split_info.name} split",
-                ):
-                    example = self.info.features.encode_example(record)
-                    writer.write(example, key)
-            finally:
-                num_examples, num_bytes = writer.finalize()
-
-        split_generator.split_info.num_examples = num_examples
-        split_generator.split_info.num_bytes = num_bytes
+        )
+        total_num_examples, total_num_bytes = 0, 0
+        try:
+            for key, record in logging.tqdm(
+                generator,
+                unit=" examples",
+                total=split_info.num_examples,
+                leave=False,
+                disable=not logging.is_progress_bar_enabled(),
+                desc=f"Generating {split_info.name} split",
+            ):
+                if max_shard_size is not None and writer._num_bytes > max_shard_size:
+                    num_examples, num_bytes = writer.finalize()
+                    total_num_examples += num_examples
+                    total_num_bytes += num_bytes
+                    shard_id += 1
+                    writer = writer_class(
+                        features=writer._features,
+                        path=fpath.replace("SSSSS", f"{shard_id:05d}"),
+                        writer_batch_size=self._writer_batch_size,
+                        hash_salt=split_info.name,
+                        check_duplicates=check_duplicate_keys,
+                        storage_options=self._fs.storage_options,
+                    )
+                example = self.info.features.encode_example(record)
+                writer.write(example, key)
+        finally:
+            num_shards = shard_id + 1
+            num_examples, num_bytes = writer.finalize()
+            total_num_examples += num_examples
+            total_num_bytes += num_bytes
 
-    def _download_and_prepare(self, dl_manager, verify_infos, file_format=None):
+        if file_format == "parquet":
+
+            def _rename_shard(shard_id: int):
+                self._rename(
+                    fpath.replace("SSSSS", f"{shard_id:05d}"),
+                    fpath.replace("SSSSS", f"{shard_id:05d}").replace("NNNNN", f"{num_shards:05d}"),
+                )
+
+            logger.debug(f"Renaming {num_shards} shards.")
+            thread_map(_rename_shard, range(num_shards), disable=True, max_workers=64)
+
+        split_generator.split_info.num_examples = total_num_examples
+        split_generator.split_info.num_bytes = total_num_bytes
+        if self.info.features is None:
+            self.info.features = writer._features
+
+    def _download_and_prepare(self, dl_manager, verify_infos, file_format=None, **kwargs):
         super()._download_and_prepare(
-            dl_manager, verify_infos, file_format=file_format, check_duplicate_keys=verify_infos
+            dl_manager, verify_infos, file_format=file_format, check_duplicate_keys=verify_infos, **kwargs
         )
 
     def _get_examples_iterable_for_split(self, split_generator: SplitGenerator) -> ExamplesIterable:
@@ -1310,26 +1355,70 @@ def _generate_tables(self, **kwargs):
         """
         raise NotImplementedError()
 
-    def _prepare_split(self, split_generator, file_format=None):
+    def _prepare_split(self, split_generator, file_format=None, max_shard_size=None):
         is_local = not is_remote_filesystem(self._fs)
         path_join = os.path.join if is_local else posixpath.join
-
         file_format = file_format or "arrow"
-        suffix = "-00000-of-00001" if file_format == "parquet" else ""
+
+        if max_shard_size is not None:
+            if file_format == "arrow":
+                raise NotImplementedError(
+                    "Writing sharded arrow files is not supported. Please don't use max_shard_size or use parquet."
+                )
+            max_shard_size = convert_file_size_to_int(max_shard_size or "500MB")
+
+        suffix = "-SSSSS-of-NNNNN" if file_format == "parquet" else ""
         fname = f"{self.name}-{split_generator.name}{suffix}.{file_format}"
         fpath = path_join(self._cache_dir, fname)
 
         generator = self._generate_tables(**split_generator.gen_kwargs)
+
         writer_class = ParquetWriter if file_format == "parquet" else ArrowWriter
-        with writer_class(features=self.info.features, path=fpath, storage_options=self._fs.storage_options) as writer:
+
+        shard_id = 0
+        writer = writer_class(
+            features=self.info.features,
+            path=fpath.replace("SSSSS", f"{shard_id:05d}"),
+            storage_options=self._fs.storage_options,
+        )
+        total_num_examples, total_num_bytes = 0, 0
+        try:
             for key, table in logging.tqdm(
-                generator, unit=" tables", leave=False, disable=(not logging.is_progress_bar_enabled())
+                generator,
+                unit=" tables",
+                leave=False,
+                disable=not logging.is_progress_bar_enabled(),
             ):
+                if max_shard_size is not None and writer._num_bytes > max_shard_size:
+                    num_examples, num_bytes = writer.finalize()
+                    total_num_examples += num_examples
+                    total_num_bytes += num_bytes
+                    shard_id += 1
+                    writer = writer_class(
+                        features=writer._features,
+                        path=fpath.replace("SSSSS", f"{shard_id:05d}"),
+                        storage_options=self._fs.storage_options,
+                    )
                 writer.write_table(table)
+        finally:
+            num_shards = shard_id + 1
             num_examples, num_bytes = writer.finalize()
+            total_num_examples += num_examples
+            total_num_bytes += num_bytes
+
+        if file_format == "parquet":
+
+            def _rename_shard(shard_id: int):
+                self._rename(
+                    fpath.replace("SSSSS", f"{shard_id:05d}"),
+                    fpath.replace("SSSSS", f"{shard_id:05d}").replace("NNNNN", f"{num_shards:05d}"),
+                )
+
+            logger.debug(f"Renaming {num_shards} shards.")
+            thread_map(_rename_shard, range(num_shards), disable=True, max_workers=64)
 
-        split_generator.split_info.num_examples = num_examples
-        split_generator.split_info.num_bytes = num_bytes
+        split_generator.split_info.num_examples = total_num_examples
+        split_generator.split_info.num_bytes = total_num_bytes
         if self.info.features is None:
             self.info.features = writer._features
 

tests/test_builder.py

@@ -952,6 +952,24 @@ def test_generator_based_builder_download_and_prepare_as_parquet(tmp_path):
     assert pq.ParquetFile(parquet_path) is not None
 
 
+def test_generator_based_builder_download_and_prepare_as_sharded_parquet(tmp_path):
+    writer_batch_size = 25
+    builder = DummyGeneratorBasedBuilder(cache_dir=tmp_path, writer_batch_size=writer_batch_size)
+    builder.download_and_prepare(file_format="parquet", max_shard_size=1)  # one batch per shard
+    expected_num_shards = 100 // writer_batch_size
+    assert builder.info.splits["train"].num_examples, 100
+    parquet_path = os.path.join(
+        tmp_path, builder.name, "default", "0.0.0", f"{builder.name}-train-00000-of-{expected_num_shards:05d}.parquet"
+    )
+    assert os.path.exists(parquet_path)
+    parquet_files = [
+        pq.ParquetFile(parquet_path)
+        for parquet_path in Path(tmp_path).rglob(f"{builder.name}-train-*-of-{expected_num_shards:05d}.parquet")
+    ]
+    assert len(parquet_files) == expected_num_shards
+    assert sum(parquet_file.metadata.num_rows for parquet_file in parquet_files) == 100
+
+
 def test_arrow_based_builder_download_and_prepare_as_parquet(tmp_path):
     builder = DummyArrowBasedBuilder(cache_dir=tmp_path)
     builder.download_and_prepare(file_format="parquet")
@@ -963,6 +981,23 @@ def test_arrow_based_builder_download_and_prepare_as_parquet(tmp_path):
     assert pq.ParquetFile(parquet_path) is not None
 
 
+def test_arrow_based_builder_download_and_prepare_as_sharded_parquet(tmp_path):
+    builder = DummyArrowBasedBuilder(cache_dir=tmp_path)
+    builder.download_and_prepare(file_format="parquet", max_shard_size=1)  # one table per shard
+    expected_num_shards = 10
+    assert builder.info.splits["train"].num_examples, 100
+    parquet_path = os.path.join(
+        tmp_path, builder.name, "default", "0.0.0", f"{builder.name}-train-00000-of-{expected_num_shards:05d}.parquet"
+    )
+    assert os.path.exists(parquet_path)
+    parquet_files = [
+        pq.ParquetFile(parquet_path)
+        for parquet_path in Path(tmp_path).rglob(f"{builder.name}-train-*-of-{expected_num_shards:05d}.parquet")
+    ]
+    assert len(parquet_files) == expected_num_shards
+    assert sum(parquet_file.metadata.num_rows for parquet_file in parquet_files) == 100
+
+
 def test_beam_based_builder_download_and_prepare_as_parquet(tmp_path):
     builder = DummyBeamBasedBuilder(cache_dir=tmp_path, beam_runner="DirectRunner")
     builder.download_and_prepare(file_format="parquet")