datasets 5.1.0__tar.gz → 5.2.0__tar.gz

{datasets-5.1.0/src/datasets.egg-info → datasets-5.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datasets
-Version: 5.1.0
+Version: 5.2.0
 Summary: HuggingFace community-driven open-source library of datasets
 Home-page: https://github.com/huggingface/datasets
 Download-URL: https://github.com/huggingface/datasets/tags
@@ -33,7 +33,7 @@ Requires-Dist: pandas
 Requires-Dist: tqdm>=4.66.3
 Requires-Dist: xxhash
 Requires-Dist: multiprocess<0.70.20
-Requires-Dist: fsspec[http]<=2026.7.0,>=2023.1.0
+Requires-Dist: fsspec<=2026.9.0,>=2023.1.0
 Requires-Dist: huggingface-hub<3.0,>=1.31.0
 Requires-Dist: packaging
 Requires-Dist: pyyaml>=5.1

{datasets-5.1.0 → datasets-5.2.0}/setup.py

@@ -124,7 +124,7 @@ REQUIRED_PKGS = [
     "multiprocess<0.70.20",  # to align with dill<0.3.9 (see above)
     # to save datasets locally or on any filesystem
     # minimum 2023.1.0 to support protocol=kwargs in fsspec's `open`, `get_fs_token_paths`, etc.: see https://github.com/fsspec/filesystem_spec/pull/1143
-    "fsspec[http]>=2023.1.0,<=2026.7.0",
+    "fsspec>=2023.1.0,<=2026.9.0",
     # To get datasets from the Datasets Hub on huggingface.co
     # (also provides the HTTP client: `from huggingface_hub.utils import httpx`)
     # minimum 1.31.0 for the `huggingface_hub.utils.httpx` re-export
@@ -252,7 +252,7 @@ EXTRAS_REQUIRE = {
 
 setup(
     name="datasets",
-    version="5.1.0",  # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
+    version="5.2.0",  # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
     description="HuggingFace community-driven open-source library of datasets",
     long_description=open("README.md", encoding="utf-8").read(),
     long_description_content_type="text/markdown",

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "5.1.0"
+__version__ = "5.2.0"
 
 from .arrow_dataset import Column, Dataset
 from .arrow_reader import ReadInstruction

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/arrow_dataset.py

@@ -855,7 +855,7 @@ class Dataset(DatasetInfoMixin, IndexableMixin, TensorflowDatasetMixin):
         table = InMemoryTable.from_buffer(buffer)
 
         if indices_buffer is not None:
-            indices_table = InMemoryTable.from_buffer(buffer)
+            indices_table = InMemoryTable.from_buffer(indices_buffer)
         else:
             indices_table = None
 
@@ -1830,14 +1830,21 @@ class Dataset(DatasetInfoMixin, IndexableMixin, TensorflowDatasetMixin):
         if self.list_indexes():
             raise ValueError("please remove all the indexes using `dataset.drop_index` before saving a dataset")
 
+        from .arrow_writer import get_writer_batch_size_from_data_size
+
+        dataset_nbytes = self._estimate_nbytes()
         if num_shards is None:
-            dataset_nbytes = self._estimate_nbytes()
             max_shard_size = convert_file_size_to_int(max_shard_size or config.MAX_SHARD_SIZE)
             num_shards = int(dataset_nbytes / max_shard_size) + 1
             num_shards = max(num_shards, num_proc or 1)
             # if we have only a few large samples, we should only create as many shards as samples
             num_shards = min(len(self.data), num_shards)
 
+        # use smaller batches for big samples (e.g. images) to stay below the 2GB limit of arrow binary arrays
+        writer_batch_size = config.DEFAULT_MAX_BATCH_SIZE
+        if dataset_nbytes > 0:
+            writer_batch_size = min(writer_batch_size, get_writer_batch_size_from_data_size(len(self), dataset_nbytes))
+
         fs: fsspec.AbstractFileSystem
         fs, _ = url_to_fs(dataset_path, **(storage_options or {}))
 
@@ -1890,6 +1897,7 @@ class Dataset(DatasetInfoMixin, IndexableMixin, TensorflowDatasetMixin):
                 "shard": self.shard(num_shards=num_shards, index=shard_idx, contiguous=True),
                 "fpath": posixpath.join(dataset_path, f"data-{shard_idx:05d}-of-{num_shards:05d}.arrow"),
                 "storage_options": storage_options,
+                "writer_batch_size": writer_batch_size,
             }
             for shard_idx in range(num_shards)
         )
@@ -1931,9 +1939,13 @@ class Dataset(DatasetInfoMixin, IndexableMixin, TensorflowDatasetMixin):
             json.dump(sorted_keys_dataset_info, dataset_info_file, indent=2)
 
     @staticmethod
-    def _save_to_disk_single(job_id: int, shard: "Dataset", fpath: str, storage_options: Optional[dict]):
-        batch_size = config.DEFAULT_MAX_BATCH_SIZE
-
+    def _save_to_disk_single(
+        job_id: int,
+        shard: "Dataset",
+        fpath: str,
+        storage_options: Optional[dict],
+        writer_batch_size: int,
+    ):
         num_examples_progress_update = 0
         writer = ArrowWriter(
             features=shard.features,
@@ -1943,7 +1955,7 @@ class Dataset(DatasetInfoMixin, IndexableMixin, TensorflowDatasetMixin):
         )
         try:
             _time = time.time()
-            for pa_table in shard.with_format("arrow").iter(batch_size):
+            for pa_table in shard.with_format("arrow").iter(writer_batch_size):
                 writer.write_table(pa_table)
                 num_examples_progress_update += len(pa_table)
                 if time.time() > _time + config.PBAR_REFRESH_TIME_INTERVAL:
@@ -3844,7 +3856,14 @@ class Dataset(DatasetInfoMixin, IndexableMixin, TensorflowDatasetMixin):
                         processed_inputs.pop(column)
             if check_same_num_examples:
                 input_num_examples = len(pa_inputs)
-                processed_inputs_num_examples = len(processed_inputs[next(iter(processed_inputs.keys()))])
+                if not isinstance(processed_inputs, Mapping):
+                    # pa.Table, pd.DataFrame or pl.DataFrame
+                    processed_inputs_num_examples = len(processed_inputs)
+                elif processed_inputs:
+                    processed_inputs_num_examples = len(next(iter(processed_inputs.values())))
+                else:
+                    # an empty dict doesn't update any column
+                    processed_inputs_num_examples = input_num_examples
                 if input_num_examples != processed_inputs_num_examples:
                     raise DatasetTransformationNotAllowedError(
                         "Using `.map` in batched mode on a dataset with attached indexes is allowed only if it doesn't create or remove existing examples. You can first run `.drop_index() to remove your index and then re-add it."
@@ -4690,7 +4709,10 @@ class Dataset(DatasetInfoMixin, IndexableMixin, TensorflowDatasetMixin):
          'text': 'the gorgeously elaborate continuation of " the lord of the rings " trilogy is so huge that a column of words cannot adequately describe co-writer/director peter jackson\'s expanded vision of j . r . r . tolkien\'s middle-earth .'}]
         ```
         """
-        return self.select(range(n))
+        # Clamp `n` to the dataset length so taking more elements than available
+        # returns the whole dataset instead of raising an IndexError, matching
+        # the behavior of `IterableDataset.take`.
+        return self.select(range(min(n, len(self))))
 
     @transmit_format
     @fingerprint_transform(inplace=False, ignore_kwargs=["load_from_cache_file", "indices_cache_file_name"])
@@ -5722,7 +5744,7 @@ class Dataset(DatasetInfoMixin, IndexableMixin, TensorflowDatasetMixin):
             k for k, v in self._info.features.items() if require_decoding(v, ignore_decode_attribute=True)
         ]
 
-        if decodable_columns:
+        if decodable_columns and len(self) > 0:
             # Approximate the space needed to store the bytes from the external files by analyzing the first 1000 examples
             extra_nbytes = 0
 
@@ -6010,12 +6032,14 @@ class Dataset(DatasetInfoMixin, IndexableMixin, TensorflowDatasetMixin):
             logger.warning(
                 f"Setting num_proc from {num_jobs} back to 1 for the {split} split to disable multiprocessing as it only contains one shard."
             )
+            num_proc = None
             num_jobs = 1
         elif num_shards < num_jobs:
             logger.warning(
                 f"Setting num_proc from {num_jobs} to {num_shards} for the {split} split as it only contains {num_shards} shards."
             )
             num_proc = num_shards
+            num_jobs = num_shards
         kwargs_iterable = [
             {
                 "self": self.shard(num_shards=num_jobs, index=job_id, contiguous=True),
@@ -7115,7 +7139,9 @@ def _concatenate_map_style_datasets(
     if info is None:
         info = DatasetInfo.from_merge([dset.info for dset in dsets])
     fingerprint = update_fingerprint(
-        "".join(dset._fingerprint for dset in dsets), _concatenate_map_style_datasets, {"info": info, "split": split}
+        "".join(dset._fingerprint for dset in dsets),
+        _concatenate_map_style_datasets,
+        {"info": info, "split": split, "axis": axis},
     )
 
     # Make final concatenated dataset

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/builder.py

@@ -1175,15 +1175,23 @@ class DatasetBuilder:
         return datasets
 
     def _as_streaming_dataset_single(
-        self,
-        splits_generator,
+        self, splits_generator, decode_num_threads: Optional[int] = None
     ) -> IterableDataset:
         ex_iterable = self._get_examples_iterable_for_split(splits_generator)
         # add auth to be able to access and decode audio/image files from private repositories.
         token_per_repo_id = {self.repo_id: self.token} if self.repo_id else {}
-        return IterableDataset(
+        ds = IterableDataset(
             ex_iterable, info=self.info, split=splits_generator.name, token_per_repo_id=token_per_repo_id
         )
+        # Read and decode media examples (Image/Audio/Video/Mesh/...) in worker threads by default: opening and
+        # reading a remote file costs network round trips rather than CPU time, so streaming remote media data is
+        # latency bound without any overlap. Datasets without decodable features (text, tabular, parquet, ...) and
+        # datasets read from the local disk are returned unchanged.
+        if decode_num_threads is None and hasattr(self.config, "decode_num_threads"):
+            decode_num_threads = self.config.decode_num_threads
+        if self.base_path and self.base_path.startswith("hf://") and decode_num_threads:
+            ds = ds.decode(num_threads=decode_num_threads)
+        return ds
 
     @abc.abstractmethod
     def _split_generators(self, dl_manager: Union[DownloadManager, StreamingDownloadManager]):

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/combine.py

@@ -180,6 +180,8 @@ def concatenate_datasets(
 
     * if axis=0, the resulting dataset's `num_shards` is the sum of each dataset's `num_shards`.
     * if axis=1, the resulting dataset has one (1) shard to not misalign data.
+    * if the datasets were split with [`datasets.distributed.split_dataset_by_node`], they must all be split the same way,
+      and the resulting dataset keeps that split.
 
     Args:
         dsets (`List[datasets.Dataset]` or `List[datasets.IterableDataset]`):

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/features/audio.py

@@ -188,6 +188,9 @@ class Audio:
         if not self.decode:
             raise RuntimeError("Decoding is disabled for this feature. Please use Audio(decode=True) instead.")
 
+        if not isinstance(value, dict):
+            return value  # decoding is idempotent: a value that is already an audio decoder is returned as is
+
         path, bytes = (value["path"], value["bytes"]) if value["bytes"] is not None else (value["path"], None)
         if path is None and bytes is None:
             raise ValueError(f"An audio sample should have one of 'path' or 'bytes' but both are None in {value}.")

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/features/features.py

@@ -1234,7 +1234,15 @@ class ClassLabel:
                 raise ValueError(
                     f"Class label {min_max['max']} greater than configured num_classes {self.num_classes}"
                 )
-        elif isinstance(storage, pa.StringArray):
+        elif (
+            pa.types.is_string(storage.type)
+            or pa.types.is_large_string(storage.type)
+            or pa.types.is_string_view(storage.type)
+        ):
+            # Check the Arrow type rather than the array class: `pa.LargeStringArray` and
+            # `pa.StringViewArray` are not subclasses of `pa.StringArray`, and pandas 3
+            # produces `large_string` for plain string columns, so an isinstance check let
+            # those fall through to an int64 cast of the label names.
             storage = pa.array(
                 [self._strval2int(label) if label is not None else None for label in storage.to_pylist()]
             )
@@ -1974,8 +1982,14 @@ class Features(dict):
         for column_name, feature in self.items():
             self[column_name] = _visit(feature, _check_old_list)
 
-    __setitem__ = keep_features_dicts_synced(dict.__setitem__)
-    __delitem__ = keep_features_dicts_synced(dict.__delitem__)
+    def __setitem__(self, column_name: str, feature: FeatureType):
+        super().__setitem__(column_name, feature)
+        self._column_requires_decoding[column_name] = require_decoding(feature)
+
+    def __delitem__(self, column_name: str):
+        super().__delitem__(column_name)
+        self._column_requires_decoding.pop(column_name, None)
+
     update = keep_features_dicts_synced(dict.update)
     setdefault = keep_features_dicts_synced(dict.setdefault)
     pop = keep_features_dicts_synced(dict.pop)

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/features/image.py

@@ -166,6 +166,9 @@ class Image:
         if token_per_repo_id is None:
             token_per_repo_id = {}
 
+        if not isinstance(value, (dict, str)):
+            return value  # decoding is idempotent: a value that is already an image is returned as is
+
         path, bytes_ = value["path"], value["bytes"]
         if bytes_ is None:
             if path is None:

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/features/mesh.py

@@ -41,7 +41,7 @@ class Mesh:
         decode (`bool`, defaults to `True`):
             Whether to decode the mesh data. If `False`,
             returns the underlying dictionary in the format `{"path": mesh_path, "bytes": mesh_bytes}`.
-            Mesh decoding uses `trimesh` and supports `.glb`, `.ply`, and `.stl` files.
+            Mesh decoding uses `trimesh` and supports `.glb`, `.gltf`, `.ply`, and `.stl` files.
     """
 
     decode: bool = True
@@ -116,6 +116,9 @@ class Mesh:
         if token_per_repo_id is None:
             token_per_repo_id = {}
 
+        if not isinstance(value, (dict, str)):
+            return value  # decoding is idempotent: a value that is already a mesh is returned as is
+
         path, bytes_ = value["path"], value["bytes"]
         if bytes_ is None:
             if path is None:
@@ -123,7 +126,7 @@ class Mesh:
             if is_local_path(path):
                 file_type = _infer_mesh_file_type(path)
                 if file_type is None:
-                    raise ValueError("A mesh path should have a .glb, .ply, or .stl extension.")
+                    raise ValueError("A mesh path should have a .glb, .gltf, .ply, or .stl extension.")
                 return trimesh.load(path, file_type=file_type)
             source_url = path.split("::")[-1]
             pattern = (
@@ -138,7 +141,7 @@ class Mesh:
         file_type = _infer_mesh_file_type(path)
         if file_type is None:
             raise ValueError(
-                "Decoding mesh bytes requires a 'path' value with a .glb, .ply, or .stl extension "
+                "Decoding mesh bytes requires a 'path' value with a .glb, .gltf, .ply, or .stl extension "
                 "to infer the mesh file type."
             )
         return trimesh.load(BytesIO(bytes_), file_type=file_type)
@@ -276,7 +279,7 @@ class Mesh:
 
 
 def _infer_mesh_file_type(path: Optional[str]) -> Optional[str]:
-    supported_file_types = {"glb", "ply", "stl"}
+    supported_file_types = {"glb", "gltf", "ply", "stl"}
     if path is None:
         return None
     path_without_archive = path.split("::", 1)[0]

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/features/nifti.py

@@ -174,6 +174,9 @@ class Nifti:
         if token_per_repo_id is None:
             token_per_repo_id = {}
 
+        if not isinstance(value, (dict, str)):
+            return value  # decoding is idempotent: a value that is already a NIfTI image is returned as is
+
         path, bytes_ = value["path"], value["bytes"]
         if bytes_ is None:
             if path is None:

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/features/video.py

@@ -184,6 +184,9 @@ class Video:
         if token_per_repo_id is None:
             token_per_repo_id = {}
 
+        if not isinstance(value, (dict, str)):
+            return value  # decoding is idempotent: a value that is already a video decoder is returned as is
+
         if isinstance(value, str):
             path, bytes_ = value, None
         else:

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/io/csv.py

@@ -115,6 +115,9 @@ class CsvDatasetWriter:
 
         Caller is responsible for opening and closing the handle.
         """
+        if len(self.dataset) == 0:
+            return file_obj.write(self._batch_csv((0, header, index, to_csv_kwargs)))
+
         written = 0
 
         if self.num_proc is None or self.num_proc == 1:

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/io/json.py

@@ -163,6 +163,9 @@ class JsonDatasetWriter:
 
         Caller is responsible for opening and closing the handle.
         """
+        if len(self.dataset) == 0 and not lines:
+            return file_obj.write(self._batch_json((0, orient, lines, to_json_kwargs)))
+
         written = 0
 
         if self.num_proc is None or self.num_proc == 1:

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/iterable_dataset.py

@@ -3,10 +3,10 @@ import concurrent.futures
 import contextlib
 import inspect
 import itertools
-import multiprocessing.pool
 import re
 import sys
 import tempfile
+import threading
 import time
 from collections import Counter
 from collections.abc import Iterable, Iterator
@@ -1407,6 +1407,10 @@ class MappedExamplesIterable(_BaseExamplesIterable):
             max_num_running_async_map_functions_in_parallel or config.MAX_NUM_RUNNING_ASYNC_MAP_FUNCTIONS_IN_PARALLEL
         )
         self.is_batch_accumulate_arrow_table_function = is_batch_accumulate_arrow_table_function
+        # Set by IterableDataset.decode when it enables threaded decoding: since examples must be decoded exactly
+        # once, a later `.decode()` call reconfigures decoding from `decoding_dataset` (the dataset `.decode()` was
+        # called on) rather than stacking a second decoding step on top of this one.
+        self.decoding_pool: Optional[_LazyThreadPoolExecutor] = None
         # sanity checks
         if formatting and formatting.is_table:
             # batch_size should match for iter_arrow
@@ -1604,16 +1608,12 @@ class MappedExamplesIterable(_BaseExamplesIterable):
                 for i, key_example in inputs_iterator:
                     indices.append(i)
                     tasks.append(loop.create_task(async_apply_function(key_example, i)))
-                    # keep the total active tasks under a certain number
+                    # keep the total active tasks under a certain number, waiting for the oldest task to complete.
+                    # Polling the whole task list with asyncio.wait can't be used for that: it returns as soon as any
+                    # task in the list is done, so as soon as one example is slower than the examples created after it
+                    # (e.g. a large remote file), every wait returns instantly and this spins at 100% CPU, taking the
+                    # GIL away from the worker threads that were doing the actual work.
                     if len(tasks) >= self.max_num_running_async_map_functions_in_parallel:
-                        done, pending = loop.run_until_complete(
-                            asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
-                        )
-                        while tasks and len(pending) >= self.max_num_running_async_map_functions_in_parallel:
-                            done, pending = loop.run_until_complete(
-                                asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
-                            )
-                    if len(tasks) >= 10 * self.max_num_running_async_map_functions_in_parallel:
                         loop.run_until_complete(tasks[0])
                     # yield finished tasks
                     while tasks and tasks[0].done():
@@ -4517,7 +4517,7 @@ class IterableDataset(DatasetInfoMixin):
             token_per_repo_id=self._token_per_repo_id,
         )
 
-    def decode(self, enable: bool = True, num_threads: int = 0) -> "IterableDataset":
+    def decode(self, enable: bool = True, num_threads: Optional[int] = None) -> "IterableDataset":
         """
         Enable or disable the dataset features decoding for audio, image, video.
 
@@ -4527,8 +4527,13 @@ class IterableDataset(DatasetInfoMixin):
         * image -> PIL.Image
         * video -> torchcodec.decoders.VideoDecoder
 
-        You can enable multithreading using `num_threads`. This is especially useful to speed up remote
-        data streaming. However it can be slower than `num_threads=0` for local data on fast disks.
+        Multithreading can be enabled with `num_threads`: every media file is then read and decoded in a worker
+        thread, so that the network round trips of consecutive examples overlap. This is especially useful to speed
+        up remote data streaming, and [`load_dataset`] already enables it by default when streaming a remote dataset
+        that has decodable features (see `HF_DATASETS_DECODE_NUM_THREADS`). Datasets read from the local disk are
+        streamed with `num_threads=0`, since multithreading can be slower than decoding in the main thread for local
+        data on fast disks. Since examples are decoded in advance, every thread in flight can keep one example in
+        memory: lower `num_threads` to use less memory.
 
         Disabling decoding is useful if you want to iterate on the paths or bytes of the media files
         without actually decoding their content. To disable decoding you can use `.decode(False)`, which
@@ -4538,8 +4543,10 @@ class IterableDataset(DatasetInfoMixin):
         Args:
             enable (`bool`, defaults to `True`):
                 Enable or disable features decoding.
-            num_threads (`int`, defaults to `0`):
-                Enable multithreading for features decoding.
+            num_threads (`int`, *optional*, defaults to `None`):
+                Enable multithreading for features decoding. Note that [`load_dataset`] already uses a few threads
+                by default when streaming a dataset on HF with decodable features.
+                By default, `num_threads` is set to 32 if at least one feature require decoding, and 0 otherwise.
 
         Returns:
             `IterableDataset`: A copy of the dataset with casted features.
@@ -4575,9 +4582,9 @@ class IterableDataset(DatasetInfoMixin):
         >>> from datasets import load_dataset
         >>> from tqdm import tqdm
         >>> ds = load_dataset("sshh12/planet-textures", split="train", streaming=True)
-        >>> num_threads = min(32, (os.cpu_count() or 1) + 4)
-        >>> ds = ds.decode(num_threads=num_threads)
-        >>> for _ in tqdm(ds):  # 20 times faster !
+        >>> # already threaded by default with 32 threads, but more threads can help going even faster
+        >>> ds = ds.decode(num_threads=64)
+        >>> for _ in tqdm(ds):  # 2 times faster !
         ...     ...
         ```
         """
@@ -4586,28 +4593,36 @@ class IterableDataset(DatasetInfoMixin):
                 "Features decoding is only available for datasets with known features, but features are Unknown. "
                 "Please set the datasets features with `ds = ds.cast(features)`."
             )
-        ds = self
-
-        def set_decoding(decode: bool, feature):
-            if hasattr(feature, "decode"):
-                feature.decode = decode
+        if num_threads is None:
+            num_threads = 32 if require_decoding(self.features) else 0
 
-        if enable and num_threads > 0:
-            disabled_decoding_features = self.features.copy()
-            enabled_decoding_features = self.features.copy()
+        ds = self
+        # If decoding is already configured, start again from the dataset it was configured on: examples must be
+        # decoded exactly once, so `.decode()` replaces the previous decoding step instead of adding a new one.
+        if isinstance(ds._ex_iterable, MappedExamplesIterable) and ds._ex_iterable.decoding_pool is not None:
+            # The original dataset may still have active iterators using its decoding pool.
+            ds = IterableDataset(
+                ds._ex_iterable.ex_iterable,
+                info=ds.info,
+                split=ds.split,
+                formatting=ds._formatting,
+                distributed=ds._distributed,
+                token_per_repo_id=ds._token_per_repo_id,
+            )
 
-            _visit(disabled_decoding_features, partial(set_decoding, False))
-            _visit(enabled_decoding_features, partial(set_decoding, True))
+        if enable and num_threads:
+            disabled_decoding_features = _features_with_decoding(self.features, False)
+            enabled_decoding_features = _features_with_decoding(self.features, True)
             ds = ds.cast(disabled_decoding_features)
-            pool = multiprocessing.pool.ThreadPool(num_threads)
-            func = partial(_apply_async, pool, enabled_decoding_features.decode_example)
+            pool = _LazyThreadPoolExecutor(num_threads)
+            decode_example = partial(enabled_decoding_features.decode_example, token_per_repo_id=ds._token_per_repo_id)
+            func = partial(_apply_async, pool, decode_example)
             ds = ds.map(func, features=enabled_decoding_features)
             assert isinstance(ds._ex_iterable, MappedExamplesIterable)
+            ds._ex_iterable.decoding_pool = pool
             ds._ex_iterable.max_num_running_async_map_functions_in_parallel = 2 * num_threads
         else:
-            features = ds.features.copy()
-            _visit(features, partial(set_decoding, enable))
-            ds = ds.cast(features)
+            ds = ds.cast(_features_with_decoding(ds.features, enable))
         return ds
 
     def _step(self, step: int, offset: int) -> "IterableDataset":
@@ -4663,10 +4678,20 @@ class IterableDataset(DatasetInfoMixin):
                 Whether to drop the last incomplete batch.
 
         Example:
+
+        Fixed batch size:
+
         ```py
-        >>> ds = load_dataset("some_dataset", streaming=True)
+        >>> ds = load_dataset("some_dataset", streaming=True, split="train")
         >>> batched_ds = ds.batch(batch_size=32)
         ```
+
+        Batches of successive samples with identical value for "my_column":
+
+        ```py
+        >>> ds = load_dataset("some_dataset", streaming=True, split="train")
+        >>> batched_ds = ds.batch(by_column="my_column")
+        ```
         """
         if batch_size is None and by_column is None:
             raise ValueError("IterableDataset.batch() misses `batch_size` or `by_column` argument.")
@@ -4682,7 +4707,7 @@ class IterableDataset(DatasetInfoMixin):
                     partial(_batch_accumulate_arrow_table_by_columns, columns=columns),
                     with_indices=True,
                     batched=True,
-                    batch_size=batch_size,
+                    batch_size=batch_size or config.DEFAULT_MAX_BATCH_SIZE,
                     drop_last_batch=drop_last_batch,
                     features=features,
                     is_batch_accumulate_arrow_table_function=True,
@@ -5532,6 +5557,11 @@ def _concatenate_iterable_datasets(
         _check_if_features_can_be_aligned([dset.features for dset in dsets])
     else:
         _check_column_names([col_name for dset in dsets for col_name in dset.features])
+    for i, dset in enumerate(dsets):
+        if dsets[0]._distributed != dset._distributed:
+            raise ValueError(
+                f"Datasets should be identically split_by_node before concatenating, but got {dsets[0]._distributed}!={dset._distributed} at index 0 and {i}"
+            )
 
     # Check format is consistent; if so, will set format for concatenated dataset
     if all(dset._formatting is None for dset in dsets):
@@ -5586,6 +5616,7 @@ def _concatenate_iterable_datasets(
         split=split,
         token_per_repo_id=token_per_repo_id,
         formatting=formatting,
+        distributed=dsets[0]._distributed,
     )
 
 
@@ -5729,13 +5760,66 @@ def _split_by_node_iterable_dataset(
     )
 
 
-async def _apply_async(pool, func, x):
-    future = pool.apply_async(func, (x,))
-    while True:
-        if future.ready():
-            return future.get()
-        else:
-            await asyncio.sleep(0)
+def _features_with_decoding(features: Features, decode: bool) -> Features:
+    """Return a copy of `features` where media features (Image/Audio/Video/Mesh/...) are set to decode or not."""
+
+    def _set_decode_flag(feature):
+        if hasattr(feature, "decode"):
+            feature.decode = decode
+
+    features = features.copy()
+    _visit(features, _set_decode_flag)
+    # Features caches which columns require decoding: re-assign every column to refresh it after the flags change
+    for column_name, feature in list(features.items()):
+        features[column_name] = feature
+    return features
+
+
+class _LazyThreadPoolExecutor:
+    """A `concurrent.futures.ThreadPoolExecutor` that is created on first use instead of upfront.
+
+    Threaded decoding holds one of these in the dataset, and datasets get pickled all the time (to send them to
+    `to_torch(num_workers=...)` worker processes for example). Creating the pool when examples are actually read
+    keeps the dataset picklable: every process that iterates it creates its own pool.
+    """
+
+    def __init__(self, max_workers: int):
+        self.max_workers = max_workers
+        self._pool: Optional[concurrent.futures.ThreadPoolExecutor] = None
+        self._lock = threading.Lock()
+
+    def submit(self, fn, *args, **kwargs):
+        with self._lock:
+            if self._pool is None:
+                self._pool = concurrent.futures.ThreadPoolExecutor(
+                    self.max_workers, thread_name_prefix="datasets_decode"
+                )
+        return self._pool.submit(fn, *args, **kwargs)
+
+    def shutdown(self, wait: bool = True, *, cancel_futures: bool = False):
+        with self._lock:
+            pool, self._pool = self._pool, None
+        if pool is not None:
+            pool.shutdown(wait=wait, cancel_futures=cancel_futures)
+
+    def __getstate__(self):
+        return {"max_workers": self.max_workers}
+
+    def __setstate__(self, state):
+        self.max_workers = state["max_workers"]
+        self._pool = None
+        self._lock = threading.Lock()
+
+
+async def _apply_async(pool: "_LazyThreadPoolExecutor", func, x):
+    """Run `func(x)` in a worker thread and await its result without blocking the event loop.
+
+    The worker future is wrapped rather than polled in a `while not future.ready(): await asyncio.sleep(0)` loop:
+    spinning wakes up the event loop thousands of times per second, which burns CPU and starves the worker threads
+    that do the actual work (mostly network reads), making threaded decoding slower than plain sequential decoding.
+    """
+
+    return await asyncio.wrap_future(pool.submit(func, x))
 
 
 def _batch_fn(unbatched):

{datasets-5.1.0 → datasets-5.2.0}/src/datasets/packaged_modules/folder_based_builder/folder_based_builder.py

@@ -34,6 +34,7 @@ class FolderBasedBuilderConfig(datasets.BuilderConfig):
     drop_metadata: bool = None
     metadata_filenames: list[str] = None
     filters: Optional[Union[ds.Expression, list[tuple], list[list[tuple]]]] = None
+    decode_num_threads: int = 32
 
     def __post_init__(self):
         super().__post_init__()