chore: import upstream snapshot with attribution

src/datasets/__init__.py

@@ -0,0 +1,47 @@
+# Copyright 2020 The HuggingFace Datasets Authors and the TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+__version__ = "5.0.1.dev0"
+
+from .arrow_dataset import Column, Dataset
+from .arrow_reader import ReadInstruction
+from .builder import ArrowBasedBuilder, BuilderConfig, DatasetBuilder, GeneratorBasedBuilder
+from .combine import concatenate_datasets, interleave_datasets
+from .dataset_dict import DatasetDict, IterableDatasetDict
+from .download import *
+from .features import *
+from .fingerprint import disable_caching, enable_caching, is_caching_enabled
+from .info import DatasetInfo
+from .inspect import (
+    get_dataset_config_info,
+    get_dataset_config_names,
+    get_dataset_default_config_name,
+    get_dataset_infos,
+    get_dataset_split_names,
+)
+from .iterable_dataset import IterableColumn, IterableDataset
+from .load import load_dataset, load_dataset_builder, load_from_disk
+from .splits import (
+    NamedSplit,
+    NamedSplitAll,
+    Split,
+    SplitBase,
+    SplitDict,
+    SplitGenerator,
+    SplitInfo,
+    SubSplitInfo,
+    percent,
+)
+from .utils import *
+from .utils import logging

src/datasets/arrow_reader.py

@@ -0,0 +1,620 @@
+# Copyright 2020 The HuggingFace Datasets Authors and the TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Lint as: python3
+"""Arrow ArrowReader."""
+
+import copy
+import math
+import os
+import re
+from dataclasses import dataclass
+from functools import partial
+from typing import TYPE_CHECKING, Optional, Union
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+from tqdm.contrib.concurrent import thread_map
+
+from .download.download_config import DownloadConfig  # noqa: F401
+from .naming import _split_re, filenames_for_dataset_split
+from .table import InMemoryTable, MemoryMappedTable, Table, concat_tables
+from .utils import logging
+from .utils import tqdm as hf_tqdm
+
+
+if TYPE_CHECKING:
+    from .info import DatasetInfo  # noqa: F401
+    from .splits import Split, SplitInfo  # noqa: F401
+
+
+logger = logging.get_logger(__name__)
+
+HF_GCP_BASE_URL = "https://storage.googleapis.com/huggingface-nlp/cache/datasets"
+
+_SUB_SPEC_RE = re.compile(
+    rf"""
+^
+ (?P<split>{_split_re[1:-1]})
+ (\[
+    ((?P<from>-?[\d_]+)
+     (?P<from_pct>%)?)?
+    :
+    ((?P<to>-?[\d_]+)
+     (?P<to_pct>%)?)?
+ \])?(\((?P<rounding>[^\)]*)\))?
+$
+""",  # remove ^ and $
+    re.X,
+)
+
+_ADDITION_SEP_RE = re.compile(r"\s*\+\s*")
+
+
+class DatasetNotOnHfGcsError(ConnectionError):
+    """When you can't get the dataset from the Hf google cloud storage"""
+
+    pass
+
+
+class MissingFilesOnHfGcsError(ConnectionError):
+    """When some files are missing on the Hf oogle cloud storage"""
+
+    pass
+
+
+@dataclass(frozen=True)
+class FileInstructions:
+    """The file instructions associated with a split ReadInstruction.
+
+    Attributes:
+        num_examples: `int`, The total number of examples
+        file_instructions: List[dict(filename, skip, take)], the files information.
+            The filenames contains the relative path, not absolute.
+            skip/take indicates which example read in the file: `ds.slice(skip, take)`
+    """
+
+    num_examples: int
+    file_instructions: list[dict]
+
+
+def make_file_instructions(
+    name: str,
+    split_infos: list["SplitInfo"],
+    instruction: Union[str, "ReadInstruction"],
+    filetype_suffix: Optional[str] = None,
+    prefix_path: Optional[str] = None,
+) -> FileInstructions:
+    """Returns instructions of the split dict.
+
+    Args:
+        name (`str`): Name of the dataset.
+        split_infos (`list` of `[SplitInfo]`): Dataset splits information.
+        instruction ([`ReadInstruction`] or `str`): Reading instruction for a dataset.
+        filetype_suffix (`str`, *optional*): Suffix of dataset files, e.g. 'arrow' or 'parquet'.
+        prefix_path (`str`, *optional*): Prefix of dataset files, e.g. directory name.
+
+    Returns:
+        [`FileInstructions`]
+    """
+    if not isinstance(name, str):
+        raise TypeError(f"Expected str 'name', but got: {type(name).__name__}")
+    elif not name:
+        raise ValueError("Expected non-empty str 'name'")
+    name2len = {info.name: info.num_examples for info in split_infos}
+    name2shard_lengths = {info.name: info.shard_lengths for info in split_infos}
+    name2filenames = {
+        info.name: filenames_for_dataset_split(
+            path=prefix_path,
+            dataset_name=name,
+            split=info.name,
+            filetype_suffix=filetype_suffix,
+            shard_lengths=name2shard_lengths[info.name],
+        )
+        for info in split_infos
+    }
+    if not isinstance(instruction, ReadInstruction):
+        instruction = ReadInstruction.from_spec(instruction)
+    # Create the absolute instruction (per split)
+    absolute_instructions = instruction.to_absolute(name2len)
+
+    # For each split, return the files instruction (skip/take)
+    file_instructions = []
+    num_examples = 0
+    for abs_instr in absolute_instructions:
+        split_length = name2len[abs_instr.splitname]
+        filenames = name2filenames[abs_instr.splitname]
+        shard_lengths = name2shard_lengths[abs_instr.splitname]
+        from_ = 0 if abs_instr.from_ is None else abs_instr.from_
+        to = split_length if abs_instr.to is None else abs_instr.to
+        if shard_lengths is None:  # not sharded
+            for filename in filenames:
+                take = to - from_
+                if take == 0:
+                    continue
+                num_examples += take
+                file_instructions.append({"filename": filename, "skip": from_, "take": take})
+        else:  # sharded
+            index_start = 0  # Beginning (included) of moving window.
+            index_end = 0  # End (excluded) of moving window.
+            for filename, shard_length in zip(filenames, shard_lengths):
+                index_end += shard_length
+                if from_ < index_end and to > index_start:  # There is something to take.
+                    skip = from_ - index_start if from_ > index_start else 0
+                    take = to - index_start - skip if to < index_end else -1
+                    if take == 0:
+                        continue
+                    file_instructions.append({"filename": filename, "skip": skip, "take": take})
+                    num_examples += shard_length - skip if take == -1 else take
+                index_start += shard_length
+    return FileInstructions(
+        num_examples=num_examples,
+        file_instructions=file_instructions,
+    )
+
+
+class BaseReader:
+    """
+    Build a Dataset object out of Instruction instance(s).
+    """
+
+    def __init__(self, path: str, info: Optional["DatasetInfo"]):
+        """Initializes ArrowReader.
+
+        Args:
+            path (str): path where tfrecords are stored.
+            info (DatasetInfo): info about the dataset.
+        """
+        self._path: str = path
+        self._info: Optional["DatasetInfo"] = info
+        self._filetype_suffix: Optional[str] = None
+
+    def _get_table_from_filename(self, filename_skip_take, in_memory=False) -> Table:
+        """Returns a Dataset instance from given (filename, skip, take)."""
+        raise NotImplementedError
+
+    def _read_files(self, files, in_memory=False) -> Table:
+        """Returns Dataset for given file instructions.
+
+        Args:
+            files: List[dict(filename, skip, take)], the files information.
+                The filenames contain the absolute path, not relative.
+                skip/take indicates which example read in the file: `ds.slice(skip, take)`
+            in_memory (bool, default False): Whether to copy the data in-memory.
+        """
+        if len(files) == 0 or not all(isinstance(f, dict) for f in files):
+            raise ValueError("please provide valid file informations")
+        files = copy.deepcopy(files)
+        for f in files:
+            f["filename"] = os.path.join(self._path, f["filename"])
+
+        pa_tables = thread_map(
+            partial(self._get_table_from_filename, in_memory=in_memory),
+            files,
+            tqdm_class=hf_tqdm,
+            desc="Loading dataset shards",
+            # set `disable=None` rather than `disable=False` by default to disable progress bar when no TTY attached
+            disable=len(files) <= 16 or None,
+        )
+        pa_tables = [t for t in pa_tables if len(t) > 0]
+        if not pa_tables and (self._info is None or self._info.features is None):
+            raise ValueError(
+                "Tried to read an empty table. Please specify at least info.features to create an empty table with the right type."
+            )
+        pa_tables = pa_tables or [InMemoryTable.from_batches([], schema=pa.schema(self._info.features.type))]
+        pa_table = concat_tables(pa_tables) if len(pa_tables) != 1 else pa_tables[0]
+        return pa_table
+
+    def get_file_instructions(self, name, instruction, split_infos):
+        """Return list of dict {'filename': str, 'skip': int, 'take': int}"""
+        file_instructions = make_file_instructions(
+            name, split_infos, instruction, filetype_suffix=self._filetype_suffix, prefix_path=self._path
+        )
+        files = file_instructions.file_instructions
+        return files
+
+    def read(
+        self,
+        name,
+        instructions,
+        split_infos,
+        in_memory=False,
+    ):
+        """Returns Dataset instance(s).
+
+        Args:
+            name (str): name of the dataset.
+            instructions (ReadInstruction): instructions to read.
+                Instruction can be string and will then be passed to the Instruction
+                constructor as it.
+            split_infos (list of SplitInfo proto): the available splits for dataset.
+            in_memory (bool, default False): Whether to copy the data in-memory.
+
+        Returns:
+             kwargs to build a single Dataset instance.
+        """
+
+        files = self.get_file_instructions(name, instructions, split_infos)
+        if not files:
+            msg = f'Instruction "{instructions}" corresponds to no data!'
+            raise ValueError(msg)
+        return self.read_files(files=files, original_instructions=instructions, in_memory=in_memory)
+
+    def read_files(
+        self,
+        files: list[dict],
+        original_instructions: Union[None, "ReadInstruction", "Split"] = None,
+        in_memory=False,
+    ):
+        """Returns single Dataset instance for the set of file instructions.
+
+        Args:
+            files: List[dict(filename, skip, take)], the files information.
+                The filenames contains the relative path, not absolute.
+                skip/take indicates which example read in the file: `ds.skip().take()`
+            original_instructions: store the original instructions used to build the dataset split in the dataset.
+            in_memory (bool, default False): Whether to copy the data in-memory.
+
+        Returns:
+            kwargs to build a Dataset instance.
+        """
+        # Prepend path to filename
+        pa_table = self._read_files(files, in_memory=in_memory)
+        # If original_instructions is not None, convert it to a human-readable NamedSplit
+        if original_instructions is not None:
+            from .splits import Split  # noqa
+
+            split = Split(str(original_instructions))
+        else:
+            split = None
+        dataset_kwargs = {"arrow_table": pa_table, "info": self._info, "split": split}
+        return dataset_kwargs
+
+
+class ArrowReader(BaseReader):
+    """
+    Build a Dataset object out of Instruction instance(s).
+    This Reader uses either memory mapping or file descriptors (in-memory) on arrow files.
+    """
+
+    def __init__(self, path: str, info: Optional["DatasetInfo"]):
+        """Initializes ArrowReader.
+
+        Args:
+            path (str): path where Arrow files are stored.
+            info (DatasetInfo): info about the dataset.
+        """
+        super().__init__(path, info)
+        self._filetype_suffix = "arrow"
+
+    def _get_table_from_filename(self, filename_skip_take, in_memory=False) -> Table:
+        """Returns a Dataset instance from given (filename, skip, take)."""
+        filename, skip, take = (
+            filename_skip_take["filename"],
+            filename_skip_take["skip"] if "skip" in filename_skip_take else None,
+            filename_skip_take["take"] if "take" in filename_skip_take else None,
+        )
+        table = ArrowReader.read_table(filename, in_memory=in_memory)
+        if take == -1:
+            take = len(table) - skip
+        # here we don't want to slice an empty table, or it may segfault
+        if skip is not None and take is not None and not (skip == 0 and take == len(table)):
+            table = table.slice(skip, take)
+        return table
+
+    @staticmethod
+    def read_table(filename, in_memory=False) -> Table:
+        """
+        Read table from file.
+
+        Args:
+            filename (str): File name of the table.
+            in_memory (bool, default=False): Whether to copy the data in-memory.
+
+        Returns:
+            pyarrow.Table
+        """
+        table_cls = InMemoryTable if in_memory else MemoryMappedTable
+        return table_cls.from_file(filename)
+
+
+class ParquetReader(BaseReader):
+    """
+    Build a Dataset object out of Instruction instance(s).
+    This Reader uses memory mapping on parquet files.
+    """
+
+    def __init__(self, path: str, info: Optional["DatasetInfo"]):
+        """Initializes ParquetReader.
+
+        Args:
+            path (str): path where tfrecords are stored.
+            info (DatasetInfo): info about the dataset.
+        """
+        super().__init__(path, info)
+        self._filetype_suffix = "parquet"
+
+    def _get_table_from_filename(self, filename_skip_take, **kwargs):
+        """Returns a Dataset instance from given (filename, skip, take)."""
+        filename, skip, take = (
+            filename_skip_take["filename"],
+            filename_skip_take["skip"] if "skip" in filename_skip_take else None,
+            filename_skip_take["take"] if "take" in filename_skip_take else None,
+        )
+        # Parquet read_table always loads data in memory, independently of memory_map
+        pa_table = pq.read_table(filename, memory_map=True)
+        # here we don't want to slice an empty table, or it may segfault
+        if skip is not None and take is not None and not (skip == 0 and take == len(pa_table)):
+            pa_table = pa_table.slice(skip, take)
+        return pa_table
+
+
+@dataclass(frozen=True)
+class _AbsoluteInstruction:
+    """A machine friendly slice: defined absolute positive boundaries."""
+
+    splitname: str
+    from_: int  # uint (starting index).
+    to: int  # uint (ending index).
+
+
+@dataclass(frozen=True)
+class _RelativeInstruction:
+    """Represents a single parsed slicing instruction, can use % and negatives."""
+
+    splitname: str
+    from_: Optional[int] = None  # int (starting index) or None if no lower boundary.
+    to: Optional[int] = None  # int (ending index) or None if no upper boundary.
+    unit: Optional[str] = None
+    rounding: Optional[str] = None
+
+    def __post_init__(self):
+        if self.unit is not None and self.unit not in ["%", "abs"]:
+            raise ValueError("unit must be either % or abs")
+        if self.rounding is not None and self.rounding not in ["closest", "pct1_dropremainder"]:
+            raise ValueError("rounding must be either closest or pct1_dropremainder")
+        if self.unit != "%" and self.rounding is not None:
+            raise ValueError("It is forbidden to specify rounding if not using percent slicing.")
+        if self.unit == "%" and self.from_ is not None and abs(self.from_) > 100:
+            raise ValueError("Percent slice boundaries must be > -100 and < 100.")
+        if self.unit == "%" and self.to is not None and abs(self.to) > 100:
+            raise ValueError("Percent slice boundaries must be > -100 and < 100.")
+        # Update via __dict__ due to instance being "frozen"
+        self.__dict__["rounding"] = "closest" if self.rounding is None and self.unit == "%" else self.rounding
+
+
+def _str_to_read_instruction(spec):
+    """Returns ReadInstruction for given string."""
+    res = _SUB_SPEC_RE.match(spec)
+    if not res:
+        raise ValueError(f"Unrecognized instruction format: {spec}")
+    unit = "%" if res.group("from_pct") or res.group("to_pct") else "abs"
+    return ReadInstruction(
+        split_name=res.group("split"),
+        rounding=res.group("rounding"),
+        from_=int(res.group("from")) if res.group("from") else None,
+        to=int(res.group("to")) if res.group("to") else None,
+        unit=unit,
+    )
+
+
+def _pct_to_abs_pct1(boundary, num_examples):
+    # Using math.trunc here, since -99.5% should give -99%, not -100%.
+    if num_examples < 100:
+        msg = (
+            'Using "pct1_dropremainder" rounding on a split with less than 100 '
+            "elements is forbidden: it always results in an empty dataset."
+        )
+        raise ValueError(msg)
+    return boundary * math.trunc(num_examples / 100.0)
+
+
+def _pct_to_abs_closest(boundary, num_examples):
+    return int(round(boundary * num_examples / 100.0))
+
+
+def _rel_to_abs_instr(rel_instr, name2len):
+    """Returns _AbsoluteInstruction instance for given RelativeInstruction.
+
+    Args:
+        rel_instr: RelativeInstruction instance.
+        name2len: dict {split_name: num_examples}.
+    """
+    pct_to_abs = _pct_to_abs_closest if rel_instr.rounding == "closest" else _pct_to_abs_pct1
+    split = rel_instr.splitname
+    if split not in name2len:
+        raise ValueError(f'Unknown split "{split}". Should be one of {list(name2len)}.')
+    num_examples = name2len[split]
+    from_ = rel_instr.from_
+    to = rel_instr.to
+    if rel_instr.unit == "%":
+        from_ = 0 if from_ is None else pct_to_abs(from_, num_examples)
+        to = num_examples if to is None else pct_to_abs(to, num_examples)
+    else:
+        from_ = 0 if from_ is None else from_
+        to = num_examples if to is None else to
+    if from_ < 0:
+        from_ = max(num_examples + from_, 0)
+    if to < 0:
+        to = max(num_examples + to, 0)
+    from_ = min(from_, num_examples)
+    to = min(to, num_examples)
+    return _AbsoluteInstruction(split, from_, to)
+
+
+class ReadInstruction:
+    """Reading instruction for a dataset.
+
+    Examples::
+
+      # The following lines are equivalent:
+      ds = datasets.load_dataset('ylecun/mnist', split='test[:33%]')
+      ds = datasets.load_dataset('ylecun/mnist', split=datasets.ReadInstruction.from_spec('test[:33%]'))
+      ds = datasets.load_dataset('ylecun/mnist', split=datasets.ReadInstruction('test', to=33, unit='%'))
+      ds = datasets.load_dataset('ylecun/mnist', split=datasets.ReadInstruction(
+          'test', from_=0, to=33, unit='%'))
+
+      # The following lines are equivalent:
+      ds = datasets.load_dataset('ylecun/mnist', split='test[:33%]+train[1:-1]')
+      ds = datasets.load_dataset('ylecun/mnist', split=datasets.ReadInstruction.from_spec(
+          'test[:33%]+train[1:-1]'))
+      ds = datasets.load_dataset('ylecun/mnist', split=(
+          datasets.ReadInstruction('test', to=33, unit='%') +
+          datasets.ReadInstruction('train', from_=1, to=-1, unit='abs')))
+
+      # The following lines are equivalent:
+      ds = datasets.load_dataset('ylecun/mnist', split='test[:33%](pct1_dropremainder)')
+      ds = datasets.load_dataset('ylecun/mnist', split=datasets.ReadInstruction.from_spec(
+          'test[:33%](pct1_dropremainder)'))
+      ds = datasets.load_dataset('ylecun/mnist', split=datasets.ReadInstruction(
+          'test', from_=0, to=33, unit='%', rounding="pct1_dropremainder"))
+
+      # 10-fold validation:
+      tests = datasets.load_dataset(
+          'ylecun/mnist',
+          [datasets.ReadInstruction('train', from_=k, to=k+10, unit='%')
+          for k in range(0, 100, 10)])
+      trains = datasets.load_dataset(
+          'ylecun/mnist',
+          [datasets.ReadInstruction('train', to=k, unit='%') + datasets.ReadInstruction('train', from_=k+10, unit='%')
+          for k in range(0, 100, 10)])
+
+    """
+
+    def _init(self, relative_instructions):
+        # Private initializer.
+        self._relative_instructions = relative_instructions
+
+    @classmethod
+    def _read_instruction_from_relative_instructions(cls, relative_instructions):
+        """Returns ReadInstruction obj initialized with relative_instructions."""
+        # Use __new__ to bypass __init__ used by public API and not conveniant here.
+        result = cls.__new__(cls)
+        result._init(relative_instructions)  # pylint: disable=protected-access
+        return result
+
+    def __init__(self, split_name, rounding=None, from_=None, to=None, unit=None):
+        """Initialize ReadInstruction.
+
+        Args:
+            split_name (str): name of the split to read. Eg: 'train'.
+            rounding (str, optional): The rounding behaviour to use when percent slicing is
+                used. Ignored when slicing with absolute indices.
+                Possible values:
+                 - 'closest' (default): The specified percentages are rounded to the
+                     closest value. Use this if you want specified percents to be as
+                     much exact as possible.
+                 - 'pct1_dropremainder': the specified percentages are treated as
+                     multiple of 1%. Use this option if you want consistency. Eg:
+                         len(5%) == 5 * len(1%).
+                     Using this option, one might not be able to use the full set of
+                     examples, if the number of those is not a multiple of 100.
+            from_ (int):
+            to (int): alternative way of specifying slicing boundaries. If any of
+                {from_, to, unit} argument is used, slicing cannot be specified as
+                string.
+            unit (str): optional, one of:
+                '%': to set the slicing unit as percents of the split size.
+                'abs': to set the slicing unit as absolute numbers.
+        """
+        # This constructor is not always called. See factory method
+        # `_read_instruction_from_relative_instructions`. Common init instructions
+        # MUST be placed in the _init method.
+        self._init([_RelativeInstruction(split_name, from_, to, unit, rounding)])
+
+    @classmethod
+    def from_spec(cls, spec):
+        """Creates a `ReadInstruction` instance out of a string spec.
+
+        Args:
+            spec (`str`):
+                Split(s) + optional slice(s) to read + optional rounding
+                if percents are used as the slicing unit. A slice can be specified,
+                using absolute numbers (`int`) or percentages (`int`).
+
+        Examples:
+
+            ```
+            test: test split.
+            test + validation: test split + validation split.
+            test[10:]: test split, minus its first 10 records.
+            test[:10%]: first 10% records of test split.
+            test[:20%](pct1_dropremainder): first 10% records, rounded with the pct1_dropremainder rounding.
+            test[:-5%]+train[40%:60%]: first 95% of test + middle 20% of train.
+            ```
+
+        Returns:
+            ReadInstruction instance.
+        """
+        spec = str(spec)  # Need to convert to str in case of NamedSplit instance.
+        subs = _ADDITION_SEP_RE.split(spec)
+        if not subs:
+            raise ValueError(f"No instructions could be built out of {spec}")
+        instruction = _str_to_read_instruction(subs[0])
+        return sum((_str_to_read_instruction(sub) for sub in subs[1:]), instruction)
+
+    def to_spec(self):
+        rel_instr_specs = []
+        for rel_instr in self._relative_instructions:
+            rel_instr_spec = rel_instr.splitname
+            if rel_instr.from_ is not None or rel_instr.to is not None:
+                from_ = rel_instr.from_
+                to = rel_instr.to
+                unit = rel_instr.unit
+                rounding = rel_instr.rounding
+                unit = unit if unit == "%" else ""
+                from_ = str(from_) + unit if from_ is not None else ""
+                to = str(to) + unit if to is not None else ""
+                slice_str = f"[{from_}:{to}]"
+                rounding_str = (
+                    f"({rounding})" if unit == "%" and rounding is not None and rounding != "closest" else ""
+                )
+                rel_instr_spec += slice_str + rounding_str
+            rel_instr_specs.append(rel_instr_spec)
+        return "+".join(rel_instr_specs)
+
+    def __add__(self, other):
+        """Returns a new ReadInstruction obj, result of appending other to self."""
+        if not isinstance(other, ReadInstruction):
+            msg = "ReadInstruction can only be added to another ReadInstruction obj."
+            raise TypeError(msg)
+        self_ris = self._relative_instructions
+        other_ris = other._relative_instructions  # pylint: disable=protected-access
+        if (
+            self_ris[0].unit != "abs"
+            and other_ris[0].unit != "abs"
+            and self._relative_instructions[0].rounding != other_ris[0].rounding
+        ):
+            raise ValueError("It is forbidden to sum ReadInstruction instances with different rounding values.")
+        return self._read_instruction_from_relative_instructions(self_ris + other_ris)
+
+    def __str__(self):
+        return self.to_spec()
+
+    def __repr__(self):
+        return f"ReadInstruction({self._relative_instructions})"
+
+    def to_absolute(self, name2len):
+        """Translate instruction into a list of absolute instructions.
+
+        Those absolute instructions are then to be added together.
+
+        Args:
+            name2len (`dict`):
+                Associating split names to number of examples.
+
+        Returns:
+            list of _AbsoluteInstruction instances (corresponds to the + in spec).
+        """
+        return [_rel_to_abs_instr(rel_instr, name2len) for rel_instr in self._relative_instructions]

src/datasets/arrow_writer.py

@@ -0,0 +1,828 @@
+# Copyright 2020 The HuggingFace Datasets Authors and the TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Lint as: python3
+"""To write records into Parquet files."""
+
+import io
+import json
+import sys
+from collections.abc import Iterable
+from typing import Any, Literal, Optional
+
+import fsspec
+import numpy as np
+import pyarrow as pa
+import pyarrow.json as paj
+import pyarrow.parquet as pq
+from fsspec.core import url_to_fs
+
+from . import config
+from .features import Audio, Features, Image, Pdf, Value, Video
+from .features.features import (
+    FeatureType,
+    List,
+    _ArrayXDExtensionType,
+    _visit,
+    cast_to_python_objects,
+    generate_from_arrow_type,
+    get_nested_type,
+    list_of_np_array_to_pyarrow_listarray,
+    numpy_to_pyarrow_listarray,
+    require_storage_embed,
+    to_pyarrow_listarray,
+)
+from .filesystems import is_remote_filesystem
+from .info import DatasetInfo
+from .table import array_cast, cast_array_to_feature, embed_table_storage, table_cast
+from .utils import logging
+from .utils.json import (
+    find_mixed_struct_types_field_paths,
+    get_json_field_path_from_pyarrow_json_error,
+    get_json_field_paths_from_feature,
+    insert_json_field_path,
+    json_encode_field,
+    json_encode_fields_in_json_lines,
+    set_json_types_in_feature,
+    ujson_dumps,
+)
+from .utils.py_utils import asdict, convert_file_size_to_int, first_non_null_non_empty_value
+
+
+logger = logging.get_logger(__name__)
+
+type_ = type  # keep python's type function
+
+
+def get_arrow_writer_batch_size_from_features(features: Optional[Features]) -> Optional[int]:
+    """
+    Get the writer_batch_size that defines the maximum record batch size in the arrow files based on configuration values.
+    The default value is 100 for image/audio datasets and 10 for videos.
+    This allows to avoid overflows in arrow buffers.
+
+    Args:
+        features (`datasets.Features` or `None`):
+            Dataset Features from `datasets`.
+    Returns:
+        writer_batch_size (`Optional[int]`):
+            Writer batch size to pass to a dataset builder.
+            If `None`, then it will use the `datasets` default, i.e. `datasets.config.DEFAULT_MAX_BATCH_SIZE`.
+    """
+    if not features:
+        return None
+
+    batch_size = np.inf
+
+    def set_batch_size(feature: FeatureType) -> None:
+        nonlocal batch_size
+        if isinstance(feature, Image) and config.ARROW_RECORD_BATCH_SIZE_FOR_IMAGE_DATASETS is not None:
+            batch_size = min(batch_size, config.ARROW_RECORD_BATCH_SIZE_FOR_IMAGE_DATASETS)
+        elif isinstance(feature, Audio) and config.ARROW_RECORD_BATCH_SIZE_FOR_AUDIO_DATASETS is not None:
+            batch_size = min(batch_size, config.ARROW_RECORD_BATCH_SIZE_FOR_AUDIO_DATASETS)
+        elif isinstance(feature, Video) and config.ARROW_RECORD_BATCH_SIZE_FOR_VIDEO_DATASETS is not None:
+            batch_size = min(batch_size, config.ARROW_RECORD_BATCH_SIZE_FOR_VIDEO_DATASETS)
+        elif (
+            isinstance(feature, Value)
+            and feature.dtype == "binary"
+            and config.ARROW_RECORD_BATCH_SIZE_FOR_BINARY_DATASETS is not None
+        ):
+            batch_size = min(batch_size, config.ARROW_RECORD_BATCH_SIZE_FOR_BINARY_DATASETS)
+
+    _visit(features, set_batch_size)
+
+    return None if batch_size is np.inf else batch_size
+
+
+def get_writer_batch_size_from_features(features: Optional[Features]) -> Optional[int]:
+    """
+    Get the writer_batch_size that defines the maximum row group size in the parquet files based on configuration values.
+    By default these are not set, but it can be helpful to hard set those values in some cases.
+    This allows to optimize random access to parquet file, since accessing 1 row requires
+    to read its entire row group.
+
+    Args:
+        features (`datasets.Features` or `None`):
+            Dataset Features from `datasets`.
+    Returns:
+        writer_batch_size (`Optional[int]`):
+            Writer batch size to pass to a parquet writer.
+            If `None`, then it will use the `datasets` default, i.e. aiming for row groups of 100MB.
+    """
+    if not features:
+        return None
+
+    batch_size = np.inf
+
+    def set_batch_size(feature: FeatureType) -> None:
+        nonlocal batch_size
+        if isinstance(feature, Image) and config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS is not None:
+            batch_size = min(batch_size, config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS)
+        elif isinstance(feature, Audio) and config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS is not None:
+            batch_size = min(batch_size, config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS)
+        elif isinstance(feature, Video) and config.PARQUET_ROW_GROUP_SIZE_FOR_VIDEO_DATASETS is not None:
+            batch_size = min(batch_size, config.PARQUET_ROW_GROUP_SIZE_FOR_VIDEO_DATASETS)
+        elif (
+            isinstance(feature, Value)
+            and feature.dtype == "binary"
+            and config.PARQUET_ROW_GROUP_SIZE_FOR_BINARY_DATASETS is not None
+        ):
+            batch_size = min(batch_size, config.PARQUET_ROW_GROUP_SIZE_FOR_BINARY_DATASETS)
+
+    _visit(features, set_batch_size)
+
+    return None if batch_size is np.inf else batch_size
+
+
+def get_writer_batch_size_from_data_size(num_rows: int, num_bytes: int) -> int:
+    """
+    Get the writer_batch_size that defines the maximum row group size in the parquet files.
+    The default in `datasets` is aiming for row groups of maximum 100MB uncompressed.
+    This allows to optimize random access to parquet file, since accessing 1 row requires
+    to read its entire row group.
+
+    This can be improved to get optimized size for querying/iterating
+    but at least it matches the dataset viewer expectations on HF.
+
+    Args:
+        num_rows (`int`):
+            Number of rows in the dataset.
+        num_bytes (`int`):
+            Number of bytes in the dataset.
+            For dataset with external files to embed (image, audio, videos), this can also be an
+            estimate from `dataset._estimate_nbytes()`.
+    Returns:
+        writer_batch_size (`Optional[int]`):
+            Writer batch size to pass to a parquet writer.
+    """
+    return max(1, num_rows * convert_file_size_to_int(config.MAX_ROW_GROUP_SIZE) // num_bytes) if num_bytes > 0 else 1
+
+
+class SchemaInferenceError(ValueError):
+    pass
+
+
+class TypedSequence:
+    """
+    This data container generalizes the typing when instantiating pyarrow arrays, tables or batches.
+
+    More specifically it adds several features:
+    - Support extension types like ``datasets.features.Array2DExtensionType``:
+        By default pyarrow arrays don't return extension arrays. One has to call
+        ``pa.ExtensionArray.from_storage(type, pa.array(data, type.storage_type))``
+        in order to get an extension array.
+    - Support for ``try_type`` parameter that can be used instead of ``type``:
+        When an array is transformed, we like to keep the same type as before if possible.
+        For example when calling :func:`datasets.Dataset.map`, we don't want to change the type
+        of each column by default.
+    - Better error message when a pyarrow array overflows.
+
+    Example::
+
+        from datasets.features import Array2D, Array2DExtensionType, Value
+        from datasets.arrow_writer import TypedSequence
+        import pyarrow as pa
+
+        arr = pa.array(TypedSequence([1, 2, 3], type=Value("int32")))
+        assert arr.type == pa.int32()
+
+        arr = pa.array(TypedSequence([1, 2, 3], try_type=Value("int32")))
+        assert arr.type == pa.int32()
+
+        arr = pa.array(TypedSequence(["foo", "bar"], try_type=Value("int32")))
+        assert arr.type == pa.string()
+
+        arr = pa.array(TypedSequence([[[1, 2, 3]]], type=Array2D((1, 3), "int64")))
+        assert arr.type == Array2DExtensionType((1, 3), "int64")
+
+        table = pa.Table.from_pydict({
+            "image": TypedSequence([[[1, 2, 3]]], type=Array2D((1, 3), "int64"))
+        })
+        assert table["image"].type == Array2DExtensionType((1, 3), "int64")
+
+    """
+
+    def __init__(
+        self,
+        data: Iterable,
+        type: Optional[FeatureType] = None,
+        try_type: Optional[FeatureType] = None,
+        optimized_int_type: Optional[FeatureType] = None,
+        on_mixed_types: Optional[Literal["use_json"]] = None,
+    ):
+        # assert type is None or try_type is None,
+        if type is not None and try_type is not None:
+            raise ValueError("You cannot specify both type and try_type")
+        # set attributes
+        self.data = data
+        self.type = type
+        self.try_type = try_type  # is ignored if it doesn't match the data
+        self.optimized_int_type = optimized_int_type
+        self.on_mixed_types = on_mixed_types
+        # when trying a type (is ignored if data is not compatible)
+        self.trying_type = self.try_type is not None
+        self.trying_int_optimization = optimized_int_type is not None and type is None and try_type is None
+        # used to get back the inferred type after __arrow_array__() is called once
+        self._inferred_type = None
+
+    def get_inferred_type(self) -> FeatureType:
+        """Return the inferred feature type.
+        This is done by converting the sequence to an Arrow array, and getting the corresponding
+        feature type.
+
+        Since building the Arrow array can be expensive, the value of the inferred type is cached
+        as soon as pa.array is called on the typed sequence.
+
+        Returns:
+            FeatureType: inferred feature type of the sequence.
+        """
+        if self._inferred_type is None:
+            pa.array(self)
+        return self._inferred_type
+
+    @staticmethod
+    def _infer_custom_type_and_encode(data: Iterable) -> tuple[Iterable, Optional[FeatureType]]:
+        """Implement type inference for custom objects like PIL.Image.Image -> Image type.
+
+        This function is only used for custom python objects that can't be directly passed to build
+        an Arrow array. In such cases is infers the feature type to use, and it encodes the data so
+        that they can be passed to an Arrow array.
+
+        Args:
+            data (Iterable): array of data to infer the type, e.g. a list of PIL images.
+
+        Returns:
+            Tuple[Iterable, Optional[FeatureType]]: a tuple with:
+                - the (possibly encoded) array, if the inferred feature type requires encoding
+                - the inferred feature type if the array is made of supported custom objects like
+                    PIL images, else None.
+        """
+        if config.PIL_AVAILABLE and "PIL" in sys.modules:
+            import PIL.Image
+
+            non_null_idx, non_null_value = first_non_null_non_empty_value(data)
+            if isinstance(non_null_value, PIL.Image.Image):
+                return [Image().encode_example(value) if value is not None else None for value in data], Image()
+            if isinstance(non_null_value, list) and isinstance(non_null_value[0], PIL.Image.Image):
+                return [
+                    [Image().encode_example(x) for x in value] if value is not None else None for value in data
+                ], List(Image())
+        if config.PDFPLUMBER_AVAILABLE and "pdfplumber" in sys.modules:
+            import pdfplumber
+
+            non_null_idx, non_null_value = first_non_null_non_empty_value(data)
+            if isinstance(non_null_value, pdfplumber.pdf.PDF):
+                return [Pdf().encode_example(value) if value is not None else None for value in data], Pdf()
+            if isinstance(non_null_value, list) and isinstance(non_null_value[0], pdfplumber.pdf.PDF):
+                return [
+                    [Pdf().encode_example(x) for x in value] if value is not None else None for value in data
+                ], List(Pdf())
+        return data, None
+
+    def __arrow_array__(self, type: Optional[pa.DataType] = None):
+        out = self._arrow_array(type=type)
+        if self._inferred_type is None:
+            self._inferred_type = generate_from_arrow_type(out.type)
+        return out
+
+    def _arrow_array(self, type: Optional[pa.DataType] = None):
+        """This function is called when calling pa.array(typed_sequence)"""
+
+        if type is not None:
+            raise ValueError("TypedSequence is supposed to be used with pa.array(typed_sequence, type=None)")
+        del type  # make sure we don't use it
+        data = self.data
+        # automatic type inference for custom objects
+        if self.type is None and self.try_type is None:
+            data, self._inferred_type = self._infer_custom_type_and_encode(data)
+        if self._inferred_type is None:
+            type = self.try_type if self.trying_type else self.type
+        else:
+            type = self._inferred_type
+        pa_type = get_nested_type(type) if type is not None else None
+        optimized_int_pa_type = (
+            get_nested_type(self.optimized_int_type) if self.optimized_int_type is not None else None
+        )
+        trying_cast_to_python_objects = False
+        json_field_paths = []
+        try:
+            # custom pyarrow types
+            if isinstance(pa_type, _ArrayXDExtensionType):
+                storage = to_pyarrow_listarray(data, pa_type)
+                return pa.ExtensionArray.from_storage(pa_type, storage)
+
+            # efficient np array to pyarrow array
+            if isinstance(data, np.ndarray):
+                out = numpy_to_pyarrow_listarray(data)
+            elif isinstance(data, list) and data and isinstance(first_non_null_non_empty_value(data)[1], np.ndarray):
+                out = list_of_np_array_to_pyarrow_listarray(data)
+            else:
+                trying_cast_to_python_objects = True
+                examples = data
+                # find fields to json-encode
+                if self.on_mixed_types == "use_json" and type is None:
+                    json_field_paths = find_mixed_struct_types_field_paths(examples, allow_root=True)
+                elif type is not None:
+                    json_field_paths = get_json_field_paths_from_feature(type)
+                # json encode if needed
+                if json_field_paths:
+                    for json_field_path in json_field_paths:
+                        examples = [json_encode_field(examples, json_field_path) for examples in examples]
+                # to arrow array
+                out = pa.array(cast_to_python_objects(examples, only_1d_for_numpy=True))
+                # cast to json type if needed
+                if json_field_paths:
+                    pa_table = pa.Table.from_arrays([out], names=["obj"])
+                    features = Features.from_arrow_schema(pa_table.schema)
+                    feature = set_json_types_in_feature(features["obj"], json_field_paths)
+                    pa_table = table_cast(pa_table, Features({"obj": feature}).arrow_schema)
+                    out = pa_table[0]  # get the "obj" column
+            # use smaller integer precisions if possible
+            if self.trying_int_optimization:
+                if pa.types.is_int64(out.type):
+                    out = out.cast(optimized_int_pa_type)
+                elif pa.types.is_list(out.type):
+                    if pa.types.is_int64(out.type.value_type):
+                        out = array_cast(out, pa.list_(optimized_int_pa_type))
+                    elif pa.types.is_list(out.type.value_type) and pa.types.is_int64(out.type.value_type.value_type):
+                        out = array_cast(out, pa.list_(pa.list_(optimized_int_pa_type)))
+            # otherwise we can finally use the user's type
+            elif type is not None:
+                # We use cast_array_to_feature to support casting to custom types like Audio and Image
+                # Also, when trying type "string", we don't want to convert integers or floats to "string".
+                # We only do it if trying_type is False - since this is what the user asks for.
+                out = cast_array_to_feature(
+                    out, type, allow_primitive_to_str=not self.trying_type, allow_decimal_to_str=not self.trying_type
+                )
+            return out
+        except (
+            TypeError,
+            pa.lib.ArrowTypeError,
+            pa.lib.ArrowInvalid,
+            pa.lib.ArrowNotImplementedError,
+        ) as e:  # handle type errors and overflows
+            # Ignore ArrowNotImplementedError caused by trying type, otherwise re-raise
+            if not self.trying_type and isinstance(e, pa.lib.ArrowNotImplementedError):
+                raise
+
+            if self.trying_type:
+                try:  # second chance
+                    if isinstance(data, np.ndarray):
+                        return numpy_to_pyarrow_listarray(data)
+                    elif isinstance(data, list) and data and any(isinstance(value, np.ndarray) for value in data):
+                        return list_of_np_array_to_pyarrow_listarray(data)
+                    else:
+                        trying_cast_to_python_objects = True
+                        return pa.array(cast_to_python_objects(data, only_1d_for_numpy=True))
+                except pa.lib.ArrowInvalid as e:
+                    if "overflow" in str(e):
+                        raise OverflowError(
+                            f"There was an overflow with type {type_(data)}. Try to reduce writer_batch_size to have batches smaller than 2GB.\n({e})"
+                        ) from None
+                    elif self.trying_int_optimization and "not in range" in str(e):
+                        optimized_int_pa_type_str = np.dtype(optimized_int_pa_type.to_pandas_dtype()).name
+                        logger.info(
+                            f"Failed to cast a sequence to {optimized_int_pa_type_str}. Falling back to int64."
+                        )
+                        return out
+                    elif trying_cast_to_python_objects and "Could not convert" in str(e):
+                        out = pa.array(
+                            cast_to_python_objects(data, only_1d_for_numpy=True, optimize_list_casting=False)
+                        )
+                        if type is not None:
+                            out = cast_array_to_feature(
+                                out, type, allow_primitive_to_str=True, allow_decimal_to_str=True
+                            )
+                        return out
+                    else:
+                        raise
+            elif "overflow" in str(e):
+                raise OverflowError(
+                    f"There was an overflow with type {type_(data)}. Try to reduce writer_batch_size to have batches smaller than 2GB.\n({e})"
+                ) from None
+            elif self.trying_int_optimization and "not in range" in str(e):
+                optimized_int_pa_type_str = np.dtype(optimized_int_pa_type.to_pandas_dtype()).name
+                logger.info(f"Failed to cast a sequence to {optimized_int_pa_type_str}. Falling back to int64.")
+                return out
+            elif trying_cast_to_python_objects and (
+                "Could not convert" in str(e) or "cannot mix struct and non-struct" in str(e) or "Expected " in str(e)
+            ):
+                try:  # third chance
+                    out = pa.array(cast_to_python_objects(data, only_1d_for_numpy=True, optimize_list_casting=False))
+                except (pa.ArrowInvalid, pa.ArrowTypeError) as ee:
+                    # in case of mixed types, we use the JSON Lines reader in pyarrow to locate them and set them to json fields
+                    if self.on_mixed_types == "use_json" and (
+                        "Could not convert " in str(ee)
+                        or "cannot mix struct and non-struct" in str(ee)
+                        or "Expected " in str(ee)
+                    ):
+                        # we use "obj" to have valid JSON Lines since data may contain lists
+                        original_batch = "\n".join([ujson_dumps({"obj": example}) for example in data]).encode()
+                        json_field_paths = [["obj"] + json_field_path for json_field_path in json_field_paths]
+                        batch = json_encode_fields_in_json_lines(original_batch, json_field_paths)
+                        pa_table = None
+                        while True:
+                            try:  # fourth chance
+                                pa_table = paj.read_json(
+                                    io.BytesIO(batch), read_options=paj.ReadOptions(use_threads=False)
+                                )
+                                break
+                            except pa.ArrowInvalid as eee:
+                                if "JSON parse error: Column(" in str(eee) and ") changed from" in str(eee):
+                                    json_field_path = get_json_field_path_from_pyarrow_json_error(str(eee))
+                                    insert_json_field_path(json_field_paths, json_field_path)
+                                    batch = json_encode_fields_in_json_lines(original_batch, json_field_paths)
+                                else:
+                                    break
+                        if pa_table is not None:
+                            features = Features.from_arrow_schema(pa_table.schema)
+                            features = set_json_types_in_feature(features, json_field_paths)
+                            pa_table = table_cast(pa_table, features.arrow_schema)
+                            out = pa_table[0]  # get the "obj" column
+                            return out
+                        else:
+                            raise
+                    else:
+                        raise
+                if type is not None:
+                    out = cast_array_to_feature(out, type, allow_primitive_to_str=True, allow_decimal_to_str=True)
+                return out
+            else:
+                raise
+
+
+class OptimizedTypedSequence(TypedSequence):
+    def __init__(
+        self,
+        data,
+        type: Optional[FeatureType] = None,
+        try_type: Optional[FeatureType] = None,
+        col: Optional[str] = None,
+        optimized_int_type: Optional[FeatureType] = None,
+        on_mixed_types: Optional[Literal["use_json"]] = None,
+    ):
+        optimized_int_type_by_col = {
+            "attention_mask": Value("int8"),  # binary tensor
+            "special_tokens_mask": Value("int8"),
+            "input_ids": Value("int32"),  # typical vocab size: 0-50k (max ~500k, never > 1M)
+            "token_type_ids": Value(
+                "int8"
+            ),  # binary mask; some (XLNetModel) use an additional token represented by a 2
+        }
+        if type is None and try_type is None:
+            optimized_int_type = optimized_int_type_by_col.get(col, None)
+        super().__init__(
+            data, type=type, try_type=try_type, optimized_int_type=optimized_int_type, on_mixed_types=on_mixed_types
+        )
+
+
+class ArrowWriter:
+    """Shuffles and writes Examples to Arrow files."""
+
+    def __init__(
+        self,
+        schema: Optional[pa.Schema] = None,
+        features: Optional[Features] = None,
+        path: Optional[str] = None,
+        stream: Optional[pa.NativeFile] = None,
+        fingerprint: Optional[str] = None,
+        writer_batch_size: Optional[int] = None,
+        disable_nullable: bool = False,
+        update_features: bool = False,
+        on_mixed_types: Optional[Literal["use_json"]] = "use_json",
+        with_metadata: bool = True,
+        unit: str = "examples",
+        embed_local_files: bool = False,
+        storage_options: Optional[dict] = None,
+    ):
+        if path is None and stream is None:
+            raise ValueError("At least one of path and stream must be provided.")
+        if features is not None:
+            self._features = features
+            self._schema = None
+        elif schema is not None:
+            self._schema: pa.Schema = schema
+            self._features = Features.from_arrow_schema(self._schema)
+        else:
+            self._features = None
+            self._schema = None
+
+        self._disable_nullable = disable_nullable
+
+        if stream is None:
+            fs, path = url_to_fs(path, **(storage_options or {}))
+            self._fs: fsspec.AbstractFileSystem = fs
+            self._path = path if not is_remote_filesystem(self._fs) else self._fs.unstrip_protocol(path)
+            self.stream = self._fs.open(path, "wb")
+            self._closable_stream = True
+        else:
+            self._fs = None
+            self._path = None
+            self.stream = stream
+            self._closable_stream = False
+
+        self.fingerprint = fingerprint
+        self.disable_nullable = disable_nullable
+        self.writer_batch_size = (
+            writer_batch_size
+            or get_arrow_writer_batch_size_from_features(self._features)
+            or config.DEFAULT_MAX_BATCH_SIZE
+        )
+        self.update_features = update_features
+        self.on_mixed_types = on_mixed_types
+        self.with_metadata = with_metadata
+        self.unit = unit
+        self.embed_local_files = embed_local_files
+
+        self._num_examples = 0
+        self._num_bytes = 0
+        self.current_examples: list[tuple[dict[str, Any], str]] = []
+        self.current_rows: list[pa.Table] = []
+        self.pa_writer: Optional[pa.RecordBatchStreamWriter] = None
+        self.hkey_record = []
+
+    def __len__(self):
+        """Return the number of writed and staged examples"""
+        return self._num_examples + len(self.current_examples) + len(self.current_rows)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    def close(self):
+        # Try closing if opened; if closed: pyarrow.lib.ArrowInvalid: Invalid operation on closed file
+        if self.pa_writer:  # it might be None
+            try:
+                self.pa_writer.close()
+            except Exception:  # pyarrow.lib.ArrowInvalid, OSError
+                pass
+        if self._closable_stream and not self.stream.closed:
+            self.stream.close()  # This also closes self.pa_writer if it is opened
+
+    def _build_schema(self, inferred_schema: pa.Schema):
+        schema = self.schema
+        features = self._features
+        inferred_features = Features.from_arrow_schema(inferred_schema)
+        if self._features is not None:
+            if self.update_features:  # keep original features it they match, or update them
+                fields = {field.name: field for field in self._features.type}
+                for inferred_field in inferred_features.type:
+                    name = inferred_field.name
+                    if name in fields:
+                        if inferred_field == fields[name]:
+                            inferred_features[name] = self._features[name]
+                features = inferred_features
+                schema: pa.Schema = inferred_schema
+        else:
+            features = inferred_features
+            schema: pa.Schema = inferred_features.arrow_schema
+
+        if self.disable_nullable:
+            schema = pa.schema(pa.field(field.name, field.type, nullable=False) for field in schema)
+        if self.with_metadata:
+            schema = schema.with_metadata(self._build_metadata(DatasetInfo(features=features), self.fingerprint))
+        else:
+            schema = schema.with_metadata({})
+
+        return schema, features
+
+    def _build_writer(self, inferred_schema: pa.Schema):
+        self._schema, self._features = self._build_schema(inferred_schema)
+        self.pa_writer = pa.RecordBatchStreamWriter(self.stream, self._schema)
+
+    @property
+    def schema(self):
+        _schema = (
+            self._schema
+            if self._schema is not None
+            else (pa.schema(self._features.type) if self._features is not None else None)
+        )
+        if self._disable_nullable and _schema is not None:
+            _schema = pa.schema(pa.field(field.name, field.type, nullable=False) for field in _schema)
+        return _schema if _schema is not None else []
+
+    @staticmethod
+    def _build_metadata(info: DatasetInfo, fingerprint: Optional[str] = None) -> dict[str, str]:
+        info_keys = ["features"]  # we can add support for more DatasetInfo keys in the future
+        info_as_dict = asdict(info)
+        metadata = {}
+        metadata["info"] = {key: info_as_dict[key] for key in info_keys}
+        if fingerprint is not None:
+            metadata["fingerprint"] = fingerprint
+        return {"huggingface": json.dumps(metadata)}
+
+    def write_examples_on_file(self):
+        """Write stored examples from the write-pool of examples. It makes a table out of the examples and write it."""
+        if not self.current_examples:
+            return
+        # preserve the order the columns
+        if self.schema:
+            schema_cols = set(self.schema.names)
+            examples_cols = self.current_examples[0][0].keys()  # .keys() preserves the order (unlike set)
+            common_cols = [col for col in self.schema.names if col in examples_cols]
+            extra_cols = [col for col in examples_cols if col not in schema_cols]
+            cols = common_cols + extra_cols
+        else:
+            cols = list(self.current_examples[0][0])
+        batch_examples = {}
+        for col in cols:
+            # We use row[0][col] since current_examples contains (example, key) tuples.
+            # Moreover, examples could be Arrow arrays of 1 element.
+            # This can happen in `.map()` when we want to re-write the same Arrow data
+            if all(isinstance(row[0][col], (pa.Array, pa.ChunkedArray)) for row in self.current_examples):
+                arrays = [row[0][col] for row in self.current_examples]
+                arrays = [
+                    chunk
+                    for array in arrays
+                    for chunk in (array.chunks if isinstance(array, pa.ChunkedArray) else [array])
+                ]
+                batch_examples[col] = pa.concat_arrays(arrays)
+            else:
+                batch_examples[col] = [
+                    row[0][col].to_pylist()[0] if isinstance(row[0][col], (pa.Array, pa.ChunkedArray)) else row[0][col]
+                    for row in self.current_examples
+                ]
+        self._write_batch(batch_examples=batch_examples)
+        self.current_examples = []
+
+    def write_rows_on_file(self):
+        """Write stored rows from the write-pool of rows. It concatenates the single-row tables and it writes the resulting table."""
+        if not self.current_rows:
+            return
+        table = pa.concat_tables(self.current_rows)
+        self._write_table(table)
+        self.current_rows = []
+
+    def write(
+        self,
+        example: dict[str, Any],
+        writer_batch_size: Optional[int] = None,
+    ):
+        """Add a given (Example,Key) pair to the write-pool of examples which is written to file.
+
+        Args:
+            example: the Example to add.
+        """
+        # Store example as a tuple so as to keep the structure of `self.current_examples` uniform
+        self.current_examples.append((example, ""))
+
+        if writer_batch_size is None:
+            writer_batch_size = self.writer_batch_size
+        if writer_batch_size is not None and len(self.current_examples) >= writer_batch_size:
+            self.write_examples_on_file()
+
+    def write_row(self, row: pa.Table, writer_batch_size: Optional[int] = None):
+        """Add a given single-row Table to the write-pool of rows which is written to file.
+
+        Args:
+            row: the row to add.
+        """
+        if len(row) != 1:
+            raise ValueError(f"Only single-row pyarrow tables are allowed but got table with {len(row)} rows.")
+        self.current_rows.append(row)
+        if writer_batch_size is None:
+            writer_batch_size = self.writer_batch_size
+        if writer_batch_size is not None and len(self.current_rows) >= writer_batch_size:
+            self.write_rows_on_file()
+
+    def write_batch(
+        self,
+        batch_examples: dict[str, list],
+        writer_batch_size: Optional[int] = None,
+        try_original_type: Optional[bool] = True,
+    ):
+        """Write a batch of Example to file.
+        Ignores the batch if it appears to be empty,
+        preventing a potential schema update of unknown types.
+
+        Args:
+            batch_examples: the batch of examples to add.
+            try_original_type: use `try_type` when instantiating OptimizedTypedSequence if `True`, otherwise `try_type = None`.
+        """
+        self.write_examples_on_file()  # in case there are buffered examples to write first
+        self._write_batch(batch_examples, writer_batch_size=writer_batch_size, try_original_type=try_original_type)
+
+    def _write_batch(
+        self,
+        batch_examples: dict[str, list],
+        writer_batch_size: Optional[int] = None,
+        try_original_type: Optional[bool] = True,
+    ):
+        if batch_examples and len(next(iter(batch_examples.values()))) == 0:
+            return
+        features = None if self.pa_writer is None and self.update_features else self._features
+        try_features = self._features if self.pa_writer is None and self.update_features else None
+        arrays = []
+        inferred_features = Features()
+        # preserve the order the columns
+        if self.schema:
+            schema_cols = set(self.schema.names)
+            batch_cols = batch_examples.keys()  # .keys() preserves the order (unlike set)
+            common_cols = [col for col in self.schema.names if col in batch_cols]
+            extra_cols = [col for col in batch_cols if col not in schema_cols]
+            cols = common_cols + extra_cols
+        else:
+            cols = list(batch_examples)
+        for col in cols:
+            col_values = batch_examples[col]
+            col_type = features[col] if features else None
+            if isinstance(col_values, (pa.Array, pa.ChunkedArray)):
+                array = cast_array_to_feature(col_values, col_type) if col_type is not None else col_values
+                arrays.append(array)
+                inferred_features[col] = generate_from_arrow_type(col_values.type)
+            else:
+                col_try_type = (
+                    try_features[col]
+                    if try_features is not None and col in try_features and try_original_type
+                    else None
+                )
+                typed_sequence = OptimizedTypedSequence(
+                    col_values, type=col_type, try_type=col_try_type, col=col, on_mixed_types=self.on_mixed_types
+                )
+                arrays.append(pa.array(typed_sequence))
+                inferred_features[col] = typed_sequence.get_inferred_type()
+        schema = inferred_features.arrow_schema if self.pa_writer is None else self.schema
+        pa_table = pa.Table.from_arrays(arrays, schema=schema)
+        self.write_table(pa_table, writer_batch_size)
+
+    def write_table(self, pa_table: pa.Table, writer_batch_size: Optional[int] = None):
+        """Write a Table to file.
+
+        Args:
+            example: the Table to add.
+        """
+        self.write_rows_on_file()  # in case there are buffered rows to write first
+        self._write_table(pa_table, writer_batch_size=writer_batch_size)
+
+    def _write_table(self, pa_table: pa.Table, writer_batch_size: Optional[int] = None):
+        if writer_batch_size is None:
+            writer_batch_size = self.writer_batch_size
+        if self.pa_writer is None:
+            self._build_writer(inferred_schema=pa_table.schema)
+        pa_table = pa_table.combine_chunks()
+        pa_table = table_cast(pa_table, self._schema)
+        if self.embed_local_files:
+            pa_table = embed_table_storage(pa_table, local_files=True, remote_files=False)
+        self._num_bytes += pa_table.nbytes
+        self._num_examples += pa_table.num_rows
+        self.pa_writer.write_table(pa_table, writer_batch_size)
+
+    def finalize(self, close_stream=True):
+        self.write_rows_on_file()
+        # In case current_examples < writer_batch_size, but user uses finalize()
+        self.write_examples_on_file()
+        # If schema is known, infer features even if no examples were written
+        if self.pa_writer is None and self.schema:
+            self._build_writer(self.schema)
+        if self.pa_writer is not None:
+            self.pa_writer.close()
+            self.pa_writer = None
+            if close_stream:
+                self.stream.close()
+        else:
+            if close_stream:
+                self.stream.close()
+            raise SchemaInferenceError("Please pass `features` or at least one example when writing data")
+        logger.debug(
+            f"Done writing {self._num_examples} {self.unit} in {self._num_bytes} bytes {self._path if self._path else ''}."
+        )
+        return self._num_examples, self._num_bytes
+
+
+class ParquetWriter(ArrowWriter):
+    def __init__(self, *args, use_content_defined_chunking=True, write_page_index=True, **kwargs):
+        super().__init__(*args, **kwargs)
+        if use_content_defined_chunking is True:
+            use_content_defined_chunking = config.DEFAULT_CDC_OPTIONS
+        self.use_content_defined_chunking = use_content_defined_chunking
+        self.write_page_index = write_page_index
+
+    def _build_writer(self, inferred_schema: pa.Schema):
+        self._schema, self._features = self._build_schema(inferred_schema)
+        self.pa_writer = pq.ParquetWriter(
+            self.stream,
+            self._schema,
+            use_content_defined_chunking=self.use_content_defined_chunking,
+            write_page_index=self.write_page_index,
+            compression={
+                col: "none" if require_storage_embed(feature) else "snappy" for col, feature in self._features.items()
+            },
+            use_dictionary=[col for col, feature in self._features.items() if not require_storage_embed(feature)],
+            column_encoding={
+                col: "PLAIN" for col, feature in self._features.items() if require_storage_embed(feature)
+            },
+        )
+        if self.use_content_defined_chunking is not False:
+            self.pa_writer.add_key_value_metadata(
+                {"content_defined_chunking": json.dumps(self.use_content_defined_chunking)}
+            )

src/datasets/combine.py

@@ -0,0 +1,232 @@
+from typing import Optional, TypeVar
+
+from .arrow_dataset import Dataset, _concatenate_map_style_datasets, _interleave_map_style_datasets
+from .dataset_dict import DatasetDict, IterableDatasetDict
+from .info import DatasetInfo
+from .iterable_dataset import IterableDataset, _concatenate_iterable_datasets, _interleave_iterable_datasets
+from .splits import NamedSplit
+from .utils import logging
+from .utils.py_utils import Literal
+
+
+logger = logging.get_logger(__name__)
+
+
+DatasetType = TypeVar("DatasetType", Dataset, IterableDataset)
+
+
+def interleave_datasets(
+    datasets: list[DatasetType],
+    probabilities: Optional[list[float]] = None,
+    seed: Optional[int] = None,
+    info: Optional[DatasetInfo] = None,
+    split: Optional[NamedSplit] = None,
+    stopping_strategy: Literal[
+        "first_exhausted", "all_exhausted", "all_exhausted_without_replacement"
+    ] = "first_exhausted",
+) -> DatasetType:
+    """
+    Interleave several datasets (sources) into a single dataset.
+    The new dataset is constructed by alternating between the sources to get the examples.
+
+    You can use this function on a list of [`Dataset`] objects, or on a list of [`IterableDataset`] objects.
+
+        - If `probabilities` is `None` (default) the new dataset is constructed by cycling between each source to get the examples.
+        - If `probabilities` is not `None`, the new dataset is constructed by getting examples from a random source at a time according to the provided probabilities.
+
+    The resulting dataset ends when one of the source datasets runs out of examples except when `oversampling` is `True`,
+    in which case, the resulting dataset ends when all datasets have ran out of examples at least one time.
+
+    Note for iterable datasets:
+
+    * The resulting dataset's `num_shards` is the minimum of each dataset's `num_shards` to ensure good parallelism.
+      If some of your datasets have a very low number of shards, you may use [`IterableDataset.reshard`].
+    * In a distributed setup or in PyTorch DataLoader workers, the stopping strategy is applied per process.
+      Therefore the "first_exhausted" strategy on an sharded iterable dataset can generate less samples in total (up to 1 missing sample per subdataset per worker).
+
+    Args:
+        datasets (`List[Dataset]` or `List[IterableDataset]`):
+            List of datasets to interleave.
+        probabilities (`List[float]`, *optional*, defaults to `None`):
+            If specified, the new dataset is constructed by sampling
+            examples from one source at a time according to these probabilities.
+        seed (`int`, *optional*, defaults to `None`):
+            The random seed used to choose a source for each example.
+        info ([`DatasetInfo`], *optional*):
+            Dataset information, like description, citation, etc.
+            <Added version="2.4.0"/>
+        split ([`NamedSplit`], *optional*):
+            Name of the dataset split.
+            <Added version="2.4.0"/>
+        stopping_strategy (`str`, defaults to `first_exhausted`):
+            Three strategies are proposed right now, `first_exhausted`, `all_exhausted` and `all_exhausted_without_replacement`.
+            By default, `first_exhausted` is an undersampling strategy, i.e the dataset construction is stopped as soon as one dataset has ran out of samples.
+            If the strategy is `all_exhausted`,  we use an oversampling strategy, i.e the dataset construction is stopped as soon as every samples of every dataset has been added at least once.
+            When strategy is `all_exhausted_without_replacement` we make sure that each sample in each dataset is sampled only once.
+            Note that if the strategy is `all_exhausted`, the interleaved dataset size can get enormous:
+            - with no probabilities, the resulting dataset will have `max_length_datasets*nb_dataset` samples.
+            - with given probabilities, the resulting dataset will have more samples if some datasets have really low probability of visiting.
+    Returns:
+        [`Dataset`] or [`IterableDataset`]: Return type depends on the input `datasets`
+        parameter. `Dataset` if the input is a list of `Dataset`, `IterableDataset` if the input is a list of
+        `IterableDataset`.
+
+    Example:
+
+        For regular datasets (map-style):
+
+        ```python
+        >>> from datasets import Dataset, interleave_datasets
+        >>> d1 = Dataset.from_dict({"a": [0, 1, 2]})
+        >>> d2 = Dataset.from_dict({"a": [10, 11, 12]})
+        >>> d3 = Dataset.from_dict({"a": [20, 21, 22]})
+        >>> dataset = interleave_datasets([d1, d2, d3], probabilities=[0.7, 0.2, 0.1], seed=42, stopping_strategy="all_exhausted")
+        >>> dataset["a"]
+        [10, 0, 11, 1, 2, 20, 12, 10, 0, 1, 2, 21, 0, 11, 1, 2, 0, 1, 12, 2, 10, 0, 22]
+        >>> dataset = interleave_datasets([d1, d2, d3], probabilities=[0.7, 0.2, 0.1], seed=42)
+        >>> dataset["a"]
+        [10, 0, 11, 1, 2]
+        >>> dataset = interleave_datasets([d1, d2, d3])
+        >>> dataset["a"]
+        [0, 10, 20, 1, 11, 21, 2, 12, 22]
+        >>> dataset = interleave_datasets([d1, d2, d3], stopping_strategy="all_exhausted")
+        >>> dataset["a"]
+        [0, 10, 20, 1, 11, 21, 2, 12, 22]
+        >>> d1 = Dataset.from_dict({"a": [0, 1, 2]})
+        >>> d2 = Dataset.from_dict({"a": [10, 11, 12, 13]})
+        >>> d3 = Dataset.from_dict({"a": [20, 21, 22, 23, 24]})
+        >>> dataset = interleave_datasets([d1, d2, d3])
+        >>> dataset["a"]
+        [0, 10, 20, 1, 11, 21, 2, 12, 22]
+        >>> dataset = interleave_datasets([d1, d2, d3], stopping_strategy="all_exhausted")
+        >>> dataset["a"]
+        [0, 10, 20, 1, 11, 21, 2, 12, 22, 0, 13, 23, 1, 10, 24]
+        >>> dataset = interleave_datasets([d1, d2, d3], probabilities=[0.7, 0.2, 0.1], seed=42)
+        >>> dataset["a"]
+        [10, 0, 11, 1, 2]
+        >>> dataset = interleave_datasets([d1, d2, d3], probabilities=[0.7, 0.2, 0.1], seed=42, stopping_strategy="all_exhausted")
+        >>> dataset["a"]
+        [10, 0, 11, 1, 2, 20, 12, 13, ..., 0, 1, 2, 0, 24]
+        For datasets in streaming mode (iterable):
+
+        >>> from datasets import interleave_datasets
+        >>> d1 = load_dataset('allenai/c4', 'es', split='train', streaming=True)
+        >>> d2 = load_dataset('allenai/c4', 'fr', split='train', streaming=True)
+        >>> dataset = interleave_datasets([d1, d2])
+        >>> iterator = iter(dataset)
+        >>> next(iterator)
+        {'text': 'Comprar Zapatillas para niña en chancla con goma por...'}
+        >>> next(iterator)
+        {'text': 'Le sacre de philippe ier, 23 mai 1059 - Compte Rendu...'
+        ```
+    """
+    from .arrow_dataset import Dataset
+    from .iterable_dataset import IterableDataset
+
+    if not datasets:
+        raise ValueError("Unable to interleave an empty list of datasets.")
+    for i, dataset in enumerate(datasets):
+        if not isinstance(dataset, (Dataset, IterableDataset)):
+            if isinstance(dataset, (DatasetDict, IterableDatasetDict)):
+                if not dataset:
+                    raise ValueError(
+                        f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
+                        "is an empty dataset dictionary."
+                    )
+                raise ValueError(
+                    f"Dataset at position {i} has at least one split: {list(dataset)}\n"
+                    f"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(dataset))}']"
+                )
+            raise ValueError(
+                f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(dataset).__name__}."
+            )
+        if i == 0:
+            dataset_type, other_type = (
+                (Dataset, IterableDataset) if isinstance(dataset, Dataset) else (IterableDataset, Dataset)
+            )
+        elif not isinstance(dataset, dataset_type):
+            raise ValueError(
+                f"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."
+            )
+    if stopping_strategy not in ["first_exhausted", "all_exhausted", "all_exhausted_without_replacement"]:
+        raise ValueError(f"{stopping_strategy} is not supported. Please enter a valid stopping_strategy.")
+    if dataset_type is Dataset:
+        return _interleave_map_style_datasets(
+            datasets, probabilities, seed, info=info, split=split, stopping_strategy=stopping_strategy
+        )
+    else:
+        return _interleave_iterable_datasets(
+            datasets,
+            probabilities,
+            seed,
+            info=info,
+            split=split,
+            stopping_strategy=stopping_strategy,
+        )
+
+
+def concatenate_datasets(
+    dsets: list[DatasetType],
+    info: Optional[DatasetInfo] = None,
+    split: Optional[NamedSplit] = None,
+    axis: int = 0,
+) -> DatasetType:
+    """
+    Concatenate several datasets (sources) into a single dataset.
+
+    Use axis=0 to concatenate vertically (default), or axis=1 to concatenate horizontally.
+
+    Note for iterable 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.
+
+    Args:
+        dsets (`List[datasets.Dataset]` or `List[datasets.IterableDataset]`):
+            List of Datasets to concatenate.
+        info (`DatasetInfo`, *optional*):
+            Dataset information, like description, citation, etc.
+        split (`NamedSplit`, *optional*):
+            Name of the dataset split.
+        axis (`{0, 1}`, defaults to `0`):
+            Axis to concatenate over, where `0` means over rows (vertically) and `1` means over columns
+            (horizontally).
+
+            <Added version="1.6.0"/>
+
+    Example:
+
+    ```py
+    >>> ds3 = concatenate_datasets([ds1, ds2])
+    ```
+    """
+
+    if not dsets:
+        raise ValueError("Unable to concatenate an empty list of datasets.")
+    for i, dataset in enumerate(dsets):
+        if not isinstance(dataset, (Dataset, IterableDataset)):
+            if isinstance(dataset, (DatasetDict, IterableDatasetDict)):
+                if not dataset:
+                    raise ValueError(
+                        f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
+                        "is an empty dataset dictionary."
+                    )
+                raise ValueError(
+                    f"Dataset at position {i} has at least one split: {list(dataset)}\n"
+                    f"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(dataset))}']"
+                )
+            raise ValueError(
+                f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(dataset).__name__}."
+            )
+        if i == 0:
+            dataset_type, other_type = (
+                (Dataset, IterableDataset) if isinstance(dataset, Dataset) else (IterableDataset, Dataset)
+            )
+        elif not isinstance(dataset, dataset_type):
+            raise ValueError(
+                f"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."
+            )
+    if dataset_type is Dataset:
+        return _concatenate_map_style_datasets(dsets, info=info, split=split, axis=axis)
+    else:
+        return _concatenate_iterable_datasets(dsets, info=info, split=split, axis=axis)

src/datasets/commands/__init__.py

@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+from argparse import ArgumentParser
+
+
+class BaseDatasetsCLICommand(ABC):
+    @staticmethod
+    @abstractmethod
+    def register_subcommand(parser: ArgumentParser):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def run(self):
+        raise NotImplementedError()

src/datasets/commands/datasets_cli.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+from argparse import ArgumentParser
+
+from datasets.commands.delete_from_hub import DeleteFromHubCommand
+from datasets.commands.env import EnvironmentCommand
+from datasets.commands.test import TestCommand
+from datasets.utils.logging import set_verbosity_info
+
+
+def parse_unknown_args(unknown_args):
+    return {key.lstrip("-"): value for key, value in zip(unknown_args[::2], unknown_args[1::2])}
+
+
+def main():
+    parser = ArgumentParser(
+        "HuggingFace Datasets CLI tool", usage="datasets-cli <command> [<args>]", allow_abbrev=False
+    )
+    commands_parser = parser.add_subparsers(help="datasets-cli command helpers")
+    set_verbosity_info()
+
+    # Register commands
+    EnvironmentCommand.register_subcommand(commands_parser)
+    TestCommand.register_subcommand(commands_parser)
+    DeleteFromHubCommand.register_subcommand(commands_parser)
+
+    # Parse args
+    args, unknown_args = parser.parse_known_args()
+    if not hasattr(args, "func"):
+        parser.print_help()
+        exit(1)
+    kwargs = parse_unknown_args(unknown_args)
+
+    # Run
+    service = args.func(args, **kwargs)
+    service.run()
+
+
+if __name__ == "__main__":
+    main()

src/datasets/commands/delete_from_hub.py

@@ -0,0 +1,42 @@
+from argparse import ArgumentParser
+from typing import Optional
+
+from datasets.commands import BaseDatasetsCLICommand
+from datasets.hub import delete_from_hub
+
+
+def _command_factory(args):
+    return DeleteFromHubCommand(
+        args.dataset_id,
+        args.config_name,
+        args.token,
+        args.revision,
+    )
+
+
+class DeleteFromHubCommand(BaseDatasetsCLICommand):
+    @staticmethod
+    def register_subcommand(parser):
+        parser: ArgumentParser = parser.add_parser("delete_from_hub", help="Delete dataset config from the Hub")
+        parser.add_argument(
+            "dataset_id", help="source dataset ID, e.g. USERNAME/DATASET_NAME or ORGANIZATION/DATASET_NAME"
+        )
+        parser.add_argument("config_name", help="config name to delete")
+        parser.add_argument("--token", help="access token to the Hugging Face Hub")
+        parser.add_argument("--revision", help="source revision")
+        parser.set_defaults(func=_command_factory)
+
+    def __init__(
+        self,
+        dataset_id: str,
+        config_name: str,
+        token: Optional[str],
+        revision: Optional[str],
+    ):
+        self._dataset_id = dataset_id
+        self._config_name = config_name
+        self._token = token
+        self._revision = revision
+
+    def run(self) -> None:
+        _ = delete_from_hub(self._dataset_id, self._config_name, revision=self._revision, token=self._token)

src/datasets/commands/env.py

@@ -0,0 +1,41 @@
+import platform
+from argparse import ArgumentParser
+
+import fsspec
+import huggingface_hub
+import pandas
+import pyarrow
+
+from datasets import __version__ as version
+from datasets.commands import BaseDatasetsCLICommand
+
+
+def info_command_factory(_):
+    return EnvironmentCommand()
+
+
+class EnvironmentCommand(BaseDatasetsCLICommand):
+    @staticmethod
+    def register_subcommand(parser: ArgumentParser):
+        download_parser = parser.add_parser("env", help="Print relevant system environment info.")
+        download_parser.set_defaults(func=info_command_factory)
+
+    def run(self):
+        info = {
+            "`datasets` version": version,
+            "Platform": platform.platform(),
+            "Python version": platform.python_version(),
+            "`huggingface_hub` version": huggingface_hub.__version__,
+            "PyArrow version": pyarrow.__version__,
+            "Pandas version": pandas.__version__,
+            "`fsspec` version": fsspec.__version__,
+        }
+
+        print("\nCopy-and-paste the text below in your GitHub issue.\n")
+        print(self.format_dict(info))
+
+        return info
+
+    @staticmethod
+    def format_dict(d):
+        return "\n".join([f"- {prop}: {val}" for prop, val in d.items()]) + "\n"

src/datasets/commands/test.py

@@ -0,0 +1,180 @@
+import logging
+import os
+from argparse import ArgumentParser
+from collections.abc import Generator
+from shutil import rmtree
+
+import datasets.config
+from datasets.builder import DatasetBuilder
+from datasets.commands import BaseDatasetsCLICommand
+from datasets.download.download_manager import DownloadMode
+from datasets.info import DatasetInfosDict
+from datasets.load import dataset_module_factory, get_dataset_builder_class
+from datasets.utils.info_utils import VerificationMode
+from datasets.utils.logging import ERROR, get_logger
+
+
+logger = get_logger(__name__)
+
+
+def _test_command_factory(args):
+    return TestCommand(
+        args.dataset,
+        args.name,
+        args.cache_dir,
+        args.data_dir,
+        args.all_configs,
+        args.save_info or args.save_infos,
+        args.ignore_verifications,
+        args.force_redownload,
+        args.clear_cache,
+        args.num_proc,
+    )
+
+
+class TestCommand(BaseDatasetsCLICommand):
+    __test__ = False  # to tell pytest it's not a test class
+
+    @staticmethod
+    def register_subcommand(parser: ArgumentParser):
+        test_parser = parser.add_parser("test", help="Test dataset loading.")
+        test_parser.add_argument("--name", type=str, default=None, help="Dataset processing name")
+        test_parser.add_argument(
+            "--cache_dir",
+            type=str,
+            default=None,
+            help="Cache directory where the datasets are stored.",
+        )
+        test_parser.add_argument(
+            "--data_dir",
+            type=str,
+            default=None,
+            help="Can be used to specify a manual directory to get the files from.",
+        )
+        test_parser.add_argument("--all_configs", action="store_true", help="Test all dataset configurations")
+        test_parser.add_argument(
+            "--save_info", action="store_true", help="Save the dataset infos in the dataset card (README.md)"
+        )
+        test_parser.add_argument(
+            "--ignore_verifications",
+            action="store_true",
+            help="Run the test without checksums and splits checks.",
+        )
+        test_parser.add_argument("--force_redownload", action="store_true", help="Force dataset redownload")
+        test_parser.add_argument(
+            "--clear_cache",
+            action="store_true",
+            help="Remove downloaded files and cached datasets after each config test",
+        )
+        test_parser.add_argument("--num_proc", type=int, default=None, help="Number of processes")
+        # aliases
+        test_parser.add_argument("--save_infos", action="store_true", help="alias to save_info")
+        test_parser.add_argument("dataset", type=str, help="Name of the dataset to download")
+        test_parser.set_defaults(func=_test_command_factory)
+
+    def __init__(
+        self,
+        dataset: str,
+        name: str,
+        cache_dir: str,
+        data_dir: str,
+        all_configs: bool,
+        save_infos: bool,
+        ignore_verifications: bool,
+        force_redownload: bool,
+        clear_cache: bool,
+        num_proc: int,
+    ):
+        self._dataset = dataset
+        self._name = name
+        self._cache_dir = cache_dir
+        self._data_dir = data_dir
+        self._all_configs = all_configs
+        self._save_infos = save_infos
+        self._ignore_verifications = ignore_verifications
+        self._force_redownload = force_redownload
+        self._clear_cache = clear_cache
+        self._num_proc = num_proc
+        if clear_cache and not cache_dir:
+            print(
+                "When --clear_cache is used, specifying a cache directory is mandatory.\n"
+                "The 'download' folder of the cache directory and the dataset builder cache will be deleted after each configuration test.\n"
+                "Please provide a --cache_dir that will be used to test the dataset."
+            )
+            exit(1)
+        if save_infos:
+            self._ignore_verifications = True
+
+    def run(self):
+        logging.getLogger("filelock").setLevel(ERROR)
+        if self._name is not None and self._all_configs:
+            print("Both parameters `config` and `all_configs` can't be used at once.")
+            exit(1)
+        path, config_name = self._dataset, self._name
+        module = dataset_module_factory(path)
+        builder_cls = get_dataset_builder_class(module)
+        n_builders = len(builder_cls.BUILDER_CONFIGS) if self._all_configs and builder_cls.BUILDER_CONFIGS else 1
+
+        def get_builders() -> Generator[DatasetBuilder, None, None]:
+            if self._all_configs and builder_cls.BUILDER_CONFIGS:
+                for i, config in enumerate(builder_cls.BUILDER_CONFIGS):
+                    if "config_name" in module.builder_kwargs:
+                        yield builder_cls(
+                            cache_dir=self._cache_dir,
+                            data_dir=self._data_dir,
+                            **module.builder_kwargs,
+                        )
+                    else:
+                        yield builder_cls(
+                            config_name=config.name,
+                            cache_dir=self._cache_dir,
+                            data_dir=self._data_dir,
+                            **module.builder_kwargs,
+                        )
+            else:
+                if "config_name" in module.builder_kwargs:
+                    yield builder_cls(cache_dir=self._cache_dir, data_dir=self._data_dir, **module.builder_kwargs)
+                else:
+                    yield builder_cls(
+                        config_name=config_name,
+                        cache_dir=self._cache_dir,
+                        data_dir=self._data_dir,
+                        **module.builder_kwargs,
+                    )
+
+        for j, builder in enumerate(get_builders()):
+            print(f"Testing builder '{builder.config.name}' ({j + 1}/{n_builders})")
+            builder._record_checksums = os.path.exists(
+                os.path.join(builder.get_imported_module_dir(), datasets.config.DATASETDICT_INFOS_FILENAME)
+            )  # record checksums only if we need to update a (deprecated) dataset_infos.json
+            builder.download_and_prepare(
+                download_mode=DownloadMode.REUSE_CACHE_IF_EXISTS
+                if not self._force_redownload
+                else DownloadMode.FORCE_REDOWNLOAD,
+                verification_mode=VerificationMode.NO_CHECKS
+                if self._ignore_verifications
+                else VerificationMode.ALL_CHECKS,
+                num_proc=self._num_proc,
+            )
+            builder.as_dataset()
+
+            # If save_infos=True, we create the dataset card (README.md)
+            # The dataset_infos are saved in the YAML part of the README.md
+            # This is to allow the user to upload them on HF afterwards.
+            if self._save_infos:
+                save_infos_dir = os.path.basename(path) if not os.path.isdir(path) else path
+                os.makedirs(save_infos_dir, exist_ok=True)
+                DatasetInfosDict(**{builder.config.name: builder.info}).write_to_directory(save_infos_dir)
+                print(f"Dataset card saved at {os.path.join(save_infos_dir, datasets.config.REPOCARD_FILENAME)}")
+
+            # If clear_cache=True, the download folder and the dataset builder cache directory are deleted
+            if self._clear_cache:
+                if os.path.isdir(builder._cache_dir):
+                    logger.warning(f"Clearing cache at {builder._cache_dir}")
+                    rmtree(builder._cache_dir)
+                download_dir = os.path.join(self._cache_dir, datasets.config.DOWNLOADED_DATASETS_DIR)
+                if os.path.isdir(download_dir):
+                    logger.warning(f"Clearing cache at {download_dir}")
+                    rmtree(download_dir)
+
+        print("Test successful.")

src/datasets/config.py

@@ -0,0 +1,277 @@
+import importlib
+import importlib.metadata
+import logging
+import os
+import platform
+from pathlib import Path
+from typing import Optional
+
+from huggingface_hub import constants
+from packaging import version
+
+
+logger = logging.getLogger(__name__.split(".", 1)[0])  # to avoid circular import from .utils.logging
+
+# Datasets
+S3_DATASETS_BUCKET_PREFIX = "https://s3.amazonaws.com/datasets.huggingface.co/datasets/datasets"
+CLOUDFRONT_DATASETS_DISTRIB_PREFIX = "https://cdn-datasets.huggingface.co/datasets/datasets"
+REPO_DATASETS_URL = "https://raw.githubusercontent.com/huggingface/datasets/{revision}/datasets/{path}/{name}"
+
+# Hub
+HF_ENDPOINT = os.environ.get("HF_ENDPOINT", "https://huggingface.co")
+HUB_DATASETS_URL = HF_ENDPOINT + "/datasets/{repo_id}/resolve/{revision}/{path}"
+HUB_DATASETS_HFFS_URL = "hf://datasets/{repo_id}@{revision}/{path}"
+HUB_DEFAULT_VERSION = "main"
+
+PY_VERSION = version.parse(platform.python_version())
+
+# General environment variables accepted values for booleans
+ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"}
+ENV_VARS_FALSE_VALUES = {"0", "OFF", "NO", "FALSE"}
+ENV_VARS_TRUE_AND_AUTO_VALUES = ENV_VARS_TRUE_VALUES.union({"AUTO"})
+ENV_VARS_FALSE_AND_AUTO_VALUES = ENV_VARS_FALSE_VALUES.union({"AUTO"})
+
+
+# Imports
+DILL_VERSION = version.parse(importlib.metadata.version("dill"))
+FSSPEC_VERSION = version.parse(importlib.metadata.version("fsspec"))
+PANDAS_VERSION = version.parse(importlib.metadata.version("pandas"))
+PYARROW_VERSION = version.parse(importlib.metadata.version("pyarrow"))
+HF_HUB_VERSION = version.parse(importlib.metadata.version("huggingface_hub"))
+
+USE_TF = os.environ.get("USE_TF", "AUTO").upper()
+USE_TORCH = os.environ.get("USE_TORCH", "AUTO").upper()
+USE_JAX = os.environ.get("USE_JAX", "AUTO").upper()
+
+TORCH_VERSION = "N/A"
+TORCH_AVAILABLE = False
+
+if USE_TORCH in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TF not in ENV_VARS_TRUE_VALUES:
+    TORCH_AVAILABLE = importlib.util.find_spec("torch") is not None
+    if TORCH_AVAILABLE:
+        try:
+            TORCH_VERSION = version.parse(importlib.metadata.version("torch"))
+            logger.debug(f"PyTorch version {TORCH_VERSION} available.")
+        except importlib.metadata.PackageNotFoundError:
+            pass
+else:
+    logger.info("Disabling PyTorch because USE_TF is set")
+
+POLARS_VERSION = "N/A"
+POLARS_AVAILABLE = importlib.util.find_spec("polars") is not None
+
+if POLARS_AVAILABLE:
+    try:
+        POLARS_VERSION = version.parse(importlib.metadata.version("polars"))
+        logger.debug(f"Polars version {POLARS_VERSION} available.")
+    except importlib.metadata.PackageNotFoundError:
+        pass
+
+
+DUCKDB_VERSION = "N/A"
+DUCKDB_AVAILABLE = importlib.util.find_spec("duckdb") is not None
+
+if DUCKDB_AVAILABLE:
+    try:
+        DUCKDB_VERSION = version.parse(importlib.metadata.version("duckdb"))
+        logger.debug(f"Duckdb version {DUCKDB_VERSION} available.")
+    except importlib.metadata.PackageNotFoundError:
+        pass
+
+TF_VERSION = "N/A"
+TF_AVAILABLE = False
+
+if USE_TF in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TORCH not in ENV_VARS_TRUE_VALUES:
+    TF_AVAILABLE = importlib.util.find_spec("tensorflow") is not None
+    if TF_AVAILABLE:
+        # For the metadata, we have to look for both tensorflow and tensorflow-cpu
+        for package in [
+            "tensorflow",
+            "tensorflow-cpu",
+            "tensorflow-gpu",
+            "tf-nightly",
+            "tf-nightly-cpu",
+            "tf-nightly-gpu",
+            "intel-tensorflow",
+            "tensorflow-rocm",
+            "tensorflow-macos",
+        ]:
+            try:
+                TF_VERSION = version.parse(importlib.metadata.version(package))
+            except importlib.metadata.PackageNotFoundError:
+                continue
+            else:
+                break
+        else:
+            TF_AVAILABLE = False
+    if TF_AVAILABLE:
+        if TF_VERSION.major < 2:
+            logger.info(f"TensorFlow found but with version {TF_VERSION}. `datasets` requires version 2 minimum.")
+            TF_AVAILABLE = False
+        else:
+            logger.info(f"TensorFlow version {TF_VERSION} available.")
+else:
+    logger.info("Disabling Tensorflow because USE_TORCH is set")
+
+
+JAX_VERSION = "N/A"
+JAX_AVAILABLE = False
+
+if USE_JAX in ENV_VARS_TRUE_AND_AUTO_VALUES:
+    JAX_AVAILABLE = importlib.util.find_spec("jax") is not None and importlib.util.find_spec("jaxlib") is not None
+    if JAX_AVAILABLE:
+        try:
+            JAX_VERSION = version.parse(importlib.metadata.version("jax"))
+            logger.info(f"JAX version {JAX_VERSION} available.")
+        except importlib.metadata.PackageNotFoundError:
+            pass
+else:
+    logger.info("Disabling JAX because USE_JAX is set to False")
+
+
+# Optional tools for data loading
+SQLALCHEMY_AVAILABLE = importlib.util.find_spec("sqlalchemy") is not None
+
+# Optional tools for file parsing and feature decoding
+PIL_AVAILABLE = importlib.util.find_spec("PIL") is not None
+IS_OPUS_SUPPORTED = True
+IS_MP3_SUPPORTED = True
+TORCHCODEC_AVAILABLE = importlib.util.find_spec("torchcodec") is not None
+TORCHVISION_AVAILABLE = importlib.util.find_spec("torchvision") is not None
+PDFPLUMBER_AVAILABLE = importlib.util.find_spec("pdfplumber") is not None
+NIBABEL_AVAILABLE = importlib.util.find_spec("nibabel") is not None
+TRIMESH_AVAILABLE = importlib.util.find_spec("trimesh") is not None
+TEICH_AVAILABLE = importlib.util.find_spec("teich") is not None
+
+# Optional compression tools
+RARFILE_AVAILABLE = importlib.util.find_spec("rarfile") is not None
+ZSTANDARD_AVAILABLE = importlib.util.find_spec("zstandard") is not None
+LZ4_AVAILABLE = importlib.util.find_spec("lz4") is not None
+PY7ZR_AVAILABLE = importlib.util.find_spec("py7zr") is not None
+
+# Cache location
+DEFAULT_XDG_CACHE_HOME = "~/.cache"
+XDG_CACHE_HOME = os.getenv("XDG_CACHE_HOME", DEFAULT_XDG_CACHE_HOME)
+DEFAULT_HF_CACHE_HOME = os.path.join(XDG_CACHE_HOME, "huggingface")
+HF_CACHE_HOME = os.path.expanduser(os.getenv("HF_HOME", DEFAULT_HF_CACHE_HOME))
+
+DEFAULT_HF_DATASETS_CACHE = os.path.join(HF_CACHE_HOME, "datasets")
+HF_DATASETS_CACHE = Path(os.getenv("HF_DATASETS_CACHE", DEFAULT_HF_DATASETS_CACHE))
+
+DEFAULT_HF_MODULES_CACHE = os.path.join(HF_CACHE_HOME, "modules")
+HF_MODULES_CACHE = Path(os.getenv("HF_MODULES_CACHE", DEFAULT_HF_MODULES_CACHE))
+
+DOWNLOADED_DATASETS_DIR = "downloads"
+DEFAULT_DOWNLOADED_DATASETS_PATH = os.path.join(HF_DATASETS_CACHE, DOWNLOADED_DATASETS_DIR)
+DOWNLOADED_DATASETS_PATH = Path(os.getenv("HF_DATASETS_DOWNLOADED_DATASETS_PATH", DEFAULT_DOWNLOADED_DATASETS_PATH))
+
+EXTRACTED_DATASETS_DIR = "extracted"
+DEFAULT_EXTRACTED_DATASETS_PATH = os.path.join(DEFAULT_DOWNLOADED_DATASETS_PATH, EXTRACTED_DATASETS_DIR)
+EXTRACTED_DATASETS_PATH = Path(os.getenv("HF_DATASETS_EXTRACTED_DATASETS_PATH", DEFAULT_EXTRACTED_DATASETS_PATH))
+
+# Cached dataset info options
+SAVE_ORIGINAL_SHARD_LENGTHS = False
+
+# Download count for the website
+HF_UPDATE_DOWNLOAD_COUNTS = (
+    os.environ.get("HF_UPDATE_DOWNLOAD_COUNTS", "AUTO").upper() in ENV_VARS_TRUE_AND_AUTO_VALUES
+)
+
+# For downloads and to check remote files metadata
+HF_DATASETS_MULTITHREADING_MAX_WORKERS = 16
+
+# Dataset viewer API
+USE_PARQUET_EXPORT = True
+
+# Batch size constants. For more info, see:
+# https://github.com/apache/arrow/blob/master/docs/source/cpp/arrays.rst#size-limitations-and-recommendations)
+DEFAULT_MAX_BATCH_SIZE = 1000
+
+DEFAULT_CDC_OPTIONS = {"min_chunk_size": 256 * 1024, "max_chunk_size": 1024 * 1024, "norm_level": 0}
+
+# Size of the preloaded record batch in `Dataset.__iter__`
+ARROW_READER_BATCH_SIZE_IN_DATASET_ITER = 10
+
+# Max uncompressed shard size in bytes (e.g. to shard parquet datasets in push_to_hub or download_and_prepare)
+MAX_SHARD_SIZE = "500MB"
+
+# Max uncompressed row group size in bytes (e.g. for parquet files in push_to_hub or download_and_prepare)
+MAX_ROW_GROUP_SIZE = "100MB"
+
+# Parquet configuration
+PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS = None
+PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS = None
+PARQUET_ROW_GROUP_SIZE_FOR_BINARY_DATASETS = None
+PARQUET_ROW_GROUP_SIZE_FOR_VIDEO_DATASETS = None
+
+# Arrow configuration
+ARROW_RECORD_BATCH_SIZE_FOR_AUDIO_DATASETS = 100
+ARROW_RECORD_BATCH_SIZE_FOR_IMAGE_DATASETS = 100
+ARROW_RECORD_BATCH_SIZE_FOR_BINARY_DATASETS = 100
+ARROW_RECORD_BATCH_SIZE_FOR_VIDEO_DATASETS = 10
+
+# Offline mode
+_offline = os.environ.get("HF_DATASETS_OFFLINE")
+HF_HUB_OFFLINE = constants.HF_HUB_OFFLINE if _offline is None else _offline.upper() in ENV_VARS_TRUE_VALUES
+HF_DATASETS_OFFLINE = HF_HUB_OFFLINE  # kept for backward-compatibility
+
+# Here, `True` will disable progress bars globally without possibility of enabling it
+# programmatically. `False` will enable them without possibility of disabling them.
+# If environment variable is not set (None), then the user is free to enable/disable
+# them programmatically.
+# TL;DR: env variable has priority over code
+__HF_DATASETS_DISABLE_PROGRESS_BARS = os.environ.get("HF_DATASETS_DISABLE_PROGRESS_BARS")
+HF_DATASETS_DISABLE_PROGRESS_BARS: Optional[bool] = (
+    __HF_DATASETS_DISABLE_PROGRESS_BARS.upper() in ENV_VARS_TRUE_VALUES
+    if __HF_DATASETS_DISABLE_PROGRESS_BARS is not None
+    else None
+)
+
+# In-memory
+DEFAULT_IN_MEMORY_MAX_SIZE = 0  # Disabled
+IN_MEMORY_MAX_SIZE = float(os.environ.get("HF_DATASETS_IN_MEMORY_MAX_SIZE", DEFAULT_IN_MEMORY_MAX_SIZE))
+
+# File names
+DATASET_ARROW_FILENAME = "dataset.arrow"
+DATASET_INDICES_FILENAME = "indices.arrow"
+DATASET_STATE_JSON_FILENAME = "state.json"
+DATASET_INFO_FILENAME = "dataset_info.json"
+DATASETDICT_INFOS_FILENAME = "dataset_infos.json"
+LICENSE_FILENAME = "LICENSE"
+DATASETDICT_JSON_FILENAME = "dataset_dict.json"
+METADATA_CONFIGS_FIELD = "configs"
+REPOCARD_FILENAME = "README.md"
+REPOYAML_FILENAME = ".huggingface.yaml"
+
+MODULE_NAME_FOR_DYNAMIC_MODULES = "datasets_modules"
+
+MAX_DATASET_CONFIG_ID_READABLE_LENGTH = 255
+
+# Temporary cache directory prefix
+TEMP_CACHE_DIR_PREFIX = "hf_datasets-"
+
+# Streaming
+STREAMING_READ_MAX_RETRIES = 20
+STREAMING_READ_RETRY_INTERVAL = 5
+STREAMING_READ_SERVER_UNAVAILABLE_RETRY_INTERVAL = 20
+STREAMING_READ_RATE_LIMIT_RETRY_INTERVAL = 60
+STREAMING_OPEN_MAX_RETRIES = 20
+STREAMING_OPEN_RETRY_INTERVAL = 5
+
+# Datasets repositories exploration
+ARCHIVES_MAX_NUMBER_FOR_MODULE_INFERENCE = 10
+
+# Async map functions
+MAX_NUM_RUNNING_ASYNC_MAP_FUNCTIONS_IN_PARALLEL = 1000
+
+# Progress bars
+PBAR_REFRESH_TIME_INTERVAL = 0.05  # 20 progress updates per sec
+
+# Maximum number of uploaded files per commit
+UPLOADS_MAX_NUMBER_PER_COMMIT = 50
+
+# Backward compatibility
+MAX_TABLE_NBYTES_FOR_PICKLING = 4 << 30
+
+# Time to let PyArrow close threads gracefully
+SLEEP_TIME_ON_THREADS_SHUTDOWN = 5

src/datasets/data_files.py

@@ -0,0 +1,829 @@
+import os
+import re
+from functools import partial
+from glob import has_magic
+from pathlib import Path, PurePath
+from typing import Callable, Optional, Union
+
+import huggingface_hub
+from fsspec.core import url_to_fs
+from huggingface_hub import HfFileSystem
+from packaging import version
+from tqdm.contrib.concurrent import thread_map
+
+from . import config
+from .download import DownloadConfig
+from .naming import _split_re
+from .splits import Split
+from .utils import logging
+from .utils import tqdm as hf_tqdm
+from .utils.file_utils import _prepare_path_and_storage_options, is_local_path, is_relative_path, xbasename, xjoin
+from .utils.py_utils import string_to_dict
+
+
+SingleOriginMetadata = Union[tuple[str, str], tuple[str], tuple[()]]
+
+
+SANITIZED_DEFAULT_SPLIT = str(Split.TRAIN)
+
+
+logger = logging.get_logger(__name__)
+
+
+class Url(str):
+    pass
+
+
+class EmptyDatasetError(FileNotFoundError):
+    pass
+
+
+SPLIT_PATTERN_SHARDED = "data/{split}-[0-9][0-9][0-9][0-9][0-9]-of-[0-9][0-9][0-9][0-9][0-9]*.*"
+
+SPLIT_KEYWORDS = {
+    Split.TRAIN: ["train", "training"],
+    Split.VALIDATION: ["validation", "valid", "dev", "val"],
+    Split.TEST: ["test", "testing", "eval", "evaluation"],
+}
+NON_WORDS_CHARS = "-._ 0-9"
+if config.FSSPEC_VERSION < version.parse("2023.9.0"):
+    KEYWORDS_IN_FILENAME_BASE_PATTERNS = ["**[{sep}/]{keyword}[{sep}]*", "{keyword}[{sep}]*"]
+    KEYWORDS_IN_DIR_NAME_BASE_PATTERNS = [
+        "{keyword}/**",
+        "{keyword}[{sep}]*/**",
+        "**[{sep}/]{keyword}/**",
+        "**[{sep}/]{keyword}[{sep}]*/**",
+    ]
+elif config.FSSPEC_VERSION < version.parse("2023.12.0"):
+    KEYWORDS_IN_FILENAME_BASE_PATTERNS = ["**/*[{sep}/]{keyword}[{sep}]*", "{keyword}[{sep}]*"]
+    KEYWORDS_IN_DIR_NAME_BASE_PATTERNS = [
+        "{keyword}/**/*",
+        "{keyword}[{sep}]*/**/*",
+        "**/*[{sep}/]{keyword}/**/*",
+        "**/*[{sep}/]{keyword}[{sep}]*/**/*",
+    ]
+else:
+    KEYWORDS_IN_FILENAME_BASE_PATTERNS = ["**/{keyword}[{sep}]*", "**/*[{sep}]{keyword}[{sep}]*"]
+    KEYWORDS_IN_DIR_NAME_BASE_PATTERNS = [
+        "**/{keyword}/**",
+        "**/{keyword}[{sep}]*/**",
+        "**/*[{sep}]{keyword}/**",
+        "**/*[{sep}]{keyword}[{sep}]*/**",
+    ]
+
+DEFAULT_SPLITS = [Split.TRAIN, Split.VALIDATION, Split.TEST]
+DEFAULT_PATTERNS_SPLIT_IN_FILENAME = {
+    split: [
+        pattern.format(keyword=keyword, sep=NON_WORDS_CHARS)
+        for keyword in SPLIT_KEYWORDS[split]
+        for pattern in KEYWORDS_IN_FILENAME_BASE_PATTERNS
+    ]
+    for split in DEFAULT_SPLITS
+}
+DEFAULT_PATTERNS_SPLIT_IN_DIR_NAME = {
+    split: [
+        pattern.format(keyword=keyword, sep=NON_WORDS_CHARS)
+        for keyword in SPLIT_KEYWORDS[split]
+        for pattern in KEYWORDS_IN_DIR_NAME_BASE_PATTERNS
+    ]
+    for split in DEFAULT_SPLITS
+}
+
+
+DEFAULT_PATTERNS_ALL = {
+    Split.TRAIN: ["**"],
+}
+
+DEFAULT_PATTERNS_LOGS = {"logs": ["**/*.eval"]}
+
+ALL_SPLIT_PATTERNS = [SPLIT_PATTERN_SHARDED]
+ALL_DEFAULT_PATTERNS = [
+    DEFAULT_PATTERNS_LOGS,
+    DEFAULT_PATTERNS_SPLIT_IN_DIR_NAME,
+    DEFAULT_PATTERNS_SPLIT_IN_FILENAME,
+    DEFAULT_PATTERNS_ALL,
+]
+WILDCARD_CHARACTERS = "*[]"
+FILES_TO_IGNORE = [
+    "README.md",
+    "config.json",
+    "dataset_info.json",
+    "dataset_infos.json",
+    "dummy_data.zip",
+    "dataset_dict.json",
+]
+
+
+def contains_wildcards(pattern: str) -> bool:
+    return any(wildcard_character in pattern for wildcard_character in WILDCARD_CHARACTERS)
+
+
+def sanitize_patterns(patterns: Union[dict, list, str]) -> dict[str, Union[list[str], "DataFilesList"]]:
+    """
+    Take the data_files patterns from the user, and format them into a dictionary.
+    Each key is the name of the split, and each value is a list of data files patterns (paths or urls).
+    The default split is "train".
+
+    Returns:
+        patterns: dictionary of split_name -> list of patterns
+    """
+    if isinstance(patterns, dict):
+        return {str(key): value if isinstance(value, list) else [value] for key, value in patterns.items()}
+    elif isinstance(patterns, str):
+        return {SANITIZED_DEFAULT_SPLIT: [patterns]}
+    elif isinstance(patterns, list):
+        if any(isinstance(pattern, dict) for pattern in patterns):
+            for pattern in patterns:
+                if not (
+                    isinstance(pattern, dict)
+                    and len(pattern) == 2
+                    and "split" in pattern
+                    and isinstance(pattern.get("path"), (str, list))
+                ):
+                    raise ValueError(
+                        "Invalid format for data_files entry. "
+                        "Each item must be a dictionary with the structure "
+                        "{'split': <split_name>, 'path': <path_or_list_of_paths>}.\n"
+                        f"Received: {pattern}"
+                    )
+            splits = [pattern["split"] for pattern in patterns]
+            if len(set(splits)) != len(splits):
+                raise ValueError(f"Some splits are duplicated in data_files: {splits}")
+            return {
+                str(pattern["split"]): pattern["path"] if isinstance(pattern["path"], list) else [pattern["path"]]
+                for pattern in patterns
+            }
+        else:
+            return {SANITIZED_DEFAULT_SPLIT: patterns}
+    else:
+        return sanitize_patterns(list(patterns))
+
+
+def _is_inside_unrequested_special_dir(matched_rel_path: str, pattern: str) -> bool:
+    """
+    When a path matches a pattern, we additionally check if it's inside a special directory
+    we ignore by default (if it starts with a double underscore).
+
+    Users can still explicitly request a filepath inside such a directory if "__pycache__" is
+    mentioned explicitly in the requested pattern.
+
+    Some examples:
+
+    base directory:
+
+        ./
+        └── __pycache__
+            └── b.txt
+
+    >>> _is_inside_unrequested_special_dir("__pycache__/b.txt", "**")
+    True
+    >>> _is_inside_unrequested_special_dir("__pycache__/b.txt", "*/b.txt")
+    True
+    >>> _is_inside_unrequested_special_dir("__pycache__/b.txt", "__pycache__/*")
+    False
+    >>> _is_inside_unrequested_special_dir("__pycache__/b.txt", "__*/*")
+    False
+    """
+    # We just need to check if every special directories from the path is present explicitly in the pattern.
+    # Since we assume that the path matches the pattern, it's equivalent to counting that both
+    # the parent path and the parent pattern have the same number of special directories.
+    data_dirs_to_ignore_in_path = [part for part in PurePath(matched_rel_path).parent.parts if part.startswith("__")]
+    data_dirs_to_ignore_in_pattern = [part for part in PurePath(pattern).parent.parts if part.startswith("__")]
+    return len(data_dirs_to_ignore_in_path) != len(data_dirs_to_ignore_in_pattern)
+
+
+def _is_unrequested_hidden_file_or_is_inside_unrequested_hidden_dir(matched_rel_path: str, pattern: str) -> bool:
+    """
+    When a path matches a pattern, we additionally check if it's a hidden file or if it's inside
+    a hidden directory we ignore by default, i.e. if the file name or a parent directory name starts with a dot.
+
+    Users can still explicitly request a filepath that is hidden or is inside a hidden directory
+    if the hidden part is mentioned explicitly in the requested pattern.
+
+    Some examples:
+
+    base directory:
+
+        ./
+        └── .hidden_file.txt
+
+    >>> _is_unrequested_hidden_file_or_is_inside_unrequested_hidden_dir(".hidden_file.txt", "**")
+    True
+    >>> _is_unrequested_hidden_file_or_is_inside_unrequested_hidden_dir(".hidden_file.txt", ".*")
+    False
+
+    base directory:
+
+        ./
+        └── .hidden_dir
+            └── a.txt
+
+    >>> _is_unrequested_hidden_file_or_is_inside_unrequested_hidden_dir(".hidden_dir/a.txt", "**")
+    True
+    >>> _is_unrequested_hidden_file_or_is_inside_unrequested_hidden_dir(".hidden_dir/a.txt", ".*/*")
+    False
+    >>> _is_unrequested_hidden_file_or_is_inside_unrequested_hidden_dir(".hidden_dir/a.txt", ".hidden_dir/*")
+    False
+
+    base directory:
+
+        ./
+        └── .hidden_dir
+            └── .hidden_file.txt
+
+    >>> _is_unrequested_hidden_file_or_is_inside_unrequested_hidden_dir(".hidden_dir/.hidden_file.txt", "**")
+    True
+    >>> _is_unrequested_hidden_file_or_is_inside_unrequested_hidden_dir(".hidden_dir/.hidden_file.txt", ".*/*")
+    True
+    >>> _is_unrequested_hidden_file_or_is_inside_unrequested_hidden_dir(".hidden_dir/.hidden_file.txt", ".*/.*")
+    False
+    >>> _is_unrequested_hidden_file_or_is_inside_unrequested_hidden_dir(".hidden_dir/.hidden_file.txt", ".hidden_dir/*")
+    True
+    >>> _is_unrequested_hidden_file_or_is_inside_unrequested_hidden_dir(".hidden_dir/.hidden_file.txt", ".hidden_dir/.*")
+    False
+    """
+    # We just need to check if every hidden part from the path is present explicitly in the pattern.
+    # Since we assume that the path matches the pattern, it's equivalent to counting that both
+    # the path and the pattern have the same number of hidden parts.
+    hidden_directories_in_path = [
+        part for part in PurePath(matched_rel_path).parts if part.startswith(".") and not set(part) == {"."}
+    ]
+    hidden_directories_in_pattern = [
+        part for part in PurePath(pattern).parts if part.startswith(".") and not set(part) == {"."}
+    ]
+    return len(hidden_directories_in_path) != len(hidden_directories_in_pattern)
+
+
+def _get_data_files_patterns(pattern_resolver: Callable[[str], list[str]]) -> dict[str, list[str]]:
+    """
+    Get the default pattern from a directory or repository by testing all the supported patterns.
+    The first patterns to return a non-empty list of data files is returned.
+
+    In order, it first tests if SPLIT_PATTERN_SHARDED works, otherwise it tests the patterns in ALL_DEFAULT_PATTERNS.
+    """
+    # first check the split patterns like data/{split}-00000-of-00001.parquet
+    for split_pattern in ALL_SPLIT_PATTERNS:
+        pattern = split_pattern.replace("{split}", "*")
+        try:
+            data_files = pattern_resolver(pattern)
+        except FileNotFoundError:
+            continue
+        if len(data_files) > 0:
+            splits: set[str] = set()
+            for p in data_files:
+                p_parts = string_to_dict(xbasename(p), xbasename(split_pattern))
+                assert p_parts is not None
+                splits.add(p_parts["split"])
+
+            if any(not re.match(_split_re, split) for split in splits):
+                raise ValueError(f"Split name should match '{_split_re}'' but got '{splits}'.")
+            sorted_splits = [str(split) for split in DEFAULT_SPLITS if split in splits] + sorted(
+                splits - {str(split) for split in DEFAULT_SPLITS}
+            )
+            return {split: [split_pattern.format(split=split)] for split in sorted_splits}
+    # then check the default patterns based on train/valid/test splits
+    for patterns_dict in ALL_DEFAULT_PATTERNS:
+        non_empty_splits = []
+        for split, patterns in patterns_dict.items():
+            for pattern in patterns:
+                try:
+                    data_files = pattern_resolver(pattern)
+                except FileNotFoundError:
+                    continue
+                if len(data_files) > 0:
+                    non_empty_splits.append(split)
+                    break
+        if non_empty_splits:
+            return {split: patterns_dict[split] for split in non_empty_splits}
+    raise FileNotFoundError(f"Couldn't resolve pattern {pattern} with resolver {pattern_resolver}")
+
+
+def resolve_pattern(
+    pattern: str,
+    base_path: str,
+    allowed_extensions: Optional[list[str]] = None,
+    download_config: Optional[DownloadConfig] = None,
+) -> list[str]:
+    """
+    Resolve the paths and URLs of the data files from the pattern passed by the user.
+
+    You can use patterns to resolve multiple local files. Here are a few examples:
+    - *.csv to match all the CSV files at the first level
+    - **.csv to match all the CSV files at any level
+    - data/* to match all the files inside "data"
+    - data/** to match all the files inside "data" and its subdirectories
+
+    The patterns are resolved using the fsspec glob. In fsspec>=2023.12.0 this is equivalent to
+    Python's glob.glob, Path.glob, Path.match and fnmatch where ** is unsupported with a prefix/suffix
+    other than a forward slash /.
+
+    More generally:
+    - '*' matches any character except a forward-slash (to match just the file or directory name)
+    - '**' matches any character including a forward-slash /
+
+    Hidden files and directories (i.e. whose names start with a dot) are ignored, unless they are explicitly requested.
+    The same applies to special directories that start with a double underscore like "__pycache__".
+    You can still include one if the pattern explicitly mentions it:
+    - to include a hidden file: "*/.hidden.txt" or "*/.*"
+    - to include a hidden directory: ".hidden/*" or ".*/*"
+    - to include a special directory: "__special__/*" or "__*/*"
+
+    Example::
+
+        >>> from datasets.data_files import resolve_pattern
+        >>> base_path = "."
+        >>> resolve_pattern("docs/**/*.py", base_path)
+        [/Users/mariosasko/Desktop/projects/datasets/docs/source/_config.py']
+
+    Args:
+        pattern (str): Unix pattern or paths or URLs of the data files to resolve.
+            The paths can be absolute or relative to base_path.
+            Remote filesystems using fsspec are supported, e.g. with the hf:// protocol.
+        base_path (str): Base path to use when resolving relative paths.
+        allowed_extensions (Optional[list], optional): White-list of file extensions to use. Defaults to None (all extensions).
+            For example: allowed_extensions=[".csv", ".json", ".txt", ".parquet"]
+        download_config ([`DownloadConfig`], *optional*): Specific download configuration parameters.
+    Returns:
+        List[str]: List of paths or URLs to the local or remote files that match the patterns.
+    """
+    if is_relative_path(pattern):
+        pattern = xjoin(base_path, pattern)
+    elif is_local_path(pattern):
+        base_path = os.path.splitdrive(pattern)[0] + os.sep
+    else:
+        base_path = ""
+    pattern, storage_options = _prepare_path_and_storage_options(pattern, download_config=download_config)
+    fs, fs_pattern = url_to_fs(pattern, **storage_options)
+    files_to_ignore = set(FILES_TO_IGNORE) - {xbasename(pattern)}
+    protocol = (
+        pattern.split("://")[0]
+        if "://" in pattern
+        else (fs.protocol if isinstance(fs.protocol, str) else fs.protocol[0])
+    )
+    protocol_prefix = protocol + "://" if protocol != "file" else ""
+    glob_kwargs = {}
+    if protocol == "hf":
+        # 10 times faster glob with detail=True (ignores costly info like lastCommit)
+        glob_kwargs["expand_info"] = False
+
+    # if the pattern contains hops like "zip://csv/*.csv::data.zip", we need to keep them after globbing
+    _, *rest_hops = pattern.split("::")
+    matched_paths = []
+    for filepath, info in fs.glob(fs_pattern, detail=True, **glob_kwargs).items():
+        if not (info["type"] == "file" or (info.get("islink") and os.path.isfile(os.path.realpath(filepath)))) or (
+            xbasename(filepath) in files_to_ignore
+        ):
+            continue
+        if _is_inside_unrequested_special_dir(filepath, fs_pattern):
+            continue
+        if _is_unrequested_hidden_file_or_is_inside_unrequested_hidden_dir(filepath, fs_pattern):
+            continue
+        filepath = filepath if "://" in filepath else protocol_prefix + filepath
+        if rest_hops:
+            filepath = "::".join([filepath] + rest_hops)
+        matched_paths.append(filepath)
+    # ignore .ipynb and __pycache__, but keep /../
+    if allowed_extensions is not None:
+        out = [
+            filepath
+            for filepath in matched_paths
+            if any("." + suffix in allowed_extensions for suffix in xbasename(filepath).split(".")[1:])
+        ]
+        if len(out) < len(matched_paths):
+            invalid_matched_files = list(set(matched_paths) - set(out))
+            logger.info(
+                f"Some files matched the pattern '{pattern}' but don't have valid data file extensions: {invalid_matched_files}"
+            )
+    else:
+        out = matched_paths
+    if not out:
+        error_msg = f"Unable to find '{pattern}'"
+        if allowed_extensions is not None:
+            error_msg += f" with any supported extension {list(allowed_extensions)}"
+        raise FileNotFoundError(error_msg)
+    return out
+
+
+def get_data_patterns(base_path: str, download_config: Optional[DownloadConfig] = None) -> dict[str, list[str]]:
+    """
+    Get the default pattern from a directory testing all the supported patterns.
+    The first patterns to return a non-empty list of data files is returned.
+
+    Some examples of supported patterns:
+
+    Input:
+
+        my_dataset_repository/
+        ├── README.md
+        └── dataset.csv
+
+    Output:
+
+        {'train': ['**']}
+
+    Input:
+
+        my_dataset_repository/
+        ├── README.md
+        ├── train.csv
+        └── test.csv
+
+        my_dataset_repository/
+        ├── README.md
+        └── data/
+            ├── train.csv
+            └── test.csv
+
+        my_dataset_repository/
+        ├── README.md
+        ├── train_0.csv
+        ├── train_1.csv
+        ├── train_2.csv
+        ├── train_3.csv
+        ├── test_0.csv
+        └── test_1.csv
+
+    Output:
+
+        {'train': ['**/train[-._ 0-9]*', '**/*[-._ 0-9]train[-._ 0-9]*', '**/training[-._ 0-9]*', '**/*[-._ 0-9]training[-._ 0-9]*'],
+         'test': ['**/test[-._ 0-9]*', '**/*[-._ 0-9]test[-._ 0-9]*', '**/testing[-._ 0-9]*', '**/*[-._ 0-9]testing[-._ 0-9]*', ...]}
+
+    Input:
+
+        my_dataset_repository/
+        ├── README.md
+        └── data/
+            ├── train/
+            │   ├── shard_0.csv
+            │   ├── shard_1.csv
+            │   ├── shard_2.csv
+            │   └── shard_3.csv
+            └── test/
+                ├── shard_0.csv
+                └── shard_1.csv
+
+    Output:
+
+        {'train': ['**/train/**', '**/train[-._ 0-9]*/**', '**/*[-._ 0-9]train/**', '**/*[-._ 0-9]train[-._ 0-9]*/**', ...],
+         'test': ['**/test/**', '**/test[-._ 0-9]*/**', '**/*[-._ 0-9]test/**', '**/*[-._ 0-9]test[-._ 0-9]*/**', ...]}
+
+    Input:
+
+        my_dataset_repository/
+        ├── README.md
+        └── data/
+            ├── train-00000-of-00003.csv
+            ├── train-00001-of-00003.csv
+            ├── train-00002-of-00003.csv
+            ├── test-00000-of-00001.csv
+            ├── random-00000-of-00003.csv
+            ├── random-00001-of-00003.csv
+            └── random-00002-of-00003.csv
+
+    Output:
+
+        {'train': ['data/train-[0-9][0-9][0-9][0-9][0-9]-of-[0-9][0-9][0-9][0-9][0-9]*.*'],
+         'test': ['data/test-[0-9][0-9][0-9][0-9][0-9]-of-[0-9][0-9][0-9][0-9][0-9]*.*'],
+         'random': ['data/random-[0-9][0-9][0-9][0-9][0-9]-of-[0-9][0-9][0-9][0-9][0-9]*.*']}
+
+    In order, it first tests if SPLIT_PATTERN_SHARDED works, otherwise it tests the patterns in ALL_DEFAULT_PATTERNS.
+    """
+    resolver = partial(resolve_pattern, base_path=base_path, download_config=download_config)
+    try:
+        return _get_data_files_patterns(resolver)
+    except FileNotFoundError:
+        raise EmptyDatasetError(f"The directory at {base_path} doesn't contain any data files") from None
+
+
+def _get_single_origin_metadata(
+    data_file: str,
+    download_config: Optional[DownloadConfig] = None,
+) -> SingleOriginMetadata:
+    if data_file.startswith(config.HF_ENDPOINT):
+        fs = HfFileSystem(endpoint=config.HF_ENDPOINT, token=download_config.token)
+        data_file = "hf://" + data_file[len(config.HF_ENDPOINT) + 1 :]
+        data_file = data_file.replace("/resolve/", "/" if data_file.startswith("hf://buckets/") else "@", 1)
+        fs_path = data_file
+    else:
+        data_file, storage_options = _prepare_path_and_storage_options(data_file, download_config=download_config)
+        fs, fs_path = url_to_fs(data_file, **storage_options)
+    if isinstance(fs, HfFileSystem):
+        resolved_path = fs.resolve_path(fs_path)
+        if hasattr(resolved_path, "revision"):  # no revision for buckets
+            return resolved_path.repo_id, resolved_path.revision
+    info = fs.info(fs_path)
+    # s3fs uses "ETag", gcsfs uses "etag", and for local we simply check mtime
+    for key in ["ETag", "etag", "mtime"]:
+        if key in info:
+            return (str(info[key]),)
+    return ()
+
+
+def _get_origin_metadata(
+    data_files: list[str],
+    download_config: Optional[DownloadConfig] = None,
+    max_workers: Optional[int] = None,
+) -> list[SingleOriginMetadata]:
+    max_workers = max_workers if max_workers is not None else config.HF_DATASETS_MULTITHREADING_MAX_WORKERS
+    if all("hf://" in data_file for data_file in data_files):
+        # No need for multithreading here since the origin metadata of HF files
+        # is (repo_id, revision) and is cached after first .info() call.
+        return [
+            _get_single_origin_metadata(data_file, download_config=download_config)
+            for data_file in hf_tqdm(
+                data_files,
+                desc="Resolving data files",
+                # set `disable=None` rather than `disable=False` by default to disable progress bar when no TTY attached
+                disable=len(data_files) <= 16 or None,
+            )
+        ]
+    return thread_map(
+        partial(_get_single_origin_metadata, download_config=download_config),
+        data_files,
+        max_workers=max_workers,
+        tqdm_class=hf_tqdm,
+        desc="Resolving data files",
+        # set `disable=None` rather than `disable=False` by default to disable progress bar when no TTY attached
+        disable=len(data_files) <= 16 or None,
+    )
+
+
+class DataFilesList(list[str]):
+    """
+    List of data files (absolute local paths or URLs).
+    It has two construction methods given the user's data files patterns:
+    - ``from_hf_repo``: resolve patterns inside a dataset repository
+    - ``from_local_or_remote``: resolve patterns from a local path
+
+    Moreover, DataFilesList has an additional attribute ``origin_metadata``.
+    It can store:
+    - the last modified time of local files
+    - ETag of remote files
+    - commit sha of a dataset repository
+
+    Thanks to this additional attribute, it is possible to hash the list
+    and get a different hash if and only if at least one file changed.
+    This is useful for caching Dataset objects that are obtained from a list of data files.
+    """
+
+    def __init__(self, data_files: list[str], origin_metadata: list[SingleOriginMetadata]) -> None:
+        super().__init__(data_files)
+        self.origin_metadata = origin_metadata
+
+    def __add__(self, other: "DataFilesList") -> "DataFilesList":
+        return DataFilesList([*self, *other], self.origin_metadata + other.origin_metadata)
+
+    @classmethod
+    def from_hf_repo(
+        cls,
+        patterns: list[str],
+        dataset_info: huggingface_hub.hf_api.DatasetInfo,
+        base_path: Optional[str] = None,
+        allowed_extensions: Optional[list[str]] = None,
+        download_config: Optional[DownloadConfig] = None,
+    ) -> "DataFilesList":
+        base_path = f"hf://datasets/{dataset_info.id}@{dataset_info.sha}/{base_path or ''}".rstrip("/")
+        return cls.from_patterns(
+            patterns, base_path=base_path, allowed_extensions=allowed_extensions, download_config=download_config
+        )
+
+    @classmethod
+    def from_local_or_remote(
+        cls,
+        patterns: list[str],
+        base_path: Optional[str] = None,
+        allowed_extensions: Optional[list[str]] = None,
+        download_config: Optional[DownloadConfig] = None,
+    ) -> "DataFilesList":
+        base_path = base_path if base_path is not None else Path().resolve().as_posix()
+        return cls.from_patterns(
+            patterns, base_path=base_path, allowed_extensions=allowed_extensions, download_config=download_config
+        )
+
+    @classmethod
+    def from_patterns(
+        cls,
+        patterns: list[str],
+        base_path: Optional[str] = None,
+        allowed_extensions: Optional[list[str]] = None,
+        download_config: Optional[DownloadConfig] = None,
+    ) -> "DataFilesList":
+        base_path = base_path if base_path is not None else Path().resolve().as_posix()
+        data_files = []
+        for pattern in patterns:
+            try:
+                data_files.extend(
+                    resolve_pattern(
+                        pattern,
+                        base_path=base_path,
+                        allowed_extensions=allowed_extensions,
+                        download_config=download_config,
+                    )
+                )
+            except FileNotFoundError:
+                if not has_magic(pattern):
+                    raise
+        origin_metadata = _get_origin_metadata(data_files, download_config=download_config)
+        return cls(data_files, origin_metadata)
+
+    def filter(
+        self, *, extensions: Optional[list[str]] = None, file_names: Optional[list[str]] = None
+    ) -> "DataFilesList":
+        patterns = []
+        if extensions:
+            ext_pattern = "|".join(re.escape(ext) for ext in extensions)
+            patterns.append(re.compile(f".*({ext_pattern})(\\..+)?$"))
+        if file_names:
+            fn_pattern = "|".join(re.escape(fn) for fn in file_names)
+            patterns.append(re.compile(rf".*[\/]?({fn_pattern})$"))
+        if patterns:
+            return DataFilesList(
+                [data_file for data_file in self if any(pattern.match(data_file) for pattern in patterns)],
+                origin_metadata=self.origin_metadata,
+            )
+        else:
+            return DataFilesList(list(self), origin_metadata=self.origin_metadata)
+
+
+class DataFilesDict(dict[str, DataFilesList]):
+    """
+    Dict of split_name -> list of data files (absolute local paths or URLs).
+    It has two construction methods given the user's data files patterns :
+    - ``from_hf_repo``: resolve patterns inside a dataset repository
+    - ``from_local_or_remote``: resolve patterns from a local path
+
+    Moreover, each list is a DataFilesList. It is possible to hash the dictionary
+    and get a different hash if and only if at least one file changed.
+    For more info, see [`DataFilesList`].
+
+    This is useful for caching Dataset objects that are obtained from a list of data files.
+
+    Changing the order of the keys of this dictionary also doesn't change its hash.
+    """
+
+    @classmethod
+    def from_local_or_remote(
+        cls,
+        patterns: dict[str, Union[list[str], DataFilesList]],
+        base_path: Optional[str] = None,
+        allowed_extensions: Optional[list[str]] = None,
+        download_config: Optional[DownloadConfig] = None,
+    ) -> "DataFilesDict":
+        out = cls()
+        for key, patterns_for_key in patterns.items():
+            out[key] = (
+                patterns_for_key
+                if isinstance(patterns_for_key, DataFilesList)
+                else DataFilesList.from_local_or_remote(
+                    patterns_for_key,
+                    base_path=base_path,
+                    allowed_extensions=allowed_extensions,
+                    download_config=download_config,
+                )
+            )
+        return out
+
+    @classmethod
+    def from_hf_repo(
+        cls,
+        patterns: dict[str, Union[list[str], DataFilesList]],
+        dataset_info: huggingface_hub.hf_api.DatasetInfo,
+        base_path: Optional[str] = None,
+        allowed_extensions: Optional[list[str]] = None,
+        download_config: Optional[DownloadConfig] = None,
+    ) -> "DataFilesDict":
+        out = cls()
+        for key, patterns_for_key in patterns.items():
+            out[key] = (
+                patterns_for_key
+                if isinstance(patterns_for_key, DataFilesList)
+                else DataFilesList.from_hf_repo(
+                    patterns_for_key,
+                    dataset_info=dataset_info,
+                    base_path=base_path,
+                    allowed_extensions=allowed_extensions,
+                    download_config=download_config,
+                )
+            )
+        return out
+
+    @classmethod
+    def from_patterns(
+        cls,
+        patterns: dict[str, Union[list[str], DataFilesList]],
+        base_path: Optional[str] = None,
+        allowed_extensions: Optional[list[str]] = None,
+        download_config: Optional[DownloadConfig] = None,
+    ) -> "DataFilesDict":
+        out = cls()
+        for key, patterns_for_key in patterns.items():
+            out[key] = (
+                patterns_for_key
+                if isinstance(patterns_for_key, DataFilesList)
+                else DataFilesList.from_patterns(
+                    patterns_for_key,
+                    base_path=base_path,
+                    allowed_extensions=allowed_extensions,
+                    download_config=download_config,
+                )
+            )
+        return out
+
+    def filter(
+        self, *, extensions: Optional[list[str]] = None, file_names: Optional[list[str]] = None
+    ) -> "DataFilesDict":
+        out = type(self)()
+        for key, data_files_list in self.items():
+            out[key] = data_files_list.filter(extensions=extensions, file_names=file_names)
+        return out
+
+
+class DataFilesPatternsList(list[str]):
+    """
+    List of data files patterns (absolute local paths or URLs).
+    For each pattern there should also be a list of allowed extensions
+    to keep, or a None ot keep all the files for the pattern.
+    """
+
+    def __init__(
+        self,
+        patterns: list[str],
+        allowed_extensions: list[Optional[list[str]]],
+    ):
+        super().__init__(patterns)
+        self.allowed_extensions = allowed_extensions
+
+    def __add__(self, other):
+        return DataFilesList([*self, *other], self.allowed_extensions + other.allowed_extensions)
+
+    @classmethod
+    def from_patterns(
+        cls, patterns: list[str], allowed_extensions: Optional[list[str]] = None
+    ) -> "DataFilesPatternsList":
+        return cls(patterns, [allowed_extensions] * len(patterns))
+
+    def resolve(
+        self,
+        base_path: str,
+        download_config: Optional[DownloadConfig] = None,
+    ) -> "DataFilesList":
+        base_path = base_path if base_path is not None else Path().resolve().as_posix()
+        data_files = []
+        for pattern, allowed_extensions in zip(self, self.allowed_extensions):
+            try:
+                data_files.extend(
+                    resolve_pattern(
+                        pattern,
+                        base_path=base_path,
+                        allowed_extensions=allowed_extensions,
+                        download_config=download_config,
+                    )
+                )
+            except FileNotFoundError:
+                if not has_magic(pattern):
+                    raise
+        origin_metadata = _get_origin_metadata(data_files, download_config=download_config)
+        return DataFilesList(data_files, origin_metadata)
+
+    def filter_extensions(self, extensions: list[str]) -> "DataFilesPatternsList":
+        return DataFilesPatternsList(
+            self, [allowed_extensions + extensions for allowed_extensions in self.allowed_extensions]
+        )
+
+
+class DataFilesPatternsDict(dict[str, DataFilesPatternsList]):
+    """
+    Dict of split_name -> list of data files patterns (absolute local paths or URLs).
+    """
+
+    @classmethod
+    def from_patterns(
+        cls, patterns: dict[str, list[str]], allowed_extensions: Optional[list[str]] = None
+    ) -> "DataFilesPatternsDict":
+        out = cls()
+        for key, patterns_for_key in patterns.items():
+            out[key] = (
+                patterns_for_key
+                if isinstance(patterns_for_key, DataFilesPatternsList)
+                else DataFilesPatternsList.from_patterns(
+                    patterns_for_key,
+                    allowed_extensions=allowed_extensions,
+                )
+            )
+        return out
+
+    def resolve(
+        self,
+        base_path: str,
+        download_config: Optional[DownloadConfig] = None,
+    ) -> "DataFilesDict":
+        out = DataFilesDict()
+        for key, data_files_patterns_list in self.items():
+            out[key] = data_files_patterns_list.resolve(base_path, download_config)
+        return out
+
+    def filter_extensions(self, extensions: list[str]) -> "DataFilesPatternsDict":
+        out = type(self)()
+        for key, data_files_patterns_list in self.items():
+            out[key] = data_files_patterns_list.filter_extensions(extensions)
+        return out

src/datasets/distributed.py

@@ -0,0 +1,43 @@
+from typing import TypeVar
+
+from .arrow_dataset import Dataset, _split_by_node_map_style_dataset
+from .iterable_dataset import IterableDataset, _split_by_node_iterable_dataset
+
+
+DatasetType = TypeVar("DatasetType", Dataset, IterableDataset)
+
+
+def split_dataset_by_node(dataset: DatasetType, rank: int, world_size: int) -> DatasetType:
+    """
+    Split a dataset for the node at rank `rank` in a pool of nodes of size `world_size`.
+
+    For map-style datasets:
+
+    Each node is assigned a chunk of data, e.g. rank 0 is given the first chunk of the dataset.
+    To maximize data loading throughput, chunks are made of contiguous data on disk if possible.
+
+    For iterable datasets:
+
+    If the dataset has a number of shards that is a factor of `world_size` (i.e. if `dataset.num_shards % world_size == 0`),
+    then the shards are evenly assigned across the nodes, which is the most optimized.
+    Otherwise, each node keeps 1 example out of `world_size`, skipping the other examples.
+
+    > [!WARNING]
+    > If you shuffle your iterable dataset in a distributed setup, make sure to set a fixed `seed` in [`IterableDataset.shuffle`]
+    so the same shuffled list of shards is used on every node to know which shards the node should skip.
+
+    Args:
+        dataset ([`Dataset`] or [`IterableDataset`]):
+            The dataset to split by node.
+        rank (`int`):
+            Rank of the current node.
+        world_size (`int`):
+            Total number of nodes.
+
+    Returns:
+        [`Dataset`] or [`IterableDataset`]: The dataset to be used on the node at rank `rank`.
+    """
+    if isinstance(dataset, Dataset):
+        return _split_by_node_map_style_dataset(dataset, rank=rank, world_size=world_size)
+    else:
+        return _split_by_node_iterable_dataset(dataset, rank=rank, world_size=world_size)

src/datasets/download/__init__.py

@@ -0,0 +1,10 @@
+__all__ = [
+    "DownloadConfig",
+    "DownloadManager",
+    "DownloadMode",
+    "StreamingDownloadManager",
+]
+
+from .download_config import DownloadConfig
+from .download_manager import DownloadManager, DownloadMode
+from .streaming_download_manager import StreamingDownloadManager

src/datasets/download/download_config.py

@@ -0,0 +1,85 @@
+import copy
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Optional, Union
+
+from .. import config
+
+
+@dataclass
+class DownloadConfig:
+    """Configuration for our cached path manager.
+
+    Attributes:
+        cache_dir (`str` or `Path`, *optional*):
+            Specify a cache directory to save the file to (overwrite the
+            default cache dir).
+        force_download (`bool`, defaults to `False`):
+            If `True`, re-download the file even if it's already cached in
+            the cache dir.
+        resume_download (`bool`, defaults to `False`):
+            If `True`, resume the download if an incompletely received file is
+            found.
+        proxies (`dict`, *optional*):
+        user_agent (`str`, *optional*):
+            Optional string or dict that will be appended to the user-agent on remote
+            requests.
+        extract_compressed_file (`bool`, defaults to `False`):
+            If `True` and the path point to a zip or tar file,
+            extract the compressed file in a folder along the archive.
+        force_extract (`bool`, defaults to `False`):
+            If `True` when `extract_compressed_file` is `True` and the archive
+            was already extracted, re-extract the archive and override the folder where it was extracted.
+        delete_extracted (`bool`, defaults to `False`):
+            Whether to delete (or keep) the extracted files.
+        extract_on_the_fly (`bool`, defaults to `False`):
+            If `True`, extract compressed files while they are being read.
+        use_etag (`bool`, defaults to `True`):
+            Whether to use the ETag HTTP response header to validate the cached files.
+        num_proc (`int`, *optional*):
+            The number of processes to launch to download the files in parallel.
+        max_retries (`int`, default to `1`):
+            The number of times to retry an HTTP request if it fails.
+        token (`str` or `bool`, *optional*):
+            Optional string or boolean to use as Bearer token
+            for remote files on the Datasets Hub. If `True`, or not specified, will get token from `~/.huggingface`.
+        storage_options (`dict`, *optional*):
+            Key/value pairs to be passed on to the dataset file-system backend, if any.
+        download_desc (`str`, *optional*):
+            A description to be displayed alongside with the progress bar while downloading the files.
+        disable_tqdm (`bool`, defaults to `False`):
+            Whether to disable the individual files download progress bar
+    """
+
+    cache_dir: Optional[Union[str, Path]] = None
+    force_download: bool = False
+    resume_download: bool = False
+    local_files_only: bool = False
+    proxies: Optional[dict] = None
+    user_agent: Optional[str] = None
+    extract_compressed_file: bool = False
+    force_extract: bool = False
+    delete_extracted: bool = False
+    extract_on_the_fly: bool = False
+    use_etag: bool = True
+    num_proc: Optional[int] = None
+    max_retries: int = 1
+    token: Optional[Union[str, bool]] = None
+    storage_options: dict[str, Any] = field(default_factory=dict)
+    download_desc: Optional[str] = None
+    disable_tqdm: bool = False
+
+    def copy(self) -> "DownloadConfig":
+        return self.__class__(**{k: copy.deepcopy(v) for k, v in self.__dict__.items()})
+
+    def __setattr__(self, name, value):
+        if name == "token" and getattr(self, "storage_options", None) is not None:
+            if "hf" not in self.storage_options:
+                self.storage_options["hf"] = {"endpoint": config.HF_ENDPOINT, "token": value}
+            else:
+                self.storage_options["hf"]["token"] = value
+        super().__setattr__(name, value)
+
+    def __post_init__(self):
+        # update storage_options
+        self.token = self.token

src/datasets/download/download_manager.py

@@ -0,0 +1,340 @@
+# Copyright 2020 The TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Lint as: python3
+"""Download manager interface."""
+
+import enum
+import io
+import multiprocessing
+import os
+from datetime import datetime
+from functools import partial
+from typing import Optional, Union
+
+import fsspec
+from fsspec.core import url_to_fs
+from tqdm.contrib.concurrent import thread_map
+
+from .. import config
+from ..utils import tqdm as hf_tqdm
+from ..utils.file_utils import (
+    ArchiveIterable,
+    FilesIterable,
+    cached_path,
+    is_relative_path,
+    stack_multiprocessing_download_progress_bars,
+    url_or_path_join,
+)
+from ..utils.info_utils import get_size_checksum_dict
+from ..utils.logging import get_logger, tqdm
+from ..utils.py_utils import NestedDataStructure, map_nested
+from ..utils.track import tracked_str
+from .download_config import DownloadConfig
+
+
+logger = get_logger(__name__)
+
+
+class DownloadMode(enum.Enum):
+    """`Enum` for how to treat pre-existing downloads and data.
+
+    The default mode is `REUSE_DATASET_IF_EXISTS`, which will reuse both
+    raw downloads and the prepared dataset if they exist.
+
+    The generations modes:
+
+    |                                     | Downloads | Dataset |
+    |-------------------------------------|-----------|---------|
+    | `REUSE_DATASET_IF_EXISTS` (default) | Reuse     | Reuse   |
+    | `REUSE_CACHE_IF_EXISTS`             | Reuse     | Fresh   |
+    | `FORCE_REDOWNLOAD`                  | Fresh     | Fresh   |
+
+    """
+
+    REUSE_DATASET_IF_EXISTS = "reuse_dataset_if_exists"
+    REUSE_CACHE_IF_EXISTS = "reuse_cache_if_exists"
+    FORCE_REDOWNLOAD = "force_redownload"
+
+
+class DownloadManager:
+    is_streaming = False
+
+    def __init__(
+        self,
+        dataset_name: Optional[str] = None,
+        data_dir: Optional[str] = None,
+        download_config: Optional[DownloadConfig] = None,
+        base_path: Optional[str] = None,
+        record_checksums=False,
+    ):
+        """Download manager constructor.
+
+        Args:
+            data_dir:
+                can be used to specify a manual directory to get the files from.
+            dataset_name (`str`):
+                name of dataset this instance will be used for. If
+                provided, downloads will contain which datasets they were used for.
+            download_config (`DownloadConfig`):
+                to specify the cache directory and other
+                download options
+            base_path (`str`):
+                base path that is used when relative paths are used to
+                download files. This can be a remote url.
+            record_checksums (`bool`, defaults to `False`):
+                Whether to record the checksums of the downloaded files. If None, the value is inferred from the builder.
+        """
+        self._dataset_name = dataset_name
+        self._data_dir = data_dir
+        self._base_path = base_path or os.path.abspath(".")
+        # To record what is being used: {url: {num_bytes: int, checksum: str}}
+        self._recorded_sizes_checksums: dict[str, dict[str, Optional[Union[int, str]]]] = {}
+        self.record_checksums = record_checksums
+        self.download_config = download_config or DownloadConfig()
+        self.downloaded_paths = {}
+        self.extracted_paths = {}
+
+    @property
+    def manual_dir(self):
+        return self._data_dir
+
+    @property
+    def downloaded_size(self):
+        """Returns the total size of downloaded files."""
+        return sum(checksums_dict["num_bytes"] for checksums_dict in self._recorded_sizes_checksums.values())
+
+    def _record_sizes_checksums(self, url_or_urls: NestedDataStructure, downloaded_path_or_paths: NestedDataStructure):
+        """Record size/checksum of downloaded files."""
+        delay = 5
+        for url, path in hf_tqdm(
+            list(zip(url_or_urls.flatten(), downloaded_path_or_paths.flatten())),
+            delay=delay,
+            desc="Computing checksums",
+        ):
+            # call str to support PathLike objects
+            self._recorded_sizes_checksums[str(url)] = get_size_checksum_dict(
+                path, record_checksum=self.record_checksums
+            )
+
+    def download(self, url_or_urls):
+        """Download given URL(s).
+
+        By default, only one process is used for download. Pass customized `download_config.num_proc` to change this behavior.
+
+        Args:
+            url_or_urls (`str` or `list` or `dict`):
+                URL or `list` or `dict` of URLs to download. Each URL is a `str`.
+
+        Returns:
+            `str` or `list` or `dict`:
+                The downloaded paths matching the given input `url_or_urls`.
+
+        Example:
+
+        ```py
+        >>> downloaded_files = dl_manager.download('https://storage.googleapis.com/seldon-datasets/sentence_polarity_v1/rt-polaritydata.tar.gz')
+        ```
+        """
+        download_config = self.download_config.copy()
+        download_config.extract_compressed_file = False
+        if download_config.download_desc is None:
+            download_config.download_desc = "Downloading data"
+
+        download_func = partial(self._download_batched, download_config=download_config)
+
+        start_time = datetime.now()
+        with stack_multiprocessing_download_progress_bars():
+            downloaded_path_or_paths = map_nested(
+                download_func,
+                url_or_urls,
+                map_tuple=True,
+                num_proc=download_config.num_proc,
+                desc="Downloading data files",
+                batched=True,
+                batch_size=-1,
+            )
+        duration = datetime.now() - start_time
+        logger.info(f"Downloading took {duration.total_seconds() // 60} min")
+        url_or_urls = NestedDataStructure(url_or_urls)
+        downloaded_path_or_paths = NestedDataStructure(downloaded_path_or_paths)
+        self.downloaded_paths.update(dict(zip(url_or_urls.flatten(), downloaded_path_or_paths.flatten())))
+
+        start_time = datetime.now()
+        self._record_sizes_checksums(url_or_urls, downloaded_path_or_paths)
+        duration = datetime.now() - start_time
+        logger.info(f"Checksum Computation took {duration.total_seconds() // 60} min")
+
+        return downloaded_path_or_paths.data
+
+    def _download_batched(
+        self,
+        url_or_filenames: list[str],
+        download_config: DownloadConfig,
+    ) -> list[str]:
+        if len(url_or_filenames) >= 16:
+            download_config = download_config.copy()
+            download_config.disable_tqdm = True
+            download_func = partial(self._download_single, download_config=download_config)
+
+            fs: fsspec.AbstractFileSystem
+            path = str(url_or_filenames[0])
+            if is_relative_path(path):
+                # append the relative path to the base_path
+                path = url_or_path_join(self._base_path, path)
+            fs, path = url_to_fs(path, **download_config.storage_options)
+            size = 0
+            try:
+                size = fs.info(path).get("size", 0)
+            except Exception:
+                pass
+            max_workers = (
+                config.HF_DATASETS_MULTITHREADING_MAX_WORKERS if size < (20 << 20) else 1
+            )  # enable multithreading if files are small
+
+            return thread_map(
+                download_func,
+                url_or_filenames,
+                desc=download_config.download_desc or "Downloading",
+                unit="files",
+                position=multiprocessing.current_process()._identity[-1]  # contains the ranks of subprocesses
+                if os.environ.get("HF_DATASETS_STACK_MULTIPROCESSING_DOWNLOAD_PROGRESS_BARS") == "1"
+                and multiprocessing.current_process()._identity
+                else None,
+                max_workers=max_workers,
+                tqdm_class=tqdm,
+            )
+        else:
+            return [
+                self._download_single(url_or_filename, download_config=download_config)
+                for url_or_filename in url_or_filenames
+            ]
+
+    def _download_single(self, url_or_filename: str, download_config: DownloadConfig) -> str:
+        url_or_filename = str(url_or_filename)
+        if is_relative_path(url_or_filename):
+            # append the relative path to the base_path
+            url_or_filename = url_or_path_join(self._base_path, url_or_filename)
+        out = cached_path(url_or_filename, download_config=download_config)
+        out = tracked_str(out)
+        out.set_origin(url_or_filename)
+        return out
+
+    def iter_archive(self, path_or_buf: Union[str, io.BufferedReader]):
+        """Iterate over files within an archive.
+
+        Args:
+            path_or_buf (`str` or `io.BufferedReader`):
+                Archive path or archive binary file object.
+
+        Yields:
+            `tuple[str, io.BufferedReader]`:
+                2-tuple (path_within_archive, file_object).
+                File object is opened in binary mode.
+
+        Example:
+
+        ```py
+        >>> archive = dl_manager.download('https://storage.googleapis.com/seldon-datasets/sentence_polarity_v1/rt-polaritydata.tar.gz')
+        >>> files = dl_manager.iter_archive(archive)
+        ```
+        """
+
+        if hasattr(path_or_buf, "read"):
+            return ArchiveIterable.from_buf(path_or_buf)
+        else:
+            return ArchiveIterable.from_urlpath(path_or_buf)
+
+    def iter_files(self, paths: Union[str, list[str]]):
+        """Iterate over file paths.
+
+        Args:
+            paths (`str` or `list` of `str`):
+                Root paths.
+
+        Yields:
+            `str`: File path.
+
+        Example:
+
+        ```py
+        >>> files = dl_manager.download_and_extract('https://huggingface.co/datasets/AI-Lab-Makerere/beans/resolve/main/data/train.zip')
+        >>> files = dl_manager.iter_files(files)
+        ```
+        """
+        return FilesIterable.from_urlpaths(paths)
+
+    def extract(self, path_or_paths):
+        """Extract given path(s).
+
+        Args:
+            path_or_paths (path or `list` or `dict`):
+                Path of file to extract. Each path is a `str`.
+
+        Returns:
+            extracted_path(s): `str`, The extracted paths matching the given input
+            path_or_paths.
+
+        Example:
+
+        ```py
+        >>> downloaded_files = dl_manager.download('https://storage.googleapis.com/seldon-datasets/sentence_polarity_v1/rt-polaritydata.tar.gz')
+        >>> extracted_files = dl_manager.extract(downloaded_files)
+        ```
+        """
+        download_config = self.download_config.copy()
+        download_config.extract_compressed_file = True
+        extract_func = partial(self._download_single, download_config=download_config)
+        extracted_paths = map_nested(
+            extract_func,
+            path_or_paths,
+            num_proc=download_config.num_proc,
+            desc="Extracting data files",
+        )
+        path_or_paths = NestedDataStructure(path_or_paths)
+        extracted_paths = NestedDataStructure(extracted_paths)
+        self.extracted_paths.update(dict(zip(path_or_paths.flatten(), extracted_paths.flatten())))
+        return extracted_paths.data
+
+    def download_and_extract(self, url_or_urls):
+        """Download and extract given `url_or_urls`.
+
+        Is roughly equivalent to:
+
+        ```
+        extracted_paths = dl_manager.extract(dl_manager.download(url_or_urls))
+        ```
+
+        Args:
+            url_or_urls (`str` or `list` or `dict`):
+                URL or `list` or `dict` of URLs to download and extract. Each URL is a `str`.
+
+        Returns:
+            extracted_path(s): `str`, extracted paths of given URL(s).
+        """
+        return self.extract(self.download(url_or_urls))
+
+    def get_recorded_sizes_checksums(self):
+        return self._recorded_sizes_checksums.copy()
+
+    def delete_extracted_files(self):
+        paths_to_delete = set(self.extracted_paths.values()) - set(self.downloaded_paths.values())
+        for key, path in list(self.extracted_paths.items()):
+            if path in paths_to_delete and os.path.isfile(path):
+                os.remove(path)
+                del self.extracted_paths[key]
+
+    def manage_extracted_files(self):
+        if self.download_config.delete_extracted:
+            self.delete_extracted_files()

src/datasets/download/streaming_download_manager.py

@@ -0,0 +1,219 @@
+import io
+import os
+from collections.abc import Iterable
+from typing import Optional, Union
+
+from ..utils.file_utils import (  # noqa: F401 # backward compatibility
+    SINGLE_FILE_COMPRESSION_PROTOCOLS,
+    ArchiveIterable,
+    FilesIterable,
+    _get_extraction_protocol,
+    _get_path_extension,
+    _prepare_path_and_storage_options,
+    is_relative_path,
+    url_or_path_join,
+    xbasename,
+    xdirname,
+    xet_parse,
+    xexists,
+    xgetsize,
+    xglob,
+    xgzip_open,
+    xisdir,
+    xisfile,
+    xjoin,
+    xlistdir,
+    xnumpy_load,
+    xopen,
+    xpandas_read_csv,
+    xpandas_read_excel,
+    xPath,
+    xpyarrow_parquet_read_table,
+    xrelpath,
+    xsio_loadmat,
+    xsplit,
+    xsplitext,
+    xwalk,
+    xxml_dom_minidom_parse,
+)
+from ..utils.logging import get_logger
+from ..utils.py_utils import map_nested
+from .download_config import DownloadConfig
+
+
+logger = get_logger(__name__)
+
+
+class StreamingDownloadManager:
+    """
+    Download manager that uses the "::" separator to navigate through (possibly remote) compressed archives.
+    Contrary to the regular `DownloadManager`, the `download` and `extract` methods don't actually download nor extract
+    data, but they rather return the path or url that could be opened using the `xopen` function which extends the
+    built-in `open` function to stream data from remote files.
+    """
+
+    is_streaming = True
+
+    def __init__(
+        self,
+        dataset_name: Optional[str] = None,
+        data_dir: Optional[str] = None,
+        download_config: Optional[DownloadConfig] = None,
+        base_path: Optional[str] = None,
+    ):
+        self._dataset_name = dataset_name
+        self._data_dir = data_dir
+        self._base_path = base_path or os.path.abspath(".")
+        self.download_config = download_config or DownloadConfig()
+        self.downloaded_size = None
+        self.record_checksums = False
+
+    @property
+    def manual_dir(self):
+        return self._data_dir
+
+    def download(self, url_or_urls):
+        """Normalize URL(s) of files to stream data from.
+        This is the lazy version of `DownloadManager.download` for streaming.
+
+        Args:
+            url_or_urls (`str` or `list` or `dict`):
+                URL(s) of files to stream data from. Each url is a `str`.
+
+        Returns:
+            url(s): (`str` or `list` or `dict`), URL(s) to stream data from matching the given input url_or_urls.
+
+        Example:
+
+        ```py
+        >>> downloaded_files = dl_manager.download('https://storage.googleapis.com/seldon-datasets/sentence_polarity_v1/rt-polaritydata.tar.gz')
+        ```
+        """
+        url_or_urls = map_nested(self._download_single, url_or_urls, map_tuple=True)
+        return url_or_urls
+
+    def _download_single(self, urlpath: str) -> str:
+        urlpath = str(urlpath)
+        if is_relative_path(urlpath):
+            # append the relative path to the base_path
+            urlpath = url_or_path_join(self._base_path, urlpath)
+        return urlpath
+
+    def extract(self, url_or_urls):
+        """Add extraction protocol for given url(s) for streaming.
+
+        This is the lazy version of `DownloadManager.extract` for streaming.
+
+        Args:
+            url_or_urls (`str` or `list` or `dict`):
+                URL(s) of files to stream data from. Each url is a `str`.
+
+        Returns:
+            url(s): (`str` or `list` or `dict`), URL(s) to stream data from matching the given input `url_or_urls`.
+
+        Example:
+
+        ```py
+        >>> downloaded_files = dl_manager.download('https://storage.googleapis.com/seldon-datasets/sentence_polarity_v1/rt-polaritydata.tar.gz')
+        >>> extracted_files = dl_manager.extract(downloaded_files)
+        ```
+        """
+        urlpaths = map_nested(self._extract, url_or_urls, map_tuple=True)
+        return urlpaths
+
+    def _extract(self, urlpath: str) -> str:
+        urlpath = str(urlpath)
+        # get inner file: zip://train-00000.json.gz::https://foo.bar/data.zip -> zip://train-00000.json.gz
+        protocol = _get_extraction_protocol(urlpath, download_config=self.download_config)
+        path = urlpath.split("::")[0]
+        extension = _get_path_extension(path)
+        if extension in ["tgz", "tar"] or path.endswith((".tar.gz", ".tar.bz2", ".tar.xz")):
+            raise NotImplementedError(
+                f"Extraction protocol for TAR archives like '{urlpath}' is not implemented in streaming mode. "
+                f"Please use `dl_manager.iter_archive` instead.\n\n"
+                f"Example usage:\n\n"
+                f"\turl = dl_manager.download(url)\n"
+                f"\ttar_archive_iterator = dl_manager.iter_archive(url)\n\n"
+                f"\tfor filename, file in tar_archive_iterator:\n"
+                f"\t\t..."
+            )
+        if protocol is None:
+            # no extraction
+            return urlpath
+        elif protocol in SINGLE_FILE_COMPRESSION_PROTOCOLS:
+            # there is one single file which is the uncompressed file
+            inner_file = os.path.basename(urlpath.split("::")[0])
+            inner_file = inner_file[: inner_file.rindex(".")] if "." in inner_file else inner_file
+            return f"{protocol}://{inner_file}::{urlpath}"
+        else:
+            return f"{protocol}://::{urlpath}"
+
+    def download_and_extract(self, url_or_urls):
+        """Prepare given `url_or_urls` for streaming (add extraction protocol).
+
+        This is the lazy version of `DownloadManager.download_and_extract` for streaming.
+
+        Is equivalent to:
+
+        ```
+        urls = dl_manager.extract(dl_manager.download(url_or_urls))
+        ```
+
+        Args:
+            url_or_urls (`str` or `list` or `dict`):
+                URL(s) to stream from data from. Each url is a `str`.
+
+        Returns:
+            url(s): (`str` or `list` or `dict`), URL(s) to stream data from matching the given input `url_or_urls`.
+        """
+        return self.extract(self.download(url_or_urls))
+
+    def iter_archive(self, urlpath_or_buf: Union[str, io.BufferedReader]) -> Iterable[tuple]:
+        """Iterate over files within an archive.
+
+        Args:
+            urlpath_or_buf (`str` or `io.BufferedReader`):
+                Archive path or archive binary file object.
+
+        Yields:
+            `tuple[str, io.BufferedReader]`:
+                2-tuple (path_within_archive, file_object).
+                File object is opened in binary mode.
+
+        Example:
+
+        ```py
+        >>> archive = dl_manager.download('https://storage.googleapis.com/seldon-datasets/sentence_polarity_v1/rt-polaritydata.tar.gz')
+        >>> files = dl_manager.iter_archive(archive)
+        ```
+        """
+
+        if hasattr(urlpath_or_buf, "read"):
+            return ArchiveIterable.from_buf(urlpath_or_buf)
+        else:
+            return ArchiveIterable.from_urlpath(urlpath_or_buf, download_config=self.download_config)
+
+    def iter_files(self, urlpaths: Union[str, list[str]]) -> Iterable[str]:
+        """Iterate over files.
+
+        Args:
+            urlpaths (`str` or `list` of `str`):
+                Root paths.
+
+        Yields:
+            str: File URL path.
+
+        Example:
+
+        ```py
+        >>> files = dl_manager.download_and_extract('https://huggingface.co/datasets/AI-Lab-Makerere/beans/resolve/main/data/train.zip')
+        >>> files = dl_manager.iter_files(files)
+        ```
+        """
+        return FilesIterable.from_urlpaths(urlpaths, download_config=self.download_config)
+
+    def manage_extracted_files(self):
+        pass
+
+    def get_recorded_sizes_checksums(self):
+        pass

src/datasets/exceptions.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2023 The HuggingFace Authors.
+from typing import Any, Optional, Union
+
+from huggingface_hub import HfFileSystem
+
+from . import config
+from .table import CastError
+from .utils.track import TrackedIterableFromGenerator, tracked_list, tracked_str
+
+
+class DatasetsError(Exception):
+    """Base class for exceptions in this library."""
+
+
+class DefunctDatasetError(DatasetsError):
+    """The dataset has been defunct."""
+
+
+class FileNotFoundDatasetsError(DatasetsError, FileNotFoundError):
+    """FileNotFoundError raised by this library."""
+
+
+class DataFilesNotFoundError(FileNotFoundDatasetsError):
+    """No (supported) data files found."""
+
+
+class DatasetNotFoundError(FileNotFoundDatasetsError):
+    """Dataset not found.
+
+    Raised when trying to access:
+    - a missing dataset, or
+    - a private/gated dataset and the user is not authenticated.
+    """
+
+
+class DatasetBuildError(DatasetsError):
+    pass
+
+
+class ManualDownloadError(DatasetBuildError):
+    pass
+
+
+class FileFormatError(DatasetBuildError):
+    pass
+
+
+class DatasetGenerationError(DatasetBuildError):
+    pass
+
+
+class DatasetGenerationCastError(DatasetGenerationError):
+    @classmethod
+    def from_cast_error(
+        cls,
+        cast_error: CastError,
+        builder_name: str,
+        gen_kwargs: dict[str, Any],
+        token: Optional[Union[bool, str]],
+    ) -> "DatasetGenerationCastError":
+        explanation_message = (
+            f"\n\nAll the data files must have the same columns, but at some point {cast_error.details()}"
+        )
+        formatted_tracked_gen_kwargs: list[str] = []
+        for gen_kwarg in gen_kwargs.values():
+            if not isinstance(gen_kwarg, (tracked_str, tracked_list, TrackedIterableFromGenerator)):
+                continue
+            while (
+                isinstance(gen_kwarg, (tracked_list, TrackedIterableFromGenerator)) and gen_kwarg.last_item is not None
+            ):
+                gen_kwarg = gen_kwarg.last_item
+            if isinstance(gen_kwarg, tracked_str):
+                gen_kwarg = gen_kwarg.get_origin()
+            if isinstance(gen_kwarg, str) and gen_kwarg.startswith("hf://"):
+                resolved_path = HfFileSystem(endpoint=config.HF_ENDPOINT, token=token).resolve_path(gen_kwarg)
+                gen_kwarg = "hf://" + resolved_path.unresolve()
+                if "@" + resolved_path.revision in gen_kwarg:
+                    gen_kwarg = (
+                        gen_kwarg.replace("@" + resolved_path.revision, "", 1)
+                        + f" (at revision {resolved_path.revision})"
+                    )
+            formatted_tracked_gen_kwargs.append(str(gen_kwarg))
+        if formatted_tracked_gen_kwargs:
+            explanation_message += f"\n\nThis happened while the {builder_name} dataset builder was generating data using\n\n{', '.join(formatted_tracked_gen_kwargs)}"
+        help_message = "\n\nPlease either edit the data files to have matching columns, or separate them into different configurations (see docs at https://hf.co/docs/hub/datasets-manual-configuration#multiple-configurations)"
+        return cls("An error occurred while generating the dataset" + explanation_message + help_message)
+
+
+class ChecksumVerificationError(DatasetsError):
+    """Error raised during checksums verifications of downloaded files."""
+
+
+class UnexpectedDownloadedFileError(ChecksumVerificationError):
+    """Some downloaded files were not expected."""
+
+
+class ExpectedMoreDownloadedFilesError(ChecksumVerificationError):
+    """Some files were supposed to be downloaded but were not."""
+
+
+class NonMatchingChecksumError(ChecksumVerificationError):
+    """The downloaded file checksum don't match the expected checksum."""
+
+
+class SplitsVerificationError(DatasetsError):
+    """Error raised during splits verifications."""
+
+
+class UnexpectedSplitsError(SplitsVerificationError):
+    """The expected splits of the downloaded file is missing."""
+
+
+class ExpectedMoreSplitsError(SplitsVerificationError):
+    """Some recorded splits are missing."""
+
+
+class NonMatchingSplitsSizesError(SplitsVerificationError):
+    """The splits sizes don't match the expected splits sizes."""

src/datasets/features/__init__.py

@@ -0,0 +1,29 @@
+__all__ = [
+    "Audio",
+    "Array2D",
+    "Array3D",
+    "Array4D",
+    "Array5D",
+    "ClassLabel",
+    "Features",
+    "Json",
+    "LargeList",
+    "List",
+    "Sequence",
+    "Value",
+    "Image",
+    "Mesh",
+    "Translation",
+    "TranslationVariableLanguages",
+    "Video",
+    "Pdf",
+    "Nifti",
+]
+from .audio import Audio
+from .features import Array2D, Array3D, Array4D, Array5D, ClassLabel, Features, Json, LargeList, List, Sequence, Value
+from .image import Image
+from .mesh import Mesh
+from .nifti import Nifti
+from .pdf import Pdf
+from .translation import Translation, TranslationVariableLanguages
+from .video import Video

src/datasets/features/_torchcodec.py

@@ -0,0 +1,15 @@
+import numpy as np
+from torchcodec.decoders import AudioDecoder as _AudioDecoder
+
+
+class AudioDecoder(_AudioDecoder):
+    def __getitem__(self, key: str):
+        if key == "array":
+            y = self.get_all_samples().data.cpu().numpy()
+            return np.mean(y, axis=tuple(range(y.ndim - 1))) if y.ndim > 1 else y
+        elif key == "sampling_rate":
+            return self.get_samples_played_in_range(0, 0).sample_rate
+        elif hasattr(super(), "__getitem__"):
+            return super().__getitem__(key)
+        else:
+            raise TypeError("'torchcodec.decoders.AudioDecoder' object is not subscriptable")

src/datasets/features/audio.py

@@ -0,0 +1,365 @@
+import os
+from dataclasses import dataclass, field
+from io import BytesIO
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
+
+import numpy as np
+import pyarrow as pa
+
+from .. import config
+from ..download.download_config import DownloadConfig
+from ..table import array_cast
+from ..utils.file_utils import is_local_path, is_remote_url, xopen
+from ..utils.py_utils import no_op_if_value_is_null, string_to_dict
+
+
+if TYPE_CHECKING:
+    from torchcodec.decoders import AudioDecoder
+
+    from .features import FeatureType
+
+
+@dataclass
+class Audio:
+    """Audio [`Feature`] to extract audio data from an audio file.
+
+    Input: The Audio feature accepts as input:
+    - A `str`: Absolute path to the audio file (i.e. random access is allowed).
+    - A `pathlib.Path`: path to the audio file (i.e. random access is allowed).
+    - A `dict` with the keys:
+
+        - `path`: String with relative path of the audio file to the archive file.
+        - `bytes`: Bytes content of the audio file.
+
+      This is useful for parquet or webdataset files which embed audio files.
+
+    - A `dict` with the keys:
+
+        - `array`: Array containing the audio sample
+        - `sampling_rate`: Integer corresponding to the sampling rate of the audio sample.
+
+    - A `torchcodec.decoders.AudioDecoder`: torchcodec audio decoder object.
+
+    Output: The Audio features output data as `torchcodec.decoders.AudioDecoder` objects, with additional keys:
+
+    - `array`: Array containing the audio sample
+    - `sampling_rate`: Integer corresponding to the sampling rate of the audio sample.
+
+    Args:
+        sampling_rate (`int`, *optional*):
+            Target sampling rate. If `None`, the native sampling rate is used.
+        num_channels (`int`, *optional*):
+             The desired number of channels of the samples. By default, the number of channels of the source is used.
+             Audio decoding will return samples with shape (num_channels, num_samples)
+             Currently `None` (number of channels of the source, default), `1` (mono) or `2` (stereo) channels are supported.
+             The `num_channels` argument is passed to `torchcodec.decoders.AudioDecoder`.
+
+             <Added version="4.4.0"/>
+        decode (`bool`, defaults to `True`):
+            Whether to decode the audio data. If `False`,
+            returns the underlying dictionary in the format `{"path": audio_path, "bytes": audio_bytes}`.
+        stream_index (`int`, *optional*):
+            The streaming index to use from the file. If `None` defaults to the "best" index.
+
+    Example:
+
+    ```py
+    >>> from datasets import load_dataset, Audio
+    >>> ds = load_dataset("PolyAI/minds14", name="en-US", split="train")
+    >>> ds = ds.cast_column("audio", Audio(sampling_rate=44100, num_channels=2))
+    >>> ds[0]["audio"]
+    <datasets.features._torchcodec.AudioDecoder object at 0x11642b6a0>
+    >>> audio = ds[0]["audio"]
+    >>> audio.get_samples_played_in_range(0, 10)
+    AudioSamples:
+        data (shape): torch.Size([2, 110592])
+        pts_seconds: 0.0
+        duration_seconds: 2.507755102040816
+        sample_rate: 44100
+    ```
+    """
+
+    sampling_rate: Optional[int] = None
+    decode: bool = True
+    num_channels: Optional[int] = None
+    stream_index: Optional[int] = None
+    id: Optional[str] = field(default=None, repr=False)
+    # Automatically constructed
+    dtype: ClassVar[str] = "dict"
+    pa_type: ClassVar[Any] = pa.struct({"bytes": pa.binary(), "path": pa.string()})
+    _type: str = field(default="Audio", init=False, repr=False)
+
+    def __call__(self):
+        return self.pa_type
+
+    def encode_example(self, value: Union[str, bytes, bytearray, dict, "AudioDecoder"]) -> dict:
+        """Encode example into a format for Arrow.
+
+        Args:
+            value (`str`, `bytes`,`bytearray`,`dict`, `AudioDecoder`):
+                Data passed as input to Audio feature.
+
+        Returns:
+            `dict`
+        """
+        try:
+            import torch
+            from torchcodec.encoders import AudioEncoder  # needed to write audio files
+        except ImportError as err:
+            raise ImportError("To support encoding audio data, please install 'torchcodec'.") from err
+
+        if value is None:
+            raise ValueError("value must be provided")
+
+        if config.TORCHCODEC_AVAILABLE:
+            from torchcodec.decoders import AudioDecoder
+
+        else:
+            AudioDecoder = None
+
+        if isinstance(value, str):
+            return {"bytes": None, "path": value}
+        elif isinstance(value, Path):
+            return {"bytes": None, "path": str(value.absolute())}
+        elif isinstance(value, (bytes, bytearray)):
+            return {"bytes": value, "path": None}
+        elif AudioDecoder is not None and isinstance(value, AudioDecoder):
+            return encode_torchcodec_audio(value)
+        elif "array" in value:
+            # convert the audio array to wav bytes
+            buffer = BytesIO()
+            AudioEncoder(
+                torch.from_numpy(value["array"].astype(np.float32)), sample_rate=value["sampling_rate"]
+            ).to_file_like(buffer, format="wav", num_channels=self.num_channels)
+            return {"bytes": buffer.getvalue(), "path": None}
+        elif value.get("path") is not None and os.path.isfile(value["path"]):
+            # we set "bytes": None to not duplicate the data if they're already available locally
+            if value["path"].endswith("pcm"):
+                # "PCM" only has raw audio bytes
+                if value.get("sampling_rate") is None:
+                    # At least, If you want to convert "PCM-byte" to "WAV-byte", you have to know sampling rate
+                    raise KeyError("To use PCM files, please specify a 'sampling_rate' in Audio object")
+                if value.get("bytes"):
+                    # If we already had PCM-byte, we don`t have to make "read file, make bytes" (just use it!)
+                    bytes_value = np.frombuffer(value["bytes"], dtype=np.int16).astype(np.float32) / 32767
+                else:
+                    bytes_value = np.memmap(value["path"], dtype="h", mode="r").astype(np.float32) / 32767
+
+                buffer = BytesIO()
+                AudioEncoder(torch.from_numpy(bytes_value), sample_rate=value["sampling_rate"]).to_file_like(
+                    buffer, format="wav", num_channels=self.num_channels
+                )
+                return {"bytes": buffer.getvalue(), "path": None}
+            else:
+                return {"bytes": None, "path": value.get("path")}
+        elif value.get("bytes") is not None or value.get("path") is not None:
+            # store the audio bytes, and path is used to infer the audio format using the file extension
+            return {"bytes": value.get("bytes"), "path": value.get("path")}
+        else:
+            raise ValueError(
+                f"An audio sample should have one of 'path' or 'bytes' but they are missing or None in {value}."
+            )
+
+    def decode_example(
+        self, value: dict, token_per_repo_id: Optional[dict[str, Union[str, bool, None]]] = None
+    ) -> "AudioDecoder":
+        """Decode example audio file into audio data.
+
+        Args:
+            value (`dict`):
+                A dictionary with keys:
+
+                - `path`: String with relative audio file path.
+                - `bytes`: Bytes of the audio file.
+            token_per_repo_id (`dict`, *optional*):
+                To access and decode
+                audio files from private repositories on the Hub, you can pass
+                a dictionary repo_id (`str`) -> token (`bool` or `str`)
+
+        Returns:
+            `torchcodec.decoders.AudioDecoder`
+        """
+        if config.TORCHCODEC_AVAILABLE:
+            from ._torchcodec import AudioDecoder
+        else:
+            raise ImportError("To support decoding audio data, please install 'torchcodec'.")
+
+        if not self.decode:
+            raise RuntimeError("Decoding is disabled for this feature. Please use Audio(decode=True) instead.")
+
+        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}.")
+
+        if bytes is None and is_local_path(path):
+            audio = AudioDecoder(
+                path, stream_index=self.stream_index, sample_rate=self.sampling_rate, num_channels=self.num_channels
+            )
+
+        elif bytes is None:
+            token_per_repo_id = token_per_repo_id or {}
+            source_url = path.split("::")[-1]
+            pattern = (
+                config.HUB_DATASETS_URL if source_url.startswith(config.HF_ENDPOINT) else config.HUB_DATASETS_HFFS_URL
+            )
+            source_url_fields = string_to_dict(source_url, pattern)
+            token = token_per_repo_id.get(source_url_fields["repo_id"]) if source_url_fields is not None else None
+
+            download_config = DownloadConfig(token=token)
+            f = xopen(path, "rb", download_config=download_config)
+            audio = AudioDecoder(
+                f, stream_index=self.stream_index, sample_rate=self.sampling_rate, num_channels=self.num_channels
+            )
+
+        else:
+            audio = AudioDecoder(
+                bytes, stream_index=self.stream_index, sample_rate=self.sampling_rate, num_channels=self.num_channels
+            )
+        audio._hf_encoded = {"path": path, "bytes": bytes}
+        audio.metadata.path = path
+        return audio
+
+    def flatten(self) -> Union["FeatureType", dict[str, "FeatureType"]]:
+        """If in the decodable state, raise an error, otherwise flatten the feature into a dictionary."""
+        from .features import Value
+
+        if self.decode:
+            raise ValueError("Cannot flatten a decoded Audio feature.")
+        return {
+            "bytes": Value("binary"),
+            "path": Value("string"),
+        }
+
+    def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray]) -> pa.StructArray:
+        """Cast an Arrow array to the Audio arrow storage type.
+        The Arrow types that can be converted to the Audio pyarrow storage type are:
+
+        - `pa.string()` - it must contain the "path" data
+        - `pa.binary()` - it must contain the audio bytes
+        - `pa.struct({"bytes": pa.binary()})`
+        - `pa.struct({"path": pa.string()})`
+        - `pa.struct({"bytes": pa.binary(), "path": pa.string()})`  - order doesn't matter
+
+        Args:
+            storage (`Union[pa.StringArray, pa.StructArray]`):
+                PyArrow array to cast.
+
+        Returns:
+            `pa.StructArray`: Array in the Audio arrow storage type, that is
+                `pa.struct({"bytes": pa.binary(), "path": pa.string()})`
+        """
+        if pa.types.is_string(storage.type):
+            bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            storage = pa.StructArray.from_arrays([bytes_array, storage], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_large_binary(storage.type):
+            storage = array_cast(
+                storage, pa.binary()
+            )  # this can fail in case of big audios, paths should be used instead
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_binary(storage.type):
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_struct(storage.type) and storage.type.get_all_field_indices("array"):
+            storage = pa.array(
+                [Audio().encode_example(x) if x is not None else None for x in storage.to_numpy(zero_copy_only=False)]
+            )
+        elif pa.types.is_struct(storage.type):
+            if storage.type.get_field_index("bytes") >= 0:
+                bytes_array = storage.field("bytes")
+            else:
+                bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            if storage.type.get_field_index("path") >= 0:
+                path_array = storage.field("path")
+            else:
+                path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
+        return array_cast(storage, self.pa_type)
+
+    def embed_storage(
+        self, storage: pa.StructArray, token_per_repo_id=None, local_files: bool = True, remote_files: bool = True
+    ) -> pa.StructArray:
+        """Embed audio files into the Arrow array.
+
+        Args:
+            storage (`pa.StructArray`):
+                PyArrow array to embed.
+            token_per_repo_id (`dict`, optional):
+                Dictionary repo_id -> token to fetch the files bytes.
+            local_files (`bool`, defaults to `True`)
+                Whether to embed local files data in the array
+
+                <Added version="4.8.5"/>
+            remote_files (`bool`, defaults to `True`)
+                Whether to embed remote files data in the array.
+                E.g. files with paths that start with hf:// or https://
+
+                <Added version="4.8.5"/>
+
+        Returns:
+            `pa.StructArray`: Array in the Audio arrow storage type, that is
+                `pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
+        """
+        if token_per_repo_id is None:
+            token_per_repo_id = {}
+
+        @no_op_if_value_is_null
+        def path_to_bytes(path):
+            source_url = path.split("::")[-1]
+            pattern = (
+                config.HUB_DATASETS_URL if source_url.startswith(config.HF_ENDPOINT) else config.HUB_DATASETS_HFFS_URL
+            )
+            source_url_fields = string_to_dict(source_url, pattern)
+            token = token_per_repo_id.get(source_url_fields["repo_id"]) if source_url_fields is not None else None
+            download_config = DownloadConfig(token=token)
+            with xopen(path, "rb", download_config=download_config) as f:
+                return f.read()
+
+        bytes_array = pa.array(
+            [
+                (
+                    path_to_bytes(x["path"])
+                    if x["bytes"] is None
+                    and ((local_files and is_local_path(x["path"])) or (remote_files and is_remote_url(x["path"])))
+                    else x["bytes"]
+                )
+                if x is not None
+                else None
+                for x in storage.to_pylist()
+            ],
+            type=pa.binary(),
+        )
+        path_array = pa.array(
+            [
+                (
+                    os.path.basename(path)
+                    if (local_files and is_local_path(path)) or (remote_files and is_remote_url(path))
+                    else path
+                )
+                if path is not None
+                else None
+                for path in storage.field("path").to_pylist()
+            ],
+            type=pa.string(),
+        )
+        storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
+        return array_cast(storage, self.pa_type)
+
+
+def encode_torchcodec_audio(audio: "AudioDecoder") -> dict:
+    if hasattr(audio, "_hf_encoded"):
+        return audio._hf_encoded
+    else:
+        try:
+            from torchcodec.encoders import AudioEncoder  # needed to write audio files
+        except ImportError as err:
+            raise ImportError("To support encoding audio data, please install 'torchcodec'.") from err
+
+        samples = audio.get_all_samples()
+        buffer = BytesIO()
+        num_channels = samples.data.shape[0]
+        AudioEncoder(samples.data.cpu(), sample_rate=samples.sample_rate).to_file_like(
+            buffer, format="wav", num_channels=num_channels
+        )
+        return {"bytes": buffer.getvalue(), "path": None}

src/datasets/features/image.py

@@ -0,0 +1,441 @@
+import os
+import sys
+import warnings
+from dataclasses import dataclass, field
+from io import BytesIO
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
+
+import numpy as np
+import pyarrow as pa
+
+from .. import config
+from ..download.download_config import DownloadConfig
+from ..table import array_cast
+from ..utils.file_utils import is_local_path, is_remote_url, xopen
+from ..utils.py_utils import first_non_null_value, no_op_if_value_is_null, string_to_dict
+
+
+if TYPE_CHECKING:
+    import PIL.Image
+
+    from .features import FeatureType
+
+
+_IMAGE_COMPRESSION_FORMATS: Optional[list[str]] = None
+_NATIVE_BYTEORDER = "<" if sys.byteorder == "little" else ">"
+# Origin: https://github.com/python-pillow/Pillow/blob/698951e19e19972aeed56df686868f1329981c12/src/PIL/Image.py#L3126 minus "|i1" which values are not preserved correctly when saving and loading an image
+_VALID_IMAGE_ARRAY_DTPYES = [
+    np.dtype("|b1"),
+    np.dtype("|u1"),
+    np.dtype("<u2"),
+    np.dtype(">u2"),
+    np.dtype("<i2"),
+    np.dtype(">i2"),
+    np.dtype("<u4"),
+    np.dtype(">u4"),
+    np.dtype("<i4"),
+    np.dtype(">i4"),
+    np.dtype("<f4"),
+    np.dtype(">f4"),
+    np.dtype("<f8"),
+    np.dtype(">f8"),
+]
+
+
+@dataclass
+class Image:
+    """Image [`Feature`] to read image data from an image file.
+
+    Input: The Image feature accepts as input:
+    - A `str`: Absolute path to the image file (i.e. random access is allowed).
+    - A `pathlib.Path`: path to the image file (i.e. random access is allowed).
+    - A `dict` with the keys:
+
+        - `path`: String with relative path of the image file to the archive file.
+        - `bytes`: Bytes of the image file.
+
+      This is useful for parquet or webdataset files which embed image files.
+
+    - An `np.ndarray`: NumPy array representing an image.
+    - A `PIL.Image.Image`: PIL image object.
+
+    Output: The Image features output data as `PIL.Image.Image` objects.
+
+    Args:
+        mode (`str`, *optional*):
+            The mode to convert the image to. If `None`, the native mode of the image is used.
+        decode (`bool`, defaults to `True`):
+            Whether to decode the image data. If `False`,
+            returns the underlying dictionary in the format `{"path": image_path, "bytes": image_bytes}`.
+
+    Examples:
+
+    ```py
+    >>> from datasets import load_dataset, Image
+    >>> ds = load_dataset("AI-Lab-Makerere/beans", split="train")
+    >>> ds.features["image"]
+    Image(decode=True, id=None)
+    >>> ds[0]["image"]
+    <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=500x500 at 0x15E52E7F0>
+    >>> ds = ds.cast_column('image', Image(decode=False))
+    {'bytes': None,
+     'path': '/root/.cache/huggingface/datasets/downloads/extracted/b0a21163f78769a2cf11f58dfc767fb458fc7cea5c05dccc0144a2c0f0bc1292/train/healthy/healthy_train.85.jpg'}
+    ```
+    """
+
+    mode: Optional[str] = None
+    decode: bool = True
+    id: Optional[str] = field(default=None, repr=False)
+    # Automatically constructed
+    dtype: ClassVar[str] = "PIL.Image.Image"
+    pa_type: ClassVar[Any] = pa.struct({"bytes": pa.binary(), "path": pa.string()})
+    _type: str = field(default="Image", init=False, repr=False)
+
+    def __call__(self):
+        return self.pa_type
+
+    def encode_example(self, value: Union[str, bytes, bytearray, dict, np.ndarray, "PIL.Image.Image"]) -> dict:
+        """Encode example into a format for Arrow.
+
+        Args:
+            value (`str`, `np.ndarray`, `PIL.Image.Image` or `dict`):
+                Data passed as input to Image feature.
+
+        Returns:
+            `dict` with "path" and "bytes" fields
+        """
+        if config.PIL_AVAILABLE:
+            import PIL.Image
+        else:
+            raise ImportError("To support encoding images, please install 'Pillow'.")
+
+        if isinstance(value, list):
+            value = np.array(value)
+
+        if isinstance(value, str):
+            return {"path": value, "bytes": None}
+        elif isinstance(value, Path):
+            return {"path": str(value.absolute()), "bytes": None}
+        elif isinstance(value, (bytes, bytearray)):
+            return {"path": None, "bytes": value}
+        elif isinstance(value, np.ndarray):
+            # convert the image array to PNG/TIFF bytes
+            return encode_np_array(value)
+        elif isinstance(value, PIL.Image.Image):
+            # convert the PIL image to bytes (default format is PNG/TIFF)
+            return encode_pil_image(value)
+        elif value.get("path") is not None and os.path.isfile(value["path"]):
+            # we set "bytes": None to not duplicate the data if they're already available locally
+            return {"bytes": None, "path": value.get("path")}
+        elif value.get("bytes") is not None or value.get("path") is not None:
+            # store the image bytes, and path is used to infer the image format using the file extension
+            return {"bytes": value.get("bytes"), "path": value.get("path")}
+        else:
+            raise ValueError(
+                f"An image sample should have one of 'path' or 'bytes' but they are missing or None in {value}."
+            )
+
+    def decode_example(self, value: dict, token_per_repo_id=None) -> "PIL.Image.Image":
+        """Decode example image file into image data.
+
+        Args:
+            value (`str` or `dict`):
+                A string with the absolute image file path, a dictionary with
+                keys:
+
+                - `path`: String with absolute or relative image file path.
+                - `bytes`: The bytes of the image file.
+            token_per_repo_id (`dict`, *optional*):
+                To access and decode
+                image files from private repositories on the Hub, you can pass
+                a dictionary repo_id (`str`) -> token (`bool` or `str`).
+
+        Returns:
+            `PIL.Image.Image`
+        """
+        if not self.decode:
+            raise RuntimeError("Decoding is disabled for this feature. Please use Image(decode=True) instead.")
+
+        if config.PIL_AVAILABLE:
+            import PIL.Image
+            import PIL.ImageOps
+        else:
+            raise ImportError("To support decoding images, please install 'Pillow'.")
+
+        if token_per_repo_id is None:
+            token_per_repo_id = {}
+
+        path, bytes_ = value["path"], value["bytes"]
+        if bytes_ is None:
+            if path is None:
+                raise ValueError(f"An image should have one of 'path' or 'bytes' but both are None in {value}.")
+            else:
+                if is_local_path(path):
+                    image = PIL.Image.open(path)
+                else:
+                    source_url = path.split("::")[-1]
+                    pattern = (
+                        config.HUB_DATASETS_URL
+                        if source_url.startswith(config.HF_ENDPOINT)
+                        else config.HUB_DATASETS_HFFS_URL
+                    )
+                    source_url_fields = string_to_dict(source_url, pattern)
+                    token = (
+                        token_per_repo_id.get(source_url_fields["repo_id"]) if source_url_fields is not None else None
+                    )
+                    download_config = DownloadConfig(token=token)
+                    with xopen(path, "rb", download_config=download_config) as f:
+                        bytes_ = BytesIO(f.read())
+                    image = PIL.Image.open(bytes_)
+        else:
+            image = PIL.Image.open(BytesIO(bytes_))
+        image.load()  # to avoid "Too many open files" errors
+        if image.getexif().get(PIL.Image.ExifTags.Base.Orientation) is not None:
+            image = PIL.ImageOps.exif_transpose(image)
+        if self.mode and self.mode != image.mode:
+            image = image.convert(self.mode)
+        return image
+
+    def flatten(self) -> Union["FeatureType", dict[str, "FeatureType"]]:
+        """If in the decodable state, return the feature itself, otherwise flatten the feature into a dictionary."""
+        from .features import Value
+
+        return (
+            self
+            if self.decode
+            else {
+                "bytes": Value("binary"),
+                "path": Value("string"),
+            }
+        )
+
+    def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray, pa.ListArray]) -> pa.StructArray:
+        """Cast an Arrow array to the Image arrow storage type.
+        The Arrow types that can be converted to the Image pyarrow storage type are:
+
+        - `pa.string()` - it must contain the "path" data
+        - `pa.large_string()` - it must contain the "path" data (will be cast to string if possible)
+        - `pa.binary()` - it must contain the image bytes
+        - `pa.struct({"bytes": pa.binary()})`
+        - `pa.struct({"path": pa.string()})`
+        - `pa.struct({"bytes": pa.binary(), "path": pa.string()})`  - order doesn't matter
+        - `pa.list(*)` - it must contain the image array data
+
+        Args:
+            storage (`Union[pa.StringArray, pa.StructArray, pa.ListArray]`):
+                PyArrow array to cast.
+
+        Returns:
+            `pa.StructArray`: Array in the Image arrow storage type, that is
+                `pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
+        """
+        if pa.types.is_large_string(storage.type):
+            try:
+                storage = storage.cast(pa.string())
+            except pa.ArrowInvalid as e:
+                raise ValueError(
+                    f"Failed to cast large_string to string for Image feature. "
+                    f"This can happen if string values exceed 2GB. "
+                    f"Original error: {e}"
+                ) from e
+        if pa.types.is_string(storage.type):
+            bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            storage = pa.StructArray.from_arrays([bytes_array, storage], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_large_binary(storage.type):
+            storage = array_cast(
+                storage, pa.binary()
+            )  # this can fail in case of big images, paths should be used instead
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_binary(storage.type):
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_struct(storage.type):
+            if storage.type.get_field_index("bytes") >= 0:
+                bytes_array = storage.field("bytes")
+            else:
+                bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            if storage.type.get_field_index("path") >= 0:
+                path_array = storage.field("path")
+            else:
+                path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_list(storage.type):
+            bytes_array = pa.array(
+                [encode_np_array(np.array(arr))["bytes"] if arr is not None else None for arr in storage.to_pylist()],
+                type=pa.binary(),
+            )
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays(
+                [bytes_array, path_array], ["bytes", "path"], mask=bytes_array.is_null()
+            )
+        return array_cast(storage, self.pa_type)
+
+    def embed_storage(
+        self, storage: pa.StructArray, token_per_repo_id=None, local_files: bool = True, remote_files: bool = True
+    ) -> pa.StructArray:
+        """Embed image files into the Arrow array.
+
+        Args:
+            storage (`pa.StructArray`):
+                PyArrow array to embed.
+            token_per_repo_id (`dict`, optional):
+                Dictionary repo_id -> token to fetch the files bytes.
+            local_files (`bool`, defaults to `True`)
+                Whether to embed local files data in the array
+
+                <Added version="4.8.5"/>
+            remote_files (`bool`, defaults to `True`)
+                Whether to embed remote files data in the array.
+                E.g. files with paths that start with hf:// or https://
+
+                <Added version="4.8.5"/>
+
+        Returns:
+            `pa.StructArray`: Array in the Image arrow storage type, that is
+                `pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
+        """
+        if token_per_repo_id is None:
+            token_per_repo_id = {}
+
+        @no_op_if_value_is_null
+        def path_to_bytes(path):
+            source_url = path.split("::")[-1]
+            pattern = (
+                config.HUB_DATASETS_URL if source_url.startswith(config.HF_ENDPOINT) else config.HUB_DATASETS_HFFS_URL
+            )
+            source_url_fields = string_to_dict(source_url, pattern)
+            token = token_per_repo_id.get(source_url_fields["repo_id"]) if source_url_fields is not None else None
+            download_config = DownloadConfig(token=token)
+            with xopen(path, "rb", download_config=download_config) as f:
+                return f.read()
+
+        bytes_array = pa.array(
+            [
+                (
+                    path_to_bytes(x["path"])
+                    if x["bytes"] is None
+                    and ((local_files and is_local_path(x["path"])) or (remote_files and is_remote_url(x["path"])))
+                    else x["bytes"]
+                )
+                if x is not None
+                else None
+                for x in storage.to_pylist()
+            ],
+            type=pa.binary(),
+        )
+        path_array = pa.array(
+            [
+                (
+                    os.path.basename(path)
+                    if (local_files and is_local_path(path)) or (remote_files and is_remote_url(path))
+                    else path
+                )
+                if path is not None
+                else None
+                for path in storage.field("path").to_pylist()
+            ],
+            type=pa.string(),
+        )
+        storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
+        return array_cast(storage, self.pa_type)
+
+
+def list_image_compression_formats() -> list[str]:
+    if config.PIL_AVAILABLE:
+        import PIL.Image
+    else:
+        raise ImportError("To support encoding images, please install 'Pillow'.")
+
+    global _IMAGE_COMPRESSION_FORMATS
+    if _IMAGE_COMPRESSION_FORMATS is None:
+        PIL.Image.init()
+        _IMAGE_COMPRESSION_FORMATS = list(set(PIL.Image.OPEN.keys()) & set(PIL.Image.SAVE.keys()))
+    return _IMAGE_COMPRESSION_FORMATS
+
+
+def image_to_bytes(image: "PIL.Image.Image") -> bytes:
+    """Convert a PIL Image object to bytes using native compression if possible, otherwise use PNG/TIFF compression."""
+    buffer = BytesIO()
+    if image.format in list_image_compression_formats():
+        format = image.format
+    else:
+        format = "PNG" if image.mode in ["1", "L", "LA", "RGB", "RGBA"] else "TIFF"
+    image.save(buffer, format=format)
+    return buffer.getvalue()
+
+
+def encode_pil_image(image: "PIL.Image.Image") -> dict:
+    if hasattr(image, "filename") and image.filename != "":
+        return {"path": image.filename, "bytes": None}
+    else:
+        return {"path": None, "bytes": image_to_bytes(image)}
+
+
+def encode_np_array(array: np.ndarray) -> dict:
+    if config.PIL_AVAILABLE:
+        import PIL.Image
+    else:
+        raise ImportError("To support encoding images, please install 'Pillow'.")
+
+    dtype = array.dtype
+    dtype_byteorder = dtype.byteorder if dtype.byteorder != "=" else _NATIVE_BYTEORDER
+    dtype_kind = dtype.kind
+    dtype_itemsize = dtype.itemsize
+
+    dest_dtype = None
+
+    # Multi-channel array case (only np.dtype("|u1") is allowed)
+    if array.shape[2:]:
+        if dtype_kind not in ["u", "i"]:
+            raise TypeError(
+                f"Unsupported array dtype {dtype} for image encoding. Only {dest_dtype} is supported for multi-channel arrays."
+            )
+        dest_dtype = np.dtype("|u1")
+        if dtype != dest_dtype:
+            warnings.warn(f"Downcasting array dtype {dtype} to {dest_dtype} to be compatible with 'Pillow'")
+    # Exact match
+    elif dtype in _VALID_IMAGE_ARRAY_DTPYES:
+        dest_dtype = dtype
+    else:  # Downcast the type within the kind (np.can_cast(from_type, to_type, casting="same_kind") doesn't behave as expected, so do it manually)
+        while dtype_itemsize >= 1:
+            dtype_str = dtype_byteorder + dtype_kind + str(dtype_itemsize)
+            if np.dtype(dtype_str) in _VALID_IMAGE_ARRAY_DTPYES:
+                dest_dtype = np.dtype(dtype_str)
+                warnings.warn(f"Downcasting array dtype {dtype} to {dest_dtype} to be compatible with 'Pillow'")
+                break
+            else:
+                dtype_itemsize //= 2
+        if dest_dtype is None:
+            raise TypeError(
+                f"Cannot downcast dtype {dtype} to a valid image dtype. Valid image dtypes: {_VALID_IMAGE_ARRAY_DTPYES}"
+            )
+
+    image = PIL.Image.fromarray(array.astype(dest_dtype))
+    return {"path": None, "bytes": image_to_bytes(image)}
+
+
+def objects_to_list_of_image_dicts(
+    objs: Union[list[str], list[dict], list[np.ndarray], list["PIL.Image.Image"]],
+) -> list[dict]:
+    """Encode a list of objects into a format suitable for creating an extension array of type `ImageExtensionType`."""
+    if config.PIL_AVAILABLE:
+        import PIL.Image
+    else:
+        raise ImportError("To support encoding images, please install 'Pillow'.")
+
+    if objs:
+        _, obj = first_non_null_value(objs)
+        if isinstance(obj, str):
+            return [{"path": obj, "bytes": None} if obj is not None else None for obj in objs]
+        if isinstance(obj, np.ndarray):
+            obj_to_image_dict_func = no_op_if_value_is_null(encode_np_array)
+            return [obj_to_image_dict_func(obj) for obj in objs]
+        elif isinstance(obj, PIL.Image.Image):
+            obj_to_image_dict_func = no_op_if_value_is_null(encode_pil_image)
+            return [obj_to_image_dict_func(obj) for obj in objs]
+        else:
+            return objs
+    else:
+        return objs

src/datasets/features/mesh.py

@@ -0,0 +1,295 @@
+import os
+from dataclasses import dataclass, field
+from io import BytesIO
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
+
+import pyarrow as pa
+
+from .. import config
+from ..download.download_config import DownloadConfig
+from ..table import array_cast
+from ..utils.file_utils import is_local_path, is_remote_url, xopen
+from ..utils.py_utils import string_to_dict
+
+
+if TYPE_CHECKING:
+    import trimesh
+
+    from .features import FeatureType
+
+
+@dataclass
+class Mesh:
+    """Mesh [`Feature`] to read 3D mesh data from a file.
+
+    Input: The Mesh feature accepts as input:
+    - A `str`: Absolute path to the mesh file (i.e. random access is allowed).
+    - A `pathlib.Path`: path to the mesh file (i.e. random access is allowed).
+    - A `dict` with the keys:
+
+        - `path`: String with relative path of the mesh file to the archive file.
+        - `bytes`: Bytes of the mesh file.
+
+      This is useful for parquet or webdataset files which embed mesh files.
+
+    - A `trimesh.Trimesh` or `trimesh.Scene`: 3D mesh or scene object.
+
+    Output: The Mesh feature outputs data as `trimesh.Trimesh` or `trimesh.Scene` objects.
+
+    Args:
+        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.
+    """
+
+    decode: bool = True
+    id: Optional[str] = field(default=None, repr=False)
+    # Automatically constructed
+    dtype: ClassVar[str] = "trimesh.Trimesh"
+    pa_type: ClassVar[Any] = pa.struct({"bytes": pa.binary(), "path": pa.string()})
+    _type: str = field(default="Mesh", init=False, repr=False)
+
+    def __call__(self):
+        return self.pa_type
+
+    def encode_example(self, value: Union[str, bytes, bytearray, dict, "trimesh.Trimesh", "trimesh.Scene"]) -> dict:
+        """Encode example into a format for Arrow.
+
+        Args:
+            value (`str`, `bytes`, `dict`, `trimesh.Trimesh`, or `trimesh.Scene`):
+                Data passed as input to Mesh feature.
+
+        Returns:
+            `dict` with "path" and "bytes" fields
+        """
+        if config.TRIMESH_AVAILABLE:
+            import trimesh
+        else:
+            trimesh = None
+
+        if isinstance(value, str):
+            return {"path": value, "bytes": None}
+        elif isinstance(value, Path):
+            return {"path": str(value.absolute()), "bytes": None}
+        elif isinstance(value, (bytes, bytearray)):
+            return {"path": None, "bytes": value}
+        elif trimesh is not None and isinstance(value, (trimesh.Trimesh, trimesh.Scene)):
+            return encode_trimesh_mesh(value)
+        elif isinstance(value, dict) and value.get("path") is not None and os.path.isfile(value["path"]):
+            # we set "bytes": None to not duplicate the data if they're already available locally
+            return {"bytes": None, "path": value.get("path")}
+        elif isinstance(value, dict) and (value.get("bytes") is not None or value.get("path") is not None):
+            # store the mesh bytes, and path is used to infer the mesh format using the file extension
+            return {"bytes": value.get("bytes"), "path": value.get("path")}
+        else:
+            raise ValueError(
+                f"A mesh sample should have one of 'path' or 'bytes' but they are missing or None in {value}."
+            )
+
+    def decode_example(self, value: dict, token_per_repo_id=None) -> Union["trimesh.Trimesh", "trimesh.Scene"]:
+        """Decode example mesh file.
+
+        Args:
+            value (`dict`):
+                A dictionary with keys:
+
+                - `path`: String with absolute or relative mesh file path.
+                - `bytes`: The bytes of the mesh file.
+            token_per_repo_id (`dict`, *optional*):
+                To access and decode
+                mesh files from private repositories on the Hub, you can pass
+                a dictionary repo_id (`str`) -> token (`bool` or `str`).
+
+        Returns:
+            `trimesh.Trimesh` or `trimesh.Scene`
+        """
+        if not self.decode:
+            raise RuntimeError("Decoding is disabled for this feature. Please use Mesh(decode=True) instead.")
+
+        if config.TRIMESH_AVAILABLE:
+            import trimesh
+        else:
+            raise ImportError("To support decoding meshes, please install 'trimesh'.")
+
+        if token_per_repo_id is None:
+            token_per_repo_id = {}
+
+        path, bytes_ = value["path"], value["bytes"]
+        if bytes_ is None:
+            if path is None:
+                raise ValueError(f"A mesh should have one of 'path' or 'bytes' but both are None in {value}.")
+            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.")
+                return trimesh.load(path, file_type=file_type)
+            source_url = path.split("::")[-1]
+            pattern = (
+                config.HUB_DATASETS_URL if source_url.startswith(config.HF_ENDPOINT) else config.HUB_DATASETS_HFFS_URL
+            )
+            source_url_fields = string_to_dict(source_url, pattern)
+            token = token_per_repo_id.get(source_url_fields["repo_id"]) if source_url_fields is not None else None
+            download_config = DownloadConfig(token=token)
+            with xopen(path, "rb", download_config=download_config) as f:
+                bytes_ = f.read()
+
+        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 "
+                "to infer the mesh file type."
+            )
+        return trimesh.load(BytesIO(bytes_), file_type=file_type)
+
+    def flatten(self) -> Union["FeatureType", dict[str, "FeatureType"]]:
+        """If in the decodable state, return the feature itself, otherwise flatten the feature into a dictionary."""
+        from .features import Value
+
+        return (
+            self
+            if self.decode
+            else {
+                "bytes": Value("binary"),
+                "path": Value("string"),
+            }
+        )
+
+    def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray]) -> pa.StructArray:
+        """Cast an Arrow array to the Mesh arrow storage type.
+        The Arrow types that can be converted to the Mesh pyarrow storage type are:
+
+        - `pa.string()` - it must contain the "path" data
+        - `pa.large_string()` - it must contain the "path" data (will be cast to string if possible)
+        - `pa.binary()` - it must contain the mesh bytes
+        - `pa.struct({"bytes": pa.binary()})`
+        - `pa.struct({"path": pa.string()})`
+        - `pa.struct({"bytes": pa.binary(), "path": pa.string()})`  - order doesn't matter
+
+        Args:
+            storage (`Union[pa.StringArray, pa.StructArray]`):
+                PyArrow array to cast.
+
+        Returns:
+            `pa.StructArray`: Array in the Mesh arrow storage type, that is
+                `pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
+        """
+        if pa.types.is_large_string(storage.type):
+            try:
+                storage = storage.cast(pa.string())
+            except pa.ArrowInvalid as e:
+                raise ValueError(
+                    f"Failed to cast large_string to string for Mesh feature. "
+                    f"This can happen if string values exceed 2GB. "
+                    f"Original error: {e}"
+                ) from e
+        if pa.types.is_string(storage.type):
+            bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            storage = pa.StructArray.from_arrays([bytes_array, storage], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_large_binary(storage.type):
+            storage = array_cast(
+                storage, pa.binary()
+            )  # this can fail in case of big meshes, paths should be used instead
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_binary(storage.type):
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_struct(storage.type):
+            if storage.type.get_field_index("bytes") >= 0:
+                bytes_array = storage.field("bytes")
+            else:
+                bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            if storage.type.get_field_index("path") >= 0:
+                path_array = storage.field("path")
+            else:
+                path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
+
+        return array_cast(storage, self.pa_type)
+
+    def embed_storage(
+        self, storage: pa.StructArray, token_per_repo_id=None, local_files: bool = True, remote_files: bool = True
+    ) -> pa.StructArray:
+        """Embed mesh files into the Arrow array.
+
+        Args:
+            storage (`pa.StructArray`):
+                PyArrow array to embed.
+            token_per_repo_id (`dict`, optional):
+                Dictionary repo_id -> token to fetch the files bytes.
+            local_files (`bool`, defaults to `True`):
+                Whether to embed local files data in the array.
+            remote_files (`bool`, defaults to `True`):
+                Whether to embed remote files data in the array.
+
+        Returns:
+            `pa.StructArray`: Array in the Mesh arrow storage type, that is
+                `pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
+        """
+        if token_per_repo_id is None:
+            token_per_repo_id = {}
+
+        def path_to_bytes(path):
+            if path is None:
+                return None
+            source_url = path.split("::")[-1]
+            pattern = (
+                config.HUB_DATASETS_URL if source_url.startswith(config.HF_ENDPOINT) else config.HUB_DATASETS_HFFS_URL
+            )
+            source_url_fields = string_to_dict(source_url, pattern)
+            token = token_per_repo_id.get(source_url_fields["repo_id"]) if source_url_fields is not None else None
+            download_config = DownloadConfig(token=token)
+            with xopen(path, "rb", download_config=download_config) as f:
+                return f.read()
+
+        bytes_array = pa.array(
+            [
+                (
+                    path_to_bytes(x["path"])
+                    if x["bytes"] is None
+                    and ((local_files and is_local_path(x["path"])) or (remote_files and is_remote_url(x["path"])))
+                    else x["bytes"]
+                )
+                if x is not None
+                else None
+                for x in storage.to_pylist()
+            ],
+            type=pa.binary(),
+        )
+        path_array = pa.array(
+            [
+                (
+                    os.path.basename(path)
+                    if (local_files and is_local_path(path)) or (remote_files and is_remote_url(path))
+                    else path
+                )
+                if path is not None
+                else None
+                for path in storage.field("path").to_pylist()
+            ],
+            type=pa.string(),
+        )
+        storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
+        return array_cast(storage, self.pa_type)
+
+
+def _infer_mesh_file_type(path: Optional[str]) -> Optional[str]:
+    supported_file_types = {"glb", "ply", "stl"}
+    if path is None:
+        return None
+    path_without_archive = path.split("::", 1)[0]
+    path_without_query = path_without_archive.split("?", 1)[0]
+    extension = os.path.splitext(path_without_query)[1].lower().lstrip(".")
+    return extension if extension in supported_file_types else None
+
+
+def encode_trimesh_mesh(mesh: Union["trimesh.Trimesh", "trimesh.Scene"]) -> dict[str, Optional[bytes | str]]:
+    """Encode a trimesh mesh or scene object into GLB bytes."""
+    metadata = getattr(mesh, "metadata", None) or {}
+    path = metadata.get("file_path") or metadata.get("file_name") if isinstance(metadata, dict) else None
+    if path is not None and os.path.isfile(path):
+        return {"path": path, "bytes": None}
+    bytes_ = mesh.export(file_type="glb")
+    return {"path": os.path.basename(path) if path else "mesh.glb", "bytes": bytes_}

src/datasets/features/nifti.py

@@ -0,0 +1,351 @@
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union
+
+import pyarrow as pa
+
+from .. import config
+from ..download.download_config import DownloadConfig
+from ..table import array_cast
+from ..utils.file_utils import is_local_path, is_remote_url, xopen
+from ..utils.py_utils import no_op_if_value_is_null, string_to_dict
+
+
+if TYPE_CHECKING:
+    import nibabel as nib
+
+    from .features import FeatureType
+
+if config.NIBABEL_AVAILABLE:
+    import nibabel as nib
+
+    class Nifti1ImageWrapper(nib.nifti1.Nifti1Image):
+        """
+        A wrapper around nibabel's Nifti1Image to customize its representation.
+        """
+
+        def __init__(self, nifti_image: nib.nifti1.Nifti1Image):
+            super().__init__(
+                dataobj=nifti_image.dataobj,
+                affine=nifti_image.affine,
+                header=nifti_image.header,
+                extra=nifti_image.extra,
+                file_map=nifti_image.file_map,
+                dtype=nifti_image.get_data_dtype(),
+            )
+            self.nifti_image = nifti_image
+
+        def _repr_html_(self):
+            from ipyniivue import NiiVue, ShowRender, SliceType, Volume
+            from IPython.display import display
+
+            bytes_ = self.nifti_image.to_bytes()
+            nv = NiiVue()
+            nv.set_slice_type(SliceType.MULTIPLANAR)
+            nv.opts.multiplanar_show_render = ShowRender.ALWAYS
+            nv.opts.show_3d_crosshair = True
+            nv.opts.multiplanar_force_render = True
+            name = None
+            if hasattr(self.nifti_image, "file_map"):
+                if (
+                    "image" in self.nifti_image.file_map
+                    and getattr(self.nifti_image.file_map["image"], "filename", None) is not None
+                ):
+                    name = self.nifti_image.file_map["image"].filename
+            if name is None:
+                name = "volume.nii.gz"
+            volume = Volume(name=name, data=bytes_)
+            nv.load_volumes([volume])
+            display(nv)
+
+
+@dataclass
+class Nifti:
+    """
+    **Experimental.**
+    Nifti [`Feature`] to read NIfTI neuroimaging files.
+
+    Input: The Nifti feature accepts as input:
+    - A `str`: Absolute path to the NIfTI file (i.e. random access is allowed).
+    - A `pathlib.Path`: path to the NIfTI file (i.e. random access is allowed).
+    - A `dict` with the keys:
+        - `path`: String with relative path of the NIfTI file in a dataset repository.
+        - `bytes`: Bytes of the NIfTI file.
+      This is useful for archived files with sequential access.
+
+    - A `nibabel` image object (e.g., `nibabel.nifti1.Nifti1Image`).
+
+    Args:
+        decode (`bool`, defaults to `True`):
+            Whether to decode the NIfTI data. If `False` a string with the bytes is returned. `decode=False` is not supported when decoding examples.
+
+    Examples:
+
+    ```py
+    >>> from datasets import Dataset, Nifti
+    >>> ds = Dataset.from_dict({"nifti": ["path/to/file.nii.gz"]}).cast_column("nifti", Nifti())
+    >>> ds.features["nifti"]
+    Nifti(decode=True, id=None)
+    >>> ds[0]["nifti"]
+    <nibabel.nifti1.Nifti1Image object at 0x7f8a1c2d8f40>
+    >>> ds = ds.cast_column("nifti", Nifti(decode=False))
+    >>> ds[0]["nifti"]
+    {'bytes': None,
+    'path': 'path/to/file.nii.gz'}
+    ```
+    """
+
+    decode: bool = True
+    id: Optional[str] = field(default=None, repr=False)
+
+    # Automatically constructed
+    dtype: ClassVar[str] = "nibabel.nifti1.Nifti1Image"
+    pa_type: ClassVar[Any] = pa.struct({"bytes": pa.binary(), "path": pa.string()})
+    _type: str = field(default="Nifti", init=False, repr=False)
+
+    def __call__(self):
+        return self.pa_type
+
+    def encode_example(self, value: Union[str, bytes, bytearray, dict, "nib.Nifti1Image"]) -> dict:
+        """Encode example into a format for Arrow.
+
+        Args:
+            value (`str`, `bytes`, `nibabel.Nifti1Image` or `dict`):
+                Data passed as input to Nifti feature.
+
+        Returns:
+            `dict` with "path" and "bytes" fields
+        """
+        if config.NIBABEL_AVAILABLE:
+            import nibabel as nib
+        else:
+            nib = None
+
+        if isinstance(value, str):
+            return {"path": value, "bytes": None}
+        elif isinstance(value, Path):
+            return {"path": str(value.absolute()), "bytes": None}
+        elif isinstance(value, (bytes, bytearray)):
+            return {"path": None, "bytes": value}
+        elif nib is not None and isinstance(value, nib.spatialimages.SpatialImage):
+            # nibabel image object - try to get path or convert to bytes
+            return encode_nibabel_image(value)
+        elif isinstance(value, dict):
+            if value.get("path") is not None and os.path.isfile(value["path"]):
+                # we set "bytes": None to not duplicate the data if they're already available locally
+                return {"bytes": None, "path": value.get("path")}
+            elif value.get("bytes") is not None or value.get("path") is not None:
+                # store the nifti bytes, and path is used to infer the format using the file extension
+                return {"bytes": value.get("bytes"), "path": value.get("path")}
+            else:
+                raise ValueError(
+                    f"A nifti sample should have one of 'path' or 'bytes' but they are missing or None in {value}."
+                )
+        else:
+            raise ValueError(
+                f"A nifti sample should be a string, bytes, Path, nibabel image, or dict, but got {type(value)}."
+            )
+
+    def decode_example(self, value: dict, token_per_repo_id=None) -> "Nifti1ImageWrapper":
+        """Decode example NIfTI file into nibabel image object.
+
+        Args:
+            value (`str` or `dict`):
+                A string with the absolute NIfTI file path, a dictionary with
+                keys:
+
+                - `path`: String with absolute or relative NIfTI file path.
+                - `bytes`: The bytes of the NIfTI file.
+
+            token_per_repo_id (`dict`, *optional*):
+                To access and decode NIfTI files from private repositories on
+                the Hub, you can pass a dictionary
+                repo_id (`str`) -> token (`bool` or `str`).
+
+        Returns:
+            `nibabel.Nifti1Image` objects
+        """
+        if config.NIBABEL_AVAILABLE:
+            import nibabel as nib
+        else:
+            raise ImportError("To support decoding NIfTI files, please install 'nibabel'.")
+
+        if token_per_repo_id is None:
+            token_per_repo_id = {}
+
+        path, bytes_ = value["path"], value["bytes"]
+        if bytes_ is None:
+            if path is None:
+                raise ValueError(f"A nifti should have one of 'path' or 'bytes' but both are None in {value}.")
+            else:
+                # gzipped files have the structure: 'gzip://T1.nii::<local_path>'
+                if path.startswith("gzip://") and is_local_path(path.split("::")[-1]):
+                    path = path.split("::")[-1]
+                if is_local_path(path):
+                    nifti = nib.load(path)
+                else:
+                    source_url = path.split("::")[-1]
+                    pattern = (
+                        config.HUB_DATASETS_URL
+                        if source_url.startswith(config.HF_ENDPOINT)
+                        else config.HUB_DATASETS_HFFS_URL
+                    )
+                    source_url_fields = string_to_dict(source_url, pattern)
+                    token = (
+                        token_per_repo_id.get(source_url_fields["repo_id"]) if source_url_fields is not None else None
+                    )
+                    download_config = DownloadConfig(token=token)
+                    with xopen(path, "rb", download_config=download_config) as f:
+                        nifti = nib.load(f)
+        else:
+            import gzip
+
+            if (
+                bytes_[:2] == b"\x1f\x8b"
+            ):  # gzip magic number, see https://stackoverflow.com/a/76055284/9534390 or "Magic number" on https://en.wikipedia.org/wiki/Gzip
+                bytes_ = gzip.decompress(bytes_)
+
+            nifti = nib.Nifti1Image.from_bytes(bytes_)
+
+        return Nifti1ImageWrapper(nifti)
+
+    def embed_storage(
+        self, storage: pa.StructArray, token_per_repo_id=None, local_files: bool = True, remote_files: bool = True
+    ) -> pa.StructArray:
+        """Embed NifTI files into the Arrow array.
+
+        Args:
+            storage (`pa.StructArray`):
+                PyArrow array to embed.
+            token_per_repo_id (`dict`, optional):
+                Dictionary repo_id -> token to fetch the files bytes.
+            local_files (`bool`, defaults to `True`)
+                Whether to embed local files data in the array
+
+                <Added version="4.8.5"/>
+            remote_files (`bool`, defaults to `True`)
+                Whether to embed remote files data in the array.
+                E.g. files with paths that start with hf:// or https://
+
+                <Added version="4.8.5"/>
+
+        Returns:
+            `pa.StructArray`: Array in the NifTI arrow storage type, that is
+                `pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
+        """
+        if token_per_repo_id is None:
+            token_per_repo_id = {}
+
+        @no_op_if_value_is_null
+        def path_to_bytes(path):
+            source_url = path.split("::")[-1]
+            pattern = (
+                config.HUB_DATASETS_URL if source_url.startswith(config.HF_ENDPOINT) else config.HUB_DATASETS_HFFS_URL
+            )
+            source_url_fields = string_to_dict(source_url, pattern)
+            token = token_per_repo_id.get(source_url_fields["repo_id"]) if source_url_fields is not None else None
+            download_config = DownloadConfig(token=token)
+            with xopen(path, "rb", download_config=download_config) as f:
+                return f.read()
+
+        bytes_array = pa.array(
+            [
+                (
+                    path_to_bytes(x["path"])
+                    if x["bytes"] is None
+                    and ((local_files and is_local_path(x["path"])) or (remote_files and is_remote_url(x["path"])))
+                    else x["bytes"]
+                )
+                if x is not None
+                else None
+                for x in storage.to_pylist()
+            ],
+            type=pa.binary(),
+        )
+        path_array = pa.array(
+            [
+                (
+                    os.path.basename(path)
+                    if (local_files and is_local_path(path)) or (remote_files and is_remote_url(path))
+                    else path
+                )
+                if path is not None
+                else None
+                for path in storage.field("path").to_pylist()
+            ],
+            type=pa.string(),
+        )
+        storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
+        return array_cast(storage, self.pa_type)
+
+    def flatten(self) -> Union["FeatureType", Dict[str, "FeatureType"]]:
+        """If in the decodable state, return the feature itself, otherwise flatten the feature into a dictionary."""
+        from .features import Value
+
+        return (
+            self
+            if self.decode
+            else {
+                "bytes": Value("binary"),
+                "path": Value("string"),
+            }
+        )
+
+    def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray, pa.BinaryArray]) -> pa.StructArray:
+        """Cast an Arrow array to the Nifti arrow storage type.
+        The Arrow types that can be converted to the Nifti pyarrow storage type are:
+
+        - `pa.string()` - it must contain the "path" data
+        - `pa.binary()` - it must contain the NIfTI bytes
+        - `pa.struct({"bytes": pa.binary()})`
+        - `pa.struct({"path": pa.string()})`
+        - `pa.struct({"bytes": pa.binary(), "path": pa.string()})`  - order doesn't matter
+
+        Args:
+            storage (`Union[pa.StringArray, pa.StructArray, pa.BinaryArray]`):
+                PyArrow array to cast.
+
+        Returns:
+            `pa.StructArray`: Array in the Nifti arrow storage type, that is
+                `pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
+        """
+        if pa.types.is_string(storage.type):
+            bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            storage = pa.StructArray.from_arrays([bytes_array, storage], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_binary(storage.type):
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_struct(storage.type):
+            if storage.type.get_field_index("bytes") >= 0:
+                bytes_array = storage.field("bytes")
+            else:
+                bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            if storage.type.get_field_index("path") >= 0:
+                path_array = storage.field("path")
+            else:
+                path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
+        return array_cast(storage, self.pa_type)
+
+
+def encode_nibabel_image(img: "nib.Nifti1Image", force_bytes: bool = False) -> dict[str, Optional[Union[str, bytes]]]:
+    """
+    Encode a nibabel image object into a dictionary.
+
+    If the image has an associated file path, returns the path. Otherwise, serializes
+    the image content into bytes.
+
+    Args:
+        img: A nibabel image object (e.g., Nifti1Image).
+        force_bytes: If `True`, always serialize to bytes even if a file path exists. Needed to upload bytes properly.
+
+    Returns:
+        dict: A dictionary with "path" or "bytes" field.
+    """
+    if hasattr(img, "file_map") and img.file_map is not None and not force_bytes:
+        filename = img.file_map["image"].filename
+        return {"path": filename, "bytes": None}
+
+    bytes_data = img.to_bytes()
+    return {"path": None, "bytes": bytes_data}

src/datasets/features/pdf.py

@@ -0,0 +1,309 @@
+import os
+from dataclasses import dataclass, field
+from io import BytesIO
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union
+
+import pyarrow as pa
+
+from .. import config
+from ..download.download_config import DownloadConfig
+from ..table import array_cast
+from ..utils.file_utils import is_local_path, is_remote_url, xopen
+from ..utils.py_utils import no_op_if_value_is_null, string_to_dict
+
+
+if TYPE_CHECKING:
+    import pdfplumber
+
+    from .features import FeatureType
+
+
+def pdf_to_bytes(pdf: "pdfplumber.pdf.PDF") -> bytes:
+    """Convert a pdfplumber.pdf.PDF object to bytes."""
+    with BytesIO() as buffer:
+        for page in pdf.pages:
+            buffer.write(page.pdf.stream)
+        return buffer.getvalue()
+
+
+@dataclass
+class Pdf:
+    """
+    **Experimental.**
+    Pdf [`Feature`] to read pdf documents from a pdf file.
+
+    Input: The Pdf feature accepts as input:
+    - A `str`: Absolute path to the pdf file (i.e. random access is allowed).
+    - A `pathlib.Path`: path to the pdf file (i.e. random access is allowed).
+    - A `dict` with the keys:
+        - `path`: String with relative path of the pdf file in a dataset repository.
+        - `bytes`: Bytes of the pdf file.
+      This is useful for archived files with sequential access.
+
+    - A `pdfplumber.pdf.PDF`: pdfplumber pdf object.
+
+    Args:
+        decode (`bool`, defaults to `True`):
+            Whether to decode the pdf data. If `False`,
+            returns the underlying dictionary in the format `{"path": pdf_path, "bytes": pdf_bytes}`.
+
+    Examples:
+
+    ```py
+    >>> from datasets import Dataset, Pdf
+    >>> ds = Dataset.from_dict({"pdf": ["path/to/pdf/file.pdf"]}).cast_column("pdf", Pdf())
+    >>> ds.features["pdf"]
+    Pdf(decode=True, id=None)
+    >>> ds[0]["pdf"]
+    <pdfplumber.pdf.PDF object at 0x7f8a1c2d8f40>
+    >>> ds = ds.cast_column("pdf", Pdf(decode=False))
+    >>> ds[0]["pdf"]
+    {'bytes': None,
+    'path': 'path/to/pdf/file.pdf'}
+    ```
+    """
+
+    decode: bool = True
+    id: Optional[str] = field(default=None, repr=False)
+
+    # Automatically constructed
+    dtype: ClassVar[str] = "pdfplumber.pdf.PDF"
+    pa_type: ClassVar[Any] = pa.struct({"bytes": pa.binary(), "path": pa.string()})
+    _type: str = field(default="Pdf", init=False, repr=False)
+
+    def __call__(self):
+        return self.pa_type
+
+    def encode_example(self, value: Union[str, bytes, bytearray, dict, "pdfplumber.pdf.PDF"]) -> dict:
+        """Encode example into a format for Arrow.
+
+        Args:
+            value (`str`, `bytes`, `pdfplumber.pdf.PDF` or `dict`):
+                Data passed as input to Pdf feature.
+
+        Returns:
+            `dict` with "path" and "bytes" fields
+        """
+        if config.PDFPLUMBER_AVAILABLE:
+            import pdfplumber
+        else:
+            pdfplumber = None
+
+        if isinstance(value, str):
+            return {"path": value, "bytes": None}
+        elif isinstance(value, Path):
+            return {"path": str(value.absolute()), "bytes": None}
+        elif isinstance(value, (bytes, bytearray)):
+            return {"path": None, "bytes": value}
+        elif pdfplumber is not None and isinstance(value, pdfplumber.pdf.PDF):
+            # convert the pdfplumber.pdf.PDF to bytes
+            return encode_pdfplumber_pdf(value)
+        elif value.get("path") is not None and os.path.isfile(value["path"]):
+            # we set "bytes": None to not duplicate the data if they're already available locally
+            return {"bytes": None, "path": value.get("path")}
+        elif value.get("bytes") is not None or value.get("path") is not None:
+            # store the pdf bytes, and path is used to infer the pdf format using the file extension
+            return {"bytes": value.get("bytes"), "path": value.get("path")}
+        else:
+            raise ValueError(
+                f"A pdf sample should have one of 'path' or 'bytes' but they are missing or None in {value}."
+            )
+
+    def decode_example(self, value: dict, token_per_repo_id=None) -> "pdfplumber.pdf.PDF":
+        """Decode example pdf file into pdf data.
+
+        Args:
+            value (`str` or `dict`):
+                A string with the absolute pdf file path, a dictionary with
+                keys:
+
+                - `path`: String with absolute or relative pdf file path.
+                - `bytes`: The bytes of the pdf file.
+
+            token_per_repo_id (`dict`, *optional*):
+                To access and decode pdf files from private repositories on
+                the Hub, you can pass a dictionary
+                repo_id (`str`) -> token (`bool` or `str`).
+
+        Returns:
+            `pdfplumber.pdf.PDF`
+        """
+        if not self.decode:
+            raise RuntimeError("Decoding is disabled for this feature. Please use Pdf(decode=True) instead.")
+
+        if config.PDFPLUMBER_AVAILABLE:
+            import pdfplumber
+        else:
+            raise ImportError("To support decoding pdfs, please install 'pdfplumber'.")
+
+        if token_per_repo_id is None:
+            token_per_repo_id = {}
+
+        path, bytes_ = value["path"], value["bytes"]
+        if bytes_ is None:
+            if path is None:
+                raise ValueError(f"A pdf should have one of 'path' or 'bytes' but both are None in {value}.")
+            else:
+                if is_local_path(path):
+                    pdf = pdfplumber.open(path)
+                else:
+                    source_url = path.split("::")[-1]
+                    pattern = (
+                        config.HUB_DATASETS_URL
+                        if source_url.startswith(config.HF_ENDPOINT)
+                        else config.HUB_DATASETS_HFFS_URL
+                    )
+                    try:
+                        repo_id = string_to_dict(source_url, pattern)["repo_id"]
+                        token = token_per_repo_id.get(repo_id)
+                    except ValueError:
+                        token = None
+                    download_config = DownloadConfig(token=token)
+                    f = xopen(path, "rb", download_config=download_config)
+                    return pdfplumber.open(f)
+        else:
+            with pdfplumber.open(BytesIO(bytes_)) as p:
+                pdf = p
+
+        return pdf
+
+    def flatten(self) -> Union["FeatureType", Dict[str, "FeatureType"]]:
+        """If in the decodable state, return the feature itself, otherwise flatten the feature into a dictionary."""
+        from .features import Value
+
+        return (
+            self
+            if self.decode
+            else {
+                "bytes": Value("binary"),
+                "path": Value("string"),
+            }
+        )
+
+    def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray, pa.ListArray]) -> pa.StructArray:
+        """Cast an Arrow array to the Pdf arrow storage type.
+        The Arrow types that can be converted to the Pdf pyarrow storage type are:
+
+        - `pa.string()` - it must contain the "path" data
+        - `pa.binary()` - it must contain the image bytes
+        - `pa.struct({"bytes": pa.binary()})`
+        - `pa.struct({"path": pa.string()})`
+        - `pa.struct({"bytes": pa.binary(), "path": pa.string()})`  - order doesn't matter
+        - `pa.list(*)` - it must contain the pdf array data
+
+        Args:
+            storage (`Union[pa.StringArray, pa.StructArray, pa.ListArray]`):
+                PyArrow array to cast.
+
+        Returns:
+            `pa.StructArray`: Array in the Pdf arrow storage type, that is
+                `pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
+        """
+        if pa.types.is_string(storage.type):
+            bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            storage = pa.StructArray.from_arrays([bytes_array, storage], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_binary(storage.type):
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_struct(storage.type):
+            if storage.type.get_field_index("bytes") >= 0:
+                bytes_array = storage.field("bytes")
+            else:
+                bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            if storage.type.get_field_index("path") >= 0:
+                path_array = storage.field("path")
+            else:
+                path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
+        return array_cast(storage, self.pa_type)
+
+    def embed_storage(
+        self, storage: pa.StructArray, token_per_repo_id=None, local_files: bool = True, remote_files: bool = True
+    ) -> pa.StructArray:
+        """Embed PDF files into the Arrow array.
+
+        Args:
+            storage (`pa.StructArray`):
+                PyArrow array to embed.
+            token_per_repo_id (`dict`, optional):
+                Dictionary repo_id -> token to fetch the files bytes.
+            local_files (`bool`, defaults to `True`)
+                Whether to embed local files data in the array
+
+                <Added version="4.8.5"/>
+            remote_files (`bool`, defaults to `True`)
+                Whether to embed remote files data in the array.
+                E.g. files with paths that start with hf:// or https://
+
+                <Added version="4.8.5"/>
+
+        Returns:
+            `pa.StructArray`: Array in the PDF arrow storage type, that is
+                `pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
+        """
+        if token_per_repo_id is None:
+            token_per_repo_id = {}
+
+        @no_op_if_value_is_null
+        def path_to_bytes(path):
+            source_url = path.split("::")[-1]
+            pattern = (
+                config.HUB_DATASETS_URL if source_url.startswith(config.HF_ENDPOINT) else config.HUB_DATASETS_HFFS_URL
+            )
+            source_url_fields = string_to_dict(source_url, pattern)
+            token = token_per_repo_id.get(source_url_fields["repo_id"]) if source_url_fields is not None else None
+            download_config = DownloadConfig(token=token)
+            with xopen(path, "rb", download_config=download_config) as f:
+                return f.read()
+
+        bytes_array = pa.array(
+            [
+                (
+                    path_to_bytes(x["path"])
+                    if x["bytes"] is None
+                    and ((local_files and is_local_path(x["path"])) or (remote_files and is_remote_url(x["path"])))
+                    else x["bytes"]
+                )
+                if x is not None
+                else None
+                for x in storage.to_pylist()
+            ],
+            type=pa.binary(),
+        )
+        path_array = pa.array(
+            [
+                (
+                    os.path.basename(path)
+                    if (local_files and is_local_path(path)) or (remote_files and is_remote_url(path))
+                    else path
+                )
+                if path is not None
+                else None
+                for path in storage.field("path").to_pylist()
+            ],
+            type=pa.string(),
+        )
+        storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
+        return array_cast(storage, self.pa_type)
+
+
+def encode_pdfplumber_pdf(pdf: "pdfplumber.pdf.PDF") -> dict:
+    """
+    Encode a pdfplumber.pdf.PDF object into a dictionary.
+
+    If the PDF has an associated file path, returns the path. Otherwise, serializes
+    the PDF content into bytes.
+
+    Args:
+        pdf (pdfplumber.pdf.PDF): A pdfplumber PDF object.
+
+    Returns:
+        dict: A dictionary with "path" or "bytes" field.
+    """
+    if hasattr(pdf, "stream") and hasattr(pdf.stream, "name") and pdf.stream.name:
+        # Return the path if the PDF has an associated file path
+        return {"path": pdf.stream.name, "bytes": None}
+    else:
+        # Convert the PDF to bytes if no path is available
+        return {"path": None, "bytes": pdf_to_bytes(pdf)}

src/datasets/features/translation.py

@@ -0,0 +1,129 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
+
+import pyarrow as pa
+
+
+if TYPE_CHECKING:
+    from .features import FeatureType
+
+
+@dataclass
+class Translation:
+    """`Feature` for translations with fixed languages per example.
+    Here for compatibility with tfds.
+
+    Args:
+        languages (`dict`):
+            A dictionary for each example mapping string language codes to string translations.
+
+    Example:
+
+    ```python
+    >>> # At construction time:
+    >>> datasets.features.Translation(languages=['en', 'fr', 'de'])
+    >>> # During data generation:
+    >>> yield {
+    ...         'en': 'the cat',
+    ...         'fr': 'le chat',
+    ...         'de': 'die katze'
+    ... }
+    ```
+    """
+
+    languages: list[str]
+    id: Optional[str] = field(default=None, repr=False)
+    # Automatically constructed
+    dtype: ClassVar[str] = "dict"
+    pa_type: ClassVar[Any] = None
+    _type: str = field(default="Translation", init=False, repr=False)
+
+    def __call__(self):
+        return pa.struct({lang: pa.string() for lang in sorted(self.languages)})
+
+    def flatten(self) -> Union["FeatureType", dict[str, "FeatureType"]]:
+        """Flatten the Translation feature into a dictionary."""
+        from .features import Value
+
+        return {k: Value("string") for k in sorted(self.languages)}
+
+
+@dataclass
+class TranslationVariableLanguages:
+    """`Feature` for translations with variable languages per example.
+    Here for compatibility with tfds.
+
+    Args:
+        languages (`dict`):
+            A dictionary for each example mapping string language codes to one or more string translations.
+            The languages present may vary from example to example.
+
+    Returns:
+        - `language` or `translation` (variable-length 1D `tf.Tensor` of `tf.string`):
+            Language codes sorted in ascending order or plain text translations, sorted to align with language codes.
+
+    Example:
+
+    ```python
+    >>> # At construction time:
+    >>> datasets.features.TranslationVariableLanguages(languages=['en', 'fr', 'de'])
+    >>> # During data generation:
+    >>> yield {
+    ...         'en': 'the cat',
+    ...         'fr': ['le chat', 'la chatte,']
+    ...         'de': 'die katze'
+    ... }
+    >>> # Tensor returned :
+    >>> {
+    ...         'language': ['en', 'de', 'fr', 'fr'],
+    ...         'translation': ['the cat', 'die katze', 'la chatte', 'le chat'],
+    ... }
+    ```
+    """
+
+    languages: Optional[list] = None
+    num_languages: Optional[int] = None
+    id: Optional[str] = field(default=None, repr=False)
+    # Automatically constructed
+    dtype: ClassVar[str] = "dict"
+    pa_type: ClassVar[Any] = None
+    _type: str = field(default="TranslationVariableLanguages", init=False, repr=False)
+
+    def __post_init__(self):
+        self.languages = sorted(set(self.languages)) if self.languages else None
+        self.num_languages = len(self.languages) if self.languages else None
+
+    def __call__(self):
+        return pa.struct({"language": pa.list_(pa.string()), "translation": pa.list_(pa.string())})
+
+    def encode_example(self, translation_dict):
+        lang_set = set(self.languages)
+        if set(translation_dict) == {"language", "translation"}:
+            return translation_dict
+        elif self.languages and set(translation_dict) - lang_set:
+            raise ValueError(
+                f"Some languages in example ({', '.join(sorted(set(translation_dict) - lang_set))}) are not in valid set ({', '.join(lang_set)})."
+            )
+
+        # Convert dictionary into tuples, splitting out cases where there are
+        # multiple translations for a single language.
+        translation_tuples = []
+        for lang, text in translation_dict.items():
+            if isinstance(text, str):
+                translation_tuples.append((lang, text))
+            else:
+                translation_tuples.extend([(lang, el) for el in text])
+
+        # Ensure translations are in ascending order by language code.
+        languages, translations = zip(*sorted(translation_tuples))
+
+        return {"language": languages, "translation": translations}
+
+    def flatten(self) -> Union["FeatureType", dict[str, "FeatureType"]]:
+        """Flatten the TranslationVariableLanguages feature into a dictionary."""
+        from .features import List, Value
+
+        return {
+            "language": List(Value("string")),
+            "translation": List(Value("string")),
+        }

src/datasets/features/video.py

@@ -0,0 +1,415 @@
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional, TypedDict, Union
+
+import numpy as np
+import pyarrow as pa
+
+from .. import config
+from ..download.download_config import DownloadConfig
+from ..table import array_cast
+from ..utils.file_utils import is_local_path, is_remote_url, xopen
+from ..utils.py_utils import no_op_if_value_is_null, string_to_dict
+
+
+if TYPE_CHECKING:
+    import torch
+    from torchcodec.decoders import VideoDecoder
+
+    from .features import FeatureType
+
+
+class Example(TypedDict):
+    path: Optional[str]
+    bytes: Optional[bytes]
+
+
+@dataclass
+class Video:
+    """
+    Video [`Feature`] to read video data from a video file.
+
+    Input: The Video feature accepts as input:
+    - A `str`: Absolute path to the video file (i.e. random access is allowed).
+    - A `pathlib.Path`: path to the video file (i.e. random access is allowed).
+    - A `dict` with the keys:
+
+        - `path`: String with relative path of the video file in a dataset repository.
+        - `bytes`: Bytes of the video file.
+
+      This is useful for parquet or webdataset files which embed video files.
+
+    - A `torchcodec.decoders.VideoDecoder`: torchcodec video decoder object.
+
+    Output: The Video features output data as `torchcodec.decoders.VideoDecoder` objects.
+
+    Args:
+        decode (`bool`, defaults to `True`):
+            Whether to decode the video data. If `False`,
+            returns the underlying dictionary in the format `{"path": video_path, "bytes": video_bytes}`.
+        stream_index (`int`, *optional*):
+            The streaming index to use from the file. If `None` defaults to the "best" index.
+        dimension_order (`str`, defaults to `NCHW`):
+            The dimension order of the decoded frames.
+            where N is the batch size, C is the number of channels,
+            H is the height, and W is the width of the frames.
+        num_ffmpeg_threads (`int`, defaults to `1`):
+            The number of threads to use for decoding the video. (Recommended to keep this at 1)
+        device (`str` or `torch.device`, defaults to `cpu`):
+            The device to use for decoding the video.
+        seek_mode (`str`, defaults to `exact`):
+            Determines if frame access will be “exact” or “approximate”.
+            Exact guarantees that requesting frame i will always return frame i, but doing so requires an initial scan of the file.
+            Approximate is faster as it avoids scanning the file, but less accurate as it uses the file's metadata to calculate where i probably is.
+            read more [here](https://docs.pytorch.org/torchcodec/stable/generated_examples/approximate_mode.html#sphx-glr-generated-examples-approximate-mode-py)
+
+    Examples:
+
+    ```py
+    >>> from datasets import Dataset, Video
+    >>> ds = Dataset.from_dict({"video":["path/to/Screen Recording.mov"]}).cast_column("video", Video())
+    >>> ds.features["video"]
+    Video(decode=True, id=None)
+    >>> ds[0]["video"]
+    <torchcodec.decoders._video_decoder.VideoDecoder object at 0x14a61e080>
+    >>> video = ds[0]["video"]
+    >>> video.get_frames_in_range(0, 10)
+    FrameBatch:
+    data (shape): torch.Size([10, 3, 50, 66])
+    pts_seconds: tensor([0.4333, 0.4333, 0.4333, 0.4333, 0.4333, 0.4333, 0.4333, 0.4333, 0.4333,
+            0.4333], dtype=torch.float64)
+    duration_seconds: tensor([0.0167, 0.0167, 0.0167, 0.0167, 0.0167, 0.0167, 0.0167, 0.0167, 0.0167,
+            0.0167], dtype=torch.float64)
+    >>> ds.cast_column('video', Video(decode=False))[0]["video]
+    {'bytes': None,
+     'path': 'path/to/Screen Recording.mov'}
+    ```
+    """
+
+    decode: bool = True
+    stream_index: Optional[int] = None
+    dimension_order: Literal["NCHW", "NHWC"] = "NCHW"
+    num_ffmpeg_threads: int = 1
+    device: Optional[Union[str, "torch.device"]] = "cpu"
+    seek_mode: Literal["exact", "approximate"] = "exact"
+    id: Optional[str] = field(default=None, repr=False)
+    # Automatically constructed
+    dtype: ClassVar[str] = "torchcodec.decoders.VideoDecoder"
+    pa_type: ClassVar[Any] = pa.struct({"bytes": pa.binary(), "path": pa.string()})
+    _type: str = field(default="Video", init=False, repr=False)
+
+    def __call__(self):
+        return self.pa_type
+
+    def encode_example(self, value: Union[str, bytes, bytearray, Example, np.ndarray, "VideoDecoder"]) -> Example:
+        """Encode example into a format for Arrow.
+
+        Args:
+            value (`str`, `np.ndarray`, `bytes`, `bytearray`, `VideoDecoder` or `dict`):
+                Data passed as input to Video feature.
+
+        Returns:
+            `dict` with "path" and "bytes" fields
+        """
+        if value is None:
+            raise ValueError("value must be provided")
+
+        if config.TORCHCODEC_AVAILABLE:
+            from torchcodec.decoders import VideoDecoder
+        else:
+            VideoDecoder = None
+
+        if isinstance(value, list):
+            value = np.array(value)
+
+        if isinstance(value, str):
+            return {"path": value, "bytes": None}
+        elif isinstance(value, Path):
+            return {"path": str(value.absolute()), "bytes": None}
+        elif isinstance(value, (bytes, bytearray)):
+            return {"path": None, "bytes": value}
+        elif isinstance(value, np.ndarray):
+            # convert the video array to bytes
+            return encode_np_array(value)
+        elif VideoDecoder is not None and isinstance(value, VideoDecoder):
+            # convert the torchcodec video decoder to bytes
+            return encode_torchcodec_video(value)
+        elif isinstance(value, dict):
+            path, bytes_ = value.get("path"), value.get("bytes")
+            if path is not None and os.path.isfile(path):
+                # we set "bytes": None to not duplicate the data if they're already available locally
+                return {"bytes": None, "path": path}
+            elif bytes_ is not None or path is not None:
+                # store the video bytes, and path is used to infer the video format using the file extension
+                return {"bytes": bytes_, "path": path}
+            else:
+                raise ValueError(
+                    f"A video sample should have one of 'path' or 'bytes' but they are missing or None in {value}."
+                )
+        else:
+            raise TypeError(f"Unsupported encode_example type: {type(value)}")
+
+    def decode_example(
+        self,
+        value: Union[str, Example],
+        token_per_repo_id: Optional[dict[str, Union[bool, str]]] = None,
+    ) -> "VideoDecoder":
+        """Decode example video file into video data.
+
+        Args:
+            value (`str` or `dict`):
+                A string with the absolute video file path, a dictionary with
+                keys:
+
+                - `path`: String with absolute or relative video file path.
+                - `bytes`: The bytes of the video file.
+            token_per_repo_id (`dict`, *optional*):
+                To access and decode
+                video files from private repositories on the Hub, you can pass
+                a dictionary repo_id (`str`) -> token (`bool` or `str`).
+
+        Returns:
+            `torchcodec.decoders.VideoDecoder`
+        """
+        if not self.decode:
+            raise RuntimeError("Decoding is disabled for this feature. Please use Video(decode=True) instead.")
+
+        if config.TORCHCODEC_AVAILABLE:
+            from torchcodec.decoders import VideoDecoder
+
+        else:
+            raise ImportError("To support decoding videos, please install 'torchcodec'.")
+
+        if token_per_repo_id is None:
+            token_per_repo_id = {}
+
+        if isinstance(value, str):
+            path, bytes_ = value, None
+        else:
+            path, bytes_ = value["path"], value["bytes"]
+
+        if bytes_ is None:
+            if path is None:
+                raise ValueError(f"A video should have one of 'path' or 'bytes' but both are None in {value}.")
+            elif is_local_path(path):
+                video = VideoDecoder(
+                    path,
+                    stream_index=self.stream_index,
+                    dimension_order=self.dimension_order,
+                    num_ffmpeg_threads=self.num_ffmpeg_threads,
+                    device=self.device,
+                    seek_mode=self.seek_mode,
+                )
+            else:
+                video = hf_video_reader(
+                    path,
+                    token_per_repo_id=token_per_repo_id,
+                    dimension_order=self.dimension_order,
+                    num_ffmpeg_threads=self.num_ffmpeg_threads,
+                    device=self.device,
+                    seek_mode=self.seek_mode,
+                )
+        else:
+            video = VideoDecoder(
+                bytes_,
+                stream_index=self.stream_index,
+                dimension_order=self.dimension_order,
+                num_ffmpeg_threads=self.num_ffmpeg_threads,
+                device=self.device,
+                seek_mode=self.seek_mode,
+            )
+        video._hf_encoded = {"path": path, "bytes": bytes_}
+        video.metadata.path = path
+        return video
+
+    def flatten(self) -> Union["FeatureType", dict[str, "FeatureType"]]:
+        """If in the decodable state, return the feature itself, otherwise flatten the feature into a dictionary."""
+        from .features import Value
+
+        return (
+            self
+            if self.decode
+            else {
+                "bytes": Value("binary"),
+                "path": Value("string"),
+            }
+        )
+
+    def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray, pa.ListArray]) -> pa.StructArray:
+        """Cast an Arrow array to the Video arrow storage type.
+        The Arrow types that can be converted to the Video pyarrow storage type are:
+
+        - `pa.string()` - it must contain the "path" data
+        - `pa.binary()` - it must contain the video bytes
+        - `pa.struct({"bytes": pa.binary()})`
+        - `pa.struct({"path": pa.string()})`
+        - `pa.struct({"bytes": pa.binary(), "path": pa.string()})`  - order doesn't matter
+        - `pa.list(*)` - it must contain the video array data
+
+        Args:
+            storage (`Union[pa.StringArray, pa.StructArray, pa.ListArray]`):
+                PyArrow array to cast.
+
+        Returns:
+            `pa.StructArray`: Array in the Video arrow storage type, that is
+                `pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
+        """
+        if pa.types.is_string(storage.type):
+            bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            storage = pa.StructArray.from_arrays([bytes_array, storage], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_large_binary(storage.type):
+            storage = array_cast(
+                storage, pa.binary()
+            )  # this can fail in case of big videos, paths should be used instead
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_binary(storage.type):
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_struct(storage.type):
+            if storage.type.get_field_index("bytes") >= 0:
+                bytes_array = storage.field("bytes")
+            else:
+                bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            if storage.type.get_field_index("path") >= 0:
+                path_array = storage.field("path")
+            else:
+                path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_list(storage.type):
+            bytes_array = pa.array(
+                [encode_np_array(np.array(arr))["bytes"] if arr is not None else None for arr in storage.to_pylist()],
+                type=pa.binary(),
+            )
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays(
+                [bytes_array, path_array], ["bytes", "path"], mask=bytes_array.is_null()
+            )
+        return array_cast(storage, self.pa_type)
+
+    def embed_storage(
+        self, storage: pa.StructArray, token_per_repo_id=None, local_files: bool = True, remote_files: bool = True
+    ) -> pa.StructArray:
+        """Embed image files into the Arrow array.
+
+        Args:
+            storage (`pa.StructArray`):
+                PyArrow array to embed.
+            token_per_repo_id (`dict`, optional):
+                Dictionary repo_id -> token to fetch the files bytes.
+            local_files (`bool`, defaults to `True`)
+                Whether to embed local files data in the array
+
+                <Added version="4.8.5"/>
+            remote_files (`bool`, defaults to `True`)
+                Whether to embed remote files data in the array.
+                E.g. files with paths that start with hf:// or https://
+
+                <Added version="4.8.5"/>
+
+        Returns:
+            `pa.StructArray`: Array in the Video arrow storage type, that is
+                `pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
+        """
+        if token_per_repo_id is None:
+            token_per_repo_id = {}
+
+        @no_op_if_value_is_null
+        def path_to_bytes(path):
+            source_url = path.split("::")[-1]
+            pattern = (
+                config.HUB_DATASETS_URL if source_url.startswith(config.HF_ENDPOINT) else config.HUB_DATASETS_HFFS_URL
+            )
+            source_url_fields = string_to_dict(source_url, pattern)
+            token = token_per_repo_id.get(source_url_fields["repo_id"]) if source_url_fields is not None else None
+            download_config = DownloadConfig(token=token)
+            with xopen(path, "rb", download_config=download_config) as f:
+                return f.read()
+
+        bytes_array = pa.array(
+            [
+                (
+                    path_to_bytes(x["path"])
+                    if x["bytes"] is None
+                    and ((local_files and is_local_path(x["path"])) or (remote_files and is_remote_url(x["path"])))
+                    else x["bytes"]
+                )
+                if x is not None
+                else None
+                for x in storage.to_pylist()
+            ],
+            type=pa.binary(),
+        )
+        path_array = pa.array(
+            [
+                (
+                    os.path.basename(path)
+                    if (local_files and is_local_path(path)) or (remote_files and is_remote_url(path))
+                    else path
+                )
+                if path is not None
+                else None
+                for path in storage.field("path").to_pylist()
+            ],
+            type=pa.string(),
+        )
+        storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
+        return array_cast(storage, self.pa_type)
+
+
+def video_to_bytes(video: "VideoDecoder") -> bytes:
+    """Convert a torchcodec Video object to bytes using native compression if possible"""
+    raise NotImplementedError()
+
+
+def encode_torchcodec_video(video: "VideoDecoder") -> Example:
+    if hasattr(video, "_hf_encoded"):
+        return video._hf_encoded
+    else:
+        raise NotImplementedError(
+            "Encoding a VideoDecoder that doesn't come from datasets.Video.decode() is not implemented"
+        )
+
+
+def encode_np_array(array: np.ndarray) -> Example:
+    raise NotImplementedError()
+
+
+# No monkey patch needed!
+# 1. store the encoded video data {"path": ..., "bytes": ...} in `video._hf_encoded``
+# 2. add support for hf:// files
+
+
+def hf_video_reader(
+    path: str,
+    token_per_repo_id: Optional[dict[str, Union[bool, str]]] = None,
+    stream: str = "video",
+    dimension_order: Literal["NCHW", "NHWC"] = "NCHW",
+    num_ffmpeg_threads: int = 1,
+    device: Optional[Union[str, "torch.device"]] = "cpu",
+    seek_mode: Literal["exact", "approximate"] = "exact",
+) -> "VideoDecoder":
+    from torchcodec.decoders import VideoDecoder
+
+    # Load the file from HF
+    if token_per_repo_id is None:
+        token_per_repo_id = {}
+    source_url = path.split("::")[-1]
+    pattern = config.HUB_DATASETS_URL if source_url.startswith(config.HF_ENDPOINT) else config.HUB_DATASETS_HFFS_URL
+    source_url_fields = string_to_dict(source_url, pattern)
+    token = token_per_repo_id.get(source_url_fields["repo_id"]) if source_url_fields is not None else None
+    download_config = DownloadConfig(token=token)
+    f = xopen(path, "rb", download_config=download_config)
+
+    # Instantiate the VideoDecoder
+    stream_id = 0 if len(stream.split(":")) == 1 else int(stream.split(":")[1])
+    vd = VideoDecoder(
+        f,
+        stream_index=stream_id,
+        dimension_order=dimension_order,
+        num_ffmpeg_threads=num_ffmpeg_threads,
+        device=device,
+        seek_mode=seek_mode,
+    )
+    return vd

src/datasets/filesystems/__init__.py

@@ -0,0 +1,47 @@
+import importlib
+import shutil
+import warnings
+from typing import List
+
+import fsspec
+import fsspec.asyn
+from fsspec.implementations.local import LocalFileSystem
+
+from . import compression
+
+
+COMPRESSION_FILESYSTEMS: list[compression.BaseCompressedFileFileSystem] = [
+    compression.Bz2FileSystem,
+    compression.GzipFileSystem,
+    compression.Lz4FileSystem,
+    compression.XzFileSystem,
+    compression.ZstdFileSystem,
+]
+
+# Register custom filesystems
+for fs_class in COMPRESSION_FILESYSTEMS:
+    if fs_class.protocol in fsspec.registry and fsspec.registry[fs_class.protocol] is not fs_class:
+        warnings.warn(f"A filesystem protocol was already set for {fs_class.protocol} and will be overwritten.")
+    fsspec.register_implementation(fs_class.protocol, fs_class, clobber=True)
+
+
+def is_remote_filesystem(fs: fsspec.AbstractFileSystem) -> bool:
+    """
+    Checks if `fs` is a remote filesystem.
+
+    Args:
+        fs (`fsspec.spec.AbstractFileSystem`):
+            An abstract super-class for pythonic file-systems, e.g. `fsspec.filesystem(\'file\')` or `s3fs.S3FileSystem`.
+    """
+    return not isinstance(fs, LocalFileSystem)
+
+
+def rename(fs: fsspec.AbstractFileSystem, src: str, dst: str):
+    """
+    Renames the file `src` in `fs` to `dst`.
+    """
+    if not is_remote_filesystem(fs):
+        # LocalFileSystem.mv does copy + rm, it is more efficient to simply move a local directory
+        shutil.move(fs._strip_protocol(src), fs._strip_protocol(dst))
+    else:
+        fs.mv(src, dst, recursive=True)

src/datasets/filesystems/compression.py

@@ -0,0 +1,127 @@
+import os
+from functools import partial
+from typing import Optional
+
+import fsspec
+from fsspec.archive import AbstractArchiveFileSystem
+
+
+class BaseCompressedFileFileSystem(AbstractArchiveFileSystem):
+    """Read contents of compressed file as a filesystem with one file inside."""
+
+    root_marker = ""
+    protocol: str = (
+        None  # protocol passed in prefix to the url. ex: "gzip", for gzip://file.txt::http://foo.bar/file.txt.gz
+    )
+    compression: str = None  # compression type in fsspec. ex: "gzip"
+    extensions: list[str] = None  # extensions of the filename to strip. ex: ".gz" to get file.txt from file.txt.gz
+
+    def __init__(
+        self, fo: str = "", target_protocol: Optional[str] = None, target_options: Optional[dict] = None, **kwargs
+    ):
+        """
+        The compressed file system can be instantiated from any compressed file.
+        It reads the contents of compressed file as a filesystem with one file inside, as if it was an archive.
+
+        The single file inside the filesystem is named after the compresssed file,
+        without the compression extension at the end of the filename.
+
+        Args:
+            fo (:obj:``str``): Path to compressed file. Will fetch file using ``fsspec.open()``
+            mode (:obj:``str``): Currently, only 'rb' accepted
+            target_protocol(:obj:``str``, optional): To override the FS protocol inferred from a URL.
+            target_options (:obj:``dict``, optional): Kwargs passed when instantiating the target FS.
+        """
+        super().__init__(self, **kwargs)
+        self.fo = fo.__fspath__() if hasattr(fo, "__fspath__") else fo
+        # always open as "rb" since fsspec can then use the TextIOWrapper to make it work for "r" mode
+        self._open_with_fsspec = partial(
+            fsspec.open,
+            self.fo,
+            mode="rb",
+            protocol=target_protocol,
+            compression=self.compression,
+            client_kwargs={
+                "requote_redirect_url": False,  # see https://github.com/huggingface/datasets/pull/5459
+                "trust_env": True,  # Enable reading proxy env variables.
+                **(target_options or {}).pop("client_kwargs", {}),  # To avoid issues if it was already passed.
+            },
+            **(target_options or {}),
+        )
+        self.compressed_name = os.path.basename(self.fo.split("::")[0])
+        self.uncompressed_name = (
+            self.compressed_name[: self.compressed_name.rindex(".")]
+            if "." in self.compressed_name
+            else self.compressed_name
+        )
+        self.dir_cache = None
+
+    @classmethod
+    def _strip_protocol(cls, path):
+        # compressed file paths are always relative to the archive root
+        return super()._strip_protocol(path).lstrip("/")
+
+    def _get_dirs(self):
+        if self.dir_cache is None:
+            f = {**self._open_with_fsspec().fs.info(self.fo), "name": self.uncompressed_name}
+            self.dir_cache = {f["name"]: f}
+
+    def cat(self, path: str):
+        with self._open_with_fsspec().open() as f:
+            return f.read()
+
+    def _open(
+        self,
+        path: str,
+        mode: str = "rb",
+        block_size=None,
+        autocommit=True,
+        cache_options=None,
+        **kwargs,
+    ):
+        path = self._strip_protocol(path)
+        if mode != "rb":
+            raise ValueError(f"Tried to read with mode {mode} on file {self.fo} opened with mode 'rb'")
+        return self._open_with_fsspec().open()
+
+
+class Bz2FileSystem(BaseCompressedFileFileSystem):
+    """Read contents of BZ2 file as a filesystem with one file inside."""
+
+    protocol = "bz2"
+    compression = "bz2"
+    extensions = [".bz2"]
+
+
+class GzipFileSystem(BaseCompressedFileFileSystem):
+    """Read contents of GZIP file as a filesystem with one file inside."""
+
+    protocol = "gzip"
+    compression = "gzip"
+    extensions = [".gz", ".gzip"]
+
+
+class Lz4FileSystem(BaseCompressedFileFileSystem):
+    """Read contents of LZ4 file as a filesystem with one file inside."""
+
+    protocol = "lz4"
+    compression = "lz4"
+    extensions = [".lz4"]
+
+
+class XzFileSystem(BaseCompressedFileFileSystem):
+    """Read contents of .xz (LZMA) file as a filesystem with one file inside."""
+
+    protocol = "xz"
+    compression = "xz"
+    extensions = [".xz"]
+
+
+class ZstdFileSystem(BaseCompressedFileFileSystem):
+    """
+    Read contents of .zstd file as a filesystem with one file inside.
+    """
+
+    protocol = "zstd"
+    compression = "zstd"
+    extensions = [".zst", ".zstd"]

src/datasets/fingerprint.py

@@ -0,0 +1,480 @@
+import inspect
+import os
+import random
+import shutil
+import tempfile
+import weakref
+from functools import wraps
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+
+import numpy as np
+import xxhash
+
+from . import config
+from .naming import INVALID_WINDOWS_CHARACTERS_IN_PATH
+from .utils._dill import dumps
+from .utils.logging import get_logger
+
+
+if TYPE_CHECKING:
+    from .arrow_dataset import Dataset
+
+
+logger = get_logger(__name__)
+
+
+# Fingerprinting allows to have one deterministic fingerprint per dataset state.
+# A dataset fingerprint is updated after each transform.
+# Re-running the same transforms on a dataset in a different session results in the same fingerprint.
+# This is possible thanks to a custom hashing function that works with most python objects.
+
+# Fingerprinting is the main mechanism that enables caching.
+# The caching mechanism allows to reload an existing cache file if it's already been computed.
+
+
+#################
+# Caching
+#################
+
+_CACHING_ENABLED = True
+_TEMP_DIR_FOR_TEMP_CACHE_FILES: Optional["_TempCacheDir"] = None
+_DATASETS_WITH_TABLE_IN_TEMP_DIR: Optional[weakref.WeakSet] = None
+
+
+class _TempCacheDir:
+    """
+    A temporary directory for storing cached Arrow files with a cleanup that frees references to the Arrow files
+    before deleting the directory itself to avoid permission errors on Windows.
+    """
+
+    def __init__(self):
+        # Check if TMPDIR is set and handle the case where it doesn't exist
+        tmpdir = os.environ.get("TMPDIR") or os.environ.get("TEMP") or os.environ.get("TMP")
+        # Normalize the path to handle any path resolution issues
+        if tmpdir:
+            tmpdir = os.path.normpath(tmpdir)
+            if not os.path.exists(tmpdir):
+                # Auto-create the directory if it doesn't exist
+                # This prevents tempfile from silently falling back to /tmp
+                try:
+                    os.makedirs(tmpdir, exist_ok=True)
+                    logger.info(f"Created TMPDIR directory: {tmpdir}")
+                except OSError as e:
+                    raise OSError(
+                        f"TMPDIR is set to '{tmpdir}' but the directory does not exist and could not be created: {e}. "
+                        "Please create it manually or unset TMPDIR to fall back to the default temporary directory."
+                    ) from e
+            # If tmpdir exists, verify it's actually a directory and writable
+            elif not os.path.isdir(tmpdir):
+                raise OSError(
+                    f"TMPDIR is set to '{tmpdir}' but it is not a directory. "
+                    "Please point TMPDIR to a writable directory or unset it to fall back to the default temporary directory."
+                )
+
+        # Explicitly pass the directory to mkdtemp to ensure TMPDIR is respected
+        # This works even if tempfile.gettempdir() was already called and cached
+        # Pass dir=None if tmpdir is None to use default temp directory
+        self.name = tempfile.mkdtemp(prefix=config.TEMP_CACHE_DIR_PREFIX, dir=tmpdir)
+        self._finalizer = weakref.finalize(self, self._cleanup)
+
+    def _cleanup(self):
+        for dset in get_datasets_with_cache_file_in_temp_dir():
+            dset.__del__()
+        if os.path.exists(self.name):
+            try:
+                shutil.rmtree(self.name)
+            except Exception as e:
+                raise OSError(
+                    f"An error occurred while trying to delete temporary cache directory {self.name}. Please delete it manually."
+                ) from e
+
+    def cleanup(self):
+        if self._finalizer.detach():
+            self._cleanup()
+
+
+def maybe_register_dataset_for_temp_dir_deletion(dataset):
+    """
+    This function registers the datasets that have cache files in _TEMP_DIR_FOR_TEMP_CACHE_FILES in order
+    to properly delete them before deleting the temporary directory.
+    The temporary directory _TEMP_DIR_FOR_TEMP_CACHE_FILES is used when caching is disabled.
+    """
+    if _TEMP_DIR_FOR_TEMP_CACHE_FILES is None:
+        return
+
+    global _DATASETS_WITH_TABLE_IN_TEMP_DIR
+    if _DATASETS_WITH_TABLE_IN_TEMP_DIR is None:
+        _DATASETS_WITH_TABLE_IN_TEMP_DIR = weakref.WeakSet()
+    if any(
+        Path(_TEMP_DIR_FOR_TEMP_CACHE_FILES.name) in Path(cache_file["filename"]).parents
+        for cache_file in dataset.cache_files
+    ):
+        _DATASETS_WITH_TABLE_IN_TEMP_DIR.add(dataset)
+
+
+def get_datasets_with_cache_file_in_temp_dir():
+    return list(_DATASETS_WITH_TABLE_IN_TEMP_DIR) if _DATASETS_WITH_TABLE_IN_TEMP_DIR is not None else []
+
+
+def enable_caching():
+    """
+    When applying transforms on a dataset, the data are stored in cache files.
+    The caching mechanism allows to reload an existing cache file if it's already been computed.
+
+    Reloading a dataset is possible since the cache files are named using the dataset fingerprint, which is updated
+    after each transform.
+
+    If disabled, the library will no longer reload cached datasets files when applying transforms to the datasets.
+    More precisely, if the caching is disabled:
+    - cache files are always recreated
+    - cache files are written to a temporary directory that is deleted when session closes
+    - cache files are named using a random hash instead of the dataset fingerprint
+    - use [`~datasets.Dataset.save_to_disk`] to save a transformed dataset or it will be deleted when session closes
+    - caching doesn't affect [`~datasets.load_dataset`]. If you want to regenerate a dataset from scratch you should use
+    the `download_mode` parameter in [`~datasets.load_dataset`].
+    """
+    global _CACHING_ENABLED
+    _CACHING_ENABLED = True
+
+
+def disable_caching():
+    """
+    When applying transforms on a dataset, the data are stored in cache files.
+    The caching mechanism allows to reload an existing cache file if it's already been computed.
+
+    Reloading a dataset is possible since the cache files are named using the dataset fingerprint, which is updated
+    after each transform.
+
+    If disabled, the library will no longer reload cached datasets files when applying transforms to the datasets.
+    More precisely, if the caching is disabled:
+    - cache files are always recreated
+    - cache files are written to a temporary directory that is deleted when session closes
+    - cache files are named using a random hash instead of the dataset fingerprint
+    - use [`~datasets.Dataset.save_to_disk`] to save a transformed dataset or it will be deleted when session closes
+    - caching doesn't affect [`~datasets.load_dataset`]. If you want to regenerate a dataset from scratch you should use
+    the `download_mode` parameter in [`~datasets.load_dataset`].
+    """
+    global _CACHING_ENABLED
+    _CACHING_ENABLED = False
+
+
+def is_caching_enabled() -> bool:
+    """
+    When applying transforms on a dataset, the data are stored in cache files.
+    The caching mechanism allows to reload an existing cache file if it's already been computed.
+
+    Reloading a dataset is possible since the cache files are named using the dataset fingerprint, which is updated
+    after each transform.
+
+    If disabled, the library will no longer reload cached datasets files when applying transforms to the datasets.
+    More precisely, if the caching is disabled:
+    - cache files are always recreated
+    - cache files are written to a temporary directory that is deleted when session closes
+    - cache files are named using a random hash instead of the dataset fingerprint
+    - use [`~datasets.Dataset.save_to_disk`]] to save a transformed dataset or it will be deleted when session closes
+    - caching doesn't affect [`~datasets.load_dataset`]. If you want to regenerate a dataset from scratch you should use
+    the `download_mode` parameter in [`~datasets.load_dataset`].
+    """
+    global _CACHING_ENABLED
+    return bool(_CACHING_ENABLED)
+
+
+def get_temporary_cache_files_directory() -> str:
+    """Return a directory that is deleted when session closes."""
+    global _TEMP_DIR_FOR_TEMP_CACHE_FILES
+    if _TEMP_DIR_FOR_TEMP_CACHE_FILES is None:
+        _TEMP_DIR_FOR_TEMP_CACHE_FILES = _TempCacheDir()
+    return _TEMP_DIR_FOR_TEMP_CACHE_FILES.name
+
+
+#################
+# Hashing
+#################
+
+
+class Hasher:
+    """Hasher that accepts python objects as inputs."""
+
+    dispatch: dict = {}
+
+    def __init__(self):
+        self.m = xxhash.xxh64()
+
+    @classmethod
+    def hash_bytes(cls, value: Union[bytes, list[bytes]]) -> str:
+        value = [value] if isinstance(value, bytes) else value
+        m = xxhash.xxh64()
+        for x in value:
+            m.update(x)
+        return m.hexdigest()
+
+    @classmethod
+    def hash(cls, value: Any) -> str:
+        return cls.hash_bytes(dumps(value))
+
+    def update(self, value: Any) -> None:
+        header_for_update = f"=={type(value)}=="
+        value_for_update = self.hash(value)
+        self.m.update(header_for_update.encode("utf8"))
+        self.m.update(value_for_update.encode("utf-8"))
+
+    def hexdigest(self) -> str:
+        return self.m.hexdigest()
+
+
+#################
+# Fingerprinting
+#################
+
+fingerprint_rng = random.Random()
+# we show a warning only once when fingerprinting fails to avoid spam
+fingerprint_warnings: dict[str, bool] = {}
+
+
+def generate_fingerprint(dataset: "Dataset") -> str:
+    state = dataset.__dict__
+    hasher = Hasher()
+    for key in sorted(state):
+        if key == "_fingerprint":
+            continue
+        hasher.update(key)
+        hasher.update(state[key])
+    # hash data files last modification timestamps as well
+    for cache_file in dataset.cache_files:
+        hasher.update(os.path.getmtime(cache_file["filename"]))
+    return hasher.hexdigest()
+
+
+def generate_random_fingerprint(nbits: int = 64) -> str:
+    return f"{fingerprint_rng.getrandbits(nbits):0{nbits // 4}x}"
+
+
+def update_fingerprint(fingerprint, transform, transform_args):
+    global fingerprint_warnings
+    hasher = Hasher()
+    hasher.update(fingerprint)
+    try:
+        hasher.update(transform)
+    except:  # noqa various errors might raise here from pickle or dill
+        if _CACHING_ENABLED:
+            if not fingerprint_warnings.get("update_fingerprint_transform_hash_failed", False):
+                logger.warning(
+                    f"Transform {transform} couldn't be hashed properly, a random hash was used instead. "
+                    "Make sure your transforms and parameters are serializable with pickle or dill for the dataset fingerprinting and caching to work. "
+                    "If you reuse this transform, the caching mechanism will consider it to be different from the previous calls and recompute everything. "
+                    "This warning is only shown once. Subsequent hashing failures won't be shown."
+                )
+                fingerprint_warnings["update_fingerprint_transform_hash_failed"] = True
+            else:
+                logger.info(f"Transform {transform} couldn't be hashed properly, a random hash was used instead.")
+        else:
+            logger.info(
+                f"Transform {transform} couldn't be hashed properly, a random hash was used instead. This doesn't affect caching since it's disabled."
+            )
+
+        return generate_random_fingerprint()
+    for key in sorted(transform_args):
+        hasher.update(key)
+        try:
+            hasher.update(transform_args[key])
+        except:  # noqa various errors might raise here from pickle or dill
+            if _CACHING_ENABLED:
+                if not fingerprint_warnings.get("update_fingerprint_transform_hash_failed", False):
+                    logger.warning(
+                        f"Parameter '{key}'={transform_args[key]} of the transform {transform} couldn't be hashed properly, a random hash was used instead. "
+                        "Make sure your transforms and parameters are serializable with pickle or dill for the dataset fingerprinting and caching to work. "
+                        "If you reuse this transform, the caching mechanism will consider it to be different from the previous calls and recompute everything. "
+                        "This warning is only shown once. Subsequent hashing failures won't be shown."
+                    )
+                    fingerprint_warnings["update_fingerprint_transform_hash_failed"] = True
+                else:
+                    logger.info(
+                        f"Parameter '{key}'={transform_args[key]} of the transform {transform} couldn't be hashed properly, a random hash was used instead."
+                    )
+            else:
+                logger.info(
+                    f"Parameter '{key}'={transform_args[key]} of the transform {transform} couldn't be hashed properly, a random hash was used instead. This doesn't affect caching since it's disabled."
+                )
+            return generate_random_fingerprint()
+    return hasher.hexdigest()
+
+
+def validate_fingerprint(fingerprint: str, max_length=64):
+    """
+    Make sure the fingerprint is a non-empty string that is not longer that max_length=64 by default,
+    so that the fingerprint can be used to name cache files without issues.
+    """
+    if not isinstance(fingerprint, str) or not fingerprint:
+        raise ValueError(f"Invalid fingerprint '{fingerprint}': it should be a non-empty string.")
+    for invalid_char in INVALID_WINDOWS_CHARACTERS_IN_PATH:
+        if invalid_char in fingerprint:
+            raise ValueError(
+                f"Invalid fingerprint. Bad characters from black list '{INVALID_WINDOWS_CHARACTERS_IN_PATH}' found in '{fingerprint}'. "
+                f"They could create issues when creating cache files."
+            )
+    if len(fingerprint) > max_length:
+        raise ValueError(
+            f"Invalid fingerprint. Maximum lenth is {max_length} but '{fingerprint}' has length {len(fingerprint)}."
+            "It could create issues when creating cache files."
+        )
+
+
+def format_transform_for_fingerprint(func: Callable, version: Optional[str] = None) -> str:
+    """
+    Format a transform to the format that will be used to update the fingerprint.
+    """
+    transform = f"{func.__module__}.{func.__qualname__}"
+    if version is not None:
+        transform += f"@{version}"
+    return transform
+
+
+def format_kwargs_for_fingerprint(
+    func: Callable,
+    args: tuple,
+    kwargs: dict[str, Any],
+    use_kwargs: Optional[list[str]] = None,
+    ignore_kwargs: Optional[list[str]] = None,
+    randomized_function: bool = False,
+) -> dict[str, Any]:
+    """
+    Format the kwargs of a transform to the format that will be used to update the fingerprint.
+    """
+    kwargs_for_fingerprint = kwargs.copy()
+    if args:
+        params = [p.name for p in inspect.signature(func).parameters.values() if p != p.VAR_KEYWORD]
+        args = args[1:]  # assume the first argument is the dataset
+        params = params[1:]
+        kwargs_for_fingerprint.update(zip(params, args))
+    else:
+        del kwargs_for_fingerprint[
+            next(iter(inspect.signature(func).parameters))
+        ]  # assume the first key is the dataset
+
+    # keep the right kwargs to be hashed to generate the fingerprint
+
+    if use_kwargs:
+        kwargs_for_fingerprint = {k: v for k, v in kwargs_for_fingerprint.items() if k in use_kwargs}
+    if ignore_kwargs:
+        kwargs_for_fingerprint = {k: v for k, v in kwargs_for_fingerprint.items() if k not in ignore_kwargs}
+    if randomized_function:  # randomized functions have `seed` and `generator` parameters
+        if kwargs_for_fingerprint.get("seed") is None and kwargs_for_fingerprint.get("generator") is None:
+            _, seed, pos, *_ = np.random.get_state()
+            seed = seed[pos] if pos < 624 else seed[0]
+            kwargs_for_fingerprint["generator"] = np.random.default_rng(seed)
+
+    # remove kwargs that are the default values
+
+    default_values = {
+        p.name: p.default for p in inspect.signature(func).parameters.values() if p.default != inspect._empty
+    }
+    for default_varname, default_value in default_values.items():
+        if default_varname in kwargs_for_fingerprint and kwargs_for_fingerprint[default_varname] == default_value:
+            kwargs_for_fingerprint.pop(default_varname)
+    return kwargs_for_fingerprint
+
+
+def fingerprint_transform(
+    inplace: bool,
+    use_kwargs: Optional[list[str]] = None,
+    ignore_kwargs: Optional[list[str]] = None,
+    fingerprint_names: Optional[list[str]] = None,
+    randomized_function: bool = False,
+    version: Optional[str] = None,
+):
+    """
+    Wrapper for dataset transforms to update the dataset fingerprint using ``update_fingerprint``
+    Args:
+        inplace (:obj:`bool`):  If inplace is True, the fingerprint of the dataset is updated inplace.
+            Otherwise, a parameter "new_fingerprint" is passed to the wrapped method that should take care of
+            setting the fingerprint of the returned Dataset.
+        use_kwargs (:obj:`List[str]`, optional): optional white list of argument names to take into account
+            to update the fingerprint to the wrapped method that should take care of
+            setting the fingerprint of the returned Dataset. By default all the arguments are used.
+        ignore_kwargs (:obj:`List[str]`, optional): optional black list of argument names to take into account
+            to update the fingerprint. Note that ignore_kwargs prevails on use_kwargs.
+        fingerprint_names (:obj:`List[str]`, optional, defaults to ["new_fingerprint"]):
+            If the dataset transforms is not inplace and returns a DatasetDict, then it can require
+            several fingerprints (one per dataset in the DatasetDict). By specifying fingerprint_names,
+            one fingerprint named after each element of fingerprint_names is going to be passed.
+        randomized_function (:obj:`bool`, defaults to False): If the dataset transform is random and has
+            optional parameters "seed" and "generator", then you can set randomized_function to True.
+            This way, even if users set "seed" and "generator" to None, then the fingerprint is
+            going to be randomly generated depending on numpy's current state. In this case, the
+            generator is set to np.random.default_rng(np.random.get_state()[1][0]).
+        version (:obj:`str`, optional): version of the transform. The version is taken into account when
+            computing the fingerprint. If a datase transform changes (or at least if the output data
+            that are cached changes), then one should increase the version. If the version stays the
+            same, then old cached data could be reused that are not compatible with the new transform.
+            It should be in the format "MAJOR.MINOR.PATCH".
+    """
+
+    if use_kwargs is not None and not isinstance(use_kwargs, list):
+        raise ValueError(f"use_kwargs is supposed to be a list, not {type(use_kwargs)}")
+
+    if ignore_kwargs is not None and not isinstance(ignore_kwargs, list):
+        raise ValueError(f"ignore_kwargs is supposed to be a list, not {type(use_kwargs)}")
+
+    if inplace and fingerprint_names:
+        raise ValueError("fingerprint_names are only used when inplace is False")
+
+    fingerprint_names = fingerprint_names if fingerprint_names is not None else ["new_fingerprint"]
+
+    def _fingerprint(func):
+        if not inplace and not all(name in func.__code__.co_varnames for name in fingerprint_names):
+            raise ValueError(f"function {func} is missing parameters {fingerprint_names} in signature")
+
+        if randomized_function:  # randomized function have seed and generator parameters
+            if "seed" not in func.__code__.co_varnames:
+                raise ValueError(f"'seed' must be in {func}'s signature")
+            if "generator" not in func.__code__.co_varnames:
+                raise ValueError(f"'generator' must be in {func}'s signature")
+        # this call has to be outside the wrapper or since __qualname__ changes in multiprocessing
+        transform = format_transform_for_fingerprint(func, version=version)
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            kwargs_for_fingerprint = format_kwargs_for_fingerprint(
+                func,
+                args,
+                kwargs,
+                use_kwargs=use_kwargs,
+                ignore_kwargs=ignore_kwargs,
+                randomized_function=randomized_function,
+            )
+
+            if args:
+                dataset: Dataset = args[0]
+                args = args[1:]
+            else:
+                dataset: Dataset = kwargs.pop(next(iter(inspect.signature(func).parameters)))
+
+            # compute new_fingerprint and add it to the args of not in-place transforms
+            if inplace:
+                new_fingerprint = update_fingerprint(dataset._fingerprint, transform, kwargs_for_fingerprint)
+            else:
+                for fingerprint_name in fingerprint_names:  # transforms like `train_test_split` have several hashes
+                    if kwargs.get(fingerprint_name) is None:
+                        kwargs_for_fingerprint["fingerprint_name"] = fingerprint_name
+                        kwargs[fingerprint_name] = update_fingerprint(
+                            dataset._fingerprint, transform, kwargs_for_fingerprint
+                        )
+                    else:
+                        validate_fingerprint(kwargs[fingerprint_name])
+
+            # Call actual function
+
+            out = func(dataset, *args, **kwargs)
+
+            # Update fingerprint of in-place transforms + update in-place history of transforms
+
+            if inplace:  # update after calling func so that the fingerprint doesn't change if the function fails
+                dataset._fingerprint = new_fingerprint
+
+            return out
+
+        wrapper._decorator_name_ = "fingerprint"
+        return wrapper
+
+    return _fingerprint

src/datasets/formatting/__init__.py

@@ -0,0 +1,136 @@
+# Copyright 2020 The HuggingFace Datasets Authors and the TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, Optional, Type
+
+from .. import config
+from ..utils import logging
+from .formatting import (
+    ArrowFormatter,
+    CustomFormatter,
+    Formatter,
+    PandasFormatter,
+    PythonFormatter,
+    TableFormatter,
+    TensorFormatter,
+    format_table,
+    query_table,
+)
+from .np_formatter import NumpyFormatter
+
+
+logger = logging.get_logger(__name__)
+
+_FORMAT_TYPES: dict[Optional[str], type[Formatter]] = {}
+_FORMAT_TYPES_ALIASES: dict[Optional[str], str] = {}
+_FORMAT_TYPES_ALIASES_UNAVAILABLE: dict[Optional[str], Exception] = {}
+
+
+def _register_formatter(
+    formatter_cls: type,
+    format_type: Optional[str],
+    aliases: Optional[list[str]] = None,
+):
+    """
+    Register a Formatter object using a name and optional aliases.
+    This function must be used on a Formatter class.
+    """
+    aliases = aliases if aliases is not None else []
+    if format_type in _FORMAT_TYPES:
+        logger.warning(
+            f"Overwriting format type '{format_type}' ({_FORMAT_TYPES[format_type].__name__} -> {formatter_cls.__name__})"
+        )
+    _FORMAT_TYPES[format_type] = formatter_cls
+    for alias in set(aliases + [format_type]):
+        if alias in _FORMAT_TYPES_ALIASES:
+            logger.warning(
+                f"Overwriting format type alias '{alias}' ({_FORMAT_TYPES_ALIASES[alias]} -> {format_type})"
+            )
+        _FORMAT_TYPES_ALIASES[alias] = format_type
+
+
+def _register_unavailable_formatter(
+    unavailable_error: Exception, format_type: Optional[str], aliases: Optional[list[str]] = None
+):
+    """
+    Register an unavailable Formatter object using a name and optional aliases.
+    This function must be used on an Exception object that is raised when trying to get the unavailable formatter.
+    """
+    aliases = aliases if aliases is not None else []
+    for alias in set(aliases + [format_type]):
+        _FORMAT_TYPES_ALIASES_UNAVAILABLE[alias] = unavailable_error
+
+
+# Here we define all the available formatting functions that can be used by `Dataset.set_format`
+_register_formatter(PythonFormatter, None, aliases=["python"])
+_register_formatter(ArrowFormatter, "arrow", aliases=["pa", "pyarrow"])
+_register_formatter(NumpyFormatter, "numpy", aliases=["np"])
+_register_formatter(PandasFormatter, "pandas", aliases=["pd"])
+_register_formatter(CustomFormatter, "custom")
+
+if config.POLARS_AVAILABLE:
+    from .polars_formatter import PolarsFormatter
+
+    _register_formatter(PolarsFormatter, "polars", aliases=["pl"])
+else:
+    _polars_error = ValueError("Polars needs to be installed to be able to return Polars dataframes.")
+    _register_unavailable_formatter(_polars_error, "polars", aliases=["pl"])
+
+if config.TORCH_AVAILABLE:
+    from .torch_formatter import TorchFormatter
+
+    _register_formatter(TorchFormatter, "torch", aliases=["pt", "pytorch"])
+else:
+    _torch_error = ValueError("PyTorch needs to be installed to be able to return PyTorch tensors.")
+    _register_unavailable_formatter(_torch_error, "torch", aliases=["pt", "pytorch"])
+
+if config.TF_AVAILABLE:
+    from .tf_formatter import TFFormatter
+
+    _register_formatter(TFFormatter, "tensorflow", aliases=["tf"])
+else:
+    _tf_error = ValueError("Tensorflow needs to be installed to be able to return Tensorflow tensors.")
+    _register_unavailable_formatter(_tf_error, "tensorflow", aliases=["tf"])
+
+if config.JAX_AVAILABLE:
+    from .jax_formatter import JaxFormatter
+
+    _register_formatter(JaxFormatter, "jax", aliases=[])
+else:
+    _jax_error = ValueError("JAX needs to be installed to be able to return JAX arrays.")
+    _register_unavailable_formatter(_jax_error, "jax", aliases=[])
+
+
+def get_format_type_from_alias(format_type: Optional[str]) -> Optional[str]:
+    """If the given format type is a known alias, then return its main type name. Otherwise return the type with no change."""
+    if format_type in _FORMAT_TYPES_ALIASES:
+        return _FORMAT_TYPES_ALIASES[format_type]
+    else:
+        return format_type
+
+
+def get_formatter(format_type: Optional[str], **format_kwargs) -> Formatter:
+    """
+    Factory function to get a Formatter given its type name and keyword arguments.
+    A formatter is an object that extracts and formats data from pyarrow table.
+    It defines the formatting for rows, columns and batches.
+    If the formatter for a given type name doesn't exist or is not available, an error is raised.
+    """
+    format_type = get_format_type_from_alias(format_type)
+    if format_type in _FORMAT_TYPES:
+        return _FORMAT_TYPES[format_type](**format_kwargs)
+    if format_type in _FORMAT_TYPES_ALIASES_UNAVAILABLE:
+        raise _FORMAT_TYPES_ALIASES_UNAVAILABLE[format_type]
+    else:
+        raise ValueError(f"Format type should be one of {list(_FORMAT_TYPES.keys())}, but got '{format_type}'")

src/datasets/formatting/formatting.py

@@ -0,0 +1,678 @@
+# Copyright 2020 The HuggingFace Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numbers
+import operator
+from collections.abc import Iterable, Mapping, MutableMapping
+from functools import partial
+
+# Lint as: python3
+from typing import Any, Callable, Generic, Optional, TypeVar, Union
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+
+from ..features import Features
+from ..features.features import _ArrayXDExtensionType, _is_zero_copy_only, decode_nested_example, pandas_types_mapper
+from ..table import Table
+from ..utils.py_utils import no_op_if_value_is_null
+
+
+T = TypeVar("T")
+
+RowFormat = TypeVar("RowFormat")
+ColumnFormat = TypeVar("ColumnFormat")
+BatchFormat = TypeVar("BatchFormat")
+
+
+def _is_range_contiguous(key: range) -> bool:
+    return key.step == 1 and key.stop >= key.start
+
+
+def _raise_bad_key_type(key: Any):
+    raise TypeError(
+        f"Wrong key type: '{key}' of type '{type(key)}'. Expected one of int, slice, range, str or Iterable."
+    )
+
+
+def _query_table_with_indices_mapping(
+    table: Table, key: Union[int, slice, range, str, Iterable], indices: Table
+) -> pa.Table:
+    """
+    Query a pyarrow Table to extract the subtable that correspond to the given key.
+    The :obj:`indices` parameter corresponds to the indices mapping in case we cant to take into
+    account a shuffling or an indices selection for example.
+    The indices table must contain one column named "indices" of type uint64.
+    """
+    if isinstance(key, int):
+        key = indices.fast_slice(key % indices.num_rows, 1).column(0)[0].as_py()
+        return _query_table(table, key)
+    if isinstance(key, slice):
+        key = range(*key.indices(indices.num_rows))
+    if isinstance(key, range):
+        if _is_range_contiguous(key) and key.start >= 0:
+            return _query_table(
+                table, [i.as_py() for i in indices.fast_slice(key.start, key.stop - key.start).column(0)]
+            )
+        else:
+            pass  # treat as an iterable
+    if isinstance(key, str):
+        table = table.select([key])
+        return _query_table(table, indices.column(0).to_pylist())
+    if isinstance(key, Iterable):
+        return _query_table(table, [indices.fast_slice(i, 1).column(0)[0].as_py() for i in key])
+
+    _raise_bad_key_type(key)
+
+
+def _query_table(table: Table, key: Union[int, slice, range, str, Iterable]) -> pa.Table:
+    """
+    Query a pyarrow Table to extract the subtable that correspond to the given key.
+    """
+    if isinstance(key, int):
+        return table.fast_slice(key % table.num_rows, 1)
+    if isinstance(key, slice):
+        key = range(*key.indices(table.num_rows))
+    if isinstance(key, range):
+        if _is_range_contiguous(key) and key.start >= 0:
+            return table.fast_slice(key.start, key.stop - key.start)
+        else:
+            pass  # treat as an iterable
+    if isinstance(key, str):
+        return table.table.drop([column for column in table.column_names if column != key])
+    if isinstance(key, Iterable):
+        key = np.fromiter(key, np.int64)
+        if len(key) == 0:
+            return table.table.slice(0, 0)
+        # don't use pyarrow.Table.take even for pyarrow >=1.0 (see https://issues.apache.org/jira/browse/ARROW-9773)
+        return table.fast_gather(key % table.num_rows)
+
+    _raise_bad_key_type(key)
+
+
+def _is_array_with_nulls(pa_array: pa.Array) -> bool:
+    return pa_array.null_count > 0
+
+
+class BaseArrowExtractor(Generic[RowFormat, ColumnFormat, BatchFormat]):
+    """
+    Arrow extractor are used to extract data from pyarrow tables.
+    It makes it possible to extract rows, columns and batches.
+    These three extractions types have to be implemented.
+    """
+
+    def extract_row(self, pa_table: pa.Table) -> RowFormat:
+        raise NotImplementedError
+
+    def extract_column(self, pa_table: pa.Table) -> ColumnFormat:
+        raise NotImplementedError
+
+    def extract_batch(self, pa_table: pa.Table) -> BatchFormat:
+        raise NotImplementedError
+
+
+def _unnest(py_dict: dict[str, list[T]]) -> dict[str, T]:
+    """Return the first element of a batch (dict) as a row (dict)"""
+    return {key: array[0] for key, array in py_dict.items()}
+
+
+class SimpleArrowExtractor(BaseArrowExtractor[pa.Table, pa.Array, pa.Table]):
+    def extract_row(self, pa_table: pa.Table) -> pa.Table:
+        return pa_table
+
+    def extract_column(self, pa_table: pa.Table) -> pa.Array:
+        return pa_table.column(0)
+
+    def extract_batch(self, pa_table: pa.Table) -> pa.Table:
+        return pa_table
+
+
+class PythonArrowExtractor(BaseArrowExtractor[dict, list, dict]):
+    def extract_row(self, pa_table: pa.Table) -> dict:
+        return _unnest(pa_table.to_pydict())
+
+    def extract_column(self, pa_table: pa.Table) -> list:
+        return pa_table.column(0).to_pylist()
+
+    def extract_batch(self, pa_table: pa.Table) -> dict:
+        return pa_table.to_pydict()
+
+
+class NumpyArrowExtractor(BaseArrowExtractor[dict, np.ndarray, dict]):
+    def __init__(self, **np_array_kwargs):
+        self.np_array_kwargs = np_array_kwargs
+
+    def extract_row(self, pa_table: pa.Table) -> dict:
+        return _unnest(self.extract_batch(pa_table))
+
+    def extract_column(self, pa_table: pa.Table) -> np.ndarray:
+        return self._arrow_array_to_numpy(pa_table[pa_table.column_names[0]])
+
+    def extract_batch(self, pa_table: pa.Table) -> dict:
+        return {col: self._arrow_array_to_numpy(pa_table[col]) for col in pa_table.column_names}
+
+    def _arrow_array_to_numpy(self, pa_array: pa.Array) -> np.ndarray:
+        if isinstance(pa_array, pa.ChunkedArray):
+            if isinstance(pa_array.type, _ArrayXDExtensionType):
+                # don't call to_pylist() to preserve dtype of the fixed-size array
+                zero_copy_only = _is_zero_copy_only(pa_array.type.storage_dtype, unnest=True)
+                array: list = [
+                    row for chunk in pa_array.chunks for row in chunk.to_numpy(zero_copy_only=zero_copy_only)
+                ]
+            else:
+                zero_copy_only = _is_zero_copy_only(pa_array.type) and all(
+                    not _is_array_with_nulls(chunk) for chunk in pa_array.chunks
+                )
+                array: list = [
+                    row for chunk in pa_array.chunks for row in chunk.to_numpy(zero_copy_only=zero_copy_only)
+                ]
+        else:
+            if isinstance(pa_array.type, _ArrayXDExtensionType):
+                # don't call to_pylist() to preserve dtype of the fixed-size array
+                zero_copy_only = _is_zero_copy_only(pa_array.type.storage_dtype, unnest=True)
+                array: list = pa_array.to_numpy(zero_copy_only=zero_copy_only)
+            else:
+                zero_copy_only = _is_zero_copy_only(pa_array.type) and not _is_array_with_nulls(pa_array)
+                array: list = pa_array.to_numpy(zero_copy_only=zero_copy_only).tolist()
+
+        if len(array) > 0:
+            if any(
+                (isinstance(x, np.ndarray) and (x.dtype == object or x.shape != array[0].shape))
+                or (isinstance(x, float) and np.isnan(x))
+                for x in array
+            ):
+                if np.lib.NumpyVersion(np.__version__) >= "2.0.0b1":
+                    return np.asarray(array, dtype=object)
+                return np.array(array, copy=False, dtype=object)
+        if np.lib.NumpyVersion(np.__version__) >= "2.0.0b1":
+            return np.asarray(array)
+        else:
+            return np.array(array, copy=False)
+
+
+class PandasArrowExtractor(BaseArrowExtractor[pd.DataFrame, pd.Series, pd.DataFrame]):
+    def extract_row(self, pa_table: pa.Table) -> pd.DataFrame:
+        return pa_table.slice(length=1).to_pandas(types_mapper=pandas_types_mapper)
+
+    def extract_column(self, pa_table: pa.Table) -> pd.Series:
+        return pa_table.select([0]).to_pandas(types_mapper=pandas_types_mapper)[pa_table.column_names[0]]
+
+    def extract_batch(self, pa_table: pa.Table) -> pd.DataFrame:
+        return pa_table.to_pandas(types_mapper=pandas_types_mapper)
+
+
+class PythonFeaturesDecoder:
+    def __init__(
+        self, features: Optional[Features], token_per_repo_id: Optional[dict[str, Union[str, bool, None]]] = None
+    ):
+        self.features = features
+        self.token_per_repo_id = token_per_repo_id
+
+    def decode_row(self, row: dict) -> dict:
+        return self.features.decode_example(row, token_per_repo_id=self.token_per_repo_id) if self.features else row
+
+    def decode_column(self, column: list, column_name: str) -> list:
+        return (
+            self.features.decode_column(column, column_name, token_per_repo_id=self.token_per_repo_id)
+            if self.features
+            else column
+        )
+
+    def decode_batch(self, batch: dict) -> dict:
+        return self.features.decode_batch(batch, token_per_repo_id=self.token_per_repo_id) if self.features else batch
+
+
+class PandasFeaturesDecoder:
+    def __init__(self, features: Optional[Features]):
+        self.features = features
+
+    def decode_row(self, row: pd.DataFrame) -> pd.DataFrame:
+        decode = (
+            {
+                column_name: no_op_if_value_is_null(partial(decode_nested_example, feature))
+                for column_name, feature in self.features.items()
+                if self.features._column_requires_decoding[column_name]
+            }
+            if self.features
+            else {}
+        )
+        if decode:
+            row[list(decode.keys())] = row.transform(decode)
+        return row
+
+    def decode_column(self, column: pd.Series, column_name: str) -> pd.Series:
+        decode = (
+            no_op_if_value_is_null(partial(decode_nested_example, self.features[column_name]))
+            if self.features and column_name in self.features and self.features._column_requires_decoding[column_name]
+            else None
+        )
+        if decode:
+            column = column.transform(decode)
+        return column
+
+    def decode_batch(self, batch: pd.DataFrame) -> pd.DataFrame:
+        return self.decode_row(batch)
+
+
+class LazyDict(MutableMapping):
+    """A dictionary backed by Arrow data. The values are formatted on-the-fly when accessing the dictionary."""
+
+    def __init__(self, pa_table: pa.Table, formatter: "Formatter"):
+        self.pa_table = pa_table
+        self.formatter = formatter
+
+        self.data = dict.fromkeys(pa_table.column_names)
+        self.keys_to_format = set(self.data.keys())
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, key):
+        value = self.data[key]
+        if key in self.keys_to_format:
+            value = self.format(key)
+            self.data[key] = value
+            self.keys_to_format.remove(key)
+        return value
+
+    def __setitem__(self, key, value):
+        if key in self.keys_to_format:
+            self.keys_to_format.remove(key)
+        self.data[key] = value
+
+    def __delitem__(self, key) -> None:
+        if key in self.keys_to_format:
+            self.keys_to_format.remove(key)
+        del self.data[key]
+
+    def __iter__(self):
+        return iter(self.data)
+
+    def __contains__(self, key):
+        return key in self.data
+
+    def __repr__(self):
+        self._format_all()
+        return repr(self.data)
+
+    def __or__(self, other):
+        if isinstance(other, LazyDict):
+            inst = self.copy()
+            other = other.copy()
+            other._format_all()
+            inst.keys_to_format -= other.data.keys()
+            inst.data = inst.data | other.data
+            return inst
+        if isinstance(other, dict):
+            inst = self.copy()
+            inst.keys_to_format -= other.keys()
+            inst.data = inst.data | other
+            return inst
+        return NotImplemented
+
+    def __ror__(self, other):
+        if isinstance(other, LazyDict):
+            inst = self.copy()
+            other = other.copy()
+            other._format_all()
+            inst.keys_to_format -= other.data.keys()
+            inst.data = other.data | inst.data
+            return inst
+        if isinstance(other, dict):
+            inst = self.copy()
+            inst.keys_to_format -= other.keys()
+            inst.data = other | inst.data
+            return inst
+        return NotImplemented
+
+    def __ior__(self, other):
+        if isinstance(other, LazyDict):
+            other = other.copy()
+            other._format_all()
+            self.keys_to_format -= other.data.keys()
+            self.data |= other.data
+        else:
+            self.keys_to_format -= other.keys()
+            self.data |= other
+        return self
+
+    def __copy__(self):
+        # Identical to `UserDict.__copy__`
+        inst = self.__class__.__new__(self.__class__)
+        inst.__dict__.update(self.__dict__)
+        # Create a copy and avoid triggering descriptors
+        inst.__dict__["data"] = self.__dict__["data"].copy()
+        inst.__dict__["keys_to_format"] = self.__dict__["keys_to_format"].copy()
+        return inst
+
+    def copy(self):
+        import copy
+
+        return copy.copy(self)
+
+    @classmethod
+    def fromkeys(cls, iterable, value=None):
+        raise NotImplementedError
+
+    def format(self, key):
+        raise NotImplementedError
+
+    def _format_all(self):
+        for key in self.keys_to_format:
+            self.data[key] = self.format(key)
+        self.keys_to_format.clear()
+
+
+class LazyRow(LazyDict):
+    def format(self, key):
+        return self.formatter.format_column(self.pa_table.select([key]))[0]
+
+
+class LazyBatch(LazyDict):
+    def format(self, key):
+        return self.formatter.format_column(self.pa_table.select([key]))
+
+
+class Formatter(Generic[RowFormat, ColumnFormat, BatchFormat]):
+    """
+    A formatter is an object that extracts and formats data from pyarrow tables.
+    It defines the formatting for rows, columns and batches.
+    """
+
+    simple_arrow_extractor = SimpleArrowExtractor
+    python_arrow_extractor = PythonArrowExtractor
+    numpy_arrow_extractor = NumpyArrowExtractor
+    pandas_arrow_extractor = PandasArrowExtractor
+
+    def __init__(
+        self,
+        features: Optional[Features] = None,
+        token_per_repo_id: Optional[dict[str, Union[str, bool, None]]] = None,
+    ):
+        self.features = features
+        self.token_per_repo_id = token_per_repo_id
+        self.python_features_decoder = PythonFeaturesDecoder(self.features, self.token_per_repo_id)
+        self.pandas_features_decoder = PandasFeaturesDecoder(self.features)
+
+    def __call__(self, pa_table: pa.Table, query_type: str) -> Union[RowFormat, ColumnFormat, BatchFormat]:
+        if query_type == "row":
+            return self.format_row(pa_table)
+        elif query_type == "column":
+            return self.format_column(pa_table)
+        elif query_type == "batch":
+            return self.format_batch(pa_table)
+
+    def format_row(self, pa_table: pa.Table) -> RowFormat:
+        raise NotImplementedError
+
+    def format_column(self, pa_table: pa.Table) -> ColumnFormat:
+        raise NotImplementedError
+
+    def format_batch(self, pa_table: pa.Table) -> BatchFormat:
+        raise NotImplementedError
+
+
+class TensorFormatter(Formatter[RowFormat, ColumnFormat, BatchFormat]):
+    def recursive_tensorize(self, data_struct: dict):
+        raise NotImplementedError
+
+
+class TableFormatter(Formatter[RowFormat, ColumnFormat, BatchFormat]):
+    table_type: str
+    column_type: str
+
+
+class ArrowFormatter(TableFormatter[pa.Table, pa.Array, pa.Table]):
+    table_type = "arrow table"
+    column_type = "arrow array"
+
+    def format_row(self, pa_table: pa.Table) -> pa.Table:
+        return self.simple_arrow_extractor().extract_row(pa_table)
+
+    def format_column(self, pa_table: pa.Table) -> pa.Array:
+        return self.simple_arrow_extractor().extract_column(pa_table)
+
+    def format_batch(self, pa_table: pa.Table) -> pa.Table:
+        return self.simple_arrow_extractor().extract_batch(pa_table)
+
+
+class PythonFormatter(Formatter[Mapping, list, Mapping]):
+    def __init__(self, features=None, lazy=False, token_per_repo_id=None):
+        super().__init__(features, token_per_repo_id)
+        self.lazy = lazy
+
+    def format_row(self, pa_table: pa.Table) -> Mapping:
+        if self.lazy:
+            return LazyRow(pa_table, self)
+        row = self.python_arrow_extractor().extract_row(pa_table)
+        row = self.python_features_decoder.decode_row(row)
+        return row
+
+    def format_column(self, pa_table: pa.Table) -> list:
+        column = self.python_arrow_extractor().extract_column(pa_table)
+        column = self.python_features_decoder.decode_column(column, pa_table.column_names[0])
+        return column
+
+    def format_batch(self, pa_table: pa.Table) -> Mapping:
+        if self.lazy:
+            return LazyBatch(pa_table, self)
+        batch = self.python_arrow_extractor().extract_batch(pa_table)
+        batch = self.python_features_decoder.decode_batch(batch)
+        return batch
+
+
+class PandasFormatter(TableFormatter[pd.DataFrame, pd.Series, pd.DataFrame]):
+    table_type = "pandas dataframe"
+    column_type = "pandas series"
+
+    def format_row(self, pa_table: pa.Table) -> pd.DataFrame:
+        row = self.pandas_arrow_extractor().extract_row(pa_table)
+        row = self.pandas_features_decoder.decode_row(row)
+        return row
+
+    def format_column(self, pa_table: pa.Table) -> pd.Series:
+        column = self.pandas_arrow_extractor().extract_column(pa_table)
+        column = self.pandas_features_decoder.decode_column(column, pa_table.column_names[0])
+        return column
+
+    def format_batch(self, pa_table: pa.Table) -> pd.DataFrame:
+        row = self.pandas_arrow_extractor().extract_batch(pa_table)
+        row = self.pandas_features_decoder.decode_batch(row)
+        return row
+
+
+class CustomFormatter(Formatter[dict, ColumnFormat, dict]):
+    """
+    A user-defined custom formatter function defined by a ``transform``.
+    The transform must take as input a batch of data extracted for an arrow table using the python extractor,
+    and return a batch.
+    If the output batch is not a dict, then output_all_columns won't work.
+    If the output batch has several fields, then querying a single column won't work since we don't know which field
+    to return.
+    """
+
+    def __init__(self, transform: Callable[[dict], dict], features=None, token_per_repo_id=None, **kwargs):
+        super().__init__(features=features, token_per_repo_id=token_per_repo_id)
+        self.transform = transform
+
+    def format_row(self, pa_table: pa.Table) -> dict:
+        formatted_batch = self.format_batch(pa_table)
+        try:
+            return _unnest(formatted_batch)
+        except Exception as exc:
+            raise TypeError(
+                f"Custom formatting function must return a dict of sequences to be able to pick a row, but got {formatted_batch}"
+            ) from exc
+
+    def format_column(self, pa_table: pa.Table) -> ColumnFormat:
+        formatted_batch = self.format_batch(pa_table)
+        if hasattr(formatted_batch, "keys"):
+            if len(formatted_batch.keys()) > 1:
+                raise TypeError(
+                    "Tried to query a column but the custom formatting function returns too many columns. "
+                    f"Only one column was expected but got columns {list(formatted_batch.keys())}."
+                )
+        else:
+            raise TypeError(
+                f"Custom formatting function must return a dict to be able to pick a row, but got {formatted_batch}"
+            )
+        try:
+            return formatted_batch[pa_table.column_names[0]]
+        except Exception as exc:
+            raise TypeError(
+                f"Custom formatting function must return a dict to be able to pick a row, but got {formatted_batch}"
+            ) from exc
+
+    def format_batch(self, pa_table: pa.Table) -> dict:
+        batch = self.python_arrow_extractor().extract_batch(pa_table)
+        batch = self.python_features_decoder.decode_batch(batch)
+        return self.transform(batch)
+
+
+def _check_valid_column_key(key: str, columns: list[str]) -> None:
+    if key not in columns:
+        raise KeyError(f"Column {key} not in the dataset. Current columns in the dataset: {columns}")
+
+
+def _check_valid_index_key(key: Union[int, slice, range, Iterable], size: int) -> None:
+    if isinstance(key, int):
+        if (key < 0 and key + size < 0) or (key >= size):
+            raise IndexError(f"Invalid key: {key} is out of bounds for size {size}")
+        return
+    elif isinstance(key, slice):
+        pass
+    elif isinstance(key, range):
+        if len(key) > 0:
+            _check_valid_index_key(max(key), size=size)
+            _check_valid_index_key(min(key), size=size)
+    elif isinstance(key, Iterable):
+        if len(key) > 0:
+            _check_valid_index_key(int(max(key)), size=size)
+            _check_valid_index_key(int(min(key)), size=size)
+    else:
+        _raise_bad_key_type(key)
+
+
+def key_to_query_type(key: Union[int, slice, range, str, Iterable]) -> str:
+    if isinstance(key, numbers.Integral):
+        return "row"
+    elif isinstance(key, str):
+        return "column"
+    elif isinstance(key, (slice, range, Iterable)):
+        return "batch"
+    _raise_bad_key_type(key)
+
+
+def query_table(
+    table: Table,
+    key: Union[int, slice, range, str, Iterable],
+    indices: Optional[Table] = None,
+) -> pa.Table:
+    """
+    Query a Table to extract the subtable that correspond to the given key.
+
+    Args:
+        table (``datasets.table.Table``): The input Table to query from
+        key (``Union[int, slice, range, str, Iterable]``): The key can be of different types:
+            - an integer i: the subtable containing only the i-th row
+            - a slice [i:j:k]: the subtable containing the rows that correspond to this slice
+            - a range(i, j, k): the subtable containing the rows that correspond to this range
+            - a string c: the subtable containing all the rows but only the column c
+            - an iterable l: the subtable that is the concatenation of all the i-th rows for all i in the iterable
+        indices (Optional ``datasets.table.Table``): If not None, it is used to re-map the given key to the table rows.
+            The indices table must contain one column named "indices" of type uint64.
+            This is used in case of shuffling or rows selection.
+
+
+    Returns:
+        ``pyarrow.Table``: the result of the query on the input table
+    """
+    # Check if key is valid
+    if not isinstance(key, (int, slice, range, str, Iterable)):
+        try:
+            key = operator.index(key)
+        except TypeError:
+            _raise_bad_key_type(key)
+    if isinstance(key, str):
+        _check_valid_column_key(key, table.column_names)
+    else:
+        size = indices.num_rows if indices is not None else table.num_rows
+        _check_valid_index_key(key, size)
+    # Query the main table
+    if indices is None:
+        pa_subtable = _query_table(table, key)
+    else:
+        pa_subtable = _query_table_with_indices_mapping(table, key, indices=indices)
+    return pa_subtable
+
+
+def format_table(
+    table: Table,
+    key: Union[int, slice, range, str, Iterable],
+    formatter: Formatter,
+    format_columns: Optional[list] = None,
+    output_all_columns=False,
+):
+    """
+    Format a Table depending on the key that was used and a Formatter object.
+
+    Args:
+        table (``datasets.table.Table``): The input Table to format
+        key (``Union[int, slice, range, str, Iterable]``): Depending on the key that was used, the formatter formats
+            the table as either a row, a column or a batch.
+        formatter (``datasets.formatting.formatting.Formatter``): Any subclass of a Formatter such as
+            PythonFormatter, NumpyFormatter, etc.
+        format_columns (:obj:`List[str]`, optional): if not None, it defines the columns that will be formatted using the
+            given formatter. Other columns are discarded (unless ``output_all_columns`` is True)
+        output_all_columns (:obj:`bool`, defaults to False). If True, the formatted output is completed using the columns
+            that are not in the ``format_columns`` list. For these columns, the PythonFormatter is used.
+
+
+    Returns:
+        A row, column or batch formatted object defined by the Formatter:
+        - the PythonFormatter returns a dictionary for a row or a batch, and a list for a column.
+        - the NumpyFormatter returns a dictionary for a row or a batch, and a np.array for a column.
+        - the PandasFormatter returns a pd.DataFrame for a row or a batch, and a pd.Series for a column.
+        - the TorchFormatter returns a dictionary for a row or a batch, and a torch.Tensor for a column.
+        - the TFFormatter returns a dictionary for a row or a batch, and a tf.Tensor for a column.
+    """
+    if isinstance(table, Table):
+        pa_table = table.table
+    else:
+        pa_table = table
+    query_type = key_to_query_type(key)
+    python_formatter = PythonFormatter(features=formatter.features)
+    if format_columns is None:
+        return formatter(pa_table, query_type=query_type)
+    elif query_type == "column":
+        if key in format_columns:
+            return formatter(pa_table, query_type)
+        else:
+            return python_formatter(pa_table, query_type=query_type)
+    else:
+        pa_table_to_format = pa_table.drop(col for col in pa_table.column_names if col not in format_columns)
+        formatted_output = formatter(pa_table_to_format, query_type=query_type)
+        if output_all_columns:
+            if isinstance(formatted_output, MutableMapping):
+                pa_table_with_remaining_columns = pa_table.drop(
+                    col for col in pa_table.column_names if col in format_columns
+                )
+                remaining_columns_dict = python_formatter(pa_table_with_remaining_columns, query_type=query_type)
+                formatted_output.update(remaining_columns_dict)
+            else:
+                raise TypeError(
+                    f"Custom formatting function must return a dict to work with output_all_columns=True, but got {formatted_output}"
+                )
+        return formatted_output

src/datasets/formatting/jax_formatter.py

@@ -0,0 +1,174 @@
+# Copyright 2021 The HuggingFace Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Lint as: python3
+import sys
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Optional
+
+import numpy as np
+import pyarrow as pa
+
+from .. import config
+from ..utils.logging import get_logger
+from ..utils.py_utils import map_nested
+from .formatting import TensorFormatter
+
+
+if TYPE_CHECKING:
+    import jax
+    import jaxlib
+
+logger = get_logger()
+
+DEVICE_MAPPING: Optional[dict] = None
+
+
+class JaxFormatter(TensorFormatter[Mapping, "jax.Array", Mapping]):
+    def __init__(self, features=None, device=None, token_per_repo_id=None, **jnp_array_kwargs):
+        super().__init__(features=features, token_per_repo_id=token_per_repo_id)
+        import jax
+        from jaxlib.xla_client import Device
+
+        if isinstance(device, Device):
+            raise ValueError(
+                f"Expected {device} to be a `str` not {type(device)}, as `jaxlib.xla_extension.Device` "
+                "is not serializable neither with `pickle` nor with `dill`. Instead you can surround "
+                "the device with `str()` to get its string identifier that will be internally mapped "
+                "to the actual `jaxlib.xla_extension.Device`."
+            )
+        self.device = device if isinstance(device, str) else str(jax.devices()[0])
+        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
+        # with `pickle` nor with `dill`, so we need to use a global variable instead
+        global DEVICE_MAPPING
+        if DEVICE_MAPPING is None:
+            DEVICE_MAPPING = self._map_devices_to_str()
+        if self.device not in list(DEVICE_MAPPING.keys()):
+            logger.warning(
+                f"Device with string identifier {self.device} not listed among the available "
+                f"devices: {list(DEVICE_MAPPING.keys())}, so falling back to the default "
+                f"device: {str(jax.devices()[0])}."
+            )
+            self.device = str(jax.devices()[0])
+        self.jnp_array_kwargs = jnp_array_kwargs
+
+    @staticmethod
+    def _map_devices_to_str() -> dict[str, "jaxlib.xla_extension.Device"]:
+        import jax
+
+        return {str(device): device for device in jax.devices()}
+
+    def _consolidate(self, column):
+        import jax
+        import jax.numpy as jnp
+
+        if isinstance(column, list) and column:
+            if all(
+                isinstance(x, jax.Array) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column
+            ):
+                return jnp.stack(column, axis=0)
+        return column
+
+    def _tensorize(self, value):
+        import jax
+        import jax.numpy as jnp
+
+        if isinstance(value, (str, bytes, type(None))):
+            return value
+        elif isinstance(value, (np.character, np.ndarray)) and np.issubdtype(value.dtype, np.character):
+            return value.tolist()
+
+        default_dtype = {}
+
+        if isinstance(value, (np.number, np.ndarray)) and np.issubdtype(value.dtype, np.integer):
+            # the default int precision depends on the jax config
+            # see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
+            if jax.config.jax_enable_x64:
+                default_dtype = {"dtype": jnp.int64}
+            else:
+                default_dtype = {"dtype": jnp.int32}
+        elif isinstance(value, (np.number, np.ndarray)) and np.issubdtype(value.dtype, np.floating):
+            default_dtype = {"dtype": jnp.float32}
+
+        if config.PIL_AVAILABLE and "PIL" in sys.modules:
+            import PIL.Image
+
+            if isinstance(value, PIL.Image.Image):
+                value = np.asarray(value)
+        if config.TORCHVISION_AVAILABLE and "torchvision" in sys.modules:
+            try:
+                from torchvision.io import VideoReader
+
+                if isinstance(value, VideoReader):
+                    return value  # TODO(QL): set output to jax arrays ?
+            except ImportError:
+                pass
+        if config.TORCHCODEC_AVAILABLE and "torchcodec" in sys.modules:
+            from torchcodec.decoders import AudioDecoder, VideoDecoder
+
+            if isinstance(value, (VideoDecoder, AudioDecoder)):
+                return value  # TODO(QL): set output to jax arrays ?
+
+        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
+        # with `pickle` nor with `dill`, so we need to use a global variable instead
+        global DEVICE_MAPPING
+        if DEVICE_MAPPING is None:
+            DEVICE_MAPPING = self._map_devices_to_str()
+
+        with jax.default_device(DEVICE_MAPPING[self.device]):
+            # calling jnp.array on a np.ndarray does copy the data
+            # see https://github.com/google/jax/issues/4486
+            return jnp.array(value, **{**default_dtype, **self.jnp_array_kwargs})
+
+    def _recursive_tensorize(self, data_struct):
+        import jax
+
+        # support for torch, tf, jax etc.
+        if config.TORCH_AVAILABLE and "torch" in sys.modules:
+            import torch
+
+            if isinstance(data_struct, torch.Tensor):
+                return self._tensorize(data_struct.detach().cpu().numpy()[()])
+        if hasattr(data_struct, "__array__") and not isinstance(data_struct, jax.Array):
+            data_struct = data_struct.__array__()
+        # support for nested types like struct of list of struct
+        if isinstance(data_struct, np.ndarray):
+            if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
+                return self._consolidate([self.recursive_tensorize(substruct) for substruct in data_struct])
+        elif isinstance(data_struct, (list, tuple)):
+            return self._consolidate([self.recursive_tensorize(substruct) for substruct in data_struct])
+        return self._tensorize(data_struct)
+
+    def recursive_tensorize(self, data_struct: dict):
+        return map_nested(self._recursive_tensorize, data_struct, map_list=False)
+
+    def format_row(self, pa_table: pa.Table) -> Mapping:
+        row = self.numpy_arrow_extractor().extract_row(pa_table)
+        row = self.python_features_decoder.decode_row(row)
+        return self.recursive_tensorize(row)
+
+    def format_column(self, pa_table: pa.Table) -> "jax.Array":
+        column = self.numpy_arrow_extractor().extract_column(pa_table)
+        column = self.python_features_decoder.decode_column(column, pa_table.column_names[0])
+        column = self.recursive_tensorize(column)
+        column = self._consolidate(column)
+        return column
+
+    def format_batch(self, pa_table: pa.Table) -> Mapping:
+        batch = self.numpy_arrow_extractor().extract_batch(pa_table)
+        batch = self.python_features_decoder.decode_batch(batch)
+        batch = self.recursive_tensorize(batch)
+        for column_name in batch:
+            batch[column_name] = self._consolidate(batch[column_name])
+        return batch

src/datasets/formatting/np_formatter.py

@@ -0,0 +1,120 @@
+# Copyright 2020 The HuggingFace Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+from collections.abc import Mapping
+
+import numpy as np
+import pyarrow as pa
+
+from .. import config
+from ..utils.py_utils import map_nested
+from .formatting import TensorFormatter
+
+
+class NumpyFormatter(TensorFormatter[Mapping, np.ndarray, Mapping]):
+    def __init__(self, features=None, token_per_repo_id=None, **np_array_kwargs):
+        super().__init__(features=features, token_per_repo_id=token_per_repo_id)
+        self.np_array_kwargs = np_array_kwargs
+
+    def _consolidate(self, column):
+        if isinstance(column, list):
+            if column and all(
+                isinstance(x, np.ndarray) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column
+            ):
+                return np.stack(column)
+            else:
+                # don't use np.array(column, dtype=object)
+                # since it fails in certain cases
+                # see https://stackoverflow.com/q/51005699
+                out = np.empty(len(column), dtype=object)
+                out[:] = column
+                return out
+        return column
+
+    def _tensorize(self, value):
+        if isinstance(value, (str, bytes, type(None))):
+            return value
+        elif isinstance(value, (np.character, np.ndarray)) and np.issubdtype(value.dtype, np.character):
+            return value
+        elif isinstance(value, np.number):
+            return value
+
+        default_dtype = {}
+
+        if isinstance(value, np.ndarray) and np.issubdtype(value.dtype, np.integer):
+            default_dtype = {"dtype": np.int64}
+        elif isinstance(value, np.ndarray) and np.issubdtype(value.dtype, np.floating):
+            default_dtype = {"dtype": np.float32}
+
+        if config.PIL_AVAILABLE and "PIL" in sys.modules:
+            import PIL.Image
+
+            if isinstance(value, PIL.Image.Image):
+                return np.asarray(value, **self.np_array_kwargs)
+        if config.TORCHVISION_AVAILABLE and "torchvision" in sys.modules:
+            try:
+                from torchvision.io import VideoReader
+
+                if isinstance(value, VideoReader):
+                    return value  # TODO(QL): set output to np arrays ?
+            except ImportError:
+                pass
+        if config.TORCHCODEC_AVAILABLE and "torchcodec" in sys.modules:
+            from torchcodec.decoders import AudioDecoder, VideoDecoder
+
+            if isinstance(value, (VideoDecoder, AudioDecoder)):
+                return value  # TODO(QL): set output to np arrays ?
+
+        return np.asarray(value, **{**default_dtype, **self.np_array_kwargs})
+
+    def _recursive_tensorize(self, data_struct):
+        # support for torch, tf, jax etc.
+        if config.TORCH_AVAILABLE and "torch" in sys.modules:
+            import torch
+
+            if isinstance(data_struct, torch.Tensor):
+                return self._tensorize(data_struct.detach().cpu().numpy()[()])
+        if hasattr(data_struct, "__array__") and not isinstance(data_struct, (np.ndarray, np.character, np.number)):
+            data_struct = data_struct.__array__()
+        # support for nested types like struct of list of struct
+        if isinstance(data_struct, np.ndarray):
+            if data_struct.dtype == object:
+                return self._consolidate([self.recursive_tensorize(substruct) for substruct in data_struct])
+        if isinstance(data_struct, (list, tuple)):
+            return self._consolidate([self.recursive_tensorize(substruct) for substruct in data_struct])
+        return self._tensorize(data_struct)
+
+    def recursive_tensorize(self, data_struct: dict):
+        return map_nested(self._recursive_tensorize, data_struct, map_list=False)
+
+    def format_row(self, pa_table: pa.Table) -> Mapping:
+        row = self.numpy_arrow_extractor().extract_row(pa_table)
+        row = self.python_features_decoder.decode_row(row)
+        return self.recursive_tensorize(row)
+
+    def format_column(self, pa_table: pa.Table) -> np.ndarray:
+        column = self.numpy_arrow_extractor().extract_column(pa_table)
+        column = self.python_features_decoder.decode_column(column, pa_table.column_names[0])
+        column = self.recursive_tensorize(column)
+        column = self._consolidate(column)
+        return column
+
+    def format_batch(self, pa_table: pa.Table) -> Mapping:
+        batch = self.numpy_arrow_extractor().extract_batch(pa_table)
+        batch = self.python_features_decoder.decode_batch(batch)
+        batch = self.recursive_tensorize(batch)
+        for column_name in batch:
+            batch[column_name] = self._consolidate(batch[column_name])
+        return batch

src/datasets/formatting/polars_formatter.py

@@ -0,0 +1,124 @@
+# Copyright 2020 The HuggingFace Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+from functools import partial
+from typing import TYPE_CHECKING, Optional
+
+import pyarrow as pa
+
+from .. import config
+from ..features import Features
+from ..features.features import decode_nested_example
+from ..utils.py_utils import no_op_if_value_is_null
+from .formatting import BaseArrowExtractor, TableFormatter
+
+
+if TYPE_CHECKING:
+    import polars as pl
+
+
+class PolarsArrowExtractor(BaseArrowExtractor["pl.DataFrame", "pl.Series", "pl.DataFrame"]):
+    def extract_row(self, pa_table: pa.Table) -> "pl.DataFrame":
+        if config.POLARS_AVAILABLE:
+            if "polars" not in sys.modules:
+                import polars
+            else:
+                polars = sys.modules["polars"]
+
+            return polars.from_arrow(pa_table.slice(length=1))
+        else:
+            raise ValueError("Polars needs to be installed to be able to return Polars dataframes.")
+
+    def extract_column(self, pa_table: pa.Table) -> "pl.Series":
+        if config.POLARS_AVAILABLE:
+            if "polars" not in sys.modules:
+                import polars
+            else:
+                polars = sys.modules["polars"]
+
+            return polars.from_arrow(pa_table.select([0]))[pa_table.column_names[0]]
+        else:
+            raise ValueError("Polars needs to be installed to be able to return Polars dataframes.")
+
+    def extract_batch(self, pa_table: pa.Table) -> "pl.DataFrame":
+        if config.POLARS_AVAILABLE:
+            if "polars" not in sys.modules:
+                import polars
+            else:
+                polars = sys.modules["polars"]
+
+            return polars.from_arrow(pa_table)
+        else:
+            raise ValueError("Polars needs to be installed to be able to return Polars dataframes.")
+
+
+class PolarsFeaturesDecoder:
+    def __init__(self, features: Optional[Features]):
+        self.features = features
+        import polars as pl  # noqa: F401 - import pl at initialization
+
+    def decode_row(self, row: "pl.DataFrame") -> "pl.DataFrame":
+        decode = (
+            {
+                column_name: no_op_if_value_is_null(partial(decode_nested_example, feature))
+                for column_name, feature in self.features.items()
+                if self.features._column_requires_decoding[column_name]
+            }
+            if self.features
+            else {}
+        )
+        if decode:
+            row[list(decode.keys())] = row.map_rows(decode)
+        return row
+
+    def decode_column(self, column: "pl.Series", column_name: str) -> "pl.Series":
+        decode = (
+            no_op_if_value_is_null(partial(decode_nested_example, self.features[column_name]))
+            if self.features and column_name in self.features and self.features._column_requires_decoding[column_name]
+            else None
+        )
+        if decode:
+            column = column.map_elements(decode)
+        return column
+
+    def decode_batch(self, batch: "pl.DataFrame") -> "pl.DataFrame":
+        return self.decode_row(batch)
+
+
+class PolarsFormatter(TableFormatter["pl.DataFrame", "pl.Series", "pl.DataFrame"]):
+    table_type = "polars dataframe"
+    column_type = "polars series"
+
+    def __init__(self, features=None, **np_array_kwargs):
+        super().__init__(features=features)
+        self.np_array_kwargs = np_array_kwargs
+        self.polars_arrow_extractor = PolarsArrowExtractor
+        self.polars_features_decoder = PolarsFeaturesDecoder(features)
+        import polars as pl  # noqa: F401 - import pl at initialization
+
+    def format_row(self, pa_table: pa.Table) -> "pl.DataFrame":
+        row = self.polars_arrow_extractor().extract_row(pa_table)
+        row = self.polars_features_decoder.decode_row(row)
+        return row
+
+    def format_column(self, pa_table: pa.Table) -> "pl.Series":
+        column = self.polars_arrow_extractor().extract_column(pa_table)
+        column = self.polars_features_decoder.decode_column(column, pa_table.column_names[0])
+        return column
+
+    def format_batch(self, pa_table: pa.Table) -> "pl.DataFrame":
+        row = self.polars_arrow_extractor().extract_batch(pa_table)
+        row = self.polars_features_decoder.decode_batch(row)
+        return row

src/datasets/formatting/tf_formatter.py

@@ -0,0 +1,129 @@
+# Copyright 2020 The HuggingFace Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Lint as: python3
+import sys
+from collections.abc import Mapping
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pyarrow as pa
+
+from .. import config
+from ..utils.py_utils import map_nested
+from .formatting import TensorFormatter
+
+
+if TYPE_CHECKING:
+    import tensorflow as tf
+
+
+class TFFormatter(TensorFormatter[Mapping, "tf.Tensor", Mapping]):
+    def __init__(self, features=None, token_per_repo_id=None, **tf_tensor_kwargs):
+        super().__init__(features=features, token_per_repo_id=token_per_repo_id)
+        self.tf_tensor_kwargs = tf_tensor_kwargs
+        import tensorflow as tf  # noqa: F401 - import tf at initialization
+
+    def _consolidate(self, column):
+        import tensorflow as tf
+
+        if isinstance(column, list) and column:
+            if all(
+                isinstance(x, tf.Tensor) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column
+            ):
+                return tf.stack(column)
+            elif all(
+                isinstance(x, (tf.Tensor, tf.RaggedTensor)) and x.ndim == 1 and x.dtype == column[0].dtype
+                for x in column
+            ):
+                # only rag 1-D tensors, otherwise some dimensions become ragged even though they were consolidated
+                return tf.ragged.stack(column)
+
+        return column
+
+    def _tensorize(self, value):
+        import tensorflow as tf
+
+        if value is None:
+            return value
+
+        default_dtype = {}
+
+        if isinstance(value, (np.number, np.ndarray)) and np.issubdtype(value.dtype, np.integer):
+            default_dtype = {"dtype": tf.int64}
+        elif isinstance(value, (np.number, np.ndarray)) and np.issubdtype(value.dtype, np.floating):
+            default_dtype = {"dtype": tf.float32}
+
+        if config.PIL_AVAILABLE and "PIL" in sys.modules:
+            import PIL.Image
+
+            if isinstance(value, PIL.Image.Image):
+                value = np.asarray(value)
+        if config.TORCHVISION_AVAILABLE and "torchvision" in sys.modules:
+            try:
+                from torchvision.io import VideoReader
+
+                if isinstance(value, VideoReader):
+                    return value  # TODO(QL): set output to tf tensors ?
+            except ImportError:
+                pass
+        if config.TORCHCODEC_AVAILABLE and "torchcodec" in sys.modules:
+            from torchcodec.decoders import AudioDecoder, VideoDecoder
+
+            if isinstance(value, (VideoDecoder, AudioDecoder)):
+                return value  # TODO(QL): set output to jax arrays ?
+
+        return tf.convert_to_tensor(value, **{**default_dtype, **self.tf_tensor_kwargs})
+
+    def _recursive_tensorize(self, data_struct):
+        import tensorflow as tf
+
+        # support for torch, tf, jax etc.
+        if config.TORCH_AVAILABLE and "torch" in sys.modules:
+            import torch
+
+            if isinstance(data_struct, torch.Tensor):
+                return self._tensorize(data_struct.detach().cpu().numpy()[()])
+        if hasattr(data_struct, "__array__") and not isinstance(data_struct, tf.Tensor):
+            data_struct = data_struct.__array__()
+        # support for nested types like struct of list of struct
+        if isinstance(data_struct, np.ndarray):
+            if data_struct.dtype == object:  # tf tensors cannot be instantied from an array of objects
+                return self._consolidate([self.recursive_tensorize(substruct) for substruct in data_struct])
+        elif isinstance(data_struct, (list, tuple)):
+            return self._consolidate([self.recursive_tensorize(substruct) for substruct in data_struct])
+        return self._tensorize(data_struct)
+
+    def recursive_tensorize(self, data_struct: dict):
+        return map_nested(self._recursive_tensorize, data_struct, map_list=False)
+
+    def format_row(self, pa_table: pa.Table) -> Mapping:
+        row = self.numpy_arrow_extractor().extract_row(pa_table)
+        row = self.python_features_decoder.decode_row(row)
+        return self.recursive_tensorize(row)
+
+    def format_column(self, pa_table: pa.Table) -> "tf.Tensor":
+        column = self.numpy_arrow_extractor().extract_column(pa_table)
+        column = self.python_features_decoder.decode_column(column, pa_table.column_names[0])
+        column = self.recursive_tensorize(column)
+        column = self._consolidate(column)
+        return column
+
+    def format_batch(self, pa_table: pa.Table) -> Mapping:
+        batch = self.numpy_arrow_extractor().extract_batch(pa_table)
+        batch = self.python_features_decoder.decode_batch(batch)
+        batch = self.recursive_tensorize(batch)
+        for column_name in batch:
+            batch[column_name] = self._consolidate(batch[column_name])
+        return batch

src/datasets/formatting/torch_formatter.py

@@ -0,0 +1,130 @@
+# Copyright 2020 The HuggingFace Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Lint as: python3
+import sys
+from collections.abc import Mapping
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pyarrow as pa
+
+from .. import config
+from ..utils.py_utils import map_nested
+from .formatting import TensorFormatter
+
+
+if TYPE_CHECKING:
+    import torch
+
+
+class TorchFormatter(TensorFormatter[Mapping, "torch.Tensor", Mapping]):
+    def __init__(self, features=None, token_per_repo_id=None, **torch_tensor_kwargs):
+        super().__init__(features=features, token_per_repo_id=token_per_repo_id)
+        self.torch_tensor_kwargs = torch_tensor_kwargs
+        import torch  # noqa import torch at initialization
+
+    def _consolidate(self, column):
+        import torch
+
+        if isinstance(column, list) and column:
+            if all(
+                isinstance(x, torch.Tensor) and x.shape == column[0].shape and x.dtype == column[0].dtype
+                for x in column
+            ):
+                return torch.stack(column)
+        return column
+
+    def _tensorize(self, value):
+        import torch
+
+        if isinstance(value, (str, bytes, type(None))):
+            return value
+        elif isinstance(value, (np.character, np.ndarray)) and np.issubdtype(value.dtype, np.character):
+            return value.tolist()
+
+        default_dtype = {}
+
+        if isinstance(value, (np.number, np.ndarray)) and np.issubdtype(value.dtype, np.integer):
+            default_dtype = {"dtype": torch.int64}
+
+            # Convert dtype to np.int64 if it's either np.uint16 or np.uint32 to ensure compatibility.
+            # np.uint64 is excluded from this conversion as there is no compatible PyTorch dtype that can handle it without loss.
+            if value.dtype in [np.uint16, np.uint32]:
+                value = value.astype(np.int64)
+
+        elif isinstance(value, (np.number, np.ndarray)) and np.issubdtype(value.dtype, np.floating):
+            default_dtype = {"dtype": torch.float32}
+
+        if config.PIL_AVAILABLE and "PIL" in sys.modules:
+            import PIL.Image
+
+            if isinstance(value, PIL.Image.Image):
+                value = np.asarray(value)
+                if value.ndim == 2:
+                    value = value[:, :, np.newaxis]
+
+                value = value.transpose((2, 0, 1))
+        if config.TORCHVISION_AVAILABLE and "torchvision" in sys.modules:
+            try:
+                from torchvision.io import VideoReader
+
+                if isinstance(value, VideoReader):
+                    return value  # TODO(QL): set output to torch tensors ?
+            except ImportError:
+                pass
+        if config.TORCHCODEC_AVAILABLE and "torchcodec" in sys.modules:
+            from torchcodec.decoders import AudioDecoder, VideoDecoder
+
+            if isinstance(value, (VideoDecoder, AudioDecoder)):
+                return value  # TODO(QL): set output to jax arrays ?
+
+        return torch.tensor(value, **{**default_dtype, **self.torch_tensor_kwargs})
+
+    def _recursive_tensorize(self, data_struct):
+        import torch
+
+        # support for torch, tf, jax etc.
+        if hasattr(data_struct, "__array__") and not isinstance(data_struct, torch.Tensor):
+            data_struct = data_struct.__array__()
+        # support for nested types like struct of list of struct
+        if isinstance(data_struct, np.ndarray):
+            if data_struct.dtype == object:  # torch tensors cannot be instantied from an array of objects
+                return self._consolidate([self.recursive_tensorize(substruct) for substruct in data_struct])
+        elif isinstance(data_struct, (list, tuple)):
+            return self._consolidate([self.recursive_tensorize(substruct) for substruct in data_struct])
+        return self._tensorize(data_struct)
+
+    def recursive_tensorize(self, data_struct: dict):
+        return map_nested(self._recursive_tensorize, data_struct, map_list=False)
+
+    def format_row(self, pa_table: pa.Table) -> Mapping:
+        row = self.numpy_arrow_extractor().extract_row(pa_table)
+        row = self.python_features_decoder.decode_row(row)
+        return self.recursive_tensorize(row)
+
+    def format_column(self, pa_table: pa.Table) -> "torch.Tensor":
+        column = self.numpy_arrow_extractor().extract_column(pa_table)
+        column = self.python_features_decoder.decode_column(column, pa_table.column_names[0])
+        column = self.recursive_tensorize(column)
+        column = self._consolidate(column)
+        return column
+
+    def format_batch(self, pa_table: pa.Table) -> Mapping:
+        batch = self.numpy_arrow_extractor().extract_batch(pa_table)
+        batch = self.python_features_decoder.decode_batch(batch)
+        batch = self.recursive_tensorize(batch)
+        for column_name in batch:
+            batch[column_name] = self._consolidate(batch[column_name])
+        return batch

src/datasets/hub.py

@@ -0,0 +1,135 @@
+from itertools import chain
+from typing import Optional, Union
+
+from huggingface_hub import (
+    CommitInfo,
+    CommitOperationAdd,
+    CommitOperationDelete,
+    DatasetCard,
+    DatasetCardData,
+    HfApi,
+    HfFileSystem,
+)
+
+import datasets.config
+from datasets import __version__
+from datasets.info import DatasetInfosDict
+from datasets.load import load_dataset_builder
+from datasets.utils.metadata import MetadataConfigs
+
+
+def delete_from_hub(
+    repo_id: str,
+    config_name: str,
+    revision: Optional[str] = None,
+    token: Optional[Union[bool, str]] = None,
+) -> CommitInfo:
+    """Delete a dataset configuration from a [data-only dataset](repository_structure) on the Hub.
+
+    Args:
+        repo_id (`str`): ID of the Hub dataset repository, in the following format: `<user>/<dataset_name>` or
+            `<org>/<dataset_name>`.
+        config_name (`str`): Name of the dataset configuration.
+        revision (`str`, *optional*): Branch to delete the configuration from. Defaults to the `"main"` branch.
+        token (`bool` or `str`, *optional*): Authentication token for the Hugging Face Hub.
+
+    Returns:
+        `huggingface_hub.CommitInfo`
+    """
+    operations = []
+    # data_files
+    fs = HfFileSystem(endpoint=datasets.config.HF_ENDPOINT, token=token)
+    builder = load_dataset_builder(repo_id, config_name, revision=revision, token=token)
+    for data_file in chain(*builder.config.data_files.values()):
+        data_file_resolved_path = fs.resolve_path(data_file)
+        if data_file_resolved_path.repo_id == repo_id:
+            operations.append(CommitOperationDelete(path_in_repo=data_file_resolved_path.path_in_repo))
+    # README.md
+    dataset_card = DatasetCard.load(repo_id)
+    # config_names
+    if dataset_card.data.get("config_names", None) and config_name in dataset_card.data["config_names"]:
+        dataset_card.data["config_names"].remove(config_name)
+    # metadata_configs
+    metadata_configs = MetadataConfigs.from_dataset_card_data(dataset_card.data)
+    if metadata_configs:
+        _ = metadata_configs.pop(config_name, None)
+        dataset_card_data = DatasetCardData()
+        metadata_configs.to_dataset_card_data(dataset_card_data)
+        if datasets.config.METADATA_CONFIGS_FIELD in dataset_card_data:
+            dataset_card.data[datasets.config.METADATA_CONFIGS_FIELD] = dataset_card_data[
+                datasets.config.METADATA_CONFIGS_FIELD
+            ]
+        else:
+            _ = dataset_card.data.pop(datasets.config.METADATA_CONFIGS_FIELD, None)
+    # dataset_info
+    dataset_infos: DatasetInfosDict = DatasetInfosDict.from_dataset_card_data(dataset_card.data)
+    if dataset_infos:
+        _ = dataset_infos.pop(config_name, None)
+        dataset_card_data = DatasetCardData()
+        dataset_infos.to_dataset_card_data(dataset_card_data)
+        if "dataset_info" in dataset_card_data:
+            dataset_card.data["dataset_info"] = dataset_card_data["dataset_info"]
+        else:
+            _ = dataset_card.data.pop("dataset_info", None)
+    # Commit
+    operations.append(
+        CommitOperationAdd(path_in_repo=datasets.config.REPOCARD_FILENAME, path_or_fileobj=str(dataset_card).encode())
+    )
+    api = HfApi(
+        endpoint=datasets.config.HF_ENDPOINT,
+        token=token,
+        library_name="datasets",
+        library_version=__version__,
+    )
+    commit_info = api.create_commit(
+        repo_id,
+        operations=operations,
+        commit_message=f"Delete '{config_name}' config",
+        commit_description=f"Delete '{config_name}' config.",
+        token=token,
+        repo_type="dataset",
+        revision=revision,
+        create_pr=True,
+    )
+    print(f"You can find your PR to delete the dataset config at: {commit_info.pr_url}")
+    return commit_info
+
+
+def _delete_files(dataset_id, revision=None, token=None):
+    hf_api = HfApi(
+        endpoint=datasets.config.HF_ENDPOINT,
+        token=token,
+        library_name="datasets",
+        library_version=__version__,
+    )
+    repo_files = hf_api.list_repo_files(
+        dataset_id,
+        repo_type="dataset",
+    )
+    if repo_files:
+        legacy_json_file = []
+        data_files = []
+        for filename in repo_files:
+            if filename in {".gitattributes", "README.md"}:
+                continue
+            elif filename == "dataset_infos.json":
+                legacy_json_file.append(filename)
+            else:
+                data_files.append(filename)
+        if legacy_json_file:
+            hf_api.delete_file(
+                "dataset_infos.json",
+                dataset_id,
+                repo_type="dataset",
+                revision=revision,
+                commit_message="Delete legacy dataset_infos.json",
+            )
+        if data_files:
+            for filename in data_files:
+                hf_api.delete_file(
+                    filename,
+                    dataset_id,
+                    repo_type="dataset",
+                    revision=revision,
+                    commit_message="Delete data file",
+                )

src/datasets/info.py

@@ -0,0 +1,440 @@
+# Copyright 2020 The HuggingFace Datasets Authors and the TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Lint as: python3
+"""DatasetInfo record information we know about a dataset.
+
+This includes things that we know about the dataset statically, i.e.:
+ - description
+ - canonical location
+ - does it have validation and tests splits
+ - size
+ - etc.
+
+This also includes the things that can and should be computed once we've
+processed the dataset as well:
+ - number of examples (in each split)
+ - etc.
+"""
+
+import copy
+import dataclasses
+import json
+import os
+import posixpath
+from dataclasses import dataclass
+from pathlib import Path
+from typing import ClassVar, Optional, Union
+
+import fsspec
+from fsspec.core import url_to_fs
+from huggingface_hub import DatasetCard, DatasetCardData
+
+from . import config
+from .features import Features
+from .splits import SplitDict
+from .utils import Version
+from .utils.logging import get_logger
+from .utils.py_utils import asdict, unique_values
+
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class SupervisedKeysData:
+    input: str = ""
+    output: str = ""
+
+
+@dataclass
+class DownloadChecksumsEntryData:
+    key: str = ""
+    value: str = ""
+
+
+class MissingCachedSizesConfigError(Exception):
+    """The expected cached sizes of the download file are missing."""
+
+
+class NonMatchingCachedSizesError(Exception):
+    """The prepared split doesn't have expected sizes."""
+
+
+@dataclass
+class PostProcessedInfo:
+    features: Optional[Features] = None
+    resources_checksums: Optional[dict] = None
+
+    def __post_init__(self):
+        # Convert back to the correct classes when we reload from dict
+        if self.features is not None and not isinstance(self.features, Features):
+            self.features = Features.from_dict(self.features)
+
+    @classmethod
+    def from_dict(cls, post_processed_info_dict: dict) -> "PostProcessedInfo":
+        field_names = {f.name for f in dataclasses.fields(cls)}
+        return cls(**{k: v for k, v in post_processed_info_dict.items() if k in field_names})
+
+
+@dataclass
+class DatasetInfo:
+    """Information about a dataset.
+
+    `DatasetInfo` documents datasets, including its name, version, and features.
+    See the constructor arguments and properties for a full list.
+
+    Not all fields are known on construction and may be updated later.
+
+    Attributes:
+        description (`str`):
+            A description of the dataset.
+        citation (`str`):
+            A BibTeX citation of the dataset.
+        homepage (`str`):
+            A URL to the official homepage for the dataset.
+        license (`str`):
+            The dataset's license. It can be the name of the license or a paragraph containing the terms of the license.
+        features ([`Features`], *optional*):
+            The features used to specify the dataset's column types.
+        post_processed (`PostProcessedInfo`, *optional*):
+            Deprecated. Information regarding the resources of a possible post-processing of a dataset. For example, it can contain the information of an index.
+        supervised_keys (`SupervisedKeysData`, *optional*):
+            Specifies the input feature and the label for supervised learning if applicable for the dataset (legacy from TFDS).
+        builder_name (`str`, *optional*):
+            The name of the `GeneratorBasedBuilder` subclass used to create the dataset. It is also the snake_case version of the dataset builder class name.
+        config_name (`str`, *optional*):
+            The name of the configuration derived from [`BuilderConfig`].
+        version (`str` or [`Version`], *optional*):
+            The version of the dataset.
+        splits (`dict`, *optional*):
+            The mapping between split name and metadata.
+        download_checksums (`dict`, *optional*):
+            The mapping between the URL to download the dataset's checksums and corresponding metadata.
+        download_size (`int`, *optional*):
+            The size of the files to download to generate the dataset, in bytes.
+        post_processing_size (`int`, *optional*):
+            Deprecated. Size of the dataset in bytes after post-processing, if any.
+        dataset_size (`int`, *optional*):
+            The combined size in bytes of the Arrow tables for all splits.
+        size_in_bytes (`int`, *optional*):
+            The combined size in bytes of all files associated with the dataset (downloaded files + Arrow files).
+        **config_kwargs (additional keyword arguments):
+            Keyword arguments to be passed to the [`BuilderConfig`] and used in the [`DatasetBuilder`].
+    """
+
+    # Set in the dataset builders
+    description: str = dataclasses.field(default_factory=str)
+    citation: str = dataclasses.field(default_factory=str)
+    homepage: str = dataclasses.field(default_factory=str)
+    license: str = dataclasses.field(default_factory=str)
+    features: Optional[Features] = None
+    post_processed: Optional[PostProcessedInfo] = None  # kept for bawkard compat
+    supervised_keys: Optional[SupervisedKeysData] = None
+
+    # Set later by the builder
+    builder_name: Optional[str] = None
+    dataset_name: Optional[str] = None  # for packaged builders, to be different from builder_name
+    config_name: Optional[str] = None
+    version: Optional[Union[str, Version]] = None
+    # Set later by `download_and_prepare`
+    splits: Optional[SplitDict] = None
+    download_checksums: Optional[dict] = None
+    download_size: Optional[int] = None
+    post_processing_size: Optional[int] = None
+    dataset_size: Optional[int] = None
+    size_in_bytes: Optional[int] = None
+
+    _INCLUDED_INFO_IN_YAML: ClassVar[list[str]] = [
+        "config_name",
+        "download_size",
+        "dataset_size",
+        "features",
+        "splits",
+    ]
+
+    def __post_init__(self):
+        # Convert back to the correct classes when we reload from dict
+        if self.features is not None and not isinstance(self.features, Features):
+            self.features = Features.from_dict(self.features)
+        if self.post_processed is not None and not isinstance(self.post_processed, PostProcessedInfo):
+            self.post_processed = PostProcessedInfo.from_dict(self.post_processed)
+        if self.version is not None and not isinstance(self.version, Version):
+            if isinstance(self.version, str):
+                self.version = Version(self.version)
+            else:
+                self.version = Version.from_dict(self.version)
+        if self.splits is not None and not isinstance(self.splits, SplitDict):
+            self.splits = SplitDict.from_split_dict(self.splits)
+        if self.supervised_keys is not None and not isinstance(self.supervised_keys, SupervisedKeysData):
+            if isinstance(self.supervised_keys, (tuple, list)):
+                self.supervised_keys = SupervisedKeysData(*self.supervised_keys)
+            else:
+                self.supervised_keys = SupervisedKeysData(**self.supervised_keys)
+
+    def write_to_directory(self, dataset_info_dir, pretty_print=False, storage_options: Optional[dict] = None):
+        """Write `DatasetInfo` and license (if present) as JSON files to `dataset_info_dir`.
+
+        Args:
+            dataset_info_dir (`str`):
+                Destination directory.
+            pretty_print (`bool`, defaults to `False`):
+                If `True`, the JSON will be pretty-printed with the indent level of 4.
+            storage_options (`dict`, *optional*):
+                Key/value pairs to be passed on to the file-system backend, if any.
+
+                <Added version="2.9.0"/>
+
+        Example:
+
+        ```py
+        >>> from datasets import load_dataset
+        >>> ds = load_dataset("cornell-movie-review-data/rotten_tomatoes", split="validation")
+        >>> ds.info.write_to_directory("/path/to/directory/")
+        ```
+        """
+        fs: fsspec.AbstractFileSystem
+        fs, *_ = url_to_fs(dataset_info_dir, **(storage_options or {}))
+        with fs.open(posixpath.join(dataset_info_dir, config.DATASET_INFO_FILENAME), "wb") as f:
+            self._dump_info(f, pretty_print=pretty_print)
+        if self.license:
+            with fs.open(posixpath.join(dataset_info_dir, config.LICENSE_FILENAME), "wb") as f:
+                self._dump_license(f)
+
+    def _dump_info(self, file, pretty_print=False):
+        """Dump info in `file` file-like object open in bytes mode (to support remote files)"""
+        file.write(json.dumps(asdict(self), indent=4 if pretty_print else None).encode("utf-8"))
+
+    def _dump_license(self, file):
+        """Dump license in `file` file-like object open in bytes mode (to support remote files)"""
+        file.write(self.license.encode("utf-8"))
+
+    @classmethod
+    def from_merge(cls, dataset_infos: list["DatasetInfo"]):
+        dataset_infos = [dset_info.copy() for dset_info in dataset_infos if dset_info is not None]
+
+        if len(dataset_infos) > 0 and all(dataset_infos[0] == dset_info for dset_info in dataset_infos):
+            # if all dataset_infos are equal we don't need to merge. Just return the first.
+            return dataset_infos[0]
+
+        description = "\n\n".join(unique_values(info.description for info in dataset_infos)).strip()
+        citation = "\n\n".join(unique_values(info.citation for info in dataset_infos)).strip()
+        homepage = "\n\n".join(unique_values(info.homepage for info in dataset_infos)).strip()
+        license = "\n\n".join(unique_values(info.license for info in dataset_infos)).strip()
+        features = None
+        supervised_keys = None
+
+        return cls(
+            description=description,
+            citation=citation,
+            homepage=homepage,
+            license=license,
+            features=features,
+            supervised_keys=supervised_keys,
+        )
+
+    @classmethod
+    def from_directory(cls, dataset_info_dir: str, storage_options: Optional[dict] = None) -> "DatasetInfo":
+        """Create [`DatasetInfo`] from the JSON file in `dataset_info_dir`.
+
+        This function updates all the dynamically generated fields (num_examples,
+        hash, time of creation,...) of the [`DatasetInfo`].
+
+        This will overwrite all previous metadata.
+
+        Args:
+            dataset_info_dir (`str`):
+                The directory containing the metadata file. This
+                should be the root directory of a specific dataset version.
+            storage_options (`dict`, *optional*):
+                Key/value pairs to be passed on to the file-system backend, if any.
+
+                <Added version="2.9.0"/>
+
+        Example:
+
+        ```py
+        >>> from datasets import DatasetInfo
+        >>> ds_info = DatasetInfo.from_directory("/path/to/directory/")
+        ```
+        """
+        fs: fsspec.AbstractFileSystem
+        fs, *_ = url_to_fs(dataset_info_dir, **(storage_options or {}))
+        logger.debug(f"Loading Dataset info from {dataset_info_dir}")
+        if not dataset_info_dir:
+            raise ValueError("Calling DatasetInfo.from_directory() with undefined dataset_info_dir.")
+        with fs.open(posixpath.join(dataset_info_dir, config.DATASET_INFO_FILENAME), "r", encoding="utf-8") as f:
+            dataset_info_dict = json.load(f)
+        return cls.from_dict(dataset_info_dict)
+
+    @classmethod
+    def from_dict(cls, dataset_info_dict: dict) -> "DatasetInfo":
+        field_names = {f.name for f in dataclasses.fields(cls)}
+        return cls(**{k: v for k, v in dataset_info_dict.items() if k in field_names})
+
+    def update(self, other_dataset_info: "DatasetInfo", ignore_none=True):
+        self_dict = self.__dict__
+        self_dict.update(
+            **{
+                k: copy.deepcopy(v)
+                for k, v in other_dataset_info.__dict__.items()
+                if (v is not None or not ignore_none)
+            }
+        )
+
+    def copy(self) -> "DatasetInfo":
+        return self.__class__(**{k: copy.deepcopy(v) for k, v in self.__dict__.items()})
+
+    def _to_yaml_dict(self) -> dict:
+        yaml_dict = {}
+        dataset_info_dict = asdict(self)
+        for key in dataset_info_dict:
+            if key in self._INCLUDED_INFO_IN_YAML:
+                value = getattr(self, key)
+                if hasattr(value, "_to_yaml_list"):  # Features, SplitDict
+                    yaml_dict[key] = value._to_yaml_list()
+                elif hasattr(value, "_to_yaml_string"):  # Version
+                    yaml_dict[key] = value._to_yaml_string()
+                else:
+                    yaml_dict[key] = value
+        return yaml_dict
+
+    @classmethod
+    def _from_yaml_dict(cls, yaml_data: dict) -> "DatasetInfo":
+        yaml_data = copy.deepcopy(yaml_data)
+        if yaml_data.get("features") is not None:
+            yaml_data["features"] = Features._from_yaml_list(yaml_data["features"])
+        if yaml_data.get("splits") is not None:
+            yaml_data["splits"] = SplitDict._from_yaml_list(yaml_data["splits"])
+        field_names = {f.name for f in dataclasses.fields(cls)}
+        return cls(**{k: v for k, v in yaml_data.items() if k in field_names})
+
+    def __repr__(self):
+        return (
+            self.__class__.__qualname__
+            + "("
+            + ", ".join(
+                [f"{f.name}={repr(getattr(self, f.name))}" for f in dataclasses.fields(self) if getattr(self, f.name)]
+            )
+            + ")"
+        )
+
+
+class DatasetInfosDict(dict[str, DatasetInfo]):
+    def write_to_directory(self, dataset_infos_dir, overwrite=False, pretty_print=False) -> None:
+        total_dataset_infos = {}
+        dataset_infos_path = os.path.join(dataset_infos_dir, config.DATASETDICT_INFOS_FILENAME)
+        dataset_readme_path = os.path.join(dataset_infos_dir, config.REPOCARD_FILENAME)
+        if not overwrite:
+            total_dataset_infos = self.from_directory(dataset_infos_dir)
+        total_dataset_infos.update(self)
+        if os.path.exists(dataset_infos_path):
+            # for backward compatibility, let's update the JSON file if it exists
+            with open(dataset_infos_path, "w", encoding="utf-8") as f:
+                dataset_infos_dict = {
+                    config_name: asdict(dset_info) for config_name, dset_info in total_dataset_infos.items()
+                }
+                json.dump(dataset_infos_dict, f, indent=4 if pretty_print else None)
+        # Dump the infos in the YAML part of the README.md file
+        if os.path.exists(dataset_readme_path):
+            dataset_card = DatasetCard.load(dataset_readme_path)
+            dataset_card_data = dataset_card.data
+        else:
+            dataset_card = None
+            dataset_card_data = DatasetCardData()
+        if total_dataset_infos:
+            total_dataset_infos.to_dataset_card_data(dataset_card_data)
+            dataset_card = (
+                DatasetCard("---\n" + str(dataset_card_data) + "\n---\n") if dataset_card is None else dataset_card
+            )
+            dataset_card.save(Path(dataset_readme_path))
+
+    @classmethod
+    def from_directory(cls, dataset_infos_dir) -> "DatasetInfosDict":
+        logger.debug(f"Loading Dataset Infos from {dataset_infos_dir}")
+        # Load the info from the YAML part of README.md
+        if os.path.exists(os.path.join(dataset_infos_dir, config.REPOCARD_FILENAME)):
+            dataset_card_data = DatasetCard.load(Path(dataset_infos_dir) / config.REPOCARD_FILENAME).data
+            if "dataset_info" in dataset_card_data:
+                return cls.from_dataset_card_data(dataset_card_data)
+        if os.path.exists(os.path.join(dataset_infos_dir, config.DATASETDICT_INFOS_FILENAME)):
+            # this is just to have backward compatibility with dataset_infos.json files
+            with open(os.path.join(dataset_infos_dir, config.DATASETDICT_INFOS_FILENAME), encoding="utf-8") as f:
+                return cls(
+                    {
+                        config_name: DatasetInfo.from_dict(dataset_info_dict)
+                        for config_name, dataset_info_dict in json.load(f).items()
+                    }
+                )
+        else:
+            return cls()
+
+    @classmethod
+    def from_dataset_card_data(cls, dataset_card_data: DatasetCardData) -> "DatasetInfosDict":
+        if isinstance(dataset_card_data.get("dataset_info"), (list, dict)):
+            if isinstance(dataset_card_data["dataset_info"], list):
+                return cls(
+                    {
+                        dataset_info_yaml_dict.get("config_name", "default"): DatasetInfo._from_yaml_dict(
+                            dataset_info_yaml_dict
+                        )
+                        for dataset_info_yaml_dict in dataset_card_data["dataset_info"]
+                    }
+                )
+            else:
+                dataset_info = DatasetInfo._from_yaml_dict(dataset_card_data["dataset_info"])
+                dataset_info.config_name = dataset_card_data["dataset_info"].get("config_name", "default")
+                return cls({dataset_info.config_name: dataset_info})
+        else:
+            return cls()
+
+    def to_dataset_card_data(self, dataset_card_data: DatasetCardData) -> None:
+        if self:
+            # first get existing metadata info
+            if "dataset_info" in dataset_card_data and isinstance(dataset_card_data["dataset_info"], dict):
+                dataset_metadata_infos = {
+                    dataset_card_data["dataset_info"].get("config_name", "default"): dataset_card_data["dataset_info"]
+                }
+            elif "dataset_info" in dataset_card_data and isinstance(dataset_card_data["dataset_info"], list):
+                dataset_metadata_infos = {
+                    config_metadata["config_name"]: config_metadata
+                    for config_metadata in dataset_card_data["dataset_info"]
+                }
+            else:
+                dataset_metadata_infos = {}
+            # update/rewrite existing metadata info with the one to dump
+            total_dataset_infos = {
+                **dataset_metadata_infos,
+                **{config_name: dset_info._to_yaml_dict() for config_name, dset_info in self.items()},
+            }
+            # the config_name from the dataset_infos_dict takes over the config_name of the DatasetInfo
+            for config_name, dset_info_yaml_dict in total_dataset_infos.items():
+                dset_info_yaml_dict["config_name"] = config_name
+            if len(total_dataset_infos) == 1:
+                # use a struct instead of a list of configurations, since there's only one
+                dataset_card_data["dataset_info"] = next(iter(total_dataset_infos.values()))
+                config_name = dataset_card_data["dataset_info"].pop("config_name", None)
+                if config_name != "default":
+                    # if config_name is not "default" preserve it and put at the first position
+                    dataset_card_data["dataset_info"] = {
+                        "config_name": config_name,
+                        **dataset_card_data["dataset_info"],
+                    }
+            else:
+                dataset_card_data["dataset_info"] = []
+                for config_name, dataset_info_yaml_dict in sorted(total_dataset_infos.items()):
+                    # add the config_name field in first position
+                    dataset_info_yaml_dict.pop("config_name", None)
+                    dataset_info_yaml_dict = {"config_name": config_name, **dataset_info_yaml_dict}
+                    dataset_card_data["dataset_info"].append(dataset_info_yaml_dict)

src/datasets/inspect.py

@@ -0,0 +1,350 @@
+# Copyright 2020 The HuggingFace Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Lint as: python3
+"""List and inspect datasets."""
+
+import os
+from collections.abc import Mapping, Sequence
+from typing import Optional, Union
+
+from .download.download_config import DownloadConfig
+from .download.download_manager import DownloadMode
+from .download.streaming_download_manager import StreamingDownloadManager
+from .info import DatasetInfo
+from .load import (
+    dataset_module_factory,
+    get_dataset_builder_class,
+    load_dataset_builder,
+)
+from .utils.logging import get_logger
+from .utils.version import Version
+
+
+logger = get_logger(__name__)
+
+
+class SplitsNotFoundError(ValueError):
+    pass
+
+
+def get_dataset_infos(
+    path: str,
+    data_files: Optional[Union[dict, list, str]] = None,
+    download_config: Optional[DownloadConfig] = None,
+    download_mode: Optional[Union[DownloadMode, str]] = None,
+    revision: Optional[Union[str, Version]] = None,
+    token: Optional[Union[bool, str]] = None,
+    **config_kwargs,
+):
+    """Get the meta information about a dataset, returned as a dict mapping config name to DatasetInfoDict.
+
+    Args:
+        path (`str`): path to the dataset repository. Can be either:
+
+            - a local path to the dataset directory containing the data files,
+                e.g. `'./dataset/squad'`
+            - a dataset identifier on the Hugging Face Hub (list all available datasets and ids with [`huggingface_hub.list_datasets`]),
+                e.g. `'rajpurkar/squad'`, `'nyu-mll/glue'` or``'openai/webtext'`
+        revision (`Union[str, datasets.Version]`, *optional*):
+            If specified, the dataset module will be loaded from the datasets repository at this version.
+            By default:
+            - it is set to the local version of the lib.
+            - it will also try to load it from the main branch if it's not available at the local version of the lib.
+            Specifying a version that is different from your local version of the lib might cause compatibility issues.
+        download_config ([`DownloadConfig`], *optional*):
+            Specific download configuration parameters.
+        download_mode ([`DownloadMode`] or `str`, defaults to `REUSE_DATASET_IF_EXISTS`):
+            Download/generate mode.
+        data_files (`Union[Dict, List, str]`, *optional*):
+            Defining the data_files of the dataset configuration.
+        token (`str` or `bool`, *optional*):
+            Optional string or boolean to use as Bearer token for remote files on the Datasets Hub.
+            If `True`, or not specified, will get token from `"~/.huggingface"`.
+        **config_kwargs (additional keyword arguments):
+            Optional attributes for builder class which will override the attributes if supplied.
+
+    Example:
+
+    ```py
+    >>> from datasets import get_dataset_infos
+    >>> get_dataset_infos('cornell-movie-review-data/rotten_tomatoes')
+    {'default': DatasetInfo(description="Movie Review Dataset.\nThis is a dataset of containing 5,331 positive and 5,331 negative processed\nsentences from Rotten Tomatoes movie reviews...), ...}
+    ```
+    """
+    config_names = get_dataset_config_names(
+        path=path,
+        revision=revision,
+        download_config=download_config,
+        download_mode=download_mode,
+        data_files=data_files,
+        token=token,
+    )
+    return {
+        config_name: get_dataset_config_info(
+            path=path,
+            config_name=config_name,
+            data_files=data_files,
+            download_config=download_config,
+            download_mode=download_mode,
+            revision=revision,
+            token=token,
+            **config_kwargs,
+        )
+        for config_name in config_names
+    }
+
+
+def get_dataset_config_names(
+    path: str,
+    revision: Optional[Union[str, Version]] = None,
+    download_config: Optional[DownloadConfig] = None,
+    download_mode: Optional[Union[DownloadMode, str]] = None,
+    data_files: Optional[Union[dict, list, str]] = None,
+    **download_kwargs,
+):
+    """Get the list of available config names for a particular dataset.
+
+    Args:
+        path (`str`): path to the dataset repository. Can be either:
+
+            - a local path to the dataset directory containing the data files,
+                e.g. `'./dataset/squad'`
+            - a dataset identifier on the Hugging Face Hub (list all available datasets and ids with [`huggingface_hub.list_datasets`]),
+                e.g. `'rajpurkar/squad'`, `'nyu-mll/glue'` or``'openai/webtext'`
+        revision (`Union[str, datasets.Version]`, *optional*):
+            If specified, the dataset module will be loaded from the datasets repository at this version.
+            By default:
+            - it is set to the local version of the lib.
+            - it will also try to load it from the main branch if it's not available at the local version of the lib.
+            Specifying a version that is different from your local version of the lib might cause compatibility issues.
+        download_config ([`DownloadConfig`], *optional*):
+            Specific download configuration parameters.
+        download_mode ([`DownloadMode`] or `str`, defaults to `REUSE_DATASET_IF_EXISTS`):
+            Download/generate mode.
+        data_files (`Union[Dict, List, str]`, *optional*):
+            Defining the data_files of the dataset configuration.
+        **download_kwargs (additional keyword arguments):
+            Optional attributes for [`DownloadConfig`] which will override the attributes in `download_config` if supplied,
+            for example `token`.
+
+    Example:
+
+    ```py
+    >>> from datasets import get_dataset_config_names
+    >>> get_dataset_config_names("nyu-mll/glue")
+    ['cola',
+     'sst2',
+     'mrpc',
+     'qqp',
+     'stsb',
+     'mnli',
+     'mnli_mismatched',
+     'mnli_matched',
+     'qnli',
+     'rte',
+     'wnli',
+     'ax']
+    ```
+    """
+    dataset_module = dataset_module_factory(
+        path,
+        revision=revision,
+        download_config=download_config,
+        download_mode=download_mode,
+        data_files=data_files,
+        **download_kwargs,
+    )
+    builder_cls = get_dataset_builder_class(dataset_module, dataset_name=os.path.basename(path))
+    return list(builder_cls.builder_configs.keys()) or [
+        dataset_module.builder_kwargs.get("config_name", builder_cls.DEFAULT_CONFIG_NAME or "default")
+    ]
+
+
+def get_dataset_default_config_name(
+    path: str,
+    revision: Optional[Union[str, Version]] = None,
+    download_config: Optional[DownloadConfig] = None,
+    download_mode: Optional[Union[DownloadMode, str]] = None,
+    data_files: Optional[Union[dict, list, str]] = None,
+    **download_kwargs,
+) -> Optional[str]:
+    """Get the default config name for a particular dataset.
+    Can return None only if the dataset has multiple configurations and no default configuration.
+
+    Args:
+        path (`str`): path to the dataset repository. Can be either:
+
+            - a local path to the dataset directory containing the data files,
+                e.g. `'./dataset/squad'`
+            - a dataset identifier on the Hugging Face Hub (list all available datasets and ids with [`huggingface_hub.list_datasets`]),
+                e.g. `'rajpurkar/squad'`, `'nyu-mll/glue'` or``'openai/webtext'`
+        revision (`Union[str, datasets.Version]`, *optional*):
+            If specified, the dataset module will be loaded from the datasets repository at this version.
+            By default:
+            - it is set to the local version of the lib.
+            - it will also try to load it from the main branch if it's not available at the local version of the lib.
+            Specifying a version that is different from your local version of the lib might cause compatibility issues.
+        download_config ([`DownloadConfig`], *optional*):
+            Specific download configuration parameters.
+        download_mode ([`DownloadMode`] or `str`, defaults to `REUSE_DATASET_IF_EXISTS`):
+            Download/generate mode.
+        data_files (`Union[Dict, List, str]`, *optional*):
+            Defining the data_files of the dataset configuration.
+        **download_kwargs (additional keyword arguments):
+            Optional attributes for [`DownloadConfig`] which will override the attributes in `download_config` if supplied,
+            for example `token`.
+
+    Returns:
+        Optional[str]: the default config name if there is one
+
+    Example:
+
+    ```py
+    >>> from datasets import get_dataset_default_config_name
+    >>> get_dataset_default_config_name("openbookqa")
+    'main'
+    ```
+    """
+    dataset_module = dataset_module_factory(
+        path,
+        revision=revision,
+        download_config=download_config,
+        download_mode=download_mode,
+        data_files=data_files,
+        **download_kwargs,
+    )
+    builder_cls = get_dataset_builder_class(dataset_module, dataset_name=os.path.basename(path))
+    builder_configs = list(builder_cls.builder_configs.keys())
+    if builder_configs:
+        default_config_name = builder_configs[0] if len(builder_configs) == 1 else None
+    else:
+        default_config_name = "default"
+    return builder_cls.DEFAULT_CONFIG_NAME or default_config_name
+
+
+def get_dataset_config_info(
+    path: str,
+    config_name: Optional[str] = None,
+    data_files: Optional[Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]] = None,
+    download_config: Optional[DownloadConfig] = None,
+    download_mode: Optional[Union[DownloadMode, str]] = None,
+    revision: Optional[Union[str, Version]] = None,
+    token: Optional[Union[bool, str]] = None,
+    **config_kwargs,
+) -> DatasetInfo:
+    """Get the meta information (DatasetInfo) about a dataset for a particular config
+
+    Args:
+        path (`str`): path to the dataset repository. Can be either:
+
+            - a local path to the dataset directory containing the data files,
+                e.g. `'./dataset/squad'`
+            - a dataset identifier on the Hugging Face Hub (list all available datasets and ids with [`huggingface_hub.list_datasets`]),
+                e.g. `'rajpurkar/squad'`, `'nyu-mll/glue'` or``'openai/webtext'`
+        config_name (:obj:`str`, optional): Defining the name of the dataset configuration.
+        data_files (:obj:`str` or :obj:`Sequence` or :obj:`Mapping`, optional): Path(s) to source data file(s).
+        download_config (:class:`~download.DownloadConfig`, optional): Specific download configuration parameters.
+        download_mode (:class:`DownloadMode` or :obj:`str`, default ``REUSE_DATASET_IF_EXISTS``): Download/generate mode.
+        revision (:class:`~utils.Version` or :obj:`str`, optional): Version of the dataset to load.
+            As datasets have their own git repository on the Datasets Hub, the default version "main" corresponds to their "main" branch.
+            You can specify a different version than the default "main" by using a commit SHA or a git tag of the dataset repository.
+        token (``str`` or :obj:`bool`, optional): Optional string or boolean to use as Bearer token for remote files on the Datasets Hub.
+            If True, or not specified, will get token from `"~/.huggingface"`.
+        **config_kwargs (additional keyword arguments): optional attributes for builder class which will override the attributes if supplied.
+
+    """
+    builder = load_dataset_builder(
+        path,
+        name=config_name,
+        data_files=data_files,
+        download_config=download_config,
+        download_mode=download_mode,
+        revision=revision,
+        token=token,
+        **config_kwargs,
+    )
+    info = builder.info
+    if info.splits is None:
+        download_config = download_config.copy() if download_config else DownloadConfig()
+        if token is not None:
+            download_config.token = token
+        try:
+            info.splits = {
+                split_generator.name: {"name": split_generator.name, "dataset_name": path}
+                for split_generator in builder._split_generators(
+                    StreamingDownloadManager(base_path=builder.base_path, download_config=download_config)
+                )
+            }
+        except Exception as err:
+            raise SplitsNotFoundError("The split names could not be parsed from the dataset config.") from err
+    return info
+
+
+def get_dataset_split_names(
+    path: str,
+    config_name: Optional[str] = None,
+    data_files: Optional[Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]] = None,
+    download_config: Optional[DownloadConfig] = None,
+    download_mode: Optional[Union[DownloadMode, str]] = None,
+    revision: Optional[Union[str, Version]] = None,
+    token: Optional[Union[bool, str]] = None,
+    **config_kwargs,
+):
+    """Get the list of available splits for a particular config and dataset.
+
+    Args:
+        path (`str`): path to the dataset repository. Can be either:
+
+            - a local path to the dataset directory containing the data files,
+                e.g. `'./dataset/squad'`
+            - a dataset identifier on the Hugging Face Hub (list all available datasets and ids with [`huggingface_hub.list_datasets`]),
+                e.g. `'rajpurkar/squad'`, `'nyu-mll/glue'` or``'openai/webtext'`
+        config_name (`str`, *optional*):
+            Defining the name of the dataset configuration.
+        data_files (`str` or `Sequence` or `Mapping`, *optional*):
+            Path(s) to source data file(s).
+        download_config ([`DownloadConfig`], *optional*):
+            Specific download configuration parameters.
+        download_mode ([`DownloadMode`] or `str`, defaults to `REUSE_DATASET_IF_EXISTS`):
+            Download/generate mode.
+        revision ([`Version`] or `str`, *optional*):
+            Version of the dataset to load.
+            As datasets have their own git repository on the Datasets Hub, the default version "main" corresponds to their "main" branch.
+            You can specify a different version than the default "main" by using a commit SHA or a git tag of the dataset repository.
+        token (`str` or `bool`, *optional*):
+            Optional string or boolean to use as Bearer token for remote files on the Datasets Hub.
+            If `True`, or not specified, will get token from `"~/.huggingface"`.
+        **config_kwargs (additional keyword arguments):
+            Optional attributes for builder class which will override the attributes if supplied.
+
+    Example:
+
+    ```py
+    >>> from datasets import get_dataset_split_names
+    >>> get_dataset_split_names('cornell-movie-review-data/rotten_tomatoes')
+    ['train', 'validation', 'test']
+    ```
+    """
+    info = get_dataset_config_info(
+        path,
+        config_name=config_name,
+        data_files=data_files,
+        download_config=download_config,
+        download_mode=download_mode,
+        revision=revision,
+        token=token,
+        **config_kwargs,
+    )
+    return list(info.splits.keys())

src/datasets/io/abc.py

@@ -0,0 +1,53 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Union
+
+from .. import Dataset, DatasetDict, Features, IterableDataset, IterableDatasetDict, NamedSplit
+from ..utils.typing import NestedDataStructureLike, PathLike
+
+
+class AbstractDatasetReader(ABC):
+    def __init__(
+        self,
+        path_or_paths: Optional[NestedDataStructureLike[PathLike]] = None,
+        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,
+    ):
+        self.path_or_paths = path_or_paths
+        self.split = split if split or isinstance(path_or_paths, dict) else "train"
+        self.features = features
+        self.cache_dir = cache_dir
+        self.keep_in_memory = keep_in_memory
+        self.streaming = streaming
+        self.num_proc = num_proc
+        self.kwargs = kwargs
+
+    @abstractmethod
+    def read(self) -> Union[Dataset, DatasetDict, IterableDataset, IterableDatasetDict]:
+        pass
+
+
+class AbstractDatasetInputStream(ABC):
+    def __init__(
+        self,
+        features: Optional[Features] = None,
+        cache_dir: str = None,
+        keep_in_memory: bool = False,
+        streaming: bool = False,
+        num_proc: Optional[int] = None,
+        **kwargs,
+    ):
+        self.features = features
+        self.cache_dir = cache_dir
+        self.keep_in_memory = keep_in_memory
+        self.streaming = streaming
+        self.num_proc = num_proc
+        self.kwargs = kwargs
+
+    @abstractmethod
+    def read(self) -> Union[Dataset, IterableDataset]:
+        pass

src/datasets/io/csv.py

@@ -0,0 +1,143 @@
+import multiprocessing
+import os
+from typing import BinaryIO, Optional, Union
+
+import fsspec
+
+from .. import Dataset, Features, NamedSplit, config
+from ..formatting import query_table
+from ..packaged_modules.csv.csv import Csv
+from ..utils import tqdm as hf_tqdm
+from ..utils.typing import NestedDataStructureLike, PathLike
+from .abc import AbstractDatasetReader
+
+
+class CsvDatasetReader(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}
+        self.builder = Csv(
+            cache_dir=cache_dir,
+            data_files=path_or_paths,
+            features=features,
+            **kwargs,
+        )
+
+    def read(self):
+        # Build iterable dataset
+        if self.streaming:
+            dataset = self.builder.as_streaming_dataset(split=self.split)
+        # Build regular (map-style) dataset
+        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, in_memory=self.keep_in_memory)
+        return dataset
+
+
+class CsvDatasetWriter:
+    def __init__(
+        self,
+        dataset: Dataset,
+        path_or_buf: Union[PathLike, BinaryIO],
+        batch_size: Optional[int] = None,
+        num_proc: Optional[int] = None,
+        storage_options: Optional[dict] = None,
+        **to_csv_kwargs,
+    ):
+        if num_proc is not None and num_proc <= 0:
+            raise ValueError(f"num_proc {num_proc} must be an integer > 0.")
+
+        self.dataset = dataset
+        self.path_or_buf = path_or_buf
+        self.batch_size = batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
+        self.num_proc = num_proc
+        self.encoding = "utf-8"
+        self.storage_options = storage_options or {}
+        self.to_csv_kwargs = to_csv_kwargs
+
+    def write(self) -> int:
+        _ = self.to_csv_kwargs.pop("path_or_buf", None)
+        header = self.to_csv_kwargs.pop("header", True)
+        index = self.to_csv_kwargs.pop("index", False)
+
+        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, header=header, index=index, **self.to_csv_kwargs)
+        else:
+            written = self._write(file_obj=self.path_or_buf, header=header, index=index, **self.to_csv_kwargs)
+        return written
+
+    def _batch_csv(self, args):
+        offset, header, index, to_csv_kwargs = args
+
+        batch = query_table(
+            table=self.dataset.data,
+            key=slice(offset, offset + self.batch_size),
+            indices=self.dataset._indices,
+        )
+        csv_str = batch.to_pandas().to_csv(
+            path_or_buf=None, header=header if (offset == 0) else False, index=index, **to_csv_kwargs
+        )
+        return csv_str.encode(self.encoding)
+
+    def _write(self, file_obj: BinaryIO, header, index, **to_csv_kwargs) -> int:
+        """Writes the pyarrow table as CSV to a binary file handle.
+
+        Caller is responsible for opening and closing the handle.
+        """
+        written = 0
+
+        if self.num_proc is None or self.num_proc == 1:
+            for offset in hf_tqdm(
+                range(0, len(self.dataset), self.batch_size),
+                unit="ba",
+                desc="Creating CSV from Arrow format",
+            ):
+                csv_str = self._batch_csv((offset, header, index, to_csv_kwargs))
+                written += file_obj.write(csv_str)
+
+        else:
+            num_rows, batch_size = len(self.dataset), self.batch_size
+            with multiprocessing.Pool(self.num_proc) as pool:
+                for csv_str in hf_tqdm(
+                    pool.imap(
+                        self._batch_csv,
+                        [(offset, header, index, to_csv_kwargs) for offset in range(0, num_rows, batch_size)],
+                    ),
+                    total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size,
+                    unit="ba",
+                    desc="Creating CSV from Arrow format",
+                ):
+                    written += file_obj.write(csv_str)
+
+        return written

src/datasets/io/generator.py

@@ -0,0 +1,62 @@
+from typing import Callable, Optional
+
+from .. import Features, NamedSplit, Split
+from ..packaged_modules.generator.generator import Generator
+from .abc import AbstractDatasetInputStream
+
+
+class GeneratorDatasetInputStream(AbstractDatasetInputStream):
+    def __init__(
+        self,
+        generator: Callable,
+        features: Optional[Features] = None,
+        cache_dir: str = None,
+        keep_in_memory: bool = False,
+        streaming: bool = False,
+        gen_kwargs: Optional[dict] = None,
+        num_proc: Optional[int] = None,
+        split: NamedSplit = Split.TRAIN,
+        fingerprint: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__(
+            features=features,
+            cache_dir=cache_dir,
+            keep_in_memory=keep_in_memory,
+            streaming=streaming,
+            num_proc=num_proc,
+            **kwargs,
+        )
+        self.builder = Generator(
+            cache_dir=cache_dir,
+            features=features,
+            generator=generator,
+            gen_kwargs=gen_kwargs,
+            split=split,
+            config_id="default-fingerprint=" + fingerprint if fingerprint else None,
+            **kwargs,
+        )
+        self.fingerprint = fingerprint
+
+    def read(self):
+        # Build iterable dataset
+        if self.streaming:
+            dataset = self.builder.as_streaming_dataset(split=self.builder.config.split)
+        # Build regular (map-style) dataset
+        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.builder.config.split, in_memory=self.keep_in_memory)
+            if self.fingerprint:
+                dataset._fingerprint = self.fingerprint
+        return dataset

src/datasets/io/json.py

@@ -0,0 +1,178 @@
+import multiprocessing
+import os
+from functools import partial
+from typing import BinaryIO, Optional, Union
+
+import fsspec
+
+from .. import Dataset, Features, NamedSplit, config
+from ..formatting import query_table
+from ..packaged_modules.json.json import Json
+from ..utils import tqdm as hf_tqdm
+from ..utils.json import get_json_field_paths_from_feature, json_decode_field
+from ..utils.typing import NestedDataStructureLike, PathLike
+from .abc import AbstractDatasetReader
+
+
+class JsonDatasetReader(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,
+        field: Optional[str] = None,
+        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,
+        )
+        self.field = field
+        path_or_paths = path_or_paths if isinstance(path_or_paths, dict) else {self.split: path_or_paths}
+        self.builder = Json(
+            cache_dir=cache_dir,
+            data_files=path_or_paths,
+            features=features,
+            field=field,
+            **kwargs,
+        )
+
+    def read(self):
+        # Build iterable dataset
+        if self.streaming:
+            dataset = self.builder.as_streaming_dataset(split=self.split)
+        # Build regular (map-style) dataset
+        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, in_memory=self.keep_in_memory)
+        return dataset
+
+
+class JsonDatasetWriter:
+    def __init__(
+        self,
+        dataset: Dataset,
+        path_or_buf: Union[PathLike, BinaryIO],
+        batch_size: Optional[int] = None,
+        num_proc: Optional[int] = None,
+        storage_options: Optional[dict] = None,
+        **to_json_kwargs,
+    ):
+        if num_proc is not None and num_proc <= 0:
+            raise ValueError(f"num_proc {num_proc} must be an integer > 0.")
+
+        self.dataset = dataset
+        self.path_or_buf = path_or_buf
+        self.batch_size = batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
+        self.num_proc = num_proc
+        self.encoding = "utf-8"
+        self.storage_options = storage_options or {}
+        self.to_json_kwargs = to_json_kwargs
+
+    def write(self) -> int:
+        _ = self.to_json_kwargs.pop("path_or_buf", None)
+        orient = self.to_json_kwargs.pop("orient", "records")
+        lines = self.to_json_kwargs.pop("lines", True if orient == "records" else False)
+        if "index" not in self.to_json_kwargs and orient in ["split", "table"]:
+            self.to_json_kwargs["index"] = False
+
+        # Determine the default compression value based on self.path_or_buf type
+        default_compression = "infer" if isinstance(self.path_or_buf, (str, bytes, os.PathLike)) else None
+        compression = self.to_json_kwargs.pop("compression", default_compression)
+
+        if compression not in [None, "infer", "gzip", "bz2", "xz"]:
+            raise NotImplementedError(f"`datasets` currently does not support {compression} compression")
+
+        if not lines and self.batch_size < self.dataset.num_rows:
+            raise NotImplementedError(
+                "Output JSON will not be formatted correctly when lines = False and batch_size < number of rows in the dataset. Use pandas.DataFrame.to_json() instead."
+            )
+
+        if isinstance(self.path_or_buf, (str, bytes, os.PathLike)):
+            with fsspec.open(
+                self.path_or_buf, "wb", compression=compression, **(self.storage_options or {})
+            ) as buffer:
+                written = self._write(file_obj=buffer, orient=orient, lines=lines, **self.to_json_kwargs)
+        else:
+            if compression:
+                raise NotImplementedError(
+                    f"The compression parameter is not supported when writing to a buffer, but compression={compression}"
+                    " was passed. Please provide a local path instead."
+                )
+            written = self._write(file_obj=self.path_or_buf, orient=orient, lines=lines, **self.to_json_kwargs)
+        return written
+
+    def _batch_json(self, args):
+        offset, orient, lines, to_json_kwargs = args
+
+        batch = query_table(
+            table=self.dataset.data,
+            key=slice(offset, offset + self.batch_size),
+            indices=self.dataset._indices,
+        )
+        batch = batch.to_pandas()
+        for json_field_path in get_json_field_paths_from_feature(self.dataset.features):
+            col, *json_field_subpath = json_field_path
+            batch[col] = batch[col].apply(partial(json_decode_field, json_field_path=json_field_subpath))
+        json_str = batch.to_json(path_or_buf=None, orient=orient, lines=lines, **to_json_kwargs)
+        if not json_str.endswith("\n"):
+            json_str += "\n"
+        return json_str.encode(self.encoding)
+
+    def _write(
+        self,
+        file_obj: BinaryIO,
+        orient,
+        lines,
+        **to_json_kwargs,
+    ) -> int:
+        """Writes the pyarrow table as JSON lines to a binary file handle.
+
+        Caller is responsible for opening and closing the handle.
+        """
+        written = 0
+
+        if self.num_proc is None or self.num_proc == 1:
+            for offset in hf_tqdm(
+                range(0, len(self.dataset), self.batch_size),
+                unit="ba",
+                desc="Creating json from Arrow format",
+            ):
+                json_str = self._batch_json((offset, orient, lines, to_json_kwargs))
+                written += file_obj.write(json_str)
+        else:
+            num_rows, batch_size = len(self.dataset), self.batch_size
+            with multiprocessing.Pool(self.num_proc) as pool:
+                for json_str in hf_tqdm(
+                    pool.imap(
+                        self._batch_json,
+                        [(offset, orient, lines, to_json_kwargs) for offset in range(0, num_rows, batch_size)],
+                    ),
+                    total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size,
+                    unit="ba",
+                    desc="Creating json from Arrow format",
+                ):
+                    written += file_obj.write(json_str)
+
+        return written

src/datasets/io/parquet.py

@@ -0,0 +1,159 @@
+import json
+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_data_size, get_writer_batch_size_from_features
+from ..features.features import require_storage_embed
+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):
+        # Build iterable dataset
+        if self.streaming:
+            dataset = self.builder.as_streaming_dataset(split=self.split)
+        # Build regular (map-style) dataset
+        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, 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,
+        use_content_defined_chunking: Union[bool, dict] = True,
+        write_page_index: bool = True,
+        **parquet_writer_kwargs,
+    ):
+        self.dataset = dataset
+        self.path_or_buf = path_or_buf
+        self.batch_size = (
+            batch_size
+            or get_writer_batch_size_from_features(dataset.features)
+            or get_writer_batch_size_from_data_size(len(dataset), dataset._estimate_nbytes())
+        )
+        self.storage_options = storage_options or {}
+        self.parquet_writer_kwargs = parquet_writer_kwargs
+        if use_content_defined_chunking is True:
+            use_content_defined_chunking = config.DEFAULT_CDC_OPTIONS
+        self.use_content_defined_chunking = use_content_defined_chunking
+        self.write_page_index = write_page_index
+
+    def write(self) -> int:
+        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=self.batch_size,
+                    **self.parquet_writer_kwargs,
+                )
+        else:
+            written = self._write(
+                file_obj=self.path_or_buf,
+                batch_size=self.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,
+            use_content_defined_chunking=self.use_content_defined_chunking,
+            write_page_index=self.write_page_index,
+            compression={
+                col: "none" if require_storage_embed(feature) else "snappy"
+                for col, feature in self.dataset.features.items()
+            },
+            use_dictionary=[
+                col for col, feature in self.dataset.features.items() if not require_storage_embed(feature)
+            ],
+            column_encoding={
+                col: "PLAIN" for col, feature in self.dataset.features.items() if require_storage_embed(feature)
+            },
+            **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
+
+        # TODO(kszucs): we may want to persist multiple parameters
+        if self.use_content_defined_chunking is not False:
+            writer.add_key_value_metadata({"content_defined_chunking": json.dumps(self.use_content_defined_chunking)})
+
+        writer.close()
+        return written

src/datasets/io/spark.py

@@ -0,0 +1,57 @@
+from typing import Optional
+
+import pyspark
+
+from .. import Features, NamedSplit
+from ..download import DownloadMode
+from ..packaged_modules.spark.spark import Spark
+from .abc import AbstractDatasetReader
+
+
+class SparkDatasetReader(AbstractDatasetReader):
+    """A dataset reader that reads from a Spark DataFrame.
+
+    When caching, cache materialization is parallelized over Spark; an NFS that is accessible to the driver must be
+    provided. Streaming is not currently supported.
+    """
+
+    def __init__(
+        self,
+        df: pyspark.sql.DataFrame,
+        split: Optional[NamedSplit] = None,
+        features: Optional[Features] = None,
+        streaming: bool = True,
+        cache_dir: str = None,
+        keep_in_memory: bool = False,
+        working_dir: str = None,
+        load_from_cache_file: bool = True,
+        file_format: str = "arrow",
+        **kwargs,
+    ):
+        super().__init__(
+            split=split,
+            features=features,
+            cache_dir=cache_dir,
+            keep_in_memory=keep_in_memory,
+            streaming=streaming,
+            **kwargs,
+        )
+        self._load_from_cache_file = load_from_cache_file
+        self._file_format = file_format
+        self.builder = Spark(
+            df=df,
+            features=features,
+            cache_dir=cache_dir,
+            working_dir=working_dir,
+            **kwargs,
+        )
+
+    def read(self):
+        if self.streaming:
+            return self.builder.as_streaming_dataset(split=self.split)
+        download_mode = None if self._load_from_cache_file else DownloadMode.FORCE_REDOWNLOAD
+        self.builder.download_and_prepare(
+            download_mode=download_mode,
+            file_format=self._file_format,
+        )
+        return self.builder.as_dataset(split=self.split)

src/datasets/io/sql.py

@@ -0,0 +1,122 @@
+import multiprocessing
+from typing import TYPE_CHECKING, Optional, Union
+
+from .. import Dataset, Features, config
+from ..formatting import query_table
+from ..packaged_modules.sql.sql import Sql
+from ..utils import tqdm as hf_tqdm
+from .abc import AbstractDatasetInputStream
+
+
+if TYPE_CHECKING:
+    import sqlite3
+
+    import sqlalchemy
+
+
+class SqlDatasetReader(AbstractDatasetInputStream):
+    def __init__(
+        self,
+        sql: Union[str, "sqlalchemy.sql.Selectable"],
+        con: Union[str, "sqlalchemy.engine.Connection", "sqlalchemy.engine.Engine", "sqlite3.Connection"],
+        features: Optional[Features] = None,
+        cache_dir: str = None,
+        keep_in_memory: bool = False,
+        **kwargs,
+    ):
+        super().__init__(features=features, cache_dir=cache_dir, keep_in_memory=keep_in_memory, **kwargs)
+        self.builder = Sql(
+            cache_dir=cache_dir,
+            features=features,
+            sql=sql,
+            con=con,
+            **kwargs,
+        )
+
+    def read(self):
+        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,
+        )
+
+        # Build dataset for splits
+        dataset = self.builder.as_dataset(split="train", in_memory=self.keep_in_memory)
+        return dataset
+
+
+class SqlDatasetWriter:
+    def __init__(
+        self,
+        dataset: Dataset,
+        name: str,
+        con: Union[str, "sqlalchemy.engine.Connection", "sqlalchemy.engine.Engine", "sqlite3.Connection"],
+        batch_size: Optional[int] = None,
+        num_proc: Optional[int] = None,
+        **to_sql_kwargs,
+    ):
+        if num_proc is not None and num_proc <= 0:
+            raise ValueError(f"num_proc {num_proc} must be an integer > 0.")
+
+        self.dataset = dataset
+        self.name = name
+        self.con = con
+        self.batch_size = batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
+        self.num_proc = num_proc
+        self.to_sql_kwargs = to_sql_kwargs
+
+    def write(self) -> int:
+        _ = self.to_sql_kwargs.pop("sql", None)
+        _ = self.to_sql_kwargs.pop("con", None)
+        index = self.to_sql_kwargs.pop("index", False)
+
+        written = self._write(index=index, **self.to_sql_kwargs)
+        return written
+
+    def _batch_sql(self, args):
+        offset, index, to_sql_kwargs = args
+        to_sql_kwargs = {**to_sql_kwargs, "if_exists": "append"} if offset > 0 else to_sql_kwargs
+        batch = query_table(
+            table=self.dataset.data,
+            key=slice(offset, offset + self.batch_size),
+            indices=self.dataset._indices,
+        )
+        df = batch.to_pandas()
+        num_rows = df.to_sql(self.name, self.con, index=index, **to_sql_kwargs)
+        return num_rows or len(df)
+
+    def _write(self, index, **to_sql_kwargs) -> int:
+        """Writes the pyarrow table as SQL to a database.
+
+        Caller is responsible for opening and closing the SQL connection.
+        """
+        written = 0
+
+        if self.num_proc is None or self.num_proc == 1:
+            for offset in hf_tqdm(
+                range(0, len(self.dataset), self.batch_size),
+                unit="ba",
+                desc="Creating SQL from Arrow format",
+            ):
+                written += self._batch_sql((offset, index, to_sql_kwargs))
+        else:
+            num_rows, batch_size = len(self.dataset), self.batch_size
+            with multiprocessing.Pool(self.num_proc) as pool:
+                for num_rows in hf_tqdm(
+                    pool.imap(
+                        self._batch_sql,
+                        [(offset, index, to_sql_kwargs) for offset in range(0, num_rows, batch_size)],
+                    ),
+                    total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size,
+                    unit="ba",
+                    desc="Creating SQL from Arrow format",
+                ):
+                    written += num_rows
+
+        return written

src/datasets/io/text.py

@@ -0,0 +1,58 @@
+from typing import Optional
+
+from .. import Features, NamedSplit
+from ..packaged_modules.text.text import Text
+from ..utils.typing import NestedDataStructureLike, PathLike
+from .abc import AbstractDatasetReader
+
+
+class TextDatasetReader(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}
+        self.builder = Text(
+            cache_dir=cache_dir,
+            data_files=path_or_paths,
+            features=features,
+            **kwargs,
+        )
+
+    def read(self):
+        # Build iterable dataset
+        if self.streaming:
+            dataset = self.builder.as_streaming_dataset(split=self.split)
+        # Build regular (map-style) dataset
+        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, in_memory=self.keep_in_memory)
+        return dataset

src/datasets/naming.py

@@ -0,0 +1,84 @@
+# Copyright 2020 The HuggingFace Datasets Authors and the TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Lint as: python3
+"""Utilities for file names."""
+
+import itertools
+import os
+import re
+
+
+_uppercase_uppercase_re = re.compile(r"([A-Z]+)([A-Z][a-z])")
+_lowercase_uppercase_re = re.compile(r"([a-z\d])([A-Z])")
+
+_single_underscore_re = re.compile(r"(?<!_)_(?!_)")
+_multiple_underscores_re = re.compile(r"(_{2,})")
+
+_split_re = r"^\w+(\.\w+)*$"
+
+INVALID_WINDOWS_CHARACTERS_IN_PATH = r"<>:/\|?*"
+
+
+def camelcase_to_snakecase(name):
+    """Convert camel-case string to snake-case."""
+    name = _uppercase_uppercase_re.sub(r"\1_\2", name)
+    name = _lowercase_uppercase_re.sub(r"\1_\2", name)
+    return name.lower()
+
+
+def snakecase_to_camelcase(name):
+    """Convert snake-case string to camel-case string."""
+    name = _single_underscore_re.split(name)
+    name = [_multiple_underscores_re.split(n) for n in name]
+    return "".join(n.capitalize() for n in itertools.chain.from_iterable(name) if n != "")
+
+
+def filename_prefix_for_name(name):
+    if os.path.basename(name) != name:
+        raise ValueError(f"Should be a dataset name, not a path: {name}")
+    return camelcase_to_snakecase(name)
+
+
+def filename_prefix_for_split(name, split):
+    if os.path.basename(name) != name:
+        raise ValueError(f"Should be a dataset name, not a path: {name}")
+    if not re.match(_split_re, split):
+        raise ValueError(f"Split name should match '{_split_re}'' but got '{split}'.")
+    return f"{filename_prefix_for_name(name)}-{split}"
+
+
+def filepattern_for_dataset_split(dataset_name, split, data_dir, filetype_suffix=None):
+    prefix = filename_prefix_for_split(dataset_name, split)
+    if filetype_suffix:
+        prefix += f".{filetype_suffix}"
+    filepath = os.path.join(data_dir, prefix)
+    return f"{filepath}*"
+
+
+def filenames_for_dataset_split(path, dataset_name, split, filetype_suffix=None, shard_lengths=None):
+    prefix = filename_prefix_for_split(dataset_name, split)
+    prefix = os.path.join(path, prefix)
+
+    if shard_lengths and len(shard_lengths) > 1:
+        num_shards = len(shard_lengths)
+        filenames = [f"{prefix}-{shard_id:05d}-of-{num_shards:05d}" for shard_id in range(num_shards)]
+        if filetype_suffix:
+            filenames = [filename + f".{filetype_suffix}" for filename in filenames]
+        return filenames
+    else:
+        filename = prefix
+        if filetype_suffix:
+            filename += f".{filetype_suffix}"
+        return [filename]

src/datasets/packaged_modules/__init__.py

@@ -0,0 +1,148 @@
+import inspect
+import re
+from typing import Dict, List, Tuple
+
+from huggingface_hub.utils import insecure_hashlib
+
+from .arrow import arrow
+from .audiofolder import audiofolder
+from .cache import cache
+from .conll import conll
+from .csv import csv
+from .eval import eval
+from .hdf5 import hdf5
+from .iceberg import iceberg
+from .imagefolder import imagefolder
+from .json import json
+from .lance import lance
+from .meshfolder import meshfolder
+from .niftifolder import niftifolder
+from .pandas import pandas
+from .parquet import parquet
+from .pdffolder import pdffolder
+from .sql import sql
+from .text import text
+from .tsfile import tsfile
+from .videofolder import videofolder
+from .webdataset import webdataset
+from .xml import xml
+
+
+def _hash_python_lines(lines: list[str]) -> str:
+    filtered_lines = []
+    for line in lines:
+        line = re.sub(r"#.*", "", line)  # remove comments
+        if line:
+            filtered_lines.append(line)
+    full_str = "\n".join(filtered_lines)
+
+    # Make a hash from all this code
+    full_bytes = full_str.encode("utf-8")
+    return insecure_hashlib.sha256(full_bytes).hexdigest()
+
+
+# get importable module names and hash for caching
+_PACKAGED_DATASETS_MODULES = {
+    "csv": (csv.__name__, _hash_python_lines(inspect.getsource(csv).splitlines())),
+    "json": (json.__name__, _hash_python_lines(inspect.getsource(json).splitlines())),
+    "pandas": (pandas.__name__, _hash_python_lines(inspect.getsource(pandas).splitlines())),
+    "parquet": (parquet.__name__, _hash_python_lines(inspect.getsource(parquet).splitlines())),
+    "arrow": (arrow.__name__, _hash_python_lines(inspect.getsource(arrow).splitlines())),
+    "text": (text.__name__, _hash_python_lines(inspect.getsource(text).splitlines())),
+    "conll": (conll.__name__, _hash_python_lines(inspect.getsource(conll).splitlines())),
+    "imagefolder": (imagefolder.__name__, _hash_python_lines(inspect.getsource(imagefolder).splitlines())),
+    "audiofolder": (audiofolder.__name__, _hash_python_lines(inspect.getsource(audiofolder).splitlines())),
+    "videofolder": (videofolder.__name__, _hash_python_lines(inspect.getsource(videofolder).splitlines())),
+    "meshfolder": (meshfolder.__name__, _hash_python_lines(inspect.getsource(meshfolder).splitlines())),
+    "pdffolder": (pdffolder.__name__, _hash_python_lines(inspect.getsource(pdffolder).splitlines())),
+    "niftifolder": (niftifolder.__name__, _hash_python_lines(inspect.getsource(niftifolder).splitlines())),
+    "webdataset": (webdataset.__name__, _hash_python_lines(inspect.getsource(webdataset).splitlines())),
+    "xml": (xml.__name__, _hash_python_lines(inspect.getsource(xml).splitlines())),
+    "hdf5": (hdf5.__name__, _hash_python_lines(inspect.getsource(hdf5).splitlines())),
+    "eval": (eval.__name__, _hash_python_lines(inspect.getsource(eval).splitlines())),
+    "lance": (lance.__name__, _hash_python_lines(inspect.getsource(lance).splitlines())),
+    "tsfile": (tsfile.__name__, _hash_python_lines(inspect.getsource(tsfile).splitlines())),
+    "iceberg": (iceberg.__name__, _hash_python_lines(inspect.getsource(iceberg).splitlines())),
+}
+
+# get importable module names and hash for caching
+_PACKAGED_DATASETS_MODULES_2_15_HASHES = {
+    "csv": "eea64c71ca8b46dd3f537ed218fc9bf495d5707789152eb2764f5c78fa66d59d",
+    "json": "8bb11242116d547c741b2e8a1f18598ffdd40a1d4f2a2872c7a28b697434bc96",
+    "pandas": "3ac4ffc4563c796122ef66899b9485a3f1a977553e2d2a8a318c72b8cc6f2202",
+    "parquet": "ca31c69184d9832faed373922c2acccec0b13a0bb5bbbe19371385c3ff26f1d1",
+    "arrow": "74f69db2c14c2860059d39860b1f400a03d11bf7fb5a8258ca38c501c878c137",
+    "text": "c4a140d10f020282918b5dd1b8a49f0104729c6177f60a6b49ec2a365ec69f34",
+    "imagefolder": "7b7ce5247a942be131d49ad4f3de5866083399a0f250901bd8dc202f8c5f7ce5",
+    "audiofolder": "d3c1655c66c8f72e4efb5c79e952975fa6e2ce538473a6890241ddbddee9071c",
+}
+
+# Used to infer the module to use based on the data files extensions
+_EXTENSION_TO_MODULE: dict[str, tuple[str, dict]] = {
+    ".csv": ("csv", {}),
+    ".tsv": ("csv", {"sep": "\t"}),
+    ".json": ("json", {}),
+    ".jsonl": ("json", {}),
+    # ndjson is no longer maintained (see: https://github.com/ndjson/ndjson-spec/issues/35#issuecomment-1285673417)
+    ".ndjson": ("json", {}),
+    ".parquet": ("parquet", {}),
+    ".geoparquet": ("parquet", {}),
+    ".gpq": ("parquet", {}),
+    ".arrow": ("arrow", {}),
+    ".txt": ("text", {}),
+    ".conll": ("conll", {}),
+    ".conllu": ("conll", {"comment_prefix": "#"}),
+    ".tar": ("webdataset", {}),
+    ".xml": ("xml", {}),
+    ".hdf5": ("hdf5", {}),
+    ".h5": ("hdf5", {}),
+    ".eval": ("eval", {}),
+    ".lance": ("lance", {}),
+    ".tsfile": ("tsfile", {}),
+}
+_EXTENSION_TO_MODULE.update({ext: ("imagefolder", {}) for ext in imagefolder.ImageFolder.EXTENSIONS})
+_EXTENSION_TO_MODULE.update({ext.upper(): ("imagefolder", {}) for ext in imagefolder.ImageFolder.EXTENSIONS})
+_EXTENSION_TO_MODULE.update({ext: ("audiofolder", {}) for ext in audiofolder.AudioFolder.EXTENSIONS})
+_EXTENSION_TO_MODULE.update({ext.upper(): ("audiofolder", {}) for ext in audiofolder.AudioFolder.EXTENSIONS})
+_EXTENSION_TO_MODULE.update({ext: ("videofolder", {}) for ext in videofolder.VideoFolder.EXTENSIONS})
+_EXTENSION_TO_MODULE.update({ext.upper(): ("videofolder", {}) for ext in videofolder.VideoFolder.EXTENSIONS})
+_EXTENSION_TO_MODULE.update({ext: ("meshfolder", {}) for ext in meshfolder.MeshFolder.EXTENSIONS})
+_EXTENSION_TO_MODULE.update({ext.upper(): ("meshfolder", {}) for ext in meshfolder.MeshFolder.EXTENSIONS})
+_EXTENSION_TO_MODULE.update({ext: ("pdffolder", {}) for ext in pdffolder.PdfFolder.EXTENSIONS})
+_EXTENSION_TO_MODULE.update({ext.upper(): ("pdffolder", {}) for ext in pdffolder.PdfFolder.EXTENSIONS})
+_EXTENSION_TO_MODULE.update({ext: ("niftifolder", {}) for ext in niftifolder.NiftiFolder.EXTENSIONS})
+_EXTENSION_TO_MODULE.update({ext.upper(): ("niftifolder", {}) for ext in niftifolder.NiftiFolder.EXTENSIONS})
+
+# Used to filter data files based on extensions given a module name
+_MODULE_TO_EXTENSIONS: dict[str, list[str]] = {}
+for _ext, (_module, _) in _EXTENSION_TO_MODULE.items():
+    _MODULE_TO_EXTENSIONS.setdefault(_module, []).append(_ext)
+
+for _module in _MODULE_TO_EXTENSIONS:
+    _MODULE_TO_EXTENSIONS[_module].append(".zip")
+
+# Used to filter data files based on file names
+_MODULE_TO_METADATA_FILE_NAMES: Dict[str, List[str]] = {}
+for _module in _MODULE_TO_EXTENSIONS:
+    _MODULE_TO_METADATA_FILE_NAMES[_module] = []
+_MODULE_TO_METADATA_FILE_NAMES["imagefolder"] = imagefolder.ImageFolder.METADATA_FILENAMES
+_MODULE_TO_METADATA_FILE_NAMES["audiofolder"] = imagefolder.ImageFolder.METADATA_FILENAMES
+_MODULE_TO_METADATA_FILE_NAMES["videofolder"] = imagefolder.ImageFolder.METADATA_FILENAMES
+_MODULE_TO_METADATA_FILE_NAMES["meshfolder"] = meshfolder.MeshFolder.METADATA_FILENAMES
+_MODULE_TO_METADATA_FILE_NAMES["pdffolder"] = imagefolder.ImageFolder.METADATA_FILENAMES
+_MODULE_TO_METADATA_FILE_NAMES["niftifolder"] = imagefolder.ImageFolder.METADATA_FILENAMES
+_MODULE_TO_METADATA_FILE_NAMES["lance"] = lance.Lance.METADATA_FILE_NAMES
+
+_MODULE_TO_METADATA_EXTENSIONS: Dict[str, List[str]] = {}
+for _module in _MODULE_TO_EXTENSIONS:
+    _MODULE_TO_METADATA_EXTENSIONS[_module] = []
+_MODULE_TO_METADATA_EXTENSIONS["lance"] = lance.Lance.METADATA_EXTENSIONS
+
+# Total
+
+_ALL_EXTENSIONS = list(_EXTENSION_TO_MODULE.keys()) + [".zip"]
+_ALL_METADATA_EXTENSIONS = sorted({_ext for _exts in _MODULE_TO_METADATA_EXTENSIONS.values() for _ext in _exts})
+_ALL_ALLOWED_EXTENSIONS = _ALL_EXTENSIONS + _ALL_METADATA_EXTENSIONS
+_ALL_METADATA_FILENAMES = sorted(
+    {file_name for file_names in _MODULE_TO_METADATA_FILE_NAMES.values() for file_name in file_names}
+)

src/datasets/packaged_modules/arrow/arrow.py

@@ -0,0 +1,77 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import pyarrow as pa
+
+import datasets
+from datasets.builder import Key
+from datasets.table import table_cast
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+@dataclass
+class ArrowConfig(datasets.BuilderConfig):
+    """BuilderConfig for Arrow."""
+
+    features: Optional[datasets.Features] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+
+
+class Arrow(datasets.ArrowBasedBuilder):
+    BUILDER_CONFIG_CLASS = ArrowConfig
+
+    def _info(self):
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        """We handle string, list and dicts in datafiles"""
+        if not self.config.data_files:
+            raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
+        data_files = dl_manager.download(self.config.data_files)
+        splits = []
+        for split_name, files in data_files.items():
+            # Infer features if they are stored in the arrow schema
+            if self.info.features is None:
+                for file in files:
+                    with open(file, "rb") as f:
+                        try:
+                            reader = pa.ipc.open_stream(f)
+                        except (OSError, pa.lib.ArrowInvalid):
+                            reader = pa.ipc.open_file(f)
+                    self.info.features = datasets.Features.from_arrow_schema(reader.schema)
+                    break
+            splits.append(datasets.SplitGenerator(name=split_name, gen_kwargs={"files": files}))
+        return splits
+
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.info.features is not None:
+            # more expensive cast to support nested features with keys in a different order
+            # allows str <-> int/float or str to Audio for example
+            pa_table = table_cast(pa_table, self.info.features.arrow_schema)
+        return pa_table
+
+    def _generate_shards(self, files):
+        yield from files
+
+    def _generate_tables(self, files):
+        for file_idx, file in enumerate(files):
+            with open(file, "rb") as f:
+                try:
+                    try:
+                        batches = pa.ipc.open_stream(f)
+                    except (OSError, pa.lib.ArrowInvalid):
+                        reader = pa.ipc.open_file(f)
+                        batches = (reader.get_batch(i) for i in range(reader.num_record_batches))
+                    for batch_idx, record_batch in enumerate(batches):
+                        pa_table = pa.Table.from_batches([record_batch])
+                        # Uncomment for debugging (will print the Arrow table size and elements)
+                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
+                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
+                        yield Key(file_idx, batch_idx), self._cast_table(pa_table)
+                except ValueError as e:
+                    logger.error(f"Failed to read file '{file}' with error {type(e)}: {e}")
+                    raise

src/datasets/packaged_modules/audiofolder/audiofolder.py

@@ -0,0 +1,86 @@
+import datasets
+
+from ..folder_based_builder import folder_based_builder
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+class AudioFolderConfig(folder_based_builder.FolderBasedBuilderConfig):
+    """Builder Config for AudioFolder."""
+
+    drop_labels: bool = None
+    drop_metadata: bool = None
+
+    def __post_init__(self):
+        super().__post_init__()
+
+
+class AudioFolder(folder_based_builder.FolderBasedBuilder):
+    BASE_FEATURE = datasets.Audio
+    BASE_COLUMN_NAME = "audio"
+    BUILDER_CONFIG_CLASS = AudioFolderConfig
+    EXTENSIONS: list[str]  # definition at the bottom of the script
+
+
+# Obtained with:
+# ```
+# import soundfile as sf
+#
+# AUDIO_EXTENSIONS = [f".{format.lower()}" for format in sf.available_formats().keys()]
+#
+# # .opus decoding is supported if libsndfile >= 1.0.31:
+# AUDIO_EXTENSIONS.extend([".opus"])
+# ```
+# We intentionally did not run this code on launch because:
+# (1) Soundfile was an optional dependency, so importing it in global namespace is not allowed
+# (2) To ensure the list of supported extensions is deterministic
+# (3) We use TorchCodec now anyways instead of Soundfile
+AUDIO_EXTENSIONS = [
+    ".aiff",
+    ".au",
+    ".avr",
+    ".caf",
+    ".flac",
+    ".htk",
+    ".svx",
+    ".mat4",
+    ".mat5",
+    ".mpc2k",
+    ".ogg",
+    ".paf",
+    ".pvf",
+    ".raw",
+    ".rf64",
+    ".sd2",
+    ".sds",
+    ".ircam",
+    ".voc",
+    ".w64",
+    ".wav",
+    ".nist",
+    ".wavex",
+    ".wve",
+    ".xi",
+    ".mp3",
+    ".opus",
+    ".3gp",
+    ".3g2",
+    ".avi",
+    ".asf",
+    ".flv",
+    ".mp4",
+    ".mov",
+    ".m4v",
+    ".mkv",
+    ".mpg",
+    ".webm",
+    ".f4v",
+    ".wmv",
+    ".wma",
+    ".ogg",
+    ".ogm",
+    ".mxf",
+    ".nut",
+]
+AudioFolder.EXTENSIONS = AUDIO_EXTENSIONS

src/datasets/packaged_modules/cache/cache.py

@@ -0,0 +1,195 @@
+import glob
+import json
+import os
+import shutil
+import time
+from pathlib import Path
+from typing import Optional, Union
+
+import pyarrow as pa
+
+import datasets
+import datasets.config
+import datasets.data_files
+from datasets.builder import Key
+from datasets.naming import camelcase_to_snakecase, filenames_for_dataset_split
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+def _get_modification_time(cached_directory_path):
+    return (Path(cached_directory_path)).stat().st_mtime
+
+
+def _find_hash_in_cache(
+    dataset_name: str,
+    config_name: Optional[str],
+    cache_dir: Optional[str],
+    config_kwargs: dict,
+    custom_features: Optional[datasets.Features],
+) -> tuple[str, str, str]:
+    if config_name or config_kwargs or custom_features:
+        config_id = datasets.BuilderConfig(config_name or "default").create_config_id(
+            config_kwargs=config_kwargs, custom_features=custom_features
+        )
+    else:
+        config_id = None
+    cache_dir = os.path.expanduser(str(cache_dir or datasets.config.HF_DATASETS_CACHE))
+    namespace_and_dataset_name = dataset_name.split("/")
+    namespace_and_dataset_name[-1] = camelcase_to_snakecase(namespace_and_dataset_name[-1])
+    cached_relative_path = "___".join(namespace_and_dataset_name)
+    cached_datasets_directory_path_root = os.path.join(cache_dir, cached_relative_path)
+    cached_directory_paths = [
+        cached_directory_path
+        for cached_directory_path in glob.glob(
+            os.path.join(cached_datasets_directory_path_root, config_id or "*", "*", "*")
+        )
+        if os.path.isdir(cached_directory_path)
+        and (
+            config_kwargs
+            or custom_features
+            or json.loads(Path(cached_directory_path, "dataset_info.json").read_text(encoding="utf-8"))["config_name"]
+            == Path(cached_directory_path).parts[-3]  # no extra params => config_id == config_name
+        )
+    ]
+    if not cached_directory_paths:
+        cached_directory_paths = [
+            cached_directory_path
+            for cached_directory_path in glob.glob(os.path.join(cached_datasets_directory_path_root, "*", "*", "*"))
+            if os.path.isdir(cached_directory_path)
+        ]
+        available_configs = sorted(
+            {Path(cached_directory_path).parts[-3] for cached_directory_path in cached_directory_paths}
+        )
+        raise ValueError(
+            f"Couldn't find cache for {dataset_name}"
+            + (f" for config '{config_id}'" if config_id else "")
+            + (f"\nAvailable configs in the cache: {available_configs}" if available_configs else "")
+        )
+    # get most recent
+    cached_directory_path = Path(sorted(cached_directory_paths, key=_get_modification_time)[-1])
+    version, hash = cached_directory_path.parts[-2:]
+    other_configs = [
+        Path(_cached_directory_path).parts[-3]
+        for _cached_directory_path in glob.glob(os.path.join(cached_datasets_directory_path_root, "*", version, hash))
+        if os.path.isdir(_cached_directory_path)
+        and (
+            config_kwargs
+            or custom_features
+            or json.loads(Path(_cached_directory_path, "dataset_info.json").read_text(encoding="utf-8"))["config_name"]
+            == Path(_cached_directory_path).parts[-3]  # no extra params => config_id == config_name
+        )
+    ]
+    if not config_id and len(other_configs) > 1:
+        raise ValueError(
+            f"There are multiple '{dataset_name}' configurations in the cache: {', '.join(other_configs)}"
+            f"\nPlease specify which configuration to reload from the cache, e.g."
+            f"\n\tload_dataset('{dataset_name}', '{other_configs[0]}')"
+        )
+    config_name = cached_directory_path.parts[-3]
+    warning_msg = (
+        f"Found the latest cached dataset configuration '{config_name}' at {cached_directory_path} "
+        f"(last modified on {time.ctime(_get_modification_time(cached_directory_path))})."
+    )
+    logger.warning(warning_msg)
+    return config_name, version, hash
+
+
+class Cache(datasets.ArrowBasedBuilder):
+    def __init__(
+        self,
+        cache_dir: Optional[str] = None,
+        dataset_name: Optional[str] = None,
+        config_name: Optional[str] = None,
+        version: Optional[str] = "0.0.0",
+        hash: Optional[str] = None,
+        base_path: Optional[str] = None,
+        info: Optional[datasets.DatasetInfo] = None,
+        features: Optional[datasets.Features] = None,
+        token: Optional[Union[bool, str]] = None,
+        repo_id: Optional[str] = None,
+        data_files: Optional[Union[str, list, dict, datasets.data_files.DataFilesDict]] = None,
+        data_dir: Optional[str] = None,
+        storage_options: Optional[dict] = None,
+        writer_batch_size: Optional[int] = None,
+        **config_kwargs,
+    ):
+        if repo_id is None and dataset_name is None:
+            raise ValueError("repo_id or dataset_name is required for the Cache dataset builder")
+        if data_files is not None:
+            config_kwargs["data_files"] = data_files
+        if data_dir is not None:
+            config_kwargs["data_dir"] = data_dir
+        if hash == "auto" and version == "auto":
+            config_name, version, hash = _find_hash_in_cache(
+                dataset_name=repo_id or dataset_name,
+                config_name=config_name,
+                cache_dir=cache_dir,
+                config_kwargs=config_kwargs,
+                custom_features=features,
+            )
+        elif hash == "auto" or version == "auto":
+            raise NotImplementedError("Pass both hash='auto' and version='auto' instead")
+        super().__init__(
+            cache_dir=cache_dir,
+            dataset_name=dataset_name,
+            config_name=config_name,
+            version=version,
+            hash=hash,
+            base_path=base_path,
+            info=info,
+            token=token,
+            repo_id=repo_id,
+            storage_options=storage_options,
+            writer_batch_size=writer_batch_size,
+        )
+
+    def _info(self) -> datasets.DatasetInfo:
+        return datasets.DatasetInfo()
+
+    def download_and_prepare(self, output_dir: Optional[str] = None, *args, **kwargs):
+        if not os.path.exists(self.cache_dir):
+            raise ValueError(f"Cache directory for {self.dataset_name} doesn't exist at {self.cache_dir}")
+        if output_dir is not None and output_dir != self.cache_dir:
+            shutil.copytree(self.cache_dir, output_dir)
+
+    def _split_generators(self, dl_manager):
+        # used to stream from cache
+        if isinstance(self.info.splits, datasets.SplitDict):
+            split_infos: list[datasets.SplitInfo] = list(self.info.splits.values())
+        else:
+            raise ValueError(f"Missing splits info for {self.dataset_name} in cache directory {self.cache_dir}")
+        return [
+            datasets.SplitGenerator(
+                name=split_info.name,
+                gen_kwargs={
+                    "files": filenames_for_dataset_split(
+                        self.cache_dir,
+                        dataset_name=self.dataset_name,
+                        split=split_info.name,
+                        filetype_suffix="arrow",
+                        shard_lengths=split_info.shard_lengths,
+                    )
+                },
+            )
+            for split_info in split_infos
+        ]
+
+    def _generate_shards(self, files):
+        yield from files
+
+    def _generate_tables(self, files):
+        # used to stream from cache
+        for file_idx, file in enumerate(files):
+            with open(file, "rb") as f:
+                try:
+                    for batch_idx, record_batch in enumerate(pa.ipc.open_stream(f)):
+                        pa_table = pa.Table.from_batches([record_batch])
+                        # Uncomment for debugging (will print the Arrow table size and elements)
+                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
+                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
+                        yield Key(file_idx, batch_idx), pa_table
+                except ValueError as e:
+                    logger.error(f"Failed to read file '{file}' with error {type(e)}: {e}")
+                    raise

src/datasets/packaged_modules/conll/conll.py

@@ -0,0 +1,155 @@
+"""CoNLL format dataset builder.
+
+Reads CoNLL-style files where each line carries one token plus its tag columns,
+columns are whitespace-separated, and empty lines mark sentence boundaries.
+Each example is one sentence, with each configured column produced as a list
+aligned with the tokens list.
+
+Supports CoNLL-2000 (chunking), CoNLL-2003 (NER), CoNLL-U (Universal
+Dependencies), and any custom column schema by overriding ``column_names``.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+import pyarrow as pa
+
+import datasets
+from datasets.builder import Key
+from datasets.features.features import require_storage_cast
+from datasets.table import table_cast
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+DEFAULT_COLUMN_NAMES = ["tokens"]
+
+
+@dataclass
+class ConllConfig(datasets.BuilderConfig):
+    """BuilderConfig for CoNLL-style files.
+
+    Args:
+        features: (`Features`, *optional*):
+            Cast the data to `features`.
+        column_names: (`list[str]`, defaults to `["tokens"]`):
+            Names for each whitespace-separated column. The first name is the
+            token column; subsequent names cover the tag columns. Common
+            schemas:
+
+            - CoNLL-2003 NER: `["tokens", "pos_tags", "chunk_tags", "ner_tags"]`
+            - CoNLL-2000 chunking: `["tokens", "pos_tags", "chunk_tags"]`
+            - CoNLL-U: `["id", "form", "lemma", "upos", "xpos", "feats",
+              "head", "deprel", "deps", "misc"]`
+        delimiter: (`str`, *optional*):
+            Column delimiter inside each line. If `None` (default), any
+            whitespace is used, matching the standard CoNLL convention.
+        encoding: (`str`, defaults to `"utf-8"`):
+            Encoding to decode the file.
+        encoding_errors: (`str`, *optional*):
+            Argument to define what to do in case of encoding error. Same as
+            the `errors` argument in `open()`.
+        skip_docstart: (`bool`, defaults to `True`):
+            Skip CoNLL-2003 `-DOCSTART-` document-boundary marker lines.
+        comment_prefix: (`str`, *optional*):
+            If set, lines starting with this prefix are skipped. CoNLL-U
+            comments start with `#`.
+    """
+
+    features: Optional[datasets.Features] = None
+    column_names: list[str] = field(default_factory=lambda: list(DEFAULT_COLUMN_NAMES))
+    delimiter: Optional[str] = None
+    encoding: str = "utf-8"
+    encoding_errors: Optional[str] = None
+    skip_docstart: bool = True
+    comment_prefix: Optional[str] = None
+
+
+class Conll(datasets.ArrowBasedBuilder):
+    BUILDER_CONFIG_CLASS = ConllConfig
+
+    def _info(self):
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        """The `data_files` kwarg in load_dataset() can be a str, List[str],
+        Dict[str,str], or Dict[str,List[str]].
+
+        If str or List[str], then the dataset returns only the 'train' split.
+        If dict, then keys should be from the `datasets.Split` enum.
+        """
+        if not self.config.data_files:
+            raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
+        dl_manager.download_config.extract_on_the_fly = True
+        base_data_files = dl_manager.download(self.config.data_files)
+        extracted_data_files = dl_manager.extract(base_data_files)
+        splits = []
+        for split_name, files in extracted_data_files.items():
+            files_iterables = [dl_manager.iter_files(file) for file in files]
+            splits.append(
+                datasets.SplitGenerator(
+                    name=split_name,
+                    gen_kwargs={
+                        "files_iterables": files_iterables,
+                        "base_files": base_data_files[split_name],
+                    },
+                )
+            )
+        return splits
+
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.config.features is not None:
+            schema = self.config.features.arrow_schema
+            if all(not require_storage_cast(feature) for feature in self.config.features.values()):
+                pa_table = pa_table.cast(schema)
+            else:
+                pa_table = table_cast(pa_table, schema)
+        return pa_table
+
+    def _generate_shards(self, base_files, files_iterables):
+        yield from base_files
+
+    def _generate_tables(self, base_files, files_iterables):
+        column_names = list(self.config.column_names)
+        if not column_names:
+            raise ValueError("ConllConfig.column_names must be a non-empty list")
+        num_cols = len(column_names)
+        delimiter = self.config.delimiter
+        skip_docstart = self.config.skip_docstart
+        comment_prefix = self.config.comment_prefix
+
+        for shard_idx, files_iterable in enumerate(files_iterables):
+            for file in files_iterable:
+                sentences: list[list[list[str]]] = []
+                current: list[list[str]] = [[] for _ in range(num_cols)]
+                with open(file, encoding=self.config.encoding, errors=self.config.encoding_errors) as f:
+                    for raw_line in f:
+                        # Strip the trailing newline only — preserve embedded whitespace
+                        line = raw_line.rstrip("\r\n")
+                        if not line.strip():
+                            # Sentence boundary
+                            if current[0]:
+                                sentences.append(current)
+                                current = [[] for _ in range(num_cols)]
+                            continue
+                        if skip_docstart and line.startswith("-DOCSTART-"):
+                            continue
+                        if comment_prefix is not None and line.startswith(comment_prefix):
+                            continue
+                        parts = line.split(delimiter) if delimiter is not None else line.split()
+                        # Pad or truncate to num_cols so column alignment is preserved
+                        if len(parts) < num_cols:
+                            parts = parts + [""] * (num_cols - len(parts))
+                        elif len(parts) > num_cols:
+                            parts = parts[:num_cols]
+                        for i, val in enumerate(parts):
+                            current[i].append(val)
+                    # Tail sentence with no trailing blank line
+                    if current[0]:
+                        sentences.append(current)
+
+                if sentences:
+                    arrays = [pa.array([sentence[col_idx] for sentence in sentences]) for col_idx in range(num_cols)]
+                    pa_table = pa.Table.from_arrays(arrays, names=column_names)
+                    yield Key(shard_idx, 0), self._cast_table(pa_table)

src/datasets/packaged_modules/csv/csv.py

@@ -0,0 +1,206 @@
+from dataclasses import dataclass
+from typing import Any, Callable, Optional, Union
+
+import pandas as pd
+import pyarrow as pa
+
+import datasets
+import datasets.config
+from datasets.builder import Key
+from datasets.features.features import require_storage_cast
+from datasets.table import table_cast
+from datasets.utils.py_utils import Literal
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+_PANDAS_READ_CSV_NO_DEFAULT_PARAMETERS = ["names", "prefix"]
+_PANDAS_READ_CSV_DEPRECATED_PARAMETERS = ["warn_bad_lines", "error_bad_lines", "mangle_dupe_cols"]
+_PANDAS_READ_CSV_NEW_1_3_0_PARAMETERS = ["encoding_errors", "on_bad_lines"]
+_PANDAS_READ_CSV_NEW_2_0_0_PARAMETERS = ["date_format"]
+_PANDAS_READ_CSV_DEPRECATED_2_2_0_PARAMETERS = ["verbose"]
+
+
+@dataclass
+class CsvConfig(datasets.BuilderConfig):
+    """BuilderConfig for CSV."""
+
+    sep: str = ","
+    delimiter: Optional[str] = None
+    header: Optional[Union[int, list[int], str]] = "infer"
+    names: Optional[list[str]] = None
+    column_names: Optional[list[str]] = None
+    index_col: Optional[Union[int, str, list[int], list[str]]] = None
+    usecols: Optional[Union[list[int], list[str]]] = None
+    prefix: Optional[str] = None
+    mangle_dupe_cols: bool = True
+    engine: Optional[Literal["c", "python", "pyarrow"]] = None
+    converters: dict[Union[int, str], Callable[[Any], Any]] = None
+    true_values: Optional[list] = None
+    false_values: Optional[list] = None
+    skipinitialspace: bool = False
+    skiprows: Optional[Union[int, list[int]]] = None
+    nrows: Optional[int] = None
+    na_values: Optional[Union[str, list[str]]] = None
+    keep_default_na: bool = True
+    na_filter: bool = True
+    verbose: bool = False
+    skip_blank_lines: bool = True
+    thousands: Optional[str] = None
+    decimal: str = "."
+    lineterminator: Optional[str] = None
+    quotechar: str = '"'
+    quoting: int = 0
+    escapechar: Optional[str] = None
+    comment: Optional[str] = None
+    encoding: Optional[str] = None
+    dialect: Optional[str] = None
+    error_bad_lines: bool = True
+    warn_bad_lines: bool = True
+    skipfooter: int = 0
+    doublequote: bool = True
+    memory_map: bool = False
+    float_precision: Optional[str] = None
+    chunksize: int = 10_000
+    features: Optional[datasets.Features] = None
+    encoding_errors: Optional[str] = "strict"
+    on_bad_lines: Literal["error", "warn", "skip"] = "error"
+    date_format: Optional[str] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.delimiter is not None:
+            self.sep = self.delimiter
+        if self.column_names is not None:
+            self.names = self.column_names
+
+    @property
+    def pd_read_csv_kwargs(self):
+        pd_read_csv_kwargs = {
+            "sep": self.sep,
+            "header": self.header,
+            "names": self.names,
+            "index_col": self.index_col,
+            "usecols": self.usecols,
+            "prefix": self.prefix,
+            "mangle_dupe_cols": self.mangle_dupe_cols,
+            "engine": self.engine,
+            "converters": self.converters,
+            "true_values": self.true_values,
+            "false_values": self.false_values,
+            "skipinitialspace": self.skipinitialspace,
+            "skiprows": self.skiprows,
+            "nrows": self.nrows,
+            "na_values": self.na_values,
+            "keep_default_na": self.keep_default_na,
+            "na_filter": self.na_filter,
+            "verbose": self.verbose,
+            "skip_blank_lines": self.skip_blank_lines,
+            "thousands": self.thousands,
+            "decimal": self.decimal,
+            "lineterminator": self.lineterminator,
+            "quotechar": self.quotechar,
+            "quoting": self.quoting,
+            "escapechar": self.escapechar,
+            "comment": self.comment,
+            "encoding": self.encoding,
+            "dialect": self.dialect,
+            "error_bad_lines": self.error_bad_lines,
+            "warn_bad_lines": self.warn_bad_lines,
+            "skipfooter": self.skipfooter,
+            "doublequote": self.doublequote,
+            "memory_map": self.memory_map,
+            "float_precision": self.float_precision,
+            "chunksize": self.chunksize,
+            "encoding_errors": self.encoding_errors,
+            "on_bad_lines": self.on_bad_lines,
+            "date_format": self.date_format,
+        }
+
+        # some kwargs must not be passed if they don't have a default value
+        # some others are deprecated and we can also not pass them if they are the default value
+        for pd_read_csv_parameter in _PANDAS_READ_CSV_NO_DEFAULT_PARAMETERS + _PANDAS_READ_CSV_DEPRECATED_PARAMETERS:
+            if pd_read_csv_kwargs[pd_read_csv_parameter] == getattr(CsvConfig(), pd_read_csv_parameter):
+                del pd_read_csv_kwargs[pd_read_csv_parameter]
+
+        # Remove 1.3 new arguments
+        if not (datasets.config.PANDAS_VERSION.major >= 1 and datasets.config.PANDAS_VERSION.minor >= 3):
+            for pd_read_csv_parameter in _PANDAS_READ_CSV_NEW_1_3_0_PARAMETERS:
+                del pd_read_csv_kwargs[pd_read_csv_parameter]
+
+        # Remove 2.0 new arguments
+        if not (datasets.config.PANDAS_VERSION.major >= 2):
+            for pd_read_csv_parameter in _PANDAS_READ_CSV_NEW_2_0_0_PARAMETERS:
+                del pd_read_csv_kwargs[pd_read_csv_parameter]
+
+        # Remove 2.2 deprecated arguments
+        if datasets.config.PANDAS_VERSION.release >= (2, 2):
+            for pd_read_csv_parameter in _PANDAS_READ_CSV_DEPRECATED_2_2_0_PARAMETERS:
+                if pd_read_csv_kwargs[pd_read_csv_parameter] == getattr(CsvConfig(), pd_read_csv_parameter):
+                    del pd_read_csv_kwargs[pd_read_csv_parameter]
+
+        return pd_read_csv_kwargs
+
+
+class Csv(datasets.ArrowBasedBuilder):
+    BUILDER_CONFIG_CLASS = CsvConfig
+
+    def _info(self):
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        """We handle string, list and dicts in datafiles"""
+        if not self.config.data_files:
+            raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
+        dl_manager.download_config.extract_on_the_fly = True
+        base_data_files = dl_manager.download(self.config.data_files)
+        extracted_data_files = dl_manager.extract(base_data_files)
+        splits = []
+        for split_name, extracted_files in extracted_data_files.items():
+            files_iterables = [dl_manager.iter_files(extracted_file) for extracted_file in extracted_files]
+            splits.append(
+                datasets.SplitGenerator(
+                    name=split_name,
+                    gen_kwargs={"files_iterables": files_iterables, "base_files": base_data_files[split_name]},
+                )
+            )
+        return splits
+
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.config.features is not None:
+            schema = self.config.features.arrow_schema
+            if all(not require_storage_cast(feature) for feature in self.config.features.values()):
+                # cheaper cast
+                pa_table = pa.Table.from_arrays([pa_table[field.name] for field in schema], schema=schema)
+            else:
+                # more expensive cast; allows str <-> int/float or str to Audio for example
+                pa_table = table_cast(pa_table, schema)
+        return pa_table
+
+    def _generate_shards(self, base_files, files_iterables):
+        yield from base_files
+
+    def _generate_tables(self, base_files, files_iterables):
+        schema = self.config.features.arrow_schema if self.config.features else None
+        # dtype allows reading an int column as str
+        dtype = (
+            {
+                name: dtype.to_pandas_dtype() if not require_storage_cast(feature) else object
+                for name, dtype, feature in zip(schema.names, schema.types, self.config.features.values())
+            }
+            if schema is not None
+            else None
+        )
+        for shard_idx, files_iterable in enumerate(files_iterables):
+            for file in files_iterable:
+                csv_file_reader = pd.read_csv(file, iterator=True, dtype=dtype, **self.config.pd_read_csv_kwargs)
+                try:
+                    for batch_idx, df in enumerate(csv_file_reader):
+                        pa_table = pa.Table.from_pandas(df)
+                        # Uncomment for debugging (will print the Arrow table size and elements)
+                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
+                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
+                        yield Key(shard_idx, batch_idx), self._cast_table(pa_table)
+                except ValueError as e:
+                    logger.error(f"Failed to read file '{file}' with error {type(e)}: {e}")
+                    raise

src/datasets/packaged_modules/eval/eval.py

@@ -0,0 +1,77 @@
+import json
+import os
+from itertools import islice
+from typing import Iterable
+
+import pyarrow as pa
+
+import datasets
+from datasets.builder import Key
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+class Eval(datasets.GeneratorBasedBuilder):
+    NUM_EXAMPLES_FOR_FEATURES_INFERENCE = 5
+
+    def _info(self):
+        return datasets.DatasetInfo()
+
+    def _split_generators(self, dl_manager):
+        """We handle string, list and dicts in datafiles"""
+        if not self.config.data_files:
+            raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
+        dl_manager.download_config.extract_on_the_fly = True
+        base_data_files = dl_manager.download(self.config.data_files)
+        extracted_data_files = dl_manager.extract(base_data_files)
+        splits = []
+        for split_name, logs in extracted_data_files.items():
+            logs_files_iterables = [dl_manager.iter_files(log) for log in logs]
+            splits.append(
+                datasets.SplitGenerator(
+                    name=split_name,
+                    gen_kwargs={
+                        "logs_files_iterables": logs_files_iterables,
+                        "base_files": base_data_files[split_name],
+                    },
+                )
+            )
+        if not self.info.features:
+            first_examples = list(
+                islice(
+                    self._iter_samples_from_log_files(logs_files_iterables[0]),
+                    self.NUM_EXAMPLES_FOR_FEATURES_INFERENCE,
+                )
+            )
+            pa_tables = [pa.Table.from_pylist([example]) for example in first_examples]
+            inferred_arrow_schema = pa.concat_tables(pa_tables, promote_options="default").schema
+            self.info.features = datasets.Features.from_arrow_schema(inferred_arrow_schema)
+
+        return splits
+
+    def _sort_samples_key(self, sample_path: str):
+        # looks like "{sample_idx}_epoch_{epoch_idx}""
+        (sample_idx_str, epoch_idx_str) = os.path.splitext(os.path.basename(sample_path))[0].split("_epoch_")
+        return (int(epoch_idx_str), int(sample_idx_str))
+
+    def _iter_samples_from_log_files(self, log_files: Iterable[str]):
+        sample_files = [log_file for log_file in log_files if os.path.basename(os.path.dirname(log_file)) == "samples"]
+        sample_files.sort(key=self._sort_samples_key)
+        for sample_file in sample_files:
+            with open(sample_file) as f:
+                sample = json.load(f)
+                for field in sample:
+                    if isinstance(sample[field], dict):
+                        sample[field] = json.dumps(sample[field])
+                    if isinstance(sample[field], list):
+                        sample[field] = [json.dumps(x) for x in sample[field]]
+                yield sample
+
+    def _generate_shards(self, base_files, logs_files_iterables):
+        yield from base_files
+
+    def _generate_examples(self, base_files, logs_files_iterables):
+        for file_idx, log_files in enumerate(logs_files_iterables):
+            for sample_idx, sample in enumerate(self._iter_samples_from_log_files(log_files)):
+                yield Key(file_idx, sample_idx), sample

src/datasets/packaged_modules/folder_based_builder/folder_based_builder.py

@@ -0,0 +1,444 @@
+import collections
+import io
+import itertools
+import os
+from dataclasses import dataclass
+from typing import Any, Callable, Iterator, Optional, Union
+
+import pandas as pd
+import pyarrow as pa
+import pyarrow.dataset as ds
+import pyarrow.json as paj
+import pyarrow.parquet as pq
+
+import datasets
+from datasets import config
+from datasets.builder import Key
+from datasets.features.features import FeatureType, _visit, _visit_with_path, _VisitPath, require_storage_cast
+from datasets.utils.file_utils import readline
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+def count_path_segments(path):
+    return path.replace("\\", "/").count("/")
+
+
+@dataclass
+class FolderBasedBuilderConfig(datasets.BuilderConfig):
+    """BuilderConfig for AutoFolder."""
+
+    features: Optional[datasets.Features] = None
+    drop_labels: bool = None
+    drop_metadata: bool = None
+    metadata_filenames: list[str] = None
+    filters: Optional[Union[ds.Expression, list[tuple], list[list[tuple]]]] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+
+
+class FolderBasedBuilder(datasets.GeneratorBasedBuilder):
+    """
+    Base class for generic data loaders for vision and image data.
+
+
+    Abstract class attributes to be overridden by a child class:
+        BASE_FEATURE: feature object to decode data (i.e. datasets.Image, datasets.Audio, ...)
+        BASE_COLUMN_NAME: string key name of a base feature (i.e. "image", "audio", ...)
+        BUILDER_CONFIG_CLASS: builder config inherited from `folder_based_builder.FolderBasedBuilderConfig`
+        EXTENSIONS: list of allowed extensions (only files with these extensions and METADATA_FILENAME files
+            will be included in a dataset)
+    """
+
+    BASE_FEATURE: type[FeatureType]
+    BASE_COLUMN_NAME: str
+    BUILDER_CONFIG_CLASS: FolderBasedBuilderConfig
+    EXTENSIONS: list[str]
+
+    METADATA_FILENAMES: list[str] = ["metadata.csv", "metadata.jsonl", "metadata.parquet"]
+
+    def _info(self):
+        if not self.config.data_dir and not self.config.data_files:
+            raise ValueError(
+                "Folder-based datasets require either `data_dir` or `data_files` to be specified. "
+                "Neither was provided."
+            )
+
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        if not self.config.data_files:
+            raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
+        dl_manager.download_config.extract_on_the_fly = True
+        # Do an early pass if:
+        # * `drop_labels` is None (default) or False, to infer the class labels
+        # * `drop_metadata` is None (default) or False, to find the metadata files
+        do_analyze = not self.config.drop_labels or not self.config.drop_metadata
+        labels, path_depths = set(), set()
+        all_metadata_files = collections.defaultdict(set)
+        metadata_filenames = self.config.metadata_filenames or self.METADATA_FILENAMES
+
+        def analyze(files_or_archives, downloaded_files_or_dirs, split):
+            if len(downloaded_files_or_dirs) == 0:
+                return
+            # The files are separated from the archives at this point, so check the first sample
+            # to see if it's a file or a directory and iterate accordingly
+            if os.path.isfile(downloaded_files_or_dirs[0]):
+                original_files, downloaded_files = files_or_archives, downloaded_files_or_dirs
+                for original_file, downloaded_file in zip(original_files, downloaded_files):
+                    original_file, downloaded_file = str(original_file), str(downloaded_file)
+                    _, original_file_ext = os.path.splitext(original_file)
+                    if original_file_ext.lower() in self.EXTENSIONS:
+                        if not self.config.drop_labels:
+                            labels.add(os.path.basename(os.path.dirname(original_file)))
+                            path_depths.add(count_path_segments(original_file))
+                    elif os.path.basename(original_file) in metadata_filenames:
+                        all_metadata_files[split].add((original_file, None, downloaded_file))
+                    else:
+                        original_file_name = os.path.basename(original_file)
+                        logger.debug(
+                            f"The file '{original_file_name}' was ignored: it is not a {self.BASE_COLUMN_NAME}, and is not {metadata_filenames} either."
+                        )
+            else:
+                archives, downloaded_dirs = files_or_archives, downloaded_files_or_dirs
+                for archive, downloaded_dir in zip(archives, downloaded_dirs):
+                    archive, downloaded_dir = str(archive), str(downloaded_dir)
+                    for downloaded_dir_file in dl_manager.iter_files(downloaded_dir):
+                        _, downloaded_dir_file_ext = os.path.splitext(downloaded_dir_file)
+                        if downloaded_dir_file_ext in self.EXTENSIONS:
+                            if not self.config.drop_labels:
+                                labels.add(os.path.basename(os.path.dirname(downloaded_dir_file)))
+                                path_depths.add(count_path_segments(downloaded_dir_file))
+                        elif os.path.basename(downloaded_dir_file) in metadata_filenames:
+                            all_metadata_files[split].add((None, downloaded_dir, downloaded_dir_file))
+                        else:
+                            archive_file_name = os.path.basename(archive)
+                            original_file_name = os.path.basename(downloaded_dir_file)
+                            logger.debug(
+                                f"The file '{original_file_name}' from the archive '{archive_file_name}' was ignored: it is not a {self.BASE_COLUMN_NAME}, and is not {metadata_filenames} either."
+                            )
+
+        data_files = self.config.data_files
+        splits = []
+        for split_name, files in data_files.items():
+            files, metadata_files, archives = self._split_files_and_metadata_and_archives(files)
+            downloaded_files = dl_manager.download(files)
+            downloaded_metadata_files = dl_manager.download(metadata_files)
+            downloaded_dirs = dl_manager.download_and_extract(archives)
+            if do_analyze:  # drop_metadata is None or False, drop_labels is None or False
+                logger.info(f"Searching for labels and/or metadata files in {split_name} data files...")
+                analyze(files, downloaded_files, split_name)
+                analyze(metadata_files, downloaded_metadata_files, split_name)
+                analyze(archives, downloaded_dirs, split_name)
+
+                if all_metadata_files:
+                    # add metadata if `all_metadata_files` are found and `drop_metadata` is None (default) or False
+                    add_metadata = not self.config.drop_metadata
+                    # if `all_metadata_files` are found, don't add labels
+                    add_labels = False
+                else:
+                    # if `all_metadata_files` are not found, don't add metadata
+                    add_metadata = False
+                    # if `all_metadata_files` are not found and `drop_labels` is None (default) -
+                    # add labels if files are on the same level in directory hierarchy and there is more than one label
+                    add_labels = (
+                        (len(labels) > 1 and len(path_depths) == 1)
+                        if self.config.drop_labels is None
+                        else not self.config.drop_labels
+                    )
+
+                if add_labels:
+                    logger.info("Adding the labels inferred from data directories to the dataset's features...")
+                if add_metadata:
+                    logger.info("Adding metadata to the dataset...")
+            else:
+                add_labels, add_metadata, all_metadata_files = False, False, {}
+
+            # files info (original_file, None, downloaded_file)
+            files = tuple(zip(files, [None] * len(files), downloaded_files))
+            # archives info (original_archive_file, downloaded_dir, downloaded_files)
+            files += tuple(
+                (archive, downloaded_dir, dl_manager.iter_files(downloaded_dir))
+                for archive, downloaded_dir in zip(archives, downloaded_dirs)
+            )
+            splits.append(
+                datasets.SplitGenerator(
+                    name=split_name,
+                    gen_kwargs={
+                        "files": files,
+                        "metadata_files": all_metadata_files.get(split_name, []),
+                        "add_labels": add_labels,
+                        "add_metadata": add_metadata,
+                    },
+                )
+            )
+
+        if add_metadata:
+            # Verify that:
+            # * all metadata files have the same set of features in each split
+            # * the `file_name` key is one of the metadata keys and is of type string
+            features_per_metadata_file: list[tuple[str, datasets.Features]] = []
+
+            # Check that all metadata files share the same format
+            metadata_ext = {
+                os.path.splitext(original_metadata_file or downloaded_metadata_file)[-1]
+                for original_metadata_file, _, downloaded_metadata_file in itertools.chain.from_iterable(
+                    all_metadata_files.values()
+                )
+            }
+            if len(metadata_ext) > 1:
+                raise ValueError(f"Found metadata files with different extensions: {list(metadata_ext)}")
+            metadata_ext = metadata_ext.pop()
+
+            for split_metadata_files in all_metadata_files.values():
+                pa_metadata_table = None
+                for _, _, downloaded_metadata_file in split_metadata_files:
+                    for pa_metadata_table in self._read_metadata(downloaded_metadata_file, metadata_ext=metadata_ext):
+                        break  # just fetch the first rows
+                    if pa_metadata_table is not None:
+                        features_per_metadata_file.append(
+                            (downloaded_metadata_file, datasets.Features.from_arrow_schema(pa_metadata_table.schema))
+                        )
+                        break  # no need to fetch all the files
+            for downloaded_metadata_file, metadata_features in features_per_metadata_file:
+                if metadata_features != features_per_metadata_file[0][1]:
+                    raise ValueError(
+                        f"Metadata files {downloaded_metadata_file} and {features_per_metadata_file[0][0]} have different features: {features_per_metadata_file[0]} != {metadata_features}"
+                    )
+            metadata_features = features_per_metadata_file[0][1]
+            feature_not_found = True
+
+            def _set_feature(feature):
+                nonlocal feature_not_found
+                if isinstance(feature, dict):
+                    out = type(feature)()
+                    for key in feature:
+                        if (key == "file_name" or key.endswith("_file_name")) and (
+                            feature[key] == datasets.Value("string") or feature[key] == datasets.Value("large_string")
+                        ):
+                            key = key[: -len("_file_name")] or self.BASE_COLUMN_NAME
+                            out[key] = self.BASE_FEATURE()
+                            feature_not_found = False
+                        elif (key == "file_names" or key.endswith("_file_names")) and (
+                            feature[key]
+                            in [datasets.List(datasets.Value("string")), datasets.List(datasets.Value("large_string"))]
+                        ):
+                            key = key[: -len("_file_names")] or (self.BASE_COLUMN_NAME + "s")
+                            out[key] = datasets.List(self.BASE_FEATURE())
+                            feature_not_found = False
+                        elif (key == "file_names" or key.endswith("_file_names")) and (
+                            feature[key] == [datasets.Value("string")]
+                            or feature[key] == [datasets.Value("large_string")]
+                        ):
+                            key = key[: -len("_file_names")] or (self.BASE_COLUMN_NAME + "s")
+                            out[key] = [self.BASE_FEATURE()]
+                            feature_not_found = False
+                        else:
+                            out[key] = feature[key]
+                    return out
+                return feature
+
+            metadata_features = _visit(metadata_features, _set_feature)
+
+            if feature_not_found:
+                raise ValueError(
+                    "`file_name`, `*_file_name`, `file_names` or `*_file_names` must be present as dictionary key in metadata files"
+                )
+        else:
+            metadata_features = None
+
+        # Normally, we would do this in _info, but we need to know the labels and/or metadata
+        # before building the features
+        if self.config.features is None:
+            if add_metadata:
+                self.info.features = metadata_features
+            elif add_labels:
+                self.info.features = datasets.Features(
+                    {
+                        self.BASE_COLUMN_NAME: self.BASE_FEATURE(),
+                        "label": datasets.ClassLabel(names=sorted(labels)),
+                    }
+                )
+            else:
+                self.info.features = datasets.Features({self.BASE_COLUMN_NAME: self.BASE_FEATURE()})
+
+        return splits
+
+    def _split_files_and_metadata_and_archives(self, data_files):
+        files, metadata_files, archives = [], [], []
+        metadata_filenames = self.config.metadata_filenames or self.METADATA_FILENAMES
+        for data_file in data_files:
+            data_file_root, data_file_ext = os.path.splitext(data_file)
+            _, second_data_file_ext = os.path.splitext(data_file_root)
+            if data_file_ext.lower() in self.EXTENSIONS or second_data_file_ext.lower() in self.EXTENSIONS:
+                files.append(data_file)
+            elif os.path.basename(data_file) in metadata_filenames:
+                metadata_files.append(data_file)
+            elif data_file_ext.lower() == ".zip":
+                archives.append(data_file)
+        return files, metadata_files, archives
+
+    def _read_metadata(self, metadata_file: str, metadata_ext: str = "") -> Iterator[pa.Table]:
+        """using the same logic as the Csv, Json and Parquet dataset builders to stream the data"""
+        if self.config.filters is not None:
+            filter_expr = (
+                pq.filters_to_expression(self.config.filters)
+                if isinstance(self.config.filters, list)
+                else self.config.filters
+            )
+        else:
+            filter_expr = None
+        if metadata_ext == ".csv":
+            chunksize = 10_000  # 10k lines
+            schema = self.config.features.arrow_schema if self.config.features else None
+            # dtype allows reading an int column as str
+            dtype = (
+                {
+                    name: dtype.to_pandas_dtype() if not require_storage_cast(feature) else object
+                    for name, dtype, feature in zip(schema.names, schema.types, self.config.features.values())
+                }
+                if schema is not None
+                else None
+            )
+            csv_file_reader = pd.read_csv(metadata_file, iterator=True, dtype=dtype, chunksize=chunksize)
+            for df in csv_file_reader:
+                pa_table = pa.Table.from_pandas(df)
+                if self.config.filters is not None:
+                    pa_table = pa_table.filter(filter_expr)
+                if len(pa_table) > 0:
+                    yield pa_table
+        elif metadata_ext == ".jsonl":
+            with open(metadata_file, "rb") as f:
+                chunksize: int = 10 << 20  # 10MB
+                # Use block_size equal to the chunk size divided by 32 to leverage multithreading
+                # Set a default minimum value of 16kB if the chunk size is really small
+                block_size = max(chunksize // 32, 16 << 10)
+                while True:
+                    batch = f.read(chunksize)
+                    if not batch:
+                        break
+                    # Finish current line
+                    try:
+                        batch += f.readline()
+                    except (AttributeError, io.UnsupportedOperation):
+                        batch += readline(f)
+                    while True:
+                        try:
+                            pa_table = paj.read_json(
+                                io.BytesIO(batch), read_options=paj.ReadOptions(block_size=block_size)
+                            )
+                            break
+                        except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
+                            if (
+                                isinstance(e, pa.ArrowInvalid)
+                                and "straddling" not in str(e)
+                                or block_size > len(batch)
+                            ):
+                                raise
+                            else:
+                                # Increase the block size in case it was too small.
+                                # The block size will be reset for the next file.
+                                logger.debug(
+                                    f"Batch of {len(batch)} bytes couldn't be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}."
+                                )
+                                block_size *= 2
+                    if self.config.filters is not None:
+                        pa_table = pa_table.filter(filter_expr)
+                    if len(pa_table) > 0:
+                        yield pa_table
+        else:
+            with open(metadata_file, "rb") as f:
+                parquet_fragment = ds.ParquetFileFormat().make_fragment(f)
+                if parquet_fragment.row_groups:
+                    batch_size = parquet_fragment.row_groups[0].num_rows
+                else:
+                    batch_size = config.DEFAULT_MAX_BATCH_SIZE
+                for record_batch in parquet_fragment.to_batches(
+                    batch_size=batch_size,
+                    filter=filter_expr,
+                    batch_readahead=0,
+                    fragment_readahead=0,
+                ):
+                    yield pa.Table.from_batches([record_batch])
+
+    def _generate_shards(self, files, metadata_files, add_metadata, add_labels):
+        if add_metadata:
+            for _, _, downloaded_metadata_file in metadata_files:
+                yield downloaded_metadata_file
+        else:
+            for _, downloaded_dir, downloaded_file in files:
+                yield downloaded_dir or downloaded_file
+
+    def _generate_examples(self, files, metadata_files, add_metadata, add_labels):
+        if add_metadata:
+            feature_paths = []
+
+            def find_feature_path(feature, feature_path):
+                nonlocal feature_paths
+                if feature_path and isinstance(feature, self.BASE_FEATURE):
+                    feature_paths.append(feature_path)
+
+            _visit_with_path(self.info.features, find_feature_path)
+
+            for shard_idx, metadata_file_info in enumerate(metadata_files):
+                if len(metadata_file_info) == 2:
+                    original_metadata_file, downloaded_metadata_file = metadata_file_info
+                else:
+                    original_metadata_file, downloaded_metadata_dir, downloaded_metadata_file = metadata_file_info
+                metadata_ext = os.path.splitext(original_metadata_file or downloaded_metadata_file)[-1]
+                downloaded_metadata_dir = os.path.dirname(downloaded_metadata_file)
+
+                def set_feature(item, feature_path: _VisitPath):
+                    if len(feature_path) == 2 and isinstance(feature_path[0], str) and feature_path[1] == 0:
+                        item[feature_path[0]] = item.pop("file_names", None) or item.pop(
+                            feature_path[0] + "_file_names", None
+                        )
+                    elif len(feature_path) == 1 and isinstance(feature_path[0], str):
+                        item[feature_path[0]] = item.pop("file_name", None) or item.pop(
+                            feature_path[0] + "_file_name", None
+                        )
+                    elif len(feature_path) == 0:
+                        if item is not None:
+                            file_relpath = os.path.normpath(item).replace("\\", "/")
+                            item = os.path.join(downloaded_metadata_dir, file_relpath)
+                    return item
+
+                for pa_metadata_table in self._read_metadata(downloaded_metadata_file, metadata_ext=metadata_ext):
+                    for sample_idx, sample in enumerate(pa_metadata_table.to_pylist()):
+                        for feature_path in feature_paths:
+                            _nested_apply(sample, feature_path, set_feature)
+                        yield Key(shard_idx, sample_idx), sample
+        else:
+            if self.config.filters is not None:
+                filter_expr = (
+                    pq.filters_to_expression(self.config.filters)
+                    if isinstance(self.config.filters, list)
+                    else self.config.filters
+                )
+            for shard_idx, (original_file, _, downloaded_files) in enumerate(files):
+                if isinstance(downloaded_files, str):
+                    downloaded_files = [downloaded_files]
+                for sample_idx, downloaded_file in enumerate(downloaded_files):
+                    sample = {self.BASE_COLUMN_NAME: downloaded_file}
+                    if add_labels:
+                        sample["label"] = os.path.basename(os.path.dirname(original_file or downloaded_file))
+                    if self.config.filters is not None:
+                        pa_table = pa.Table.from_pylist([sample]).filter(filter_expr)
+                        if len(pa_table) == 0:
+                            continue
+                    yield Key(shard_idx, sample_idx), sample
+
+
+def _nested_apply(item: Any, feature_path: _VisitPath, func: Callable[[Any, _VisitPath], Any]):
+    # see _visit_with_path() to see how feature paths are constructed
+    item = func(item, feature_path)
+    if feature_path:
+        key = feature_path[0]
+        if key == 0:
+            for i in range(len(item)):
+                item[i] = _nested_apply(item[i], feature_path[1:], func)
+        else:
+            item[key] = _nested_apply(item[key], feature_path[1:], func)
+    return item

src/datasets/packaged_modules/generator/generator.py

@@ -0,0 +1,38 @@
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+import datasets
+from datasets.builder import Key
+from datasets.utils.sharding import _number_of_shards_in_gen_kwargs, _split_gen_kwargs
+
+
+@dataclass
+class GeneratorConfig(datasets.BuilderConfig):
+    generator: Optional[Callable] = None
+    gen_kwargs: Optional[dict] = None
+    features: Optional[datasets.Features] = None
+    split: datasets.NamedSplit = datasets.Split.TRAIN
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.generator is None:
+            raise ValueError("generator must be specified")
+
+        if self.gen_kwargs is None:
+            self.gen_kwargs = {}
+
+
+class Generator(datasets.GeneratorBasedBuilder):
+    BUILDER_CONFIG_CLASS = GeneratorConfig
+
+    def _info(self):
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        return [datasets.SplitGenerator(name=self.config.split, gen_kwargs=self.config.gen_kwargs)]
+
+    def _generate_examples(self, **gen_kwargs):
+        num_shards = _number_of_shards_in_gen_kwargs(gen_kwargs)
+        for shard_idx, shard_gen_kwargs in enumerate(_split_gen_kwargs(gen_kwargs, max_num_jobs=num_shards)):
+            for sample_idx, sample in enumerate(self.config.generator(**shard_gen_kwargs)):
+                yield Key(shard_idx, sample_idx), sample

src/datasets/packaged_modules/hdf5/hdf5.py

@@ -0,0 +1,389 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Optional
+
+import numpy as np
+import pyarrow as pa
+
+import datasets
+from datasets.builder import Key
+from datasets.features.features import (
+    Array2D,
+    Array3D,
+    Array4D,
+    Array5D,
+    Features,
+    LargeList,
+    List,
+    Value,
+    _ArrayXD,
+    _arrow_to_datasets_dtype,
+)
+from datasets.table import cast_table_to_features
+
+
+if TYPE_CHECKING:
+    import h5py
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+EXTENSIONS = [".h5", ".hdf5"]
+
+
+@dataclass
+class HDF5Config(datasets.BuilderConfig):
+    """BuilderConfig for HDF5."""
+
+    batch_size: Optional[int] = None
+    features: Optional[datasets.Features] = None
+
+
+class HDF5(datasets.ArrowBasedBuilder):
+    """ArrowBasedBuilder that converts HDF5 files to Arrow tables using the HF extension types."""
+
+    BUILDER_CONFIG_CLASS = HDF5Config
+
+    def _info(self):
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        import h5py
+
+        if not self.config.data_files:
+            raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
+        data_files = dl_manager.download(self.config.data_files)
+        splits = []
+        for split_name, files in data_files.items():
+            # Infer features from first file
+            if self.info.features is None:
+                for first_file in files:
+                    with open(first_file, "rb") as f:
+                        with h5py.File(f, "r") as h5:
+                            self.info.features = _recursive_infer_features(h5)
+                    break
+            splits.append(datasets.SplitGenerator(name=split_name, gen_kwargs={"files": files}))
+        return splits
+
+    def _generate_shards(self, files):
+        yield from files
+
+    def _generate_tables(self, files):
+        import h5py
+
+        batch_size_cfg = self.config.batch_size
+        for file_idx, file in enumerate(files):
+            try:
+                with open(file, "rb") as f:
+                    with h5py.File(f, "r") as h5:
+                        # Infer features and lengths from first file
+                        if self.info.features is None:
+                            self.info.features = _recursive_infer_features(h5)
+                        num_rows = _check_dataset_lengths(h5, self.info.features)
+                        if num_rows is None:
+                            logger.warning(f"File {file} contains no data, skipping...")
+                            continue
+                        effective_batch = batch_size_cfg or self._writer_batch_size or num_rows
+                        for batch_idx, start in enumerate(range(0, num_rows, effective_batch)):
+                            end = min(start + effective_batch, num_rows)
+                            pa_table = _recursive_load_arrays(h5, self.info.features, start, end)
+                            if pa_table is None:
+                                logger.warning(f"File {file} contains no data, skipping...")
+                                continue
+                            yield Key(file_idx, batch_idx), cast_table_to_features(pa_table, self.info.features)
+            except ValueError as e:
+                logger.error(f"Failed to read file '{file}' with error {type(e)}: {e}")
+                raise
+
+
+# ┌───────────┐
+# │  Complex  │
+# └───────────┘
+
+
+def _is_complex_dtype(dtype: np.dtype) -> bool:
+    if dtype.kind == "c":
+        return True
+    if dtype.subdtype is not None:
+        return _is_complex_dtype(dtype.subdtype[0])
+    return False
+
+
+def _create_complex_features(dset) -> Features:
+    if dset.dtype.subdtype is not None:
+        dtype, data_shape = dset.dtype.subdtype
+    else:
+        data_shape = dset.shape[1:]
+        dtype = dset.dtype
+
+    if dtype == np.complex64:
+        # two float32s
+        value_type = Value("float32")
+    elif dtype == np.complex128:
+        # two float64s
+        value_type = Value("float64")
+    else:
+        logger.warning(f"Found complex dtype {dtype} that is not supported. Converting to float64...")
+        value_type = Value("float64")
+
+    return Features(
+        {
+            "real": _create_sized_feature_impl(data_shape, value_type),
+            "imag": _create_sized_feature_impl(data_shape, value_type),
+        }
+    )
+
+
+def _convert_complex_to_nested(arr: np.ndarray) -> pa.StructArray:
+    data = {
+        "real": datasets.features.features.numpy_to_pyarrow_listarray(arr.real),
+        "imag": datasets.features.features.numpy_to_pyarrow_listarray(arr.imag),
+    }
+    return pa.StructArray.from_arrays([data["real"], data["imag"]], names=["real", "imag"])
+
+
+# ┌────────────┐
+# │  Compound  │
+# └────────────┘
+
+
+def _is_compound_dtype(dtype: np.dtype) -> bool:
+    return dtype.kind == "V"
+
+
+@dataclass
+class _CompoundGroup:
+    dset: "h5py.Dataset"
+    data: np.ndarray = None
+
+    def items(self):
+        for field_name in self.dset.dtype.names:
+            field_dtype = self.dset.dtype[field_name]
+            yield field_name, _CompoundField(self.data, field_name, field_dtype)
+
+
+@dataclass
+class _CompoundField:
+    data: Optional[np.ndarray]
+    name: str
+    dtype: np.dtype
+    shape: tuple[int, ...] = field(init=False)
+
+    def __post_init__(self):
+        self.shape = (len(self.data) if self.data is not None else 0,) + self.dtype.shape
+
+    def __getitem__(self, key):
+        return self.data[key][self.name]
+
+
+def _create_compound_features(dset) -> Features:
+    mock_group = _CompoundGroup(dset)
+    return _recursive_infer_features(mock_group)
+
+
+def _convert_compound_to_nested(arr, dset) -> pa.StructArray:
+    mock_group = _CompoundGroup(dset, data=arr)
+    features = _create_compound_features(dset)
+    return _recursive_load_arrays(mock_group, features, 0, len(arr))
+
+
+# ┌───────────────────┐
+# │  Variable-Length  │
+# └───────────────────┘
+
+
+def _is_vlen_dtype(dtype: np.dtype) -> bool:
+    if dtype.metadata and "vlen" in dtype.metadata:
+        return True
+    return False
+
+
+def _create_vlen_features(dset) -> Features:
+    vlen_dtype = dset.dtype.metadata["vlen"]
+    if vlen_dtype in (str, bytes):
+        return Value("string")
+    inner_feature = _np_to_pa_to_hf_value(vlen_dtype)
+    return List(inner_feature)
+
+
+def _convert_vlen_to_array(arr: np.ndarray) -> pa.Array:
+    return datasets.features.features.numpy_to_pyarrow_listarray(arr)
+
+
+# ┌───────────┐
+# │  Generic  │
+# └───────────┘
+
+
+def _recursive_infer_features(h5_obj) -> Features:
+    features_dict = {}
+    for path, dset in h5_obj.items():
+        if _is_group(dset):
+            features = _recursive_infer_features(dset)
+            if features:
+                features_dict[path] = features
+        elif _is_dataset(dset):
+            features = _infer_feature(dset)
+            if features:
+                features_dict[path] = features
+
+    return Features(features_dict)
+
+
+def _infer_feature(dset):
+    if _is_complex_dtype(dset.dtype):
+        return _create_complex_features(dset)
+    elif _is_compound_dtype(dset.dtype) or dset.dtype.kind == "V":
+        return _create_compound_features(dset)
+    elif _is_vlen_dtype(dset.dtype):
+        return _create_vlen_features(dset)
+    return _create_sized_feature(dset)
+
+
+def _load_array(dset, path: str, start: int, end: int) -> pa.Array:
+    arr = dset[start:end]
+
+    if _is_vlen_dtype(dset.dtype):
+        return _convert_vlen_to_array(arr)
+    elif _is_complex_dtype(dset.dtype):
+        return _convert_complex_to_nested(arr)
+    elif _is_compound_dtype(dset.dtype):
+        return _convert_compound_to_nested(arr, dset)
+    elif dset.dtype.kind == "O":
+        raise ValueError(
+            f"Object dtype dataset '{path}' is not supported. "
+            f"For variable-length data, please use h5py.vlen_dtype() "
+            f"when creating the HDF5 file. "
+            f"See: https://docs.h5py.org/en/stable/special.html#variable-length-strings"
+        )
+    else:
+        # If any non-batch dimension is zero, emit an unsized pa.list_
+        # to avoid creating FixedSizeListArray with list_size=0.
+        if any(dim == 0 for dim in dset.shape[1:]):
+            inner_type = pa.from_numpy_dtype(dset.dtype)
+            return pa.array([[] for _ in arr], type=pa.list_(inner_type))
+        else:
+            return datasets.features.features.numpy_to_pyarrow_listarray(arr)
+
+
+def _recursive_load_arrays(h5_obj, features: Features, start: int, end: int):
+    batch_dict = {}
+    for path, dset in h5_obj.items():
+        if path not in features:
+            continue
+        if _is_group(dset):
+            arr = _recursive_load_arrays(dset, features[path], start, end)
+        elif _is_dataset(dset):
+            arr = _load_array(dset, path, start, end)
+        else:
+            raise ValueError(f"Unexpected type {type(dset)}")
+
+        if arr is not None:
+            batch_dict[path] = arr
+
+    if _is_file(h5_obj):
+        return pa.Table.from_pydict(batch_dict)
+
+    if batch_dict:
+        should_chunk, keys, values = False, [], []
+        for k, v in batch_dict.items():
+            if isinstance(v, pa.ChunkedArray):
+                should_chunk = True
+                v = v.combine_chunks()
+            keys.append(k)
+            values.append(v)
+
+        sarr = pa.StructArray.from_arrays(values, names=keys)
+        return pa.chunked_array(sarr) if should_chunk else sarr
+
+
+# ┌─────────────┐
+# │  Utilities  │
+# └─────────────┘
+
+
+def _create_sized_feature(dset):
+    dset_shape = dset.shape[1:]
+    value_feature = _np_to_pa_to_hf_value(dset.dtype)
+    return _create_sized_feature_impl(dset_shape, value_feature)
+
+
+def _create_sized_feature_impl(dset_shape, value_feature):
+    dtype_str = value_feature.dtype
+    if any(dim == 0 for dim in dset_shape):
+        logger.warning(
+            f"HDF5 to Arrow: Found a dataset with shape {dset_shape} and dtype {dtype_str} that has a dimension with size 0. Shape information will be lost in the conversion to List({value_feature})."
+        )
+        return List(value_feature)
+
+    rank = len(dset_shape)
+    if rank == 0:
+        return value_feature
+    elif rank == 1:
+        return List(value_feature, length=dset_shape[0])
+    elif rank <= 5:
+        return _sized_arrayxd(rank)(shape=dset_shape, dtype=dtype_str)
+    else:
+        raise TypeError(f"Array{rank}D not supported. Maximum 5 dimensions allowed.")
+
+
+def _sized_arrayxd(rank: int):
+    return {2: Array2D, 3: Array3D, 4: Array4D, 5: Array5D}[rank]
+
+
+def _np_to_pa_to_hf_value(numpy_dtype: np.dtype) -> Value:
+    return Value(dtype=_arrow_to_datasets_dtype(pa.from_numpy_dtype(numpy_dtype)))
+
+
+def _first_dataset(h5_obj, features: Features, prefix=""):
+    for path, dset in h5_obj.items():
+        if path not in features:
+            continue
+        if _is_group(dset):
+            found = _first_dataset(dset, features[path], prefix=f"{prefix}{path}/")
+            if found is not None:
+                return found
+        elif _is_dataset(dset):
+            return f"{prefix}{path}"
+
+
+def _check_dataset_lengths(h5_obj, features: Features) -> int:
+    first_path = _first_dataset(h5_obj, features)
+    if first_path is None:
+        return None
+
+    num_rows = h5_obj[first_path].shape[0]
+    for path, dset in h5_obj.items():
+        if path not in features:
+            continue
+        if _is_dataset(dset):
+            if dset.shape[0] != num_rows:
+                raise ValueError(f"Dataset '{path}' has length {dset.shape[0]} but expected {num_rows}")
+    return num_rows
+
+
+def _is_group(h5_obj) -> bool:
+    import h5py
+
+    return isinstance(h5_obj, h5py.Group) or isinstance(h5_obj, _CompoundGroup)
+
+
+def _is_dataset(h5_obj) -> bool:
+    import h5py
+
+    return isinstance(h5_obj, h5py.Dataset) or isinstance(h5_obj, _CompoundField)
+
+
+def _is_file(h5_obj) -> bool:
+    import h5py
+
+    return isinstance(h5_obj, h5py.File)
+
+
+def _has_zero_dimensions(feature):
+    if isinstance(feature, _ArrayXD):
+        return any(dim == 0 for dim in feature.shape)
+    elif isinstance(feature, List):
+        return feature.length == 0 or _has_zero_dimensions(feature.feature)
+    elif isinstance(feature, LargeList):
+        return _has_zero_dimensions(feature.feature)
+    else:
+        return False

src/datasets/packaged_modules/iceberg/iceberg.py

@@ -0,0 +1,173 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+import pyarrow as pa
+
+import datasets
+from datasets.builder import Key
+from datasets.features import Features
+from datasets.table import table_cast
+
+
+if TYPE_CHECKING:
+    from pyiceberg.catalog import Catalog
+    from pyiceberg.expressions import BooleanExpression
+    from pyiceberg.table import FileScanTask
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+@dataclass
+class IcebergConfig(datasets.BuilderConfig):
+    """BuilderConfig for Apache Iceberg format.
+
+    Args:
+        catalog (`pyiceberg.catalog.Catalog`):
+            A pre-configured pyiceberg Catalog object.
+        table (`str` or `Dict[str, str]`):
+            Iceberg table identifier, e.g. ``"db.my_table"``.
+            Pass a dict to map split names to table identifiers,
+            e.g. ``{"train": "db.train", "test": "db.test"}``.
+        features (`Features`, *optional*):
+            Cast the data to these features.
+        columns (`List[str]`, *optional*):
+            List of columns to load; others are ignored.
+        filters (`str` or `BooleanExpression`, *optional*):
+            Row filter with predicate pushdown. Accepts a SQL-style string
+            (``"col > 1 AND col2 == 'foo'"``), or a pyiceberg
+            ``BooleanExpression`` object. Parsed by pyiceberg internally.
+        batch_size (`int`, defaults to ``131072``):
+            Number of rows per RecordBatch when reading.
+        snapshot_id (`int`, *optional*):
+            Load a specific snapshot for time-travel queries.
+    """
+
+    catalog: Optional["Catalog"] = None
+    table: Optional[Union[str, Dict[str, str]]] = None
+    features: Optional[datasets.Features] = None
+    columns: Optional[List[str]] = None
+    filters: Optional[Union[str, "BooleanExpression"]] = None
+    batch_size: int = 131072
+    snapshot_id: Optional[int] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.catalog is None:
+            raise ValueError("`catalog` must be a pyiceberg Catalog object, but got None.")
+        if self.table is None:
+            raise ValueError("`table` must be specified, e.g. table='db.my_table'")
+        # Normalize table to Dict[split_name, table_identifier]
+        if isinstance(self.table, str):
+            self.table = {"train": self.table}
+        # Generate a stable config name for caching
+        if self.name == "default":
+            catalog_id = f"{self.catalog.__class__.__name__}_{self.catalog.name}"
+            table_id = "_".join(sorted(self.table.values()))
+            self.name = f"{catalog_id}_{table_id}"
+
+    def create_config_id(
+        self,
+        config_kwargs: dict,
+        custom_features: Optional[Features] = None,
+    ) -> str:
+        # The catalog object is not picklable (contains SQLAlchemy engines, etc.),
+        # so we replace it with a hashable string representation before the
+        # parent class hashes config_kwargs via dill.
+        config_kwargs = config_kwargs.copy()
+        catalog = config_kwargs.pop("catalog", None)
+        if catalog is not None:
+            config_kwargs["_catalog_id"] = f"{catalog.__class__.__name__}_{catalog.name}"
+        # filters may contain pyiceberg Expression objects that are not picklable
+        filters = config_kwargs.pop("filters", None)
+        if filters is not None:
+            config_kwargs["_filters_repr"] = repr(filters)
+        return super().create_config_id(config_kwargs, custom_features=custom_features)
+
+
+class Iceberg(datasets.ArrowBasedBuilder, datasets.builder._CountableBuilderMixin):
+    BUILDER_CONFIG_CLASS = IcebergConfig
+
+    def _info(self):
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        splits = []
+        for split_name, table_id in self.config.table.items():
+            iceberg_table = self.config.catalog.load_table(table_id)
+
+            scan_kwargs = {}
+            if self.config.filters is not None:
+                scan_kwargs["row_filter"] = self.config.filters
+            if self.config.columns:
+                scan_kwargs["selected_fields"] = tuple(self.config.columns)
+            if self.config.snapshot_id is not None:
+                scan_kwargs["snapshot_id"] = self.config.snapshot_id
+
+            scan = iceberg_table.scan(**scan_kwargs)
+
+            # Infer features from Arrow schema if not user-provided
+            if self.info.features is None:
+                arrow_schema = scan.projection().as_arrow()
+                self.info.features = datasets.Features.from_arrow_schema(arrow_schema)
+
+            # Plan files for parallel processing: passing a list in gen_kwargs
+            # enables _split_gen_kwargs to distribute tasks across num_proc workers.
+            tasks = list(scan.plan_files())
+
+            # Extract picklable scan context for multiprocessing compatibility.
+            # The scan object itself is not picklable (holds catalog connections),
+            # but these components are individually serializable.
+            scan_context = (
+                scan.table_metadata,
+                scan.io,
+                scan.projection(),
+                scan.row_filter,
+                scan.case_sensitive,
+                scan.limit,
+            )
+
+            splits.append(
+                datasets.SplitGenerator(
+                    name=split_name,
+                    gen_kwargs={"tasks": tasks, "scan_context": scan_context},
+                )
+            )
+
+        # Drop the catalog reference so the builder becomes picklable for num_proc > 1.
+        # All data needed for reading has been extracted into scan_context above.
+        self.config.catalog = None
+        self.config_kwargs.pop("catalog", None)
+
+        return splits
+
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.info.features is not None:
+            # More expensive cast to support nested features with keys in a different order
+            # allows str <-> int/float or str to Audio for example
+            pa_table = table_cast(pa_table, self.info.features.arrow_schema)
+        return pa_table
+
+    def _generate_shards(self, tasks: List["FileScanTask"], scan_context):
+        for task in tasks:
+            yield task.file.file_path
+
+    def _generate_num_examples(self, tasks: List["FileScanTask"], scan_context):
+        for task in tasks:
+            yield task.file.record_count
+
+    def _generate_tables(self, tasks: List["FileScanTask"], scan_context):
+        from pyiceberg.io.pyarrow import ArrowScan
+
+        table_metadata, io, projected_schema, row_filter, case_sensitive, limit = scan_context
+        arrow_scan = ArrowScan(
+            table_metadata,
+            io,
+            projected_schema,
+            row_filter,
+            case_sensitive=case_sensitive,
+            limit=limit,
+        )
+        for task_idx, task in enumerate(tasks):
+            for batch_idx, batch in enumerate(arrow_scan.to_record_batches([task])):
+                pa_table = pa.Table.from_batches([batch])
+                yield Key(task_idx, batch_idx), self._cast_table(pa_table)

src/datasets/packaged_modules/imagefolder/imagefolder.py

@@ -0,0 +1,103 @@
+import datasets
+
+from ..folder_based_builder import folder_based_builder
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+class ImageFolderConfig(folder_based_builder.FolderBasedBuilderConfig):
+    """BuilderConfig for ImageFolder."""
+
+    drop_labels: bool = None
+    drop_metadata: bool = None
+
+    def __post_init__(self):
+        super().__post_init__()
+
+
+class ImageFolder(folder_based_builder.FolderBasedBuilder):
+    BASE_FEATURE = datasets.Image
+    BASE_COLUMN_NAME = "image"
+    BUILDER_CONFIG_CLASS = ImageFolderConfig
+    EXTENSIONS: list[str]  # definition at the bottom of the script
+
+
+# Obtained with:
+# ```
+# import PIL.Image
+# IMAGE_EXTENSIONS = []
+# PIL.Image.init()
+# for ext, format in PIL.Image.EXTENSION.items():
+#     if format in PIL.Image.OPEN:
+#         IMAGE_EXTENSIONS.append(ext[1:])
+# ```
+# We intentionally do not run this code on launch because:
+# (1) Pillow is an optional dependency, so importing Pillow in global namespace is not allowed
+# (2) To ensure the list of supported extensions is deterministic
+IMAGE_EXTENSIONS = [
+    ".blp",
+    ".bmp",
+    ".dib",
+    ".bufr",
+    ".cur",
+    ".pcx",
+    ".dcx",
+    ".dds",
+    ".ps",
+    ".eps",
+    ".fit",
+    ".fits",
+    ".fli",
+    ".flc",
+    ".ftc",
+    ".ftu",
+    ".gbr",
+    ".gif",
+    ".grib",
+    # ".h5",   # may contain zero or several images
+    # ".hdf",  # may contain zero or several images
+    ".png",
+    ".apng",
+    ".jp2",
+    ".j2k",
+    ".jpc",
+    ".jpf",
+    ".jpx",
+    ".j2c",
+    ".icns",
+    ".ico",
+    ".im",
+    ".iim",
+    ".tif",
+    ".tiff",
+    ".jfif",
+    ".jpe",
+    ".jpg",
+    ".jpeg",
+    ".mpg",
+    ".mpeg",
+    ".msp",
+    ".pcd",
+    ".pxr",
+    ".pbm",
+    ".pgm",
+    ".ppm",
+    ".pnm",
+    ".psd",
+    ".bw",
+    ".rgb",
+    ".rgba",
+    ".sgi",
+    ".ras",
+    ".tga",
+    ".icb",
+    ".vda",
+    ".vst",
+    ".webp",
+    ".wmf",
+    ".emf",
+    ".xbm",
+    ".xpm",
+]
+ImageFolder.EXTENSIONS = IMAGE_EXTENSIONS

src/datasets/packaged_modules/json/json.py

@@ -0,0 +1,536 @@
+import codecs
+import io
+import os
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Literal, Optional
+
+import pandas as pd
+import pyarrow as pa
+import pyarrow.json as paj
+
+import datasets
+import datasets.config
+from datasets import List, Value
+from datasets.builder import Key
+from datasets.table import table_cast
+from datasets.utils.file_utils import readline
+from datasets.utils.json import (
+    find_mixed_struct_types_field_paths,
+    get_json_field_path_from_pyarrow_json_error,
+    get_json_field_paths_from_feature,
+    insert_json_field_path,
+    json_encode_field,
+    json_encode_fields_in_json_lines,
+    set_json_types_in_feature,
+    ujson_dumps,
+    ujson_loads,
+)
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+def pandas_read_json(path_or_buf, **kwargs):
+    if datasets.config.PANDAS_VERSION.major >= 2:
+        kwargs["dtype_backend"] = "pyarrow"
+    return pd.read_json(path_or_buf, **kwargs)
+
+
+class FullReadDisallowed(Exception):
+    pass
+
+
+@dataclass
+class JsonConfig(datasets.BuilderConfig):
+    """BuilderConfig for JSON."""
+
+    features: Optional[datasets.Features] = None
+    encoding: str = "utf-8"
+    encoding_errors: Optional[str] = None
+    field: Optional[str] = None
+    use_threads: bool = True  # deprecated
+    block_size: Optional[int] = None  # deprecated
+    chunksize: int = 10 << 20  # 10MB
+    newlines_in_values: Optional[bool] = None
+    on_mixed_types: Optional[Literal["use_json"]] = "use_json"
+    parse_agent_traces: bool = True
+
+    def __post_init__(self):
+        super().__post_init__()
+
+
+class Json(datasets.ArrowBasedBuilder):
+    BUILDER_CONFIG_CLASS = JsonConfig
+
+    def _info(self):
+        if self.config.block_size is not None:
+            logger.warning("The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead")
+            self.config.chunksize = self.config.block_size
+        if self.config.use_threads is not True:
+            logger.warning(
+                "The JSON loader parameter `use_threads` is deprecated and doesn't have any effect anymore."
+            )
+        if self.config.newlines_in_values is not None:
+            raise ValueError("The JSON loader parameter `newlines_in_values` is no longer supported")
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        """We handle string, list and dicts in datafiles"""
+        if not self.config.data_files:
+            raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
+        dl_manager.download_config.extract_on_the_fly = True
+        base_data_files = dl_manager.download(self.config.data_files)
+        extracted_data_files = dl_manager.extract(base_data_files)
+        splits = []
+        for split_name, extracted_files in extracted_data_files.items():
+            files_iterables = [dl_manager.iter_files(extracted_file) for extracted_file in extracted_files]
+            splits.append(
+                datasets.SplitGenerator(
+                    name=split_name,
+                    gen_kwargs={
+                        "files_iterables": files_iterables,
+                        "base_files": base_data_files[split_name],
+                        "original_files": self.config.data_files[split_name],
+                    },
+                )
+            )
+        if self.info.features is None:
+            try:
+                pa_table = next(iter(self._generate_tables(**splits[0].gen_kwargs, allow_full_read=False)))[1]
+                self.info.features = datasets.Features.from_arrow_schema(pa_table.schema)
+                if self.config.parse_agent_traces and has_agent_traces_markers(self.info.features):
+                    self.info.features = AGENT_TRACES_FEATURES
+            except FullReadDisallowed:
+                pass
+        return splits
+
+    def _cast_table(self, pa_table: pa.Table, json_field_paths=()) -> pa.Table:
+        if self.info.features is not None:
+            # adding missing columns
+            for column_name in set(self.info.features) - set(pa_table.column_names):
+                type = self.info.features.arrow_schema.field(column_name).type
+                pa_table = pa_table.append_column(column_name, pa.array([None] * len(pa_table), type=type))
+            # convert to string when needed
+            for i, column_name in enumerate(pa_table.column_names):
+                if pa.types.is_struct(pa_table[column_name].type) and self.info.features.get(
+                    column_name, None
+                ) == Value("string"):
+                    jsonl = (
+                        pa_table[column_name]
+                        .to_pandas(types_mapper=pd.ArrowDtype)
+                        .to_json(orient="records", lines=True)
+                    )
+                    string_array = pa.array(
+                        (None if x.strip() == "null" else x.strip() for x in jsonl.split("\n") if x.strip()),
+                        type=pa.string(),
+                    )
+                    pa_table = pa_table.set_column(i, column_name, string_array)
+            # more expensive cast to support nested structures with keys in a different order
+            # allows str <-> int/float or str to Audio for example
+            pa_table = table_cast(pa_table, self.info.features.arrow_schema)
+        elif json_field_paths:
+            features = datasets.Features.from_arrow_schema(pa_table.schema)
+            features = set_json_types_in_feature(features, json_field_paths)
+            pa_table = table_cast(pa_table, features.arrow_schema)
+        return pa_table
+
+    def _generate_shards(self, base_files, files_iterables, original_files):
+        yield from base_files
+
+    def _generate_tables(self, base_files, files_iterables, original_files, allow_full_read=True):
+        json_field_paths = []
+        is_agent_traces = False
+
+        if self.info.features is not None:
+            if self.info.features == AGENT_TRACES_FEATURES:
+                is_agent_traces = True
+                if datasets.config.TEICH_AVAILABLE:
+                    from teich import convert_traces_to_training_data
+                else:
+                    raise ImportError("To support decoding agent traces, please install 'teich'.")
+            json_field_paths = get_json_field_paths_from_feature(self.info.features)
+
+        for shard_idx, files_iterable in enumerate(files_iterables):
+            for file in files_iterable:
+                # If the file is one json object and if we need to look at the items in one specific field
+                if self.config.field is not None:
+                    if not allow_full_read:
+                        raise FullReadDisallowed()
+                    with open(file, encoding=self.config.encoding, errors=self.config.encoding_errors) as f:
+                        dataset = ujson_loads(f.read())
+                    # We keep only the field we are interested in
+                    dataset = dataset[self.config.field]
+                    df = pandas_read_json(io.StringIO(ujson_dumps(dataset)))
+                    if df.columns.tolist() == [0]:
+                        df.columns = list(self.config.features) if self.config.features else ["text"]
+                    pa_table = pa.Table.from_pandas(df, preserve_index=False)
+                    yield Key(shard_idx, 0), self._cast_table(pa_table)
+
+                # If the files are agent traces (one row = one file except for hermes which can have multiple sessions per file)
+                elif is_agent_traces:
+                    trace_file = Path(file)
+                    trace = trace_file.read_text(encoding="utf-8")
+                    lines = trace.splitlines()
+                    file_path = original_files[shard_idx]
+                    if self.base_path is not None and file_path.startswith(self.base_path):
+                        file_path = os.path.relpath(file_path, self.base_path)
+                    training_examples = convert_traces_to_training_data(trace_file)
+                    examples = []
+                    for i, training_example in enumerate(training_examples):
+                        if training_example["metadata"]["trace_type"] == "hermes":
+                            timestamp = ujson_loads(lines[i])["started_at"]
+                            milliseconds = timestamp if timestamp <= 10_000_000_000 else timestamp * 1_000
+                            sent_at = (
+                                datetime.fromtimestamp(milliseconds, tz=timezone.utc)
+                                .isoformat(timespec="milliseconds")
+                                .replace("+00:00", "Z")
+                            )
+                            bonus_fields = {
+                                "harness": training_example["metadata"]["trace_type"],
+                                "session_id": training_example["metadata"]["session_id"],
+                                "sent_at": sent_at,
+                                "num_user_messages": training_example["metadata"]["turn_count"],
+                                "num_tool_calls": training_example["metadata"]["tool_call_count"],
+                                "trace": lines[i],
+                            }
+                        else:
+                            harness, session_id, prompt, sent_at, num_user_messages, num_tool_calls = (
+                                parse_traces_info(lines)
+                            )
+                            bonus_fields = {
+                                "harness": harness,
+                                "session_id": session_id,
+                                "prompt": prompt,
+                                "sent_at": sent_at,
+                                "num_user_messages": num_user_messages,
+                                "num_tool_calls": num_tool_calls,
+                                "trace": trace,
+                            }
+                        example = {
+                            **dict.fromkeys(AGENT_TRACES_FEATURES),
+                            **training_example,
+                            **bonus_fields,
+                            "file_path": file_path,
+                        }
+                        for json_field_path in json_field_paths:
+                            example = json_encode_field(example, json_field_path)
+                        examples.append(example)
+                    pa_table = pa.Table.from_pylist(examples)
+                    yield Key(shard_idx, 0), self._cast_table(pa_table)
+
+                # If the file has one json object per line
+                else:
+                    with open(file, "rb") as f:
+                        batch_idx = 0
+                        # Use block_size equal to the chunk size divided by 32 to leverage multithreading
+                        # Set a default minimum value of 16kB if the chunk size is really small
+                        block_size = max(self.config.chunksize // 32, 16 << 10)
+                        encoding_errors = (
+                            self.config.encoding_errors if self.config.encoding_errors is not None else "strict"
+                        )
+                        while True:
+                            batch = f.read(self.config.chunksize)
+                            if not batch:
+                                break
+                            # A leading UTF-8 BOM makes the ujson pre-scan below raise
+                            # ValueError, silently skipping mixed-struct detection, while
+                            # PyArrow tolerates the BOM -- so the inferred schema differed
+                            # depending on whether the file started with a BOM. Strip it so
+                            # both see the same bytes (a BOM only appears at the start of the file).
+                            if shard_idx == 0 and batch_idx == 0 and batch.startswith(codecs.BOM_UTF8):
+                                batch = batch[len(codecs.BOM_UTF8) :]
+                            if batch.startswith(b"["):
+                                if not allow_full_read:
+                                    raise FullReadDisallowed()
+                                else:
+                                    # convert to JSON Lines
+                                    full_data = batch + f.read()
+                                    if b"{" in batch[:100].split(b'"', 1)[0]:  # list of objects
+                                        batch = "\n".join(ujson_dumps(x) for x in ujson_loads(full_data)).encode()
+                                    else:  # list of strings
+                                        batch = "\n".join(
+                                            ujson_dumps({"text": x}) for x in ujson_loads(full_data)
+                                        ).encode()
+                            # Finish current line
+                            try:
+                                batch += f.readline()
+                            except (AttributeError, io.UnsupportedOperation):
+                                batch += readline(f)
+                            # PyArrow only accepts utf-8 encoded bytes
+                            if self.config.encoding != "utf-8":
+                                batch = batch.decode(self.config.encoding, errors=encoding_errors).encode("utf-8")
+                            # On first batch we check for lists of objects with arbitrary fields
+                            if (
+                                shard_idx == 0
+                                and batch_idx == 0
+                                and self.info.features is None
+                                and self.config.on_mixed_types == "use_json"
+                            ):
+                                try:
+                                    examples = [ujson_loads(line) for line in batch.splitlines()]
+                                except ValueError:
+                                    # the file is likely not JSON Lines and may contain one single multi-line JSON object
+                                    pass
+                                else:
+                                    json_field_paths += find_mixed_struct_types_field_paths(examples)
+                            # Re-encode JSON fields
+                            original_batch = batch
+                            if json_field_paths:
+                                examples = [ujson_loads(line) for line in batch.splitlines()]
+                                for json_field_path in json_field_paths:
+                                    examples = [json_encode_field(example, json_field_path) for example in examples]
+                                batch = "\n".join(ujson_dumps(example) for example in examples).encode()
+                            # Disable parallelism if block size is ~ len(batch) to avoid segfault
+                            block_size = len(batch) if len(batch) // 8 > block_size else block_size
+                            try:
+                                while True:
+                                    try:
+                                        pa_table = paj.read_json(
+                                            io.BytesIO(batch), read_options=paj.ReadOptions(block_size=block_size)
+                                        )
+                                        break
+                                    except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
+                                        if batch.startswith(b"["):  # paj.read_json only supports json lines
+                                            raise
+                                        elif self.config.on_mixed_types == "use_json" and (
+                                            isinstance(e, pa.ArrowInvalid)
+                                            and "JSON parse error: Column(" in str(e)
+                                            and ") changed from" in str(e)
+                                        ):
+                                            json_field_path = get_json_field_path_from_pyarrow_json_error(str(e))
+                                            insert_json_field_path(json_field_paths, json_field_path)
+                                            batch = json_encode_fields_in_json_lines(original_batch, json_field_paths)
+                                        elif (
+                                            "straddling" in str(e) or "JSON conversion to" in str(e)
+                                        ) and block_size < len(batch):
+                                            # Increase the block size in case it was too small.
+                                            # The block size will be reset for the next file.
+                                            # this is needed in case of "stradding" or for some JSON conversions (see https://github.com/huggingface/datasets/issues/2799)
+                                            logger.debug(
+                                                f"Batch of {len(batch)} bytes couldn't be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}."
+                                            )
+                                            block_size *= 2
+                                        else:
+                                            raise
+                            except pa.ArrowInvalid as e:
+                                if not allow_full_read:
+                                    raise FullReadDisallowed()
+                                try:
+                                    with open(
+                                        file, encoding=self.config.encoding, errors=self.config.encoding_errors
+                                    ) as f:
+                                        df = pandas_read_json(f)
+                                except ValueError:
+                                    logger.error(f"Failed to load JSON from file '{file}' with error {type(e)}: {e}")
+                                    raise e
+                                if df.columns.tolist() == [0]:
+                                    df.columns = list(self.config.features) if self.config.features else ["text"]
+                                try:
+                                    pa_table = pa.Table.from_pandas(df, preserve_index=False)
+                                except pa.ArrowInvalid as e:
+                                    logger.error(
+                                        f"Failed to convert pandas DataFrame to Arrow Table from file '{file}' with error {type(e)}: {e}"
+                                    )
+                                    raise ValueError(
+                                        f"Failed to convert pandas DataFrame to Arrow Table from file {file}."
+                                    ) from None
+                                yield Key(shard_idx, 0), self._cast_table(pa_table)
+                                break
+                            yield (
+                                Key(shard_idx, batch_idx),
+                                self._cast_table(pa_table, json_field_paths=json_field_paths),
+                            )
+                            batch_idx += 1
+
+
+AGENT_TRACES_TYPES_VALUES = {
+    "claude_code": ["user", "assistant", "system"],
+    "pi": ["session", "message"],
+    "codex": ["session_meta", "turn_context", "response_item", "event_msg"],
+    # droid message events share pi's "message" type, but droid traces always start with a session_start event
+    "droid": ["session_start"],
+}
+AGENT_TRACES_TYPE_TO_HARNESS = {}
+for _harness, _trace_types in AGENT_TRACES_TYPES_VALUES.items():
+    for _trace_type in _trace_types:
+        AGENT_TRACES_TYPE_TO_HARNESS[_trace_type] = _harness
+
+
+AGENT_TRACES_FEATURES_MARKERS = {
+    "claude_code_or_pi_or_openclaw": datasets.Features(
+        {
+            "type": lambda f: f == Value("string"),
+            "message": lambda f: f == datasets.Json(),
+        }
+    ),
+    "codex": datasets.Features(
+        {
+            "type": lambda f: f == Value("string"),
+            "payload": lambda f: f == datasets.Json(),
+        }
+    ),
+    "hermes": datasets.Features(
+        {
+            "id": lambda f: f == Value("string"),
+            "source": lambda f: f == datasets.Value("string"),
+            "model": lambda f: f == datasets.Value("string"),
+            "system_prompt": lambda f: f == datasets.Value("string"),
+            "messages": lambda f: isinstance(f, (datasets.List, datasets.Json)),
+        }
+    ),
+    "droid": datasets.Features(
+        {
+            "type": lambda f: f == Value("string"),
+            "id": lambda f: f == Value("string"),
+            "version": lambda f: f == Value("int64"),
+            "cwd": lambda f: f == Value("string"),
+        }
+    ),
+}
+
+AGENT_TRACES_FEATURES = datasets.Features(
+    {
+        # basic features
+        "harness": Value("string"),
+        "session_id": Value("string"),
+        # teich features
+        "prompt": Value("string"),
+        "messages": List(datasets.Json()),
+        "tools": List(datasets.Json()),
+        "metadata": datasets.Json(),
+        # bonus features
+        "sent_at": Value("string"),
+        "num_user_messages": Value("int64"),
+        "num_tool_calls": Value("int64"),
+        "trace": datasets.Json(),
+        "file_path": Value("string"),
+    }
+)
+
+
+def has_agent_traces_markers(features: datasets.Features) -> bool:
+    for agent_traces_features_marker in AGENT_TRACES_FEATURES_MARKERS.values():
+        if all(feature_marker(features.get(key)) for key, feature_marker in agent_traces_features_marker.items()):
+            return True
+    return False
+
+
+def parse_traces_info(
+    trace_events: list[str],
+) -> tuple[Optional[str], Optional[str], Optional[str], Optional[str], int, int]:
+    harness, session_id, prompt, sent_at = None, None, None, None
+    # prompt/sent_at describe the first user message; the counters summarize the whole trace file.
+    # Codex response_item user messages can include context files (for example AGENTS.md), so only event_msg
+    # user_message records are treated as Codex user messages.
+    num_user_messages = 0
+    num_tool_calls = 0
+    for event in trace_events:
+        decoded_event = ujson_loads(event)
+        if harness is None:
+            if "type" in decoded_event and isinstance(decoded_event["type"], str):
+                harness = AGENT_TRACES_TYPE_TO_HARNESS.get(decoded_event["type"])
+        if session_id is None:
+            session_id = get_session_id(decoded_event)
+            if (
+                session_id is not None
+                and decoded_event.get("type") == "session"
+                and isinstance(decoded_event.get("cwd"), str)
+                and "/.openclaw/" in decoded_event["cwd"]
+            ):
+                harness = "openclaw"
+        user_prompt = get_user_prompt(decoded_event)
+        if user_prompt is not None:
+            num_user_messages += 1
+            if prompt is None:
+                prompt = user_prompt
+                sent_at = get_trace_event_timestamp(decoded_event)
+        num_tool_calls += get_tool_call_count(decoded_event)
+    return harness, session_id, prompt, sent_at, num_user_messages, num_tool_calls
+
+
+def get_session_id(trace: dict) -> Optional[str]:
+    # claude
+    if isinstance(trace.get("sessionId"), str):
+        return trace["sessionId"]
+    # claude (not sure but this format does exist online)
+    if isinstance(trace.get("session_id"), str):
+        return trace["session_id"]
+    # codex
+    if isinstance(trace.get("payload"), dict) and isinstance(trace["payload"].get("id"), str):
+        return trace["payload"]["id"]
+    # pi / openclaw on "session" (openclaw embeds pi-agent; distinguish via cwd), droid on "session_start"
+    if trace.get("type") in ("session", "session_start") and isinstance(trace.get("id"), str):
+        return trace["id"]
+    return None
+
+
+def get_user_prompt(trace_event: dict) -> Optional[str]:
+    if trace_event.get("type") == "user" and isinstance(trace_event.get("message"), dict):
+        message = trace_event["message"]
+        if message.get("role") == "user":
+            return get_content_text(message.get("content"))
+
+    if trace_event.get("type") == "message":
+        if isinstance(trace_event.get("message"), dict):
+            message = trace_event["message"]
+            # droid marks injected context as llm_only and local-only notes as user_only, neither is a real user prompt
+            if message.get("role") == "user" and message.get("visibility") not in ("llm_only", "user_only"):
+                return get_content_text(message.get("content"))
+        if trace_event.get("role") == "user":
+            return get_content_text(trace_event.get("content"))
+
+    if trace_event.get("type") == "event_msg" and isinstance(trace_event.get("payload"), dict):
+        payload = trace_event["payload"]
+        if payload.get("type") == "user_message":
+            return get_content_text(payload.get("message"))
+
+    return None
+
+
+def get_tool_call_count(trace_event: dict) -> int:
+    trace_type = trace_event.get("type")
+    if trace_type == "response_item":
+        payload = trace_event.get("payload")
+        if isinstance(payload, dict) and payload.get("type") == "function_call":
+            return 1
+        return 0
+
+    if trace_type not in {"assistant", "message"} or not isinstance(trace_event.get("message"), dict):
+        return 0
+    message = trace_event["message"]
+    if message.get("role") != "assistant":
+        return 0
+    content = message.get("content")
+    if not isinstance(content, list):
+        return 0
+    return sum(
+        1
+        for content_part in content
+        if isinstance(content_part, dict) and content_part.get("type") in {"tool_use", "toolCall"}
+    )
+
+
+def get_trace_event_timestamp(trace_event: dict) -> Optional[str]:
+    timestamp = trace_event.get("timestamp")
+    if isinstance(timestamp, str):
+        return timestamp
+    return None
+
+
+def get_content_text(content) -> Optional[str]:
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        content_parts = []
+        for content_part in content:
+            if isinstance(content_part, str):
+                content_parts.append(content_part)
+            elif isinstance(content_part, dict):
+                text = content_part.get("text")
+                if isinstance(text, str):
+                    content_parts.append(text)
+        if content_parts:
+            return "\n".join(content_parts)
+    return None

src/datasets/packaged_modules/lance/lance.py

@@ -0,0 +1,240 @@
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+import pyarrow as pa
+from huggingface_hub import HfApi, get_token
+
+import datasets
+from datasets import Audio, Image, Video
+from datasets.builder import Key
+from datasets.table import table_cast
+from datasets.utils.file_utils import is_local_path
+
+
+if TYPE_CHECKING:
+    import lance
+    import lance.file
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+MAGIC_BYTES_EXTENSION_AND_FEATURE_TYPES = [
+    ("1A 45 DF A3", ".mkv", Video()),
+    ("66 74 79 70 69 73 6F 6D", ".mp4", Video()),
+    ("66 74 79 70 4D 53 4E 56", ".mp4", Video()),
+    ("52 49 46 46", ".avi", Video()),
+    ("00 00 01 BA", ".mpeg", Video()),
+    ("00 00 01 BA", ".mpeg", Video()),
+    ("00 00 01 B3", ".mov", Video()),
+    ("89 50 4E 47", ".png", Image()),
+    ("FF D8", ".jpg", Image()),
+    ("49 49", ".tif", Image()),
+    ("47 49 46 38", ".gif", Image()),
+    ("52 49 46 46", ".wav", Audio()),
+    ("49 44 33", ".mp3", Audio()),
+    ("66 4C 61 43", ".flac", Audio()),
+]
+
+
+@dataclass
+class LanceConfig(datasets.BuilderConfig):
+    """
+    BuilderConfig for Lance format.
+
+    Args:
+        features: (`Features`, *optional*):
+            Cast the data to `features`.
+        columns: (`List[str]`, *optional*):
+            List of columns to load, the other ones are ignored.
+        batch_size: (`int`, *optional*):
+            Size of the RecordBatches to iterate on. Default to 256.
+        token: (`str`, *optional*):
+            Optional HF token to use to download datasets.
+    """
+
+    features: Optional[datasets.Features] = None
+    columns: Optional[List[str]] = None
+    batch_size: Optional[int] = 256
+    token: Optional[str] = None
+
+
+def resolve_dataset_uris(files: List[str]) -> Dict[str, List[str]]:
+    dataset_uris = set()
+    for file_path in files:
+        path = Path(file_path)
+        if path.parent.name in {"_transactions", "_indices", "_versions"}:
+            dataset_root = path.parent.parent
+            dataset_uris.add(str(dataset_root))
+    return list(dataset_uris)
+
+
+def _fix_hf_uri(uri: str) -> str:
+    # replace the revision tag from hf uri
+    if "@" in uri:
+        matched = re.match(r"(hf://.+?)(@[0-9a-f]+)(/.*)", uri)
+        if matched:
+            uri = matched.group(1) + matched.group(3)
+    return uri
+
+
+def _fix_local_version_file(uri: str) -> str:
+    # replace symlinks with real files for _version
+    if "/_versions/" in uri and is_local_path(uri):
+        path = Path(uri)
+        if path.is_symlink():
+            data = path.read_bytes()
+            path.unlink()
+            path.write_bytes(data)
+    return uri
+
+
+class Lance(datasets.ArrowBasedBuilder, datasets.builder._CountableBuilderMixin):
+    BUILDER_CONFIG_CLASS = LanceConfig
+    METADATA_EXTENSIONS = [".idx", ".txn", ".manifest"]
+    METADATA_FILE_NAMES = ["latest_version_hint.json"]
+
+    def _info(self):
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        import lance
+        import lance.file
+
+        if not self.config.data_files:
+            raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
+        if self.repo_id:
+            api = HfApi(**dl_manager.download_config.storage_options.get("hf", {}))
+            dataset_sha = api.dataset_info(self.repo_id).sha
+            if dataset_sha != self.hash:
+                raise NotImplementedError(
+                    f"lance doesn't support loading other revisions than 'main' yet, but got {self.hash}"
+                )
+        data_files = dl_manager.download(self.config.data_files)
+
+        # TODO: remove once Lance supports HF links with revisions
+        data_files = {split: [_fix_hf_uri(file) for file in files] for split, files in data_files.items()}
+        # TODO: remove once Lance supports symlinks for _version files
+        data_files = {split: [_fix_local_version_file(file) for file in files] for split, files in data_files.items()}
+
+        splits: list[datasets.SplitGenerator] = []
+        for split_name, files in data_files.items():
+            protocol = files[0].split("://", 1)[0]
+            storage_options = dict(dl_manager.download_config.storage_options.get(protocol, {}))
+            # lance doesn't allow "token": None for hf and expects a string
+            if protocol == "hf" and storage_options.get("token") is None:
+                storage_options["token"] = get_token()
+
+            lance_dataset_uris = resolve_dataset_uris(files)
+            if lance_dataset_uris:
+                lance_datasets = [lance.dataset(uri, storage_options=storage_options) for uri in lance_dataset_uris]
+                fragments = [frag for lance_dataset in lance_datasets for frag in lance_dataset.get_fragments()]
+                if self.info.features is None:
+                    pa_schema = fragments[0]._ds.schema
+                    first_row_first_bytes = {}
+                    for field in pa_schema:
+                        if self.config.columns is not None and field.name not in self.config.columns:
+                            continue
+                        if pa.types.is_binary(field.type) or pa.types.is_large_binary(field.type):
+                            try:
+                                first_row_first_bytes[field.name] = (
+                                    lance_datasets[0].take_blobs(field.name, [0])[0].read(16)
+                                )
+                            except ValueError:
+                                first_row_first_bytes[field.name] = (
+                                    lance_datasets[0].take([0], [field.name]).to_pylist()[0][field.name][:16]
+                                )
+                splits.append(
+                    datasets.SplitGenerator(
+                        name=split_name,
+                        gen_kwargs={"fragments": fragments, "lance_files_paths": None, "lance_files": None},
+                    )
+                )
+            else:
+                lance_files = [
+                    lance.file.LanceFileReader(file, storage_options=storage_options, columns=self.config.columns)
+                    for file in files
+                ]
+                if self.info.features is None:
+                    pa_schema = lance_files[0].metadata().schema
+                    first_row_first_bytes = {
+                        field_name: value[:16]
+                        for field_name, value in lance_files[0].take_rows([0]).to_table().to_pylist()[0].items()
+                        if isinstance(value, bytes)
+                    }
+                splits.append(
+                    datasets.SplitGenerator(
+                        name=split_name,
+                        gen_kwargs={"fragments": None, "lance_files_paths": files, "lance_files": lance_files},
+                    )
+                )
+            if self.info.features is None:
+                if self.config.columns:
+                    fields = [
+                        pa_schema.field(name) for name in self.config.columns if pa_schema.get_field_index(name) != -1
+                    ]
+                    pa_schema = pa.schema(fields)
+                features = datasets.Features.from_arrow_schema(pa_schema)
+                for field_name, first_bytes in first_row_first_bytes.items():
+                    for magic_bytes_hex, _, feature_type in MAGIC_BYTES_EXTENSION_AND_FEATURE_TYPES:
+                        magic_bytes = bytes.fromhex(magic_bytes_hex)
+                        if magic_bytes in first_bytes[: len(magic_bytes) * 2]:  # allow some padding
+                            features[field_name] = feature_type
+                            break
+                self.info.features = features
+
+        return splits
+
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.info.features is not None:
+            # more expensive cast to support nested features with keys in a different order
+            # allows str <-> int/float or str to Audio for example
+            pa_table = table_cast(pa_table, self.info.features.arrow_schema)
+        return pa_table
+
+    def _generate_shards(
+        self,
+        fragments: Optional[List["lance.LanceFragment"]],
+        lance_files_paths: Optional[list[str]],
+        lance_files: Optional[List["lance.file.LanceFileReader"]],
+    ):
+        if fragments:
+            for fragment in fragments:
+                paths = [data_file.path for data_file in fragment.metadata.data_files()]
+                yield paths[0] if len(paths) == 1 else {"fragment_data_files": paths}
+        else:
+            yield from lance_files_paths
+
+    def _generate_num_examples(
+        self,
+        fragments: Optional[List["lance.LanceFragment"]],
+        lance_files_paths: Optional[list[str]],
+        lance_files: Optional[List["lance.file.LanceFileReader"]],
+    ):
+        if fragments:
+            for fragment in fragments:
+                yield fragment.count_rows()
+        else:
+            for lance_file in lance_files:
+                yield lance_file.num_rows()
+
+    def _generate_tables(
+        self,
+        fragments: Optional[List["lance.LanceFragment"]],
+        lance_files_paths: Optional[list[str]],
+        lance_files: Optional[List["lance.file.LanceFileReader"]],
+    ):
+        if fragments:
+            for frag_idx, fragment in enumerate(fragments):
+                for batch_idx, batch in enumerate(
+                    fragment.to_batches(
+                        columns=self.config.columns, batch_size=self.config.batch_size, blob_handling="all_binary"
+                    )
+                ):
+                    table = pa.Table.from_batches([batch])
+                    yield Key(frag_idx, batch_idx), self._cast_table(table)
+        else:
+            for file_idx, lance_file in enumerate(lance_files):
+                for batch_idx, batch in enumerate(lance_file.read_all(batch_size=self.config.batch_size).to_batches()):
+                    table = pa.Table.from_batches([batch])
+                    yield Key(file_idx, batch_idx), self._cast_table(table)

src/datasets/packaged_modules/meshfolder/meshfolder.py

@@ -0,0 +1,31 @@
+import datasets
+
+from ..folder_based_builder import folder_based_builder
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+class MeshFolderConfig(folder_based_builder.FolderBasedBuilderConfig):
+    """BuilderConfig for MeshFolder."""
+
+    drop_labels: bool = None
+    drop_metadata: bool = None
+
+    def __post_init__(self):
+        super().__post_init__()
+
+
+class MeshFolder(folder_based_builder.FolderBasedBuilder):
+    BASE_FEATURE = datasets.Mesh
+    BASE_COLUMN_NAME = "mesh"
+    BUILDER_CONFIG_CLASS = MeshFolderConfig
+    EXTENSIONS: list[str]  # definition at the bottom of the script
+
+
+MESH_EXTENSIONS = [
+    ".glb",
+    ".ply",
+    ".stl",
+]
+MeshFolder.EXTENSIONS = MESH_EXTENSIONS

src/datasets/packaged_modules/niftifolder/niftifolder.py

@@ -0,0 +1,23 @@
+import datasets
+
+from ..folder_based_builder import folder_based_builder
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+class NiftiFolderConfig(folder_based_builder.FolderBasedBuilderConfig):
+    """BuilderConfig for NiftiFolder."""
+
+    drop_labels: bool = None
+    drop_metadata: bool = None
+
+    def __post_init__(self):
+        super().__post_init__()
+
+
+class NiftiFolder(folder_based_builder.FolderBasedBuilder):
+    BASE_FEATURE = datasets.Nifti
+    BASE_COLUMN_NAME = "nifti"
+    BUILDER_CONFIG_CLASS = NiftiFolderConfig
+    EXTENSIONS: list[str] = [".nii"]

src/datasets/packaged_modules/pandas/pandas.py

@@ -0,0 +1,57 @@
+import warnings
+from dataclasses import dataclass
+from typing import Optional
+
+import pandas as pd
+import pyarrow as pa
+
+import datasets
+from datasets.builder import Key
+from datasets.table import table_cast
+
+
+@dataclass
+class PandasConfig(datasets.BuilderConfig):
+    """BuilderConfig for Pandas."""
+
+    features: Optional[datasets.Features] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+
+
+class Pandas(datasets.ArrowBasedBuilder):
+    BUILDER_CONFIG_CLASS = PandasConfig
+
+    def _info(self):
+        warnings.warn(
+            "The Pandas builder is deprecated and will be removed in the next major version of datasets.",
+            FutureWarning,
+        )
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        """We handle string, list and dicts in datafiles"""
+        if not self.config.data_files:
+            raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
+        data_files = dl_manager.download(self.config.data_files)
+        splits = []
+        for split_name, files in data_files.items():
+            splits.append(datasets.SplitGenerator(name=split_name, gen_kwargs={"files": files}))
+        return splits
+
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.config.features is not None:
+            # more expensive cast to support nested features with keys in a different order
+            # allows str <-> int/float or str to Audio for example
+            pa_table = table_cast(pa_table, self.config.features.arrow_schema)
+        return pa_table
+
+    def _generate_shards(self, files):
+        yield from files
+
+    def _generate_tables(self, files):
+        for i, file in enumerate(files):
+            with open(file, "rb") as f:
+                pa_table = pa.Table.from_pandas(pd.read_pickle(f))
+                yield Key(i, 0), self._cast_table(pa_table)

src/datasets/packaged_modules/parquet/parquet.py

@@ -0,0 +1,234 @@
+import gc
+from dataclasses import dataclass
+from typing import Literal, Optional, Union
+
+import pyarrow as pa
+import pyarrow.dataset as ds
+import pyarrow.parquet as pq
+from packaging import version
+
+import datasets
+import datasets.config
+from datasets.builder import Key
+from datasets.table import table_cast
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+@dataclass
+class ParquetConfig(datasets.BuilderConfig):
+    """
+    BuilderConfig for Parquet.
+
+    Args:
+        batch_size (`int`, *optional*):
+            Size of the RecordBatches to iterate on.
+            The default is the row group size (defined by the first row group).
+        columns (`list[str]`, *optional*)
+            List of columns to load, the other ones are ignored.
+            All columns are loaded by default.
+        features: (`Features`, *optional*):
+            Cast the data to `features`.
+        filters (`Union[pyarrow.dataset.Expression, list[tuple], list[list[tuple]]]`, *optional*):
+            Return only the rows matching the filter.
+            If possible the predicate will be pushed down to exploit the partition information
+            or internal metadata found in the data source, e.g. Parquet statistics.
+            Otherwise filters the loaded RecordBatches before yielding them.
+        fragment_scan_options (`pyarrow.dataset.ParquetFragmentScanOptions`, *optional*)
+            Scan-specific options for Parquet fragments.
+            This is especially useful to configure buffering and caching.
+
+            <Added version="4.2.0"/>
+        on_bad_files (`Literal["error", "warn", "skip"]`, *optional*, defaults to "error")
+            Specify what to do upon encountering a bad file (a file that can't be read). Allowed values are :
+            * 'error', raise an Exception when a bad file is encountered.
+            * 'warn', raise a warning when a bad file is encountered and skip that file.
+            * 'skip', skip bad files without raising or warning when they are encountered.
+
+            <Added version="4.2.0"/>
+
+    Example:
+
+    Load a subset of columns:
+
+    ```python
+    >>> ds = load_dataset(parquet_dataset_id, columns=["col_0", "col_1"])
+    ```
+
+    Stream data and efficiently filter data, possibly skipping entire files or row groups:
+
+    ```python
+    >>> filters = [("col_0", "==", 0)]
+    >>> ds = load_dataset(parquet_dataset_id, streaming=True, filters=filters)
+    ```
+
+    Increase the minimum request size when streaming from 32MiB (default) to 128MiB and enable prefetching:
+
+    ```python
+    >>> import pyarrow
+    >>> import pyarrow.dataset
+    >>> fragment_scan_options = pyarrow.dataset.ParquetFragmentScanOptions(
+    ...     cache_options=pyarrow.CacheOptions(
+    ...         prefetch_limit=1,
+    ...         range_size_limit=128 << 20
+    ...     ),
+    ... )
+    >>> ds = load_dataset(parquet_dataset_id, streaming=True, fragment_scan_options=fragment_scan_options)
+    ```
+
+    """
+
+    batch_size: Optional[int] = None
+    columns: Optional[list[str]] = None
+    features: Optional[datasets.Features] = None
+    filters: Optional[Union[ds.Expression, list[tuple], list[list[tuple]]]] = None
+    fragment_scan_options: Optional[ds.ParquetFragmentScanOptions] = None
+    on_bad_files: Literal["error", "warn", "skip"] = "error"
+
+    def __post_init__(self):
+        super().__post_init__()
+
+
+class Parquet(datasets.ArrowBasedBuilder):
+    BUILDER_CONFIG_CLASS = ParquetConfig
+    SLEEP_ON_THREADS_SHUTDOWNS = True  # Related to https://github.com/apache/arrow/issues/45214
+
+    def _info(self):
+        if (
+            self.config.columns is not None
+            and self.config.features is not None
+            and set(self.config.columns) != set(self.config.features)
+        ):
+            if any(col not in self.config.features for col in self.config.columns):
+                raise ValueError(
+                    "The columns and features argument must match, but got ",
+                    f"{self.config.columns} and {self.config.features}",
+                )
+            else:
+                features = datasets.Features({col: self.config.features[col] for col in self.config.columns})
+        else:
+            features = self.config.features
+        return datasets.DatasetInfo(features=features)
+
+    def _split_generators(self, dl_manager):
+        """We handle string, list and dicts in datafiles"""
+        if not self.config.data_files:
+            raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
+        dl_manager.download_config.extract_on_the_fly = True
+        data_files = dl_manager.download(self.config.data_files)
+        splits = []
+        for split_name, files in data_files.items():
+            # Infer features if they are stored in the arrow schema
+            if self.info.features is None:
+                for file in files:
+                    try:
+                        with open(file, "rb") as f:
+                            self.info.features = datasets.Features.from_arrow_schema(pq.read_schema(f))
+                            break
+                    except pa.ArrowInvalid as e:
+                        if self.config.on_bad_files == "error":
+                            logger.error(f"Failed to read schema from '{file}' with error {type(e).__name__}: {e}")
+                            raise
+                        elif self.config.on_bad_files == "warn":
+                            logger.warning(f"Skipping bad schema from '{file}'. {type(e).__name__}: {e}`")
+                        else:
+                            logger.debug(f"Skipping bad schema from '{file}'. {type(e).__name__}: {e}`")
+            if self.info.features is None:
+                raise ValueError(
+                    f"At least one valid data file must be specified, all the data_files are invalid: {self.config.data_files}"
+                )
+            splits.append(
+                datasets.SplitGenerator(
+                    name=split_name, gen_kwargs={"files": files, "row_groups_list": [None] * len(files)}
+                )
+            )
+        if self.config.columns is not None and set(self.config.columns) != set(self.info.features):
+            self.info.features = datasets.Features(
+                {col: feat for col, feat in self.info.features.items() if col in self.config.columns}
+            )
+        return splits
+
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.info.features is not None:
+            # more expensive cast to support nested features with keys in a different order
+            # allows str <-> int/float or str to Audio for example
+            pa_table = table_cast(pa_table, self.info.features.arrow_schema)
+        return pa_table
+
+    def _generate_shards(self, files, row_groups_list):
+        if not row_groups_list:
+            yield from files
+        else:
+            for file, row_groups in zip(files, row_groups_list):
+                yield {
+                    "fragment_data_file": file,
+                    "fragment_row_groups": row_groups,
+                }
+
+    def _generate_more_gen_kwargs(self, files, row_groups_list):
+        if not row_groups_list or any(row_group is None for row_group in row_groups_list):
+            parquet_file_format = ds.ParquetFileFormat(default_fragment_scan_options=self.config.fragment_scan_options)
+            for file in files:
+                with open(file, "rb") as f:
+                    parquet_fragment = parquet_file_format.make_fragment(f)
+                    yield {
+                        "files": [file] * parquet_fragment.num_row_groups,
+                        "row_groups_list": [
+                            (row_group_id,) for row_group_id in range(parquet_fragment.num_row_groups)
+                        ],
+                    }
+        else:
+            for file, row_groups in zip(files, row_groups_list):
+                yield {"files": [file], "row_groups_list": [row_groups]}
+
+    def _generate_tables(self, files, row_groups_list):
+        if self.config.features is not None and self.config.columns is not None:
+            if sorted(field.name for field in self.info.features.arrow_schema) != sorted(self.config.columns):
+                raise ValueError(
+                    f"Tried to load parquet data with columns '{self.config.columns}' with mismatching features '{self.info.features}'"
+                )
+        filter_expr = (
+            pq.filters_to_expression(self.config.filters)
+            if isinstance(self.config.filters, list)
+            else self.config.filters
+        )
+        parquet_file_format = ds.ParquetFileFormat(default_fragment_scan_options=self.config.fragment_scan_options)
+        for file_idx, (file, row_groups) in enumerate(zip(files, row_groups_list)):
+            try:
+                with open(file, "rb") as f:
+                    parquet_fragment = parquet_file_format.make_fragment(f)
+                    fragment_is_closed = False
+                    try:
+                        if row_groups is not None:
+                            parquet_fragment = parquet_fragment.subset(row_group_ids=row_groups)
+                        if parquet_fragment.row_groups:
+                            batch_size = self.config.batch_size or parquet_fragment.row_groups[0].num_rows
+                            for batch_idx, record_batch in enumerate(
+                                parquet_fragment.to_batches(
+                                    batch_size=batch_size,
+                                    columns=self.config.columns,
+                                    filter=filter_expr,
+                                    batch_readahead=0,
+                                    fragment_readahead=0,
+                                )
+                            ):
+                                pa_table = pa.Table.from_batches([record_batch])
+                                # Uncomment for debugging (will print the Arrow table size and elements)
+                                # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
+                                # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
+                                yield Key(file_idx, batch_idx), self._cast_table(pa_table)
+                            fragment_is_closed = True
+                    finally:
+                        # Fix for https://github.com/apache/arrow/issues/45214
+                        if not fragment_is_closed and datasets.config.PYARROW_VERSION <= version.parse("24.0.0"):
+                            del parquet_fragment
+                            gc.collect()
+            except (pa.ArrowInvalid, ValueError) as e:
+                if self.config.on_bad_files == "error":
+                    logger.error(f"Failed to read file '{file}' with error {type(e).__name__}: {e}")
+                    raise
+                elif self.config.on_bad_files == "warn":
+                    logger.warning(f"Skipping bad file '{file}'. {type(e).__name__}: {e}`")
+                else:
+                    logger.debug(f"Skipping bad file '{file}'. {type(e).__name__}: {e}`")

src/datasets/packaged_modules/pdffolder/pdffolder.py

@@ -0,0 +1,23 @@
+import datasets
+
+from ..folder_based_builder import folder_based_builder
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+class PdfFolderConfig(folder_based_builder.FolderBasedBuilderConfig):
+    """BuilderConfig for ImageFolder."""
+
+    drop_labels: bool = None
+    drop_metadata: bool = None
+
+    def __post_init__(self):
+        super().__post_init__()
+
+
+class PdfFolder(folder_based_builder.FolderBasedBuilder):
+    BASE_FEATURE = datasets.Pdf
+    BASE_COLUMN_NAME = "pdf"
+    BUILDER_CONFIG_CLASS = PdfFolderConfig
+    EXTENSIONS: list[str] = [".pdf"]

src/datasets/packaged_modules/spark/spark.py

@@ -0,0 +1,367 @@
+import os
+import posixpath
+import uuid
+from collections.abc import Iterable
+from dataclasses import dataclass
+from itertools import islice
+from typing import TYPE_CHECKING, Optional, Union
+
+import numpy as np
+import pyarrow as pa
+
+import datasets
+from datasets.arrow_writer import ArrowWriter, ParquetWriter
+from datasets.config import MAX_SHARD_SIZE
+from datasets.filesystems import (
+    is_remote_filesystem,
+    rename,
+)
+from datasets.iterable_dataset import _BaseExamplesIterable
+from datasets.utils import experimental
+from datasets.utils.py_utils import convert_file_size_to_int
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+if TYPE_CHECKING:
+    import pyspark
+    import pyspark.sql
+
+
+@dataclass
+class SparkConfig(datasets.BuilderConfig):
+    """BuilderConfig for Spark."""
+
+    features: Optional[datasets.Features] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+
+
+def _reorder_dataframe_by_partition(df: "pyspark.sql.DataFrame", new_partition_order: list[int]):
+    df_combined = df.select("*").where(f"part_id = {new_partition_order[0]}")
+    for partition_id in new_partition_order[1:]:
+        partition_df = df.select("*").where(f"part_id = {partition_id}")
+        df_combined = df_combined.union(partition_df)
+    return df_combined
+
+
+def _generate_iterable_examples(
+    df: "pyspark.sql.DataFrame",
+    partition_order: list[int],
+    state_dict: Optional[dict] = None,
+):
+    import pyspark
+
+    df_with_partition_id = df.select("*", pyspark.sql.functions.spark_partition_id().alias("part_id"))
+    partition_idx_start = state_dict["partition_idx"] if state_dict else 0
+    partition_df = _reorder_dataframe_by_partition(df_with_partition_id, partition_order[partition_idx_start:])
+    # pipeline next partition in parallel to hide latency
+    rows = partition_df.toLocalIterator(prefetchPartitions=True)
+    curr_partition = None
+    row_id = state_dict["partition_example_idx"] if state_dict else 0
+    for row in islice(rows, row_id, None):
+        row_as_dict = row.asDict()
+        part_id = row_as_dict["part_id"]
+        row_as_dict.pop("part_id")
+        if curr_partition != part_id:
+            if state_dict and curr_partition is not None:
+                state_dict["partition_idx"] += 1
+            curr_partition = part_id
+            row_id = 0
+        if state_dict:
+            state_dict["partition_example_idx"] = row_id + 1
+        yield (part_id, row_id), row_as_dict
+        row_id += 1
+
+
+class SparkExamplesIterable(_BaseExamplesIterable):
+    def __init__(
+        self,
+        df: "pyspark.sql.DataFrame",
+        partition_order=None,
+    ):
+        super().__init__()
+        self.df = df
+        self.partition_order = partition_order or range(self.df.rdd.getNumPartitions())
+
+    def _init_state_dict(self) -> dict:
+        self._state_dict = {"partition_idx": 0, "partition_example_idx": 0}
+        return self._state_dict
+
+    @experimental
+    def load_state_dict(self, state_dict: dict) -> dict:
+        return super().load_state_dict(state_dict)
+
+    def __iter__(self):
+        yield from _generate_iterable_examples(self.df, self.partition_order, self._state_dict)
+
+    def shuffle_data_sources(self, generator: np.random.Generator) -> "SparkExamplesIterable":
+        partition_order = list(range(self.df.rdd.getNumPartitions()))
+        generator.shuffle(partition_order)
+        return SparkExamplesIterable(self.df, partition_order=partition_order)
+
+    def shard_data_sources(self, num_shards: int, index: int, contiguous=True) -> "SparkExamplesIterable":
+        partition_order = self.split_shard_indices_by_worker(num_shards=num_shards, index=index, contiguous=contiguous)
+        return SparkExamplesIterable(self.df, partition_order=partition_order)
+
+    @property
+    def num_shards(self) -> int:
+        return len(self.partition_order)
+
+
+class Spark(datasets.DatasetBuilder):
+    BUILDER_CONFIG_CLASS = SparkConfig
+
+    def __init__(
+        self,
+        df: "pyspark.sql.DataFrame",
+        cache_dir: str = None,
+        working_dir: str = None,
+        **config_kwargs,
+    ):
+        import pyspark
+
+        self._spark = pyspark.sql.SparkSession.builder.getOrCreate()
+        self.df = df
+        self._working_dir = working_dir
+
+        super().__init__(
+            cache_dir=cache_dir,
+            config_name=str(self.df.semanticHash()),
+            **config_kwargs,
+        )
+
+    def _validate_cache_dir(self):
+        # Define this so that we don't reference self in create_cache_and_write_probe, which will result in a pickling
+        # error due to pickling the SparkContext.
+        cache_dir = self._cache_dir
+
+        # Returns the path of the created file.
+        def create_cache_and_write_probe(context):
+            # makedirs with exist_ok will recursively create the directory. It will not throw an error if directories
+            # already exist.
+            os.makedirs(cache_dir, exist_ok=True)
+            probe_file = os.path.join(cache_dir, "fs_test" + uuid.uuid4().hex)
+            # Opening the file in append mode will create a new file unless it already exists, in which case it will not
+            # change the file contents.
+            open(probe_file, "a")
+            return [probe_file]
+
+        if self._spark.conf.get("spark.master", "").startswith("local"):
+            return
+
+        # If the cluster is multi-node, make sure that the user provided a cache_dir and that it is on an NFS
+        # accessible to the driver.
+        # TODO: Stream batches to the driver using ArrowCollectSerializer instead of throwing an error.
+        if self._cache_dir:
+            probe = (
+                self._spark.sparkContext.parallelize(range(1), 1).mapPartitions(create_cache_and_write_probe).collect()
+            )
+            if os.path.isfile(probe[0]):
+                return
+
+        raise ValueError(
+            "When using Dataset.from_spark on a multi-node cluster, the driver and all workers should be able to access cache_dir"
+        )
+
+    def _info(self):
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager: datasets.download.download_manager.DownloadManager):
+        return [datasets.SplitGenerator(name=datasets.Split.TRAIN)]
+
+    def _repartition_df_if_needed(self, max_shard_size):
+        import pyspark
+
+        def get_arrow_batch_size(it):
+            for batch in it:
+                yield pa.RecordBatch.from_pydict({"batch_bytes": [batch.nbytes]})
+
+        df_num_rows = self.df.count()
+        sample_num_rows = df_num_rows if df_num_rows <= 100 else 100
+        # Approximate the size of each row (in Arrow format) by averaging over a max-100-row sample.
+        approx_bytes_per_row = (
+            self.df.limit(sample_num_rows)
+            .repartition(1)
+            .mapInArrow(get_arrow_batch_size, "batch_bytes: long")
+            .agg(pyspark.sql.functions.sum("batch_bytes").alias("sample_bytes"))
+            .collect()[0]
+            .sample_bytes
+            / sample_num_rows
+        )
+        approx_total_size = approx_bytes_per_row * df_num_rows
+        if approx_total_size > max_shard_size:
+            # Make sure there is at least one row per partition.
+            new_num_partitions = min(df_num_rows, int(approx_total_size / max_shard_size))
+            self.df = self.df.repartition(new_num_partitions)
+
+    def _prepare_split_single(
+        self,
+        fpath: str,
+        file_format: str,
+        max_shard_size: int,
+    ) -> Iterable[tuple[int, bool, Union[int, tuple]]]:
+        import pyspark
+
+        writer_class = ParquetWriter if file_format == "parquet" else ArrowWriter
+        working_fpath = os.path.join(self._working_dir, os.path.basename(fpath)) if self._working_dir else fpath
+        embed_local_files = file_format == "parquet"
+
+        # Define these so that we don't reference self in write_arrow, which will result in a pickling error due to
+        # pickling the SparkContext.
+        features = self.config.features
+        writer_batch_size = self._writer_batch_size
+        storage_options = self._fs.storage_options
+
+        def write_arrow(it):
+            # Within the same SparkContext, no two task attempts will share the same attempt ID.
+            task_id = pyspark.TaskContext().taskAttemptId()
+            first_batch = next(it, None)
+            if first_batch is None:
+                # Some partitions might not receive any data.
+                return pa.RecordBatch.from_arrays(
+                    [[task_id], [0], [0]],
+                    names=["task_id", "num_examples", "num_bytes"],
+                )
+            shard_id = 0
+            writer = writer_class(
+                features=features,
+                path=working_fpath.replace("SSSSS", f"{shard_id:05d}").replace("TTTTT", f"{task_id:05d}"),
+                writer_batch_size=writer_batch_size,
+                storage_options=storage_options,
+                embed_local_files=embed_local_files,
+            )
+            table = pa.Table.from_batches([first_batch])
+            writer.write_table(table)
+            for batch in it:
+                if max_shard_size is not None and writer._num_bytes >= max_shard_size:
+                    num_examples, num_bytes = writer.finalize()
+                    writer.close()
+                    yield pa.RecordBatch.from_arrays(
+                        [[task_id], [num_examples], [num_bytes]],
+                        names=["task_id", "num_examples", "num_bytes"],
+                    )
+                    shard_id += 1
+                    writer = writer_class(
+                        features=writer._features,
+                        path=working_fpath.replace("SSSSS", f"{shard_id:05d}").replace("TTTTT", f"{task_id:05d}"),
+                        writer_batch_size=writer_batch_size,
+                        storage_options=storage_options,
+                        embed_local_files=embed_local_files,
+                    )
+                table = pa.Table.from_batches([batch])
+                writer.write_table(table)
+
+            if writer._num_bytes > 0:
+                num_examples, num_bytes = writer.finalize()
+                writer.close()
+                yield pa.RecordBatch.from_arrays(
+                    [[task_id], [num_examples], [num_bytes]],
+                    names=["task_id", "num_examples", "num_bytes"],
+                )
+
+            if working_fpath != fpath:
+                for file in os.listdir(os.path.dirname(working_fpath)):
+                    dest = os.path.join(os.path.dirname(fpath), os.path.basename(file))
+                    shutil.move(file, dest)
+
+        stats = (
+            self.df.mapInArrow(write_arrow, "task_id: long, num_examples: long, num_bytes: long")
+            .groupBy("task_id")
+            .agg(
+                pyspark.sql.functions.sum("num_examples").alias("total_num_examples"),
+                pyspark.sql.functions.sum("num_bytes").alias("total_num_bytes"),
+                pyspark.sql.functions.count("num_bytes").alias("num_shards"),
+                pyspark.sql.functions.collect_list("num_examples").alias("shard_lengths"),
+            )
+            .collect()
+        )
+        for row in stats:
+            yield row.task_id, (row.total_num_examples, row.total_num_bytes, row.num_shards, row.shard_lengths)
+
+    def _prepare_split(
+        self,
+        split_generator: "datasets.SplitGenerator",
+        file_format: str = "arrow",
+        max_shard_size: Optional[Union[str, int]] = None,
+        num_proc: Optional[int] = None,
+        **kwargs,
+    ):
+        self._validate_cache_dir()
+
+        max_shard_size = convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE)
+        self._repartition_df_if_needed(max_shard_size)
+        is_local = not is_remote_filesystem(self._fs)
+        path_join = os.path.join if is_local else posixpath.join
+
+        SUFFIX = "-TTTTT-SSSSS-of-NNNNN"
+        fname = f"{self.name}-{split_generator.name}{SUFFIX}.{file_format}"
+        fpath = path_join(self._output_dir, fname)
+
+        total_num_examples = 0
+        total_num_bytes = 0
+        total_shards = 0
+        task_id_and_num_shards = []
+        all_shard_lengths = []
+
+        for task_id, content in self._prepare_split_single(fpath, file_format, max_shard_size):
+            (
+                num_examples,
+                num_bytes,
+                num_shards,
+                shard_lengths,
+            ) = content
+            if num_bytes > 0:
+                total_num_examples += num_examples
+                total_num_bytes += num_bytes
+                total_shards += num_shards
+                task_id_and_num_shards.append((task_id, num_shards))
+                all_shard_lengths.extend(shard_lengths)
+
+        split_generator.split_info.num_examples = total_num_examples
+        split_generator.split_info.num_bytes = total_num_bytes
+
+        # should rename everything at the end
+        logger.debug(f"Renaming {total_shards} shards.")
+        if total_shards > 1:
+            split_generator.split_info.shard_lengths = all_shard_lengths
+
+            # Define fs outside of _rename_shard so that we don't reference self in the function, which will result in a
+            # pickling error due to pickling the SparkContext.
+            fs = self._fs
+
+            # use the -SSSSS-of-NNNNN pattern
+            def _rename_shard(
+                task_id: int,
+                shard_id: int,
+                global_shard_id: int,
+            ):
+                rename(
+                    fs,
+                    fpath.replace("SSSSS", f"{shard_id:05d}").replace("TTTTT", f"{task_id:05d}"),
+                    fpath.replace("TTTTT-SSSSS", f"{global_shard_id:05d}").replace("NNNNN", f"{total_shards:05d}"),
+                )
+
+            args = []
+            global_shard_id = 0
+            for i in range(len(task_id_and_num_shards)):
+                task_id, num_shards = task_id_and_num_shards[i]
+                for shard_id in range(num_shards):
+                    args.append([task_id, shard_id, global_shard_id])
+                    global_shard_id += 1
+            self._spark.sparkContext.parallelize(args, len(args)).map(lambda args: _rename_shard(*args)).collect()
+        else:
+            # don't use any pattern
+            shard_id = 0
+            task_id = task_id_and_num_shards[0][0]
+            self._rename(
+                fpath.replace("SSSSS", f"{shard_id:05d}").replace("TTTTT", f"{task_id:05d}"),
+                fpath.replace(SUFFIX, ""),
+            )
+
+    def _get_examples_iterable_for_split(
+        self,
+        split_generator: "datasets.SplitGenerator",
+    ) -> SparkExamplesIterable:
+        return SparkExamplesIterable(self.df)

src/datasets/packaged_modules/sql/sql.py

@@ -0,0 +1,120 @@
+import sys
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional, Union
+
+import pandas as pd
+import pyarrow as pa
+
+import datasets
+import datasets.config
+from datasets.builder import Key
+from datasets.features.features import require_storage_cast
+from datasets.table import table_cast
+
+
+if TYPE_CHECKING:
+    import sqlite3
+
+    import sqlalchemy
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+@dataclass
+class SqlConfig(datasets.BuilderConfig):
+    """BuilderConfig for SQL."""
+
+    sql: Union[str, "sqlalchemy.sql.Selectable"] = None
+    con: Union[str, "sqlalchemy.engine.Connection", "sqlalchemy.engine.Engine", "sqlite3.Connection"] = None
+    index_col: Optional[Union[str, list[str]]] = None
+    coerce_float: bool = True
+    params: Optional[Union[list, tuple, dict]] = None
+    parse_dates: Optional[Union[list, dict]] = None
+    columns: Optional[list[str]] = None
+    chunksize: Optional[int] = 10_000
+    features: Optional[datasets.Features] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.sql is None:
+            raise ValueError("sql must be specified")
+        if self.con is None:
+            raise ValueError("con must be specified")
+
+    def create_config_id(
+        self,
+        config_kwargs: dict,
+        custom_features: Optional[datasets.Features] = None,
+    ) -> str:
+        config_kwargs = config_kwargs.copy()
+        # We need to stringify the Selectable object to make its hash deterministic
+
+        # The process of stringifying is explained here: http://docs.sqlalchemy.org/en/latest/faq/sqlexpressions.html
+        sql = config_kwargs["sql"]
+        if not isinstance(sql, str):
+            if datasets.config.SQLALCHEMY_AVAILABLE and "sqlalchemy" in sys.modules:
+                import sqlalchemy
+
+                if isinstance(sql, sqlalchemy.sql.Selectable):
+                    engine = sqlalchemy.create_engine(config_kwargs["con"].split("://")[0] + "://")
+                    sql_str = str(sql.compile(dialect=engine.dialect))
+                    config_kwargs["sql"] = sql_str
+                else:
+                    raise TypeError(
+                        f"Supported types for 'sql' are string and sqlalchemy.sql.Selectable but got {type(sql)}: {sql}"
+                    )
+            else:
+                raise TypeError(
+                    f"Supported types for 'sql' are string and sqlalchemy.sql.Selectable but got {type(sql)}: {sql}"
+                )
+        con = config_kwargs["con"]
+        if not isinstance(con, str):
+            config_kwargs["con"] = id(con)
+            logger.info(
+                f"SQL connection 'con' of type {type(con)} couldn't be hashed properly. To enable hashing, specify 'con' as URI string instead."
+            )
+
+        return super().create_config_id(config_kwargs, custom_features=custom_features)
+
+    @property
+    def pd_read_sql_kwargs(self):
+        pd_read_sql_kwargs = {
+            "index_col": self.index_col,
+            "columns": self.columns,
+            "params": self.params,
+            "coerce_float": self.coerce_float,
+            "parse_dates": self.parse_dates,
+        }
+        return pd_read_sql_kwargs
+
+
+class Sql(datasets.ArrowBasedBuilder):
+    BUILDER_CONFIG_CLASS = SqlConfig
+
+    def _info(self):
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        return [datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={})]
+
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.config.features is not None:
+            schema = self.config.features.arrow_schema
+            if all(not require_storage_cast(feature) for feature in self.config.features.values()):
+                # cheaper cast
+                pa_table = pa.Table.from_arrays([pa_table[field.name] for field in schema], schema=schema)
+            else:
+                # more expensive cast; allows str <-> int/float or str to Audio for example
+                pa_table = table_cast(pa_table, schema)
+        return pa_table
+
+    def _generate_tables(self):
+        chunksize = self.config.chunksize
+        sql_reader = pd.read_sql(
+            self.config.sql, self.config.con, chunksize=chunksize, **self.config.pd_read_sql_kwargs
+        )
+        sql_reader = [sql_reader] if chunksize is None else sql_reader
+        for chunk_idx, df in enumerate(sql_reader):
+            pa_table = pa.Table.from_pandas(df)
+            yield Key(0, chunk_idx), self._cast_table(pa_table)

src/datasets/packaged_modules/text/text.py

@@ -0,0 +1,137 @@
+from dataclasses import dataclass
+from io import StringIO
+from typing import Literal, Optional
+
+import pyarrow as pa
+
+import datasets
+from datasets.builder import Key
+from datasets.features.features import require_storage_cast
+from datasets.table import table_cast
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+@dataclass
+class TextConfig(datasets.BuilderConfig):
+    """BuilderConfig for text files.
+
+    Args:
+        features: (`Features`, *optional*):
+            Cast the data to `features`.
+        encoding: (`str`, defaults to "utf-8"):
+            Encoding to decode the file.
+        encoding_errors: (`str`, *optional*):
+            Argument to define what to do in case of encoding error.
+            This is the same as the `error` argument in `open()`.
+        chunksize: (`Features`, *optional*, defaults to "10MB"):
+            Chunk size to read the data.
+        keep_linebreaks: (`bool`, defaults to False):
+            Whether to keep line breaks.
+        sample_by (`Literal["line", "paragraph", "document"]`, defaults to "line"):
+            Whether to load data per line, praragraph or document.
+            By default one row in the dataset = one line.
+    """
+
+    features: Optional[datasets.Features] = None
+    encoding: str = "utf-8"
+    encoding_errors: Optional[str] = None
+    chunksize: int = 10 << 20  # 10MB
+    keep_linebreaks: bool = False
+    sample_by: Literal["line", "paragraph", "document"] = "line"
+
+
+class Text(datasets.ArrowBasedBuilder):
+    BUILDER_CONFIG_CLASS = TextConfig
+
+    def _info(self):
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        """The `data_files` kwarg in load_dataset() can be a str, List[str], Dict[str,str], or Dict[str,List[str]].
+
+        If str or List[str], then the dataset returns only the 'train' split.
+        If dict, then keys should be from the `datasets.Split` enum.
+        """
+        if not self.config.data_files:
+            raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
+        dl_manager.download_config.extract_on_the_fly = True
+        base_data_files = dl_manager.download(self.config.data_files)
+        extracted_data_files = dl_manager.extract(base_data_files)
+        splits = []
+        for split_name, files in extracted_data_files.items():
+            files_iterables = [dl_manager.iter_files(file) for file in files]
+            splits.append(
+                datasets.SplitGenerator(
+                    name=split_name,
+                    gen_kwargs={"files_iterables": files_iterables, "base_files": base_data_files[split_name]},
+                )
+            )
+        return splits
+
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.config.features is not None:
+            schema = self.config.features.arrow_schema
+            if all(not require_storage_cast(feature) for feature in self.config.features.values()):
+                # cheaper cast
+                pa_table = pa_table.cast(schema)
+            else:
+                # more expensive cast; allows str <-> int/float or str to Audio for example
+                pa_table = table_cast(pa_table, schema)
+            return pa_table
+        else:
+            return pa_table.cast(pa.schema({"text": pa.string()}))
+
+    def _generate_shards(self, base_files, files_iterables):
+        yield from base_files
+
+    def _generate_tables(self, base_files, files_iterables):
+        pa_table_names = list(self.config.features) if self.config.features is not None else ["text"]
+        for shard_idx, files_iterable in enumerate(files_iterables):
+            for file in files_iterable:
+                # open in text mode, by default translates universal newlines ("\n", "\r\n" and "\r") into "\n"
+                with open(file, encoding=self.config.encoding, errors=self.config.encoding_errors) as f:
+                    if self.config.sample_by == "line":
+                        batch_idx = 0
+                        while True:
+                            batch = f.read(self.config.chunksize)
+                            if not batch:
+                                break
+                            batch += f.readline()  # finish current line
+                            # StringIO.readlines, by default splits only on "\n" (and keeps line breaks)
+                            batch = StringIO(batch).readlines()
+                            if not self.config.keep_linebreaks:
+                                batch = [line.rstrip("\n") for line in batch]
+                            pa_table = pa.Table.from_arrays([pa.array(batch)], names=pa_table_names)
+                            # Uncomment for debugging (will print the Arrow table size and elements)
+                            # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
+                            # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
+                            yield Key(shard_idx, batch_idx), self._cast_table(pa_table)
+                            batch_idx += 1
+                    elif self.config.sample_by == "paragraph":
+                        batch_idx = 0
+                        batch = ""
+                        while True:
+                            new_batch = f.read(self.config.chunksize)
+                            if not new_batch:
+                                break
+                            batch += new_batch
+                            batch += f.readline()  # finish current line
+                            batch = batch.split("\n\n")
+                            pa_table = pa.Table.from_arrays(
+                                [pa.array([example for example in batch[:-1] if example])], names=pa_table_names
+                            )
+                            # Uncomment for debugging (will print the Arrow table size and elements)
+                            # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
+                            # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
+                            yield Key(shard_idx, batch_idx), self._cast_table(pa_table)
+                            batch_idx += 1
+                            batch = batch[-1]
+                        if batch:
+                            pa_table = pa.Table.from_arrays([pa.array([batch])], names=pa_table_names)
+                            yield (shard_idx, batch_idx), self._cast_table(pa_table)
+                    elif self.config.sample_by == "document":
+                        text = f.read()
+                        pa_table = pa.Table.from_arrays([pa.array([text])], names=pa_table_names)
+                        yield Key(shard_idx, 0), self._cast_table(pa_table)

src/datasets/packaged_modules/tsfile/tsfile.py

@@ -0,0 +1,773 @@
+"""TsFile (table model) packaged builder — per-device wide format.
+
+Each output row corresponds to a single device (identified by its TAG values).
+The ``time`` column and every FIELD column are Arrow ``list<...>`` columns
+holding the entire time series for that device. When the same device is
+present in multiple TsFiles within a split, its data is merged across files
+and the resulting lists are sorted in ascending time order.
+
+Output schema layout::
+
+    <tag1>:    string
+    <tag2>:    string                   (one column per TAG)
+    ...
+    time:      list<timestamp[unit, tz]>
+    <field1>:  list<original_type>      (one column per FIELD)
+    <field2>:  list<original_type>
+    ...
+
+Reading model
+-------------
+Data is fetched **per device** via ``TsFileReader.query_table`` with a
+push-down ``tag_filter``. For each split the builder:
+
+1. Opens every input file once, calls ``get_all_devices`` to enumerate the
+   ``(tag-tuple) → [files]`` index across all shards.
+2. Iterates the index in stable order. For each device, streams Arrow
+   batches from every contributing file, concatenates and sorts by time,
+   and emits one wide row.
+
+Peak memory is bounded by **one device's** total payload across the split,
+not by the split's total size.
+"""
+
+from __future__ import annotations
+
+import datetime as _dt
+from dataclasses import dataclass
+from typing import Any, Literal, Optional
+
+import numpy as np
+import pyarrow as pa
+
+import datasets
+from datasets.builder import Key
+from datasets.table import table_cast
+from datasets.utils.tqdm import tqdm
+
+
+logger = datasets.utils.logging.get_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Type helpers
+# ---------------------------------------------------------------------------
+
+
+def _arrow_type(ts_dtype, *, unit: str, tz: Optional[str]) -> pa.DataType:
+    """Map a tsfile ``TSDataType`` to its Arrow representation."""
+    from tsfile.constants import TSDataType
+
+    return {
+        TSDataType.BOOLEAN: pa.bool_(),
+        TSDataType.INT32: pa.int32(),
+        TSDataType.INT64: pa.int64(),
+        TSDataType.FLOAT: pa.float32(),
+        TSDataType.DOUBLE: pa.float64(),
+        TSDataType.TEXT: pa.string(),
+        TSDataType.STRING: pa.string(),
+        TSDataType.TIMESTAMP: pa.timestamp(unit, tz=tz),
+        TSDataType.DATE: pa.date32(),
+        TSDataType.BLOB: pa.binary(),
+    }.get(ts_dtype, pa.string())
+
+
+def _promote_tsdatatype(a, b):
+    """Return the widest of two ``TSDataType`` values.
+
+    Mirrors IoTDB's ``ALTER COLUMN ... SET DATA TYPE`` rules:
+
+    - ``INT32 → INT64 → DOUBLE``
+    - ``INT32 → FLOAT → DOUBLE``
+
+    ``INT64`` and ``FLOAT`` cannot widen losslessly into either, so their
+    join is ``DOUBLE``. Non-numeric or otherwise unrelated pairs raise.
+    """
+    if a == b:
+        return a
+
+    from tsfile.constants import TSDataType
+
+    table = {
+        (TSDataType.INT32, TSDataType.INT64): TSDataType.INT64,
+        (TSDataType.INT32, TSDataType.FLOAT): TSDataType.FLOAT,
+        (TSDataType.INT32, TSDataType.DOUBLE): TSDataType.DOUBLE,
+        (TSDataType.INT64, TSDataType.FLOAT): TSDataType.DOUBLE,
+        (TSDataType.INT64, TSDataType.DOUBLE): TSDataType.DOUBLE,
+        (TSDataType.FLOAT, TSDataType.DOUBLE): TSDataType.DOUBLE,
+    }
+    if (a, b) in table:
+        return table[(a, b)]
+    if (b, a) in table:
+        return table[(b, a)]
+    raise ValueError(
+        f"Incompatible column types across files: {a.name} vs {b.name}. "
+        "Only numeric widening (INT32→INT64→DOUBLE, INT32→FLOAT→DOUBLE) is supported."
+    )
+
+
+def _to_epoch(value: Any, unit: str) -> int:
+    """Coerce a timestamp boundary to an integer epoch in ``unit``.
+
+    Accepts ``int`` (raw epoch in ``unit``), ``datetime``/``date``,
+    ISO-8601 ``str``, or any ``pa.Scalar`` of timestamp type.
+    """
+    if isinstance(value, bool):  # bool is a subclass of int; reject explicitly
+        raise TypeError(f"start_time/end_time must be a timestamp, got bool: {value!r}")
+    if isinstance(value, int):
+        return value
+    try:
+        # Normalize the various input shapes into something pa.scalar() can absorb
+        # under a `timestamp[unit]` target type.
+        if isinstance(value, _dt.datetime):
+            if value.tzinfo is not None:
+                value = value.astimezone(_dt.timezone.utc).replace(tzinfo=None)
+        elif isinstance(value, _dt.date):
+            value = _dt.datetime(value.year, value.month, value.day)
+        elif isinstance(value, str):
+            value = _dt.datetime.fromisoformat(value)
+        return pa.scalar(value, type=pa.timestamp(unit)).value
+    except (pa.ArrowInvalid, pa.ArrowTypeError, TypeError, ValueError) as e:
+        raise TypeError(
+            f"start_time/end_time must be a datetime, date, pa.TimestampScalar, "
+            f"ISO-8601 str, or int epoch; got {type(value).__name__}: {value!r}"
+        ) from e
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class TsFileConfig(datasets.BuilderConfig):
+    """BuilderConfig for TsFile (table model) — per-device wide format.
+
+    Args:
+        table_name (`str`, *optional*):
+            Name of the table to read. When unset, the first table found in
+            the first valid file is used. Lookups are case-insensitive.
+        columns (`list[str]`, *optional*):
+            Subset of FIELD columns to keep. TAG columns and the TIME column
+            are *always* returned (they identify the device / its timeline
+            and cannot be excluded). Names that refer to TAG or TIME columns,
+            or to fields absent from every file, resolve quietly: TAGs/TIME
+            are emitted as usual, and never-seen fields become all-null list
+            columns. When unset, all FIELDs are returned.
+        start_time, end_time (`datetime`, `date`, `pa.TimestampScalar`, ISO-8601 `str`, or `int`, *optional*):
+            Inclusive timestamp range. Either bound may be omitted.
+            ``datetime`` values are taken in their own tz (UTC if naive);
+            ``int`` is interpreted as a raw epoch in ``timestamp_unit``.
+        input_batch_size (`int`, *optional*, defaults to 65_536):
+            Maximum number of rows fetched per Arrow batch from
+            ``TsFileReader.query_table``. Controls peak memory while
+            streaming a single device.
+        output_batch_size (`int`, *optional*, defaults to 32):
+            Number of devices (output dataset rows) packed into each Arrow
+            record batch yielded to the writer. Also the granularity at
+            which the dataset progress bar advances; smaller values give
+            more responsive feedback on slow per-device reads, larger ones
+            reduce per-batch overhead.
+        features (`Features`, *optional*):
+            Final Features schema. When provided, the metadata scan over
+            input files is skipped.
+        on_bad_files (`Literal["error", "warn", "skip"]`, *optional*, defaults to "error"):
+            What to do if a file cannot be opened or lacks the requested table.
+        timestamp_unit (`Literal["s", "ms", "us", "ns"]`, *optional*, defaults to "ms"):
+            Time unit for the timestamp column. IoTDB defaults to milliseconds.
+        timestamp_tz (`str`, *optional*):
+            Time zone for the timestamp column. ``None`` means timezone-naive.
+    """
+
+    table_name: Optional[str] = None
+    columns: Optional[list[str]] = None
+    start_time: Optional[Any] = None
+    end_time: Optional[Any] = None
+    input_batch_size: int = 65_536
+    output_batch_size: int = 32
+    features: Optional[datasets.Features] = None
+    on_bad_files: Literal["error", "warn", "skip"] = "error"
+    timestamp_unit: Literal["s", "ms", "us", "ns"] = "ms"
+    timestamp_tz: Optional[str] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.input_batch_size is None or self.input_batch_size <= 0:
+            raise ValueError(f"`input_batch_size` must be a positive integer, got {self.input_batch_size}")
+        if self.output_batch_size is None or self.output_batch_size <= 0:
+            raise ValueError(f"`output_batch_size` must be a positive integer, got {self.output_batch_size}")
+        if self.columns is not None and len(self.columns) == 0:
+            raise ValueError("`columns` must be a non-empty list when provided.")
+        if self.timestamp_unit not in ("s", "ms", "us", "ns"):
+            raise ValueError(f"`timestamp_unit` must be one of 's', 'ms', 'us', 'ns', got {self.timestamp_unit!r}")
+        if self.on_bad_files not in ("error", "warn", "skip"):
+            raise ValueError(f"`on_bad_files` must be one of 'error', 'warn', 'skip', got {self.on_bad_files!r}")
+        if self.start_time is not None:
+            self.start_time = _to_epoch(self.start_time, self.timestamp_unit)
+        if self.end_time is not None:
+            self.end_time = _to_epoch(self.end_time, self.timestamp_unit)
+
+
+# ---------------------------------------------------------------------------
+# Internal sentinels
+# ---------------------------------------------------------------------------
+
+
+class _SkipSplit(Exception):
+    """Raised internally to abort emitting a split entirely."""
+
+
+class _MissingTableError(ValueError):
+    def __init__(self, table: Optional[str], available):
+        super().__init__(f"Table {table!r} not found in file. Available tables: {available}")
+
+
+_TSFILE_MAGIC = b"TsFile"
+
+
+# ---------------------------------------------------------------------------
+# Builder
+# ---------------------------------------------------------------------------
+
+
+class TsFile(datasets.ArrowBasedBuilder):
+    """Per-device wide-format builder for TsFile (table model)."""
+
+    BUILDER_CONFIG_CLASS = TsFileConfig
+
+    # ----- builder hooks ------------------------------------------------
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._table: Optional[str] = None
+        self._time_col: str = "time"
+        self._tag_cols: list[str] = []
+        self._field_inner: dict[str, pa.DataType] = {}
+        self._requested_fields: Optional[list[str]] = None  # lowercased
+
+    def _info(self):
+        if (
+            self.config.columns is not None
+            and self.config.features is not None
+            and not set(self.config.columns).issubset(set(self.config.features))
+        ):
+            raise ValueError(
+                "Every entry in `columns` must also appear in `features`, but got "
+                f"columns={self.config.columns} and features={list(self.config.features)}"
+            )
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        if not self.config.data_files:
+            raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
+        dl_manager.download_config.extract_on_the_fly = True
+        data_files = dl_manager.download(self.config.data_files)
+
+        # Lowercase user-facing names to match tsfile's case-insensitive convention.
+        self._table = self.config.table_name.lower() if self.config.table_name else None
+        self._requested_fields = [c.lower() for c in self.config.columns] if self.config.columns else None
+
+        all_files = [f for files in data_files.values() for f in files]
+        scan = self._scan_metadata(all_files)
+        if scan is None:
+            raise ValueError(
+                "Could not infer schema from any of the provided files. "
+                "Set `features` explicitly or check the input files."
+            )
+        self._table = scan["table"]
+        self._time_col = scan["time_col"]
+        self._tag_cols = scan["tag_cols"]
+        self._field_inner = scan["field_inner"]
+
+        if self.info.features is None:
+            self.info.features = self._build_features()
+
+        return [
+            datasets.SplitGenerator(name=split, gen_kwargs={"files": list(files)})
+            for split, files in data_files.items()
+        ]
+
+    def _generate_shards(self, files):
+        yield from files
+
+    def _generate_tables(self, files):
+        target_schema = self.info.features.arrow_schema
+        try:
+            yield from self._fold_split(files, target_schema)
+        except _SkipSplit:
+            return
+
+    # ----- metadata scan ------------------------------------------------
+
+    def _scan_metadata(self, files) -> Optional[dict]:
+        """Walk every file and unify table name, TAG columns, FIELD types."""
+        from tsfile.constants import TIME_COLUMN, ColumnCategory
+
+        wanted_table = self._table
+        wanted_fields = set(self._requested_fields) if self._requested_fields is not None else None
+
+        table: Optional[str] = wanted_table
+        time_col: Optional[str] = None
+        tag_cols: list[str] = []
+        tag_seen: set[str] = set()
+        # Per-field widest TSDataType seen so far (we map to Arrow at the end).
+        field_widest: dict = {}
+
+        for file in files:
+            try:
+                with self._open_reader(file) as reader:
+                    schemas = self._schemas_by_lc(reader)
+                    self._require_table_model(file, schemas)
+                    if table is None:
+                        table = next(iter(schemas))
+                    if table not in schemas:
+                        raise _MissingTableError(table, list(schemas))
+                    for col in schemas[table].get_columns():
+                        name = col.get_column_name()
+                        cat = col.get_category()
+                        ts_dtype = col.get_data_type()
+                        if cat == ColumnCategory.TIME:
+                            time_col = name
+                        elif cat == ColumnCategory.TAG:
+                            if name not in tag_seen:
+                                tag_seen.add(name)
+                                tag_cols.append(name)
+                        else:  # FIELD
+                            if wanted_fields is not None and name not in wanted_fields:
+                                continue
+                            prev = field_widest.get(name)
+                            field_widest[name] = ts_dtype if prev is None else _promote_tsdatatype(prev, ts_dtype)
+            except Exception as e:
+                if self._should_reraise(file, e):
+                    raise
+                continue
+
+        if table is None:
+            return None
+
+        unit = self.config.timestamp_unit
+        tz = self.config.timestamp_tz
+
+        if self._requested_fields is not None:
+            # Honor user order; silently drop names that turned out to be TAGs
+            # or the TIME column (TAGs are emitted as their own scalar columns
+            # and TIME is always emitted as a list column — neither may also
+            # appear as a list-typed field, which would collide on schema name).
+            reserved = tag_seen | {time_col} if time_col is not None else tag_seen
+            field_inner: dict[str, pa.DataType] = {}
+            for name in self._requested_fields:
+                if name in reserved:
+                    continue
+                ts_dtype = field_widest.get(name)
+                if ts_dtype is not None:
+                    field_inner[name] = _arrow_type(ts_dtype, unit=unit, tz=tz)
+                else:
+                    # Field never appeared in any file — keep as a nullable
+                    # float64 list, fully filled with nulls at read time.
+                    field_inner[name] = pa.float64()
+        else:
+            field_inner = {n: _arrow_type(d, unit=unit, tz=tz) for n, d in field_widest.items()}
+
+        return {
+            "table": table,
+            "time_col": time_col or TIME_COLUMN,
+            "tag_cols": tag_cols,
+            "field_inner": field_inner,
+        }
+
+    def _build_features(self) -> datasets.Features:
+        unit = self.config.timestamp_unit
+        tz = self.config.timestamp_tz
+        fields: list[pa.Field] = [pa.field(t, pa.string()) for t in self._tag_cols]
+        fields.append(pa.field(self._time_col, pa.list_(pa.timestamp(unit, tz=tz))))
+        for name, inner in self._field_inner.items():
+            fields.append(pa.field(name, pa.list_(inner)))
+        return datasets.Features.from_arrow_schema(pa.schema(fields))
+
+    # ----- per-split folding -------------------------------------------
+
+    def _fold_split(self, files, target_schema: pa.Schema):
+        """Stream every device in this split via per-device tag-filter pushdown.
+
+        Open one ``TsFileReader`` per file, build a cross-file device index
+        keyed by ``(tag-tuple)``, then iterate devices in stable order. For
+        each device, ``query_table(tag_filter=...)`` reads only that device's
+        rows from each contributing file, so peak memory is bounded by one
+        device's payload across the split — never the split's total size.
+        """
+        if self._table is None:
+            raise _SkipSplit
+
+        readers: dict[str, Any] = {}
+        try:
+            for file in files:
+                try:
+                    readers[file] = self._open_reader(file)
+                except Exception as e:
+                    if self._should_reraise(file, e):
+                        raise
+                    continue
+
+            device_index, file_meta = self._build_device_index(readers)
+            if not device_index:
+                return
+
+            yield from self._iter_device_batches(device_index, file_meta, readers, target_schema)
+        finally:
+            for reader in readers.values():
+                try:
+                    reader.close()
+                except Exception:
+                    pass
+
+    def _build_device_index(self, readers: dict):
+        """Walk every open reader and build the cross-file device index.
+
+        Returns ``(device_index, file_meta)``:
+
+        - ``device_index``: list of ``(device_key, [file_path, ...])`` pairs
+          in stable first-seen order. ``device_key`` is a tuple aligned to
+          ``self._tag_cols`` (the unified tag-column order).
+        - ``file_meta``: maps each readable file to its per-file context
+          (``tag_cols``, ``field_cols``, ``time_col``).
+        """
+        from tsfile.constants import ColumnCategory
+
+        device_to_files: dict[tuple, list[str]] = {}
+        device_order: list[tuple] = []
+        file_meta: dict[str, dict] = {}
+        # ``self._table`` was lowercased either by user-input normalization in
+        # ``_split_generators`` or by ``_schemas_by_lc`` during auto-detect.
+        table_lc = self._table
+
+        files_iter = tqdm(
+            readers.items(),
+            total=len(readers),
+            desc="Indexing TsFile devices",
+            unit="file",
+        )
+
+        for file, reader in files_iter:
+            try:
+                schemas = self._schemas_by_lc(reader)
+                self._require_table_model(file, schemas)
+                if table_lc not in schemas:
+                    raise _MissingTableError(table_lc, list(schemas))
+                schema = schemas[table_lc]
+
+                file_tag_cols: list[str] = []
+                file_field_cols: set[str] = set()
+                time_col = self._time_col
+                for col in schema.get_columns():
+                    name = col.get_column_name()
+                    cat = col.get_category()
+                    if cat == ColumnCategory.TIME:
+                        time_col = name
+                    elif cat == ColumnCategory.TAG:
+                        file_tag_cols.append(name)
+                    elif cat == ColumnCategory.FIELD:
+                        file_field_cols.add(name)
+
+                file_meta[file] = {
+                    "tag_cols": file_tag_cols,
+                    "field_cols": file_field_cols,
+                    "time_col": time_col,
+                }
+
+                for device in reader.get_all_devices():
+                    if device.table_name is None or device.table_name.lower() != table_lc:
+                        continue
+                    file_tag_values = list(device.segments[1 : 1 + len(file_tag_cols)])
+                    file_tag_dict = dict(zip(file_tag_cols, file_tag_values))
+                    unified_key = tuple(file_tag_dict.get(c) for c in self._tag_cols)
+                    if any(v is None for v in unified_key):
+                        raise ValueError(
+                            f"Device in file '{file}' has missing tag values: "
+                            f"{dict(zip(self._tag_cols, unified_key))}. "
+                            "Schema-evolution devices with NULL tag values are not "
+                            "supported because tsfile lacks an IS NULL tag filter."
+                        )
+                    if unified_key not in device_to_files:
+                        device_to_files[unified_key] = []
+                        device_order.append(unified_key)
+                    device_to_files[unified_key].append(file)
+            except Exception as e:
+                if self._should_reraise(file, e):
+                    raise
+                file_meta.pop(file, None)
+                continue
+
+        device_index = [(key, device_to_files[key]) for key in device_order]
+        return device_index, file_meta
+
+    def _iter_device_batches(self, device_index, file_meta, readers, target_schema: pa.Schema):
+        """Materialize devices in order and emit packed Arrow tables."""
+        field_names = list(self._field_inner.keys())
+        rows: list[dict] = []
+        batch_idx = 0
+        for device_key, contributing_files in device_index:
+            time_chunks: list[np.ndarray] = []
+            field_chunks: dict[str, list] = {f: [] for f in field_names}
+            for file in contributing_files:
+                if file not in file_meta:
+                    continue
+                try:
+                    ts_arr, vals = self._read_device_from_file(readers[file], file_meta[file], device_key)
+                except Exception as e:
+                    if self._should_reraise(file, e):
+                        raise
+                    continue
+                if len(ts_arr) == 0:
+                    continue
+                time_chunks.append(ts_arr)
+                for f in field_names:
+                    field_chunks[f].append(vals.get(f))  # None → all-null contribution
+            if not time_chunks:
+                continue  # device produced no rows in time range → skip
+
+            row = self._finalize_device(device_key, time_chunks, field_chunks, field_names)
+            rows.append(row)
+            if len(rows) >= self.config.output_batch_size:
+                yield Key(0, batch_idx), self._rows_to_table(rows, target_schema)
+                rows = []
+                batch_idx += 1
+        if rows:
+            yield Key(0, batch_idx), self._rows_to_table(rows, target_schema)
+
+    def _read_device_from_file(self, reader, meta: dict, device_key: tuple) -> tuple:
+        """Stream one device's rows from one file via ``query_table`` pushdown.
+
+        Returns ``(timestamps, {field_name: values})``. The dict only includes
+        field columns that this file owns *and* that the builder requested;
+        callers fill missing fields with all-null contributions.
+        """
+        from tsfile import tag_eq
+
+        file_tag_cols: list[str] = meta["tag_cols"]
+        file_field_cols: set[str] = meta["field_cols"]
+        time_col: str = meta["time_col"]
+
+        # Build the tag filter only over this file's tag columns. The unified
+        # device key carries one value per builder tag; map back by name.
+        unified_to_value = dict(zip(self._tag_cols, device_key))
+        tag_filter = None
+        for c in file_tag_cols:
+            v = unified_to_value.get(c)
+            if v is None:
+                # Caught earlier in _build_device_index; defensive guard.
+                return np.array([], dtype=np.int64), {}
+            expr = tag_eq(c, str(v))
+            tag_filter = expr if tag_filter is None else tag_filter & expr
+
+        # Project: requested fields ∩ this file's fields. ``query_table``
+        # always returns the time column, but it requires at least one
+        # non-time column — fall back to any owned field if the user's
+        # selection has nothing in this file.
+        requested = list(self._field_inner.keys())
+        fields_to_query = [f for f in requested if f in file_field_cols]
+        fallback_only = False
+        if not fields_to_query:
+            if not file_field_cols:
+                return np.array([], dtype=np.int64), {}
+            fields_to_query = [next(iter(file_field_cols))]
+            fallback_only = True
+
+        kwargs: dict = {"tag_filter": tag_filter, "batch_size": self.config.input_batch_size}
+        if self.config.start_time is not None:
+            kwargs["start_time"] = self.config.start_time
+        if self.config.end_time is not None:
+            kwargs["end_time"] = self.config.end_time
+
+        ts_parts: list[np.ndarray] = []
+        field_parts: dict[str, list] = {f: [] for f in fields_to_query}
+        with reader.query_table(self._table, fields_to_query, **kwargs) as rs:
+            while True:
+                batch = rs.read_arrow_batch()
+                if batch is None:
+                    break
+                if batch.num_rows == 0:
+                    continue
+                ts_parts.append(np.asarray(batch.column(time_col).to_numpy(), dtype=np.int64))
+                for f in fields_to_query:
+                    col = batch.column(f)
+                    # tsfile's arrow reader tags TIMESTAMP / DATE field columns
+                    # with a fixed unit (e.g. ``timestamp[ns]``) regardless of
+                    # the value's original write unit. Reinterpret as raw
+                    # int64/int32 ticks so the downstream
+                    # ``pa.array(type=timestamp[<our unit>])`` treats them as
+                    # ticks in the unit declared by our schema, instead of
+                    # cross-unit casting (which would raise on data loss).
+                    if pa.types.is_timestamp(col.type):
+                        field_parts[f].append(col.cast(pa.int64()).to_numpy(zero_copy_only=False))
+                    elif pa.types.is_date(col.type):
+                        field_parts[f].append(col.cast(pa.int32()).to_numpy(zero_copy_only=False))
+                    else:
+                        field_parts[f].append(col.to_numpy(zero_copy_only=False))
+
+        if not ts_parts:
+            return np.array([], dtype=np.int64), {}
+        ts_full = np.concatenate(ts_parts) if len(ts_parts) > 1 else ts_parts[0]
+        vals_full = {f: (np.concatenate(parts) if len(parts) > 1 else parts[0]) for f, parts in field_parts.items()}
+
+        # Defensive boundary mask: native query paths may emit rows just
+        # outside the requested window in some chunk-boundary cases.
+        if self.config.start_time is not None or self.config.end_time is not None:
+            lo = self.config.start_time if self.config.start_time is not None else np.iinfo(np.int64).min
+            hi = self.config.end_time if self.config.end_time is not None else np.iinfo(np.int64).max
+            mask = (ts_full >= lo) & (ts_full <= hi)
+            if not mask.all():
+                ts_full = ts_full[mask]
+                vals_full = {f: arr[mask] for f, arr in vals_full.items()}
+
+        # Drop the fallback "pick one" column from the user-visible payload.
+        if fallback_only:
+            vals_full = {}
+        return ts_full, vals_full
+
+    def _finalize_device(
+        self,
+        device_key: tuple,
+        time_chunks: list,
+        field_chunks: dict,
+        field_names: list[str],
+    ) -> dict:
+        """Concatenate per-file chunks, sort by time, and return one row.
+
+        Raises ``ValueError`` if the same timestamp appears more than once
+        for a device (within or across files) — tsfile's per-device timeline
+        is required to be unique-by-timestamp.
+        """
+        time_arr = np.concatenate(time_chunks) if time_chunks else np.array([], dtype=np.int64)
+        n_total = len(time_arr)
+
+        if n_total > 0:
+            sort_idx = np.argsort(time_arr, kind="stable")
+            time_sorted = time_arr[sort_idx]
+            if n_total > 1:
+                dup_mask = time_sorted[1:] == time_sorted[:-1]
+                if dup_mask.any():
+                    dup_ts = int(time_sorted[1:][dup_mask][0])
+                    raise ValueError(
+                        f"Duplicate timestamp {dup_ts} for device "
+                        f"{dict(zip(self._tag_cols, device_key))}. "
+                        "Cross-file or within-file duplicate timestamps are not supported."
+                    )
+        else:
+            sort_idx = None
+            time_sorted = time_arr
+
+        row: dict = {}
+        for tag_name, tag_val in zip(self._tag_cols, device_key):
+            row[tag_name] = None if tag_val is None else str(tag_val)
+        row[self._time_col] = time_sorted
+
+        for fname in field_names:
+            chunks = field_chunks.get(fname, [])
+            materialized: list = []
+            for tchunk, fchunk in zip(time_chunks, chunks):
+                if fchunk is None:
+                    materialized.append(np.full(len(tchunk), None, dtype=object))
+                else:
+                    materialized.append(fchunk)
+            arr = np.concatenate(materialized) if materialized else np.array([], dtype=object)
+            if sort_idx is not None and len(arr) == n_total:
+                arr = arr[sort_idx]
+            row[fname] = arr
+        return row
+
+    # ----- arrow assembly ----------------------------------------------
+
+    def _rows_to_table(self, rows: list[dict], target_schema: pa.Schema) -> pa.Table:
+        """Convert a batch of row dicts into an Arrow table matching ``target_schema``."""
+        arrays: list[pa.Array] = []
+        for f in target_schema:
+            values = [r[f.name] for r in rows]
+            if pa.types.is_list(f.type):
+                arrays.append(self._build_list_array(values, f.type))
+            else:
+                arrays.append(pa.array(values, type=f.type))
+        pa_table = pa.Table.from_arrays(arrays, names=[f.name for f in target_schema])
+        return table_cast(pa_table, target_schema)
+
+    @staticmethod
+    def _build_list_array(values: list, list_type: pa.ListType) -> pa.ListArray:
+        """Build a ``ListArray`` from a list of per-row 1D arrays / sequences."""
+        inner_type = list_type.value_type
+        offsets = [0]
+        flat_chunks: list = []
+        total = 0
+        for v in values:
+            if v is None:
+                length = 0
+            else:
+                length = len(v)
+                flat_chunks.append(v)
+            total += length
+            offsets.append(total)
+
+        if flat_chunks:
+            try:
+                flat = np.concatenate([np.asarray(c) for c in flat_chunks])
+            except ValueError:
+                # Heterogeneous shapes / dtypes → fall back to a Python list.
+                flat = []
+                for c in flat_chunks:
+                    flat.extend(list(c))
+            flat_arr = pa.array(flat, type=inner_type, from_pandas=True)
+        else:
+            flat_arr = pa.array([], type=inner_type)
+
+        return pa.ListArray.from_arrays(pa.array(offsets, type=pa.int32()), flat_arr)
+
+    # ----- file / error handling ---------------------------------------
+
+    @staticmethod
+    def _open_reader(file: str):
+        """Open a file as a ``TsFileReader`` after verifying its magic header.
+
+        The C library's ``TsFileReader`` constructor silently returns an
+        invalid handle for non-tsfile inputs, and any subsequent call on it
+        segfaults. The 6-byte ``TsFile`` magic header is checked first to
+        bail out cleanly.
+        """
+        from tsfile import TsFileReader
+
+        try:
+            with open(file, "rb") as fh:
+                header = fh.read(len(_TSFILE_MAGIC))
+        except OSError as e:
+            raise ValueError(f"Cannot open file {file!r}: {e}") from e
+        if header != _TSFILE_MAGIC:
+            raise ValueError(f"File {file!r} is not a valid TsFile (bad magic header).")
+        return TsFileReader(file)
+
+    @staticmethod
+    def _require_table_model(file: str, schemas) -> None:
+        if not schemas:
+            raise ValueError(
+                f"File {file!r} is a tree-model TsFile, which is not supported. "
+                "Only table-model TsFiles can be loaded."
+            )
+
+    @staticmethod
+    def _schemas_by_lc(reader) -> dict:
+        """Return ``get_all_table_schemas()`` keyed by lowercased table name.
+
+        TsFile / IoTDB treat table names case-insensitively, but the Python
+        binding's ``get_all_table_schemas()`` returns a dict keyed by whatever
+        casing the file was written with. Lowercasing the keys here lets all
+        downstream lookups use a single canonical form.
+        """
+        return {name.lower(): schema for name, schema in reader.get_all_table_schemas().items()}
+
+    def _should_reraise(self, file: str, exc: BaseException) -> bool:
+        """Apply ``on_bad_files`` policy. Returns True iff the caller should re-raise."""
+        mode = self.config.on_bad_files
+        if mode == "error":
+            logger.error(f"Failed to read file '{file}' with error {type(exc).__name__}: {exc}")
+            return True
+        if mode == "warn":
+            logger.warning(f"Skipping bad file '{file}'. {type(exc).__name__}: {exc}")
+        else:
+            logger.debug(f"Skipping bad file '{file}'. {type(exc).__name__}: {exc}")
+        return False