chore: import upstream snapshot with attribution

tests/packaged_modules/test_arrow.py

@@ -0,0 +1,61 @@
+import pyarrow as pa
+import pytest
+
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesList
+from datasets.packaged_modules.arrow.arrow import Arrow, ArrowConfig
+
+
+@pytest.fixture
+def arrow_file_streaming_format(tmp_path):
+    filename = tmp_path / "stream.arrow"
+    testdata = [[1, 1, 1], [0, 100, 6], [1, 90, 900]]
+
+    schema = pa.schema([pa.field("input_ids", pa.list_(pa.int32()))])
+    array = pa.array(testdata, type=pa.list_(pa.int32()))
+    table = pa.Table.from_arrays([array], schema=schema)
+    with open(filename, "wb") as f:
+        with pa.ipc.new_stream(f, schema) as writer:
+            writer.write_table(table)
+    return str(filename)
+
+
+@pytest.fixture
+def arrow_file_file_format(tmp_path):
+    filename = tmp_path / "file.arrow"
+    testdata = [[1, 1, 1], [0, 100, 6], [1, 90, 900]]
+
+    schema = pa.schema([pa.field("input_ids", pa.list_(pa.int32()))])
+    array = pa.array(testdata, type=pa.list_(pa.int32()))
+    table = pa.Table.from_arrays([array], schema=schema)
+    with open(filename, "wb") as f:
+        with pa.ipc.new_file(f, schema) as writer:
+            writer.write_table(table)
+    return str(filename)
+
+
+@pytest.mark.parametrize(
+    "file_fixture, config_kwargs",
+    [
+        ("arrow_file_streaming_format", {}),
+        ("arrow_file_file_format", {}),
+    ],
+)
+def test_arrow_generate_tables(file_fixture, config_kwargs, request):
+    arrow = Arrow(**config_kwargs)
+    generator = arrow._generate_tables([request.getfixturevalue(file_fixture)])
+    pa_table = pa.concat_tables([table for _, table in generator])
+
+    expected = {"input_ids": [[1, 1, 1], [0, 100, 6], [1, 90, 900]]}
+    assert pa_table.to_pydict() == expected
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = ArrowConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = ArrowConfig(name="name", data_files=data_files)

tests/packaged_modules/test_audiofolder.py

@@ -0,0 +1,394 @@
+import shutil
+import textwrap
+
+import numpy as np
+import pytest
+
+from datasets import Audio, ClassLabel, Features
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesDict, DataFilesList, get_data_patterns
+from datasets.download.streaming_download_manager import StreamingDownloadManager
+from datasets.packaged_modules.audiofolder.audiofolder import AudioFolder, AudioFolderConfig
+
+from ..utils import require_torchcodec
+
+
+@pytest.fixture
+def cache_dir(tmp_path):
+    return str(tmp_path / "audiofolder_cache_dir")
+
+
+@pytest.fixture
+def data_files_with_labels_no_metadata(tmp_path, audio_file):
+    data_dir = tmp_path / "data_files_with_labels_no_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    subdir_class_0 = data_dir / "fr"
+    subdir_class_0.mkdir(parents=True, exist_ok=True)
+    subdir_class_1 = data_dir / "uk"
+    subdir_class_1.mkdir(parents=True, exist_ok=True)
+
+    audio_filename = subdir_class_0 / "audio_fr.wav"
+    shutil.copyfile(audio_file, audio_filename)
+    audio_filename2 = subdir_class_1 / "audio_uk.wav"
+    shutil.copyfile(audio_file, audio_filename2)
+
+    data_files_with_labels_no_metadata = DataFilesDict.from_patterns(
+        get_data_patterns(str(data_dir)), data_dir.as_posix()
+    )
+
+    return data_files_with_labels_no_metadata
+
+
+@pytest.fixture
+def audio_file_with_metadata(tmp_path, audio_file):
+    audio_filename = tmp_path / "audio_file.wav"
+    shutil.copyfile(audio_file, audio_filename)
+    audio_metadata_filename = tmp_path / "metadata.jsonl"
+    audio_metadata = textwrap.dedent(
+        """\
+        {"file_name": "audio_file.wav", "text": "Audio transcription"}
+        """
+    )
+    with open(audio_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(audio_metadata)
+    return str(audio_filename), str(audio_metadata_filename)
+
+
+@pytest.fixture
+def data_files_with_one_split_and_metadata(tmp_path, audio_file):
+    data_dir = tmp_path / "audiofolder_data_dir_with_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    subdir = data_dir / "subdir"
+    subdir.mkdir(parents=True, exist_ok=True)
+
+    audio_filename = data_dir / "audio_file.wav"
+    shutil.copyfile(audio_file, audio_filename)
+    audio_filename2 = data_dir / "audio_file2.wav"
+    shutil.copyfile(audio_file, audio_filename2)
+    audio_filename3 = subdir / "audio_file3.wav"  # in subdir
+    shutil.copyfile(audio_file, audio_filename3)
+
+    audio_metadata_filename = data_dir / "metadata.jsonl"
+    audio_metadata = textwrap.dedent(
+        """\
+        {"file_name": "audio_file.wav", "text": "First audio transcription"}
+        {"file_name": "audio_file2.wav", "text": "Second audio transcription"}
+        {"file_name": "subdir/audio_file3.wav", "text": "Third audio transcription (in subdir)"}
+        """
+    )
+    with open(audio_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(audio_metadata)
+    data_files_with_one_split_and_metadata = DataFilesDict.from_patterns(
+        get_data_patterns(str(data_dir)), data_dir.as_posix()
+    )
+    assert len(data_files_with_one_split_and_metadata) == 1
+    assert len(data_files_with_one_split_and_metadata["train"]) == 4
+    return data_files_with_one_split_and_metadata
+
+
+@pytest.fixture(params=["jsonl", "csv"])
+def data_files_with_two_splits_and_metadata(request, tmp_path, audio_file):
+    data_dir = tmp_path / "audiofolder_data_dir_with_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    train_dir = data_dir / "train"
+    train_dir.mkdir(parents=True, exist_ok=True)
+    test_dir = data_dir / "test"
+    test_dir.mkdir(parents=True, exist_ok=True)
+
+    audio_filename = train_dir / "audio_file.wav"  # train audio
+    shutil.copyfile(audio_file, audio_filename)
+    audio_filename2 = train_dir / "audio_file2.wav"  # train audio
+    shutil.copyfile(audio_file, audio_filename2)
+    audio_filename3 = test_dir / "audio_file3.wav"  # test audio
+    shutil.copyfile(audio_file, audio_filename3)
+
+    train_audio_metadata_filename = train_dir / f"metadata.{request.param}"
+    audio_metadata = (
+        textwrap.dedent(
+            """\
+        {"file_name": "audio_file.wav", "text": "First train audio transcription"}
+        {"file_name": "audio_file2.wav", "text": "Second train audio transcription"}
+        """
+        )
+        if request.param == "jsonl"
+        else textwrap.dedent(
+            """\
+        file_name,text
+        audio_file.wav,First train audio transcription
+        audio_file2.wav,Second train audio transcription
+        """
+        )
+    )
+    with open(train_audio_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(audio_metadata)
+    test_audio_metadata_filename = test_dir / f"metadata.{request.param}"
+    audio_metadata = (
+        textwrap.dedent(
+            """\
+        {"file_name": "audio_file3.wav", "text": "Test audio transcription"}
+        """
+        )
+        if request.param == "jsonl"
+        else textwrap.dedent(
+            """\
+        file_name,text
+        audio_file3.wav,Test audio transcription
+        """
+        )
+    )
+    with open(test_audio_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(audio_metadata)
+    data_files_with_two_splits_and_metadata = DataFilesDict.from_patterns(
+        get_data_patterns(str(data_dir)), data_dir.as_posix()
+    )
+    assert len(data_files_with_two_splits_and_metadata) == 2
+    assert len(data_files_with_two_splits_and_metadata["train"]) == 3
+    assert len(data_files_with_two_splits_and_metadata["test"]) == 2
+    return data_files_with_two_splits_and_metadata
+
+
+@pytest.fixture
+def data_files_with_zip_archives(tmp_path, audio_file_44100, audio_file_16000):
+    data_dir = tmp_path / "audiofolder_data_dir_with_zip_archives"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    archive_dir = data_dir / "archive"
+    archive_dir.mkdir(parents=True, exist_ok=True)
+    subdir = archive_dir / "subdir"
+    subdir.mkdir(parents=True, exist_ok=True)
+
+    audio_filename = archive_dir / "audio_file.mp3"
+    shutil.copyfile(audio_file_44100, audio_filename)
+    audio_filename2 = subdir / "audio_file2.mp3"  # in subdir
+    shutil.copyfile(audio_file_16000, audio_filename2)
+
+    audio_metadata_filename = archive_dir / "metadata.jsonl"
+    audio_metadata = textwrap.dedent(
+        """\
+        {"file_name": "audio_file.mp3", "text": "First audio transcription"}
+        {"file_name": "subdir/audio_file2.mp3", "text": "Second audio transcription (in subdir)"}
+        """
+    )
+
+    with open(audio_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(audio_metadata)
+
+    shutil.make_archive(str(archive_dir), "zip", archive_dir)
+    shutil.rmtree(str(archive_dir))
+
+    data_files_with_zip_archives = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+
+    assert len(data_files_with_zip_archives) == 1
+    assert len(data_files_with_zip_archives["train"]) == 1
+    return data_files_with_zip_archives
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = AudioFolderConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = AudioFolderConfig(name="name", data_files=data_files)
+
+
+@require_torchcodec
+# check that labels are inferred correctly from dir names
+def test_generate_examples_with_labels(data_files_with_labels_no_metadata, cache_dir):
+    # there are no metadata.jsonl files in this test case
+    audiofolder = AudioFolder(data_files=data_files_with_labels_no_metadata, cache_dir=cache_dir, drop_labels=False)
+    audiofolder.download_and_prepare()
+    assert audiofolder.info.features == Features({"audio": Audio(), "label": ClassLabel(names=["fr", "uk"])})
+    dataset = list(audiofolder.as_dataset()["train"])
+    label_feature = audiofolder.info.features["label"]
+
+    assert dataset[0]["label"] == label_feature._str2int["fr"]
+    assert dataset[1]["label"] == label_feature._str2int["uk"]
+
+
+@require_torchcodec
+@pytest.mark.parametrize("drop_metadata", [None, True, False])
+@pytest.mark.parametrize("drop_labels", [None, True, False])
+def test_generate_examples_drop_labels(data_files_with_labels_no_metadata, drop_metadata, drop_labels):
+    audiofolder = AudioFolder(
+        drop_metadata=drop_metadata, drop_labels=drop_labels, data_files=data_files_with_labels_no_metadata
+    )
+    gen_kwargs = audiofolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    # removing the labels explicitly requires drop_labels=True
+    assert gen_kwargs["add_labels"] is not bool(drop_labels)
+    assert gen_kwargs["add_metadata"] is False  # metadata files is not present in this case
+    generator = audiofolder._generate_examples(**gen_kwargs)
+    if not drop_labels:
+        assert all(
+            example.keys() == {"audio", "label"} and all(val is not None for val in example.values())
+            for _, example in generator
+        )
+    else:
+        assert all(
+            example.keys() == {"audio"} and all(val is not None for val in example.values())
+            for _, example in generator
+        )
+
+
+@require_torchcodec
+@pytest.mark.parametrize("drop_metadata", [None, True, False])
+@pytest.mark.parametrize("drop_labels", [None, True, False])
+def test_generate_examples_drop_metadata(audio_file_with_metadata, drop_metadata, drop_labels):
+    audio_file, audio_metadata_file = audio_file_with_metadata
+    audiofolder = AudioFolder(
+        drop_metadata=drop_metadata, drop_labels=drop_labels, data_files={"train": [audio_file, audio_metadata_file]}
+    )
+    gen_kwargs = audiofolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    # since the dataset has metadata, removing the metadata explicitly requires drop_metadata=True
+    assert gen_kwargs["add_metadata"] is not bool(drop_metadata)
+    # since the dataset has metadata, adding the labels explicitly requires drop_labels=False
+    assert gen_kwargs["add_labels"] is False
+    generator = audiofolder._generate_examples(**gen_kwargs)
+    expected_columns = {"audio"}
+    if gen_kwargs["add_metadata"]:
+        expected_columns.add("text")
+    if gen_kwargs["add_labels"]:
+        expected_columns.add("label")
+    result = [example for _, example in generator]
+    assert len(result) == 1
+    example = result[0]
+    assert example.keys() == expected_columns
+    for column in expected_columns:
+        assert example[column] is not None
+
+
+@require_torchcodec
+@pytest.mark.parametrize("streaming", [False, True])
+def test_data_files_with_metadata_and_single_split(streaming, cache_dir, data_files_with_one_split_and_metadata):
+    data_files = data_files_with_one_split_and_metadata
+    audiofolder = AudioFolder(data_files=data_files, cache_dir=cache_dir)
+    audiofolder.download_and_prepare()
+    datasets = audiofolder.as_streaming_dataset() if streaming else audiofolder.as_dataset()
+    for split, data_files in data_files.items():
+        expected_num_of_audios = len(data_files) - 1  # don't count the metadata file
+        assert split in datasets
+        dataset = list(datasets[split])
+        assert len(dataset) == expected_num_of_audios
+        # make sure each sample has its own audio and metadata
+        assert len({example["audio"].metadata.path for example in dataset}) == expected_num_of_audios
+        assert len({example["text"] for example in dataset}) == expected_num_of_audios
+        assert all(example["text"] is not None for example in dataset)
+
+
+@require_torchcodec
+@pytest.mark.parametrize("streaming", [False, True])
+def test_data_files_with_metadata_and_multiple_splits(streaming, cache_dir, data_files_with_two_splits_and_metadata):
+    data_files = data_files_with_two_splits_and_metadata
+    audiofolder = AudioFolder(data_files=data_files, cache_dir=cache_dir)
+    audiofolder.download_and_prepare()
+    datasets = audiofolder.as_streaming_dataset() if streaming else audiofolder.as_dataset()
+    for split, data_files in data_files.items():
+        expected_num_of_audios = len(data_files) - 1  # don't count the metadata file
+        assert split in datasets
+        dataset = list(datasets[split])
+        assert len(dataset) == expected_num_of_audios
+        # make sure each sample has its own audio and metadata
+        assert len({example["audio"].metadata.path for example in dataset}) == expected_num_of_audios
+        assert len({example["text"] for example in dataset}) == expected_num_of_audios
+        assert all(example["text"] is not None for example in dataset)
+
+
+@require_torchcodec
+@pytest.mark.parametrize("streaming", [False, True])
+def test_data_files_with_metadata_and_archives(streaming, cache_dir, data_files_with_zip_archives):
+    audiofolder = AudioFolder(data_files=data_files_with_zip_archives, cache_dir=cache_dir)
+    audiofolder.download_and_prepare()
+    datasets = audiofolder.as_streaming_dataset() if streaming else audiofolder.as_dataset()
+    for split, data_files in data_files_with_zip_archives.items():
+        num_of_archives = len(data_files)  # the metadata file is inside the archive
+        expected_num_of_audios = 2 * num_of_archives
+        assert split in datasets
+        dataset = list(datasets[split])
+        assert len(dataset) == expected_num_of_audios
+        # make sure each sample has its own audio (all arrays are different) and metadata
+        assert (
+            sum(
+                np.array_equal(
+                    dataset[0]["audio"].get_all_samples().data.numpy(), example["audio"].get_all_samples().data.numpy()
+                )
+                for example in dataset[1:]
+            )
+            == 0
+        )
+        assert len({example["text"] for example in dataset}) == expected_num_of_audios
+        assert all(example["text"] is not None for example in dataset)
+
+
+@require_torchcodec
+def test_data_files_with_wrong_metadata_file_name(cache_dir, tmp_path, audio_file):
+    data_dir = tmp_path / "data_dir_with_bad_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(audio_file, data_dir / "audio_file.wav")
+    audio_metadata_filename = data_dir / "bad_metadata.jsonl"  # bad file
+    audio_metadata = textwrap.dedent(
+        """\
+        {"file_name": "audio_file.wav", "text": "Audio transcription"}
+        """
+    )
+    with open(audio_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(audio_metadata)
+
+    data_files_with_bad_metadata = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+    audiofolder = AudioFolder(data_files=data_files_with_bad_metadata, cache_dir=cache_dir)
+    audiofolder.download_and_prepare()
+    dataset = audiofolder.as_dataset(split="train")
+    # check that there are no metadata, since the metadata file name doesn't have the right name
+    assert "text" not in dataset.column_names
+
+
+@require_torchcodec
+def test_data_files_with_custom_audio_file_name_column_in_metadata_file(cache_dir, tmp_path, audio_file):
+    data_dir = tmp_path / "data_dir_with_custom_file_name_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(audio_file, data_dir / "audio_file.wav")
+    audio_metadata_filename = data_dir / "metadata.jsonl"
+    audio_metadata = textwrap.dedent(  # with bad column "bad_file_name" instead of "file_name"
+        """\
+        {"speech_file_name": "audio_file.wav", "text": "Audio transcription"}
+        """
+    )
+    with open(audio_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(audio_metadata)
+
+    data_files_with_bad_metadata = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+    audiofolder = AudioFolder(data_files=data_files_with_bad_metadata, cache_dir=cache_dir)
+    audiofolder.download_and_prepare()
+    dataset = audiofolder.as_dataset(split="train")
+    assert "speech" in dataset.features
+    assert "speech_file_name" not in dataset.features
+
+
+@require_torchcodec
+def test_data_files_with_with_metadata_in_different_formats(cache_dir, tmp_path, audio_file):
+    data_dir = tmp_path / "data_dir_with_metadata_in_different_format"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(audio_file, data_dir / "audio_file.wav")
+    audio_metadata_filename_jsonl = data_dir / "metadata.jsonl"
+    audio_metadata_jsonl = textwrap.dedent(
+        """\
+        {"file_name": "audio_file.wav", "text": "Audio transcription"}
+        """
+    )
+    with open(audio_metadata_filename_jsonl, "w", encoding="utf-8") as f:
+        f.write(audio_metadata_jsonl)
+    audio_metadata_filename_csv = data_dir / "metadata.csv"
+    audio_metadata_csv = textwrap.dedent(
+        """\
+        file_name,text
+        audio_file.wav,Audio transcription
+        """
+    )
+    with open(audio_metadata_filename_csv, "w", encoding="utf-8") as f:
+        f.write(audio_metadata_csv)
+
+    data_files_with_bad_metadata = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+    audiofolder = AudioFolder(data_files=data_files_with_bad_metadata, cache_dir=cache_dir)
+    with pytest.raises(ValueError) as exc_info:
+        audiofolder.download_and_prepare()
+    assert "metadata files with different extensions" in str(exc_info.value)

tests/packaged_modules/test_cache.py

@@ -0,0 +1,158 @@
+from pathlib import Path
+
+import pytest
+
+from datasets import load_dataset
+from datasets.packaged_modules.cache.cache import Cache
+
+
+SAMPLE_DATASET_SINGLE_CONFIG_IN_METADATA = "hf-internal-testing/audiofolder_single_config_in_metadata"
+SAMPLE_DATASET_TWO_CONFIG_IN_METADATA = "hf-internal-testing/audiofolder_two_configs_in_metadata"
+SAMPLE_DATASET_CAPITAL_LETTERS_IN_NAME = "hf-internal-testing/DatasetWithCapitalLetters"
+
+
+def test_cache(text_dir: Path, tmp_path: Path):
+    cache_dir = tmp_path / "test_cache"
+    ds = load_dataset(str(text_dir), cache_dir=str(cache_dir))
+    hash = Path(ds["train"].cache_files[0]["filename"]).parts[-2]
+    cache = Cache(cache_dir=str(cache_dir), dataset_name=text_dir.name, hash=hash)
+    reloaded = cache.as_dataset()
+    assert list(ds) == list(reloaded)
+    assert list(ds["train"]) == list(reloaded["train"])
+
+
+def test_cache_streaming(text_dir: Path, tmp_path: Path):
+    cache_dir = tmp_path / "test_cache_streaming"
+    ds = load_dataset(str(text_dir), cache_dir=str(cache_dir))
+    hash = Path(ds["train"].cache_files[0]["filename"]).parts[-2]
+    cache = Cache(cache_dir=str(cache_dir), dataset_name=text_dir.name, hash=hash)
+    reloaded = cache.as_streaming_dataset()
+    assert list(ds) == list(reloaded)
+    assert list(ds["train"]) == list(reloaded["train"])
+
+
+def test_cache_auto_hash(text_dir: Path, tmp_path: Path):
+    cache_dir = tmp_path / "test_cache_auto_hash"
+    ds = load_dataset(str(text_dir), cache_dir=str(cache_dir))
+    cache = Cache(cache_dir=str(cache_dir), dataset_name=text_dir.name, version="auto", hash="auto")
+    reloaded = cache.as_dataset()
+    assert list(ds) == list(reloaded)
+    assert list(ds["train"]) == list(reloaded["train"])
+
+
+def test_cache_auto_hash_with_custom_config(text_dir: Path, tmp_path: Path):
+    cache_dir = tmp_path / "test_cache_auto_hash_with_custom_config"
+    ds = load_dataset(str(text_dir), sample_by="paragraph", cache_dir=str(cache_dir))
+    another_ds = load_dataset(str(text_dir), cache_dir=str(cache_dir))
+    cache = Cache(
+        cache_dir=str(cache_dir), dataset_name=text_dir.name, version="auto", hash="auto", sample_by="paragraph"
+    )
+    another_cache = Cache(cache_dir=str(cache_dir), dataset_name=text_dir.name, version="auto", hash="auto")
+    assert cache.config_id.endswith("paragraph")
+    assert not another_cache.config_id.endswith("paragraph")
+    reloaded = cache.as_dataset()
+    another_reloaded = another_cache.as_dataset()
+    assert list(ds) == list(reloaded)
+    assert list(ds["train"]) == list(reloaded["train"])
+    assert list(another_ds) == list(another_reloaded)
+    assert list(another_ds["train"]) == list(another_reloaded["train"])
+
+
+def test_cache_missing(text_dir: Path, tmp_path: Path):
+    cache_dir = tmp_path / "test_cache_missing"
+    load_dataset(str(text_dir), cache_dir=str(cache_dir))
+    Cache(cache_dir=str(cache_dir), dataset_name=text_dir.name, version="auto", hash="auto").download_and_prepare()
+    with pytest.raises(ValueError):
+        Cache(cache_dir=str(cache_dir), dataset_name="missing", version="auto", hash="auto").download_and_prepare()
+    with pytest.raises(ValueError):
+        Cache(cache_dir=str(cache_dir), dataset_name=text_dir.name, hash="missing").download_and_prepare()
+    with pytest.raises(ValueError):
+        Cache(
+            cache_dir=str(cache_dir), dataset_name=text_dir.name, config_name="missing", version="auto", hash="auto"
+        ).download_and_prepare()
+
+
+@pytest.mark.integration
+def test_cache_multi_configs(tmp_path: Path):
+    cache_dir = tmp_path / "test_cache_multi_configs"
+    repo_id = SAMPLE_DATASET_TWO_CONFIG_IN_METADATA
+    dataset_name = repo_id.split("/")[-1]
+    config_name = "v1"
+    ds = load_dataset(repo_id, config_name, cache_dir=str(cache_dir))
+    cache = Cache(
+        cache_dir=str(cache_dir),
+        dataset_name=dataset_name,
+        repo_id=repo_id,
+        config_name=config_name,
+        version="auto",
+        hash="auto",
+    )
+    reloaded = cache.as_dataset()
+    assert list(ds) == list(reloaded)
+    assert len(ds["train"]) == len(reloaded["train"])
+    with pytest.raises(ValueError) as excinfo:
+        Cache(
+            cache_dir=str(cache_dir),
+            dataset_name=dataset_name,
+            repo_id=repo_id,
+            config_name="missing",
+            version="auto",
+            hash="auto",
+        )
+    assert config_name in str(excinfo.value)
+
+
+@pytest.mark.integration
+def test_cache_single_config(tmp_path: Path):
+    cache_dir = tmp_path / "test_cache_single_config"
+    repo_id = SAMPLE_DATASET_SINGLE_CONFIG_IN_METADATA
+    dataset_name = repo_id.split("/")[-1]
+    config_name = "custom"
+    ds = load_dataset(repo_id, cache_dir=str(cache_dir))
+    cache = Cache(cache_dir=str(cache_dir), dataset_name=dataset_name, repo_id=repo_id, version="auto", hash="auto")
+    reloaded = cache.as_dataset()
+    assert list(ds) == list(reloaded)
+    assert len(ds["train"]) == len(reloaded["train"])
+    cache = Cache(
+        cache_dir=str(cache_dir),
+        dataset_name=dataset_name,
+        config_name=config_name,
+        repo_id=repo_id,
+        version="auto",
+        hash="auto",
+    )
+    reloaded = cache.as_dataset()
+    assert list(ds) == list(reloaded)
+    assert len(ds["train"]) == len(reloaded["train"])
+    with pytest.raises(ValueError) as excinfo:
+        Cache(
+            cache_dir=str(cache_dir),
+            dataset_name=dataset_name,
+            repo_id=repo_id,
+            config_name="missing",
+            version="auto",
+            hash="auto",
+        )
+    assert config_name in str(excinfo.value)
+
+
+@pytest.mark.integration
+def test_cache_capital_letters(tmp_path: Path):
+    cache_dir = tmp_path / "test_cache_capital_letters"
+    repo_id = SAMPLE_DATASET_CAPITAL_LETTERS_IN_NAME
+    dataset_name = repo_id.split("/")[-1]
+    ds = load_dataset(repo_id, cache_dir=str(cache_dir))
+    cache = Cache(cache_dir=str(cache_dir), dataset_name=dataset_name, repo_id=repo_id, version="auto", hash="auto")
+    reloaded = cache.as_dataset()
+    assert list(ds) == list(reloaded)
+    assert len(ds["train"]) == len(reloaded["train"])
+    cache = Cache(
+        cache_dir=str(cache_dir),
+        dataset_name=dataset_name,
+        repo_id=repo_id,
+        version="auto",
+        hash="auto",
+    )
+    reloaded = cache.as_dataset()
+    assert list(ds) == list(reloaded)
+    assert len(ds["train"]) == len(reloaded["train"])

tests/packaged_modules/test_conll.py

@@ -0,0 +1,213 @@
+import textwrap
+
+import pyarrow as pa
+import pytest
+
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesList
+from datasets.packaged_modules.conll.conll import Conll, ConllConfig
+
+
+@pytest.fixture
+def conll2003_file(tmp_path):
+    """Minimal CoNLL-2003-style NER file: token POS chunk NER."""
+    filename = tmp_path / "sample.conll"
+    data = textwrap.dedent(
+        """\
+        -DOCSTART- -X- -X- O
+
+        EU NNP B-NP B-ORG
+        rejects VBZ B-VP O
+        German JJ B-NP B-MISC
+        call NN I-NP O
+        . . O O
+
+        Peter NNP B-NP B-PER
+        Blackburn NNP I-NP I-PER
+        """
+    )
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def conll_no_trailing_blank_file(tmp_path):
+    """Sentence with no trailing blank line — should still flush."""
+    filename = tmp_path / "no_trailing.conll"
+    data = "Hello NN\nworld NN\n. ."
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def conll_tab_delimited_file(tmp_path):
+    """Tab-delimited CoNLL — fields can contain spaces."""
+    filename = tmp_path / "tab.conll"
+    data = "tok1\tNN\nphrase with spaces\tNNP\n\ntok2\tVB\n"
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def conllu_file(tmp_path):
+    """CoNLL-U with # comment lines."""
+    filename = tmp_path / "sample.conllu"
+    data = textwrap.dedent(
+        """\
+        # sent_id = test-1
+        # text = The cat sat.
+        1\tThe\tthe\tDET
+        2\tcat\tcat\tNOUN
+        3\tsat\tsit\tVERB
+        4\t.\t.\tPUNCT
+
+        # sent_id = test-2
+        1\tHello\thello\tINTJ
+        """
+    )
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write(data)
+    return str(filename)
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = ConllConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = ConllConfig(name="name", data_files=data_files)
+
+
+def test_config_default_column_names():
+    c = ConllConfig(name="test")
+    assert c.column_names == ["tokens"]
+
+
+def test_generate_tables_conll2003_with_full_schema(conll2003_file):
+    builder = Conll(
+        column_names=["tokens", "pos_tags", "chunk_tags", "ner_tags"],
+    )
+    generator = builder._generate_tables(base_files=[conll2003_file], files_iterables=[[conll2003_file]])
+    pa_table = pa.concat_tables([table for _, table in generator])
+    data = pa_table.to_pydict()
+
+    assert data["tokens"] == [
+        ["EU", "rejects", "German", "call", "."],
+        ["Peter", "Blackburn"],
+    ]
+    assert data["pos_tags"] == [
+        ["NNP", "VBZ", "JJ", "NN", "."],
+        ["NNP", "NNP"],
+    ]
+    assert data["chunk_tags"] == [
+        ["B-NP", "B-VP", "B-NP", "I-NP", "O"],
+        ["B-NP", "I-NP"],
+    ]
+    assert data["ner_tags"] == [
+        ["B-ORG", "O", "B-MISC", "O", "O"],
+        ["B-PER", "I-PER"],
+    ]
+
+
+def test_generate_tables_skips_docstart_by_default(conll2003_file):
+    builder = Conll(column_names=["tokens", "pos_tags", "chunk_tags", "ner_tags"])
+    generator = builder._generate_tables(base_files=[conll2003_file], files_iterables=[[conll2003_file]])
+    pa_table = pa.concat_tables([table for _, table in generator])
+    # DOCSTART line should not appear as a token
+    all_tokens = sum(pa_table.to_pydict()["tokens"], [])
+    assert "-DOCSTART-" not in all_tokens
+
+
+def test_generate_tables_keeps_docstart_when_disabled(conll2003_file):
+    builder = Conll(
+        column_names=["tokens", "pos_tags", "chunk_tags", "ner_tags"],
+        skip_docstart=False,
+    )
+    generator = builder._generate_tables(base_files=[conll2003_file], files_iterables=[[conll2003_file]])
+    pa_table = pa.concat_tables([table for _, table in generator])
+    all_tokens = sum(pa_table.to_pydict()["tokens"], [])
+    assert "-DOCSTART-" in all_tokens
+
+
+def test_generate_tables_flushes_tail_sentence_without_trailing_blank(
+    conll_no_trailing_blank_file,
+):
+    builder = Conll(column_names=["tokens", "pos_tags"])
+    generator = builder._generate_tables(
+        base_files=[conll_no_trailing_blank_file],
+        files_iterables=[[conll_no_trailing_blank_file]],
+    )
+    pa_table = pa.concat_tables([table for _, table in generator])
+    assert pa_table.to_pydict()["tokens"] == [["Hello", "world", "."]]
+    assert pa_table.to_pydict()["pos_tags"] == [["NN", "NN", "."]]
+
+
+def test_generate_tables_with_tab_delimiter(conll_tab_delimited_file):
+    builder = Conll(column_names=["tokens", "pos_tags"], delimiter="\t")
+    generator = builder._generate_tables(
+        base_files=[conll_tab_delimited_file],
+        files_iterables=[[conll_tab_delimited_file]],
+    )
+    pa_table = pa.concat_tables([table for _, table in generator])
+    # 'phrase with spaces' must stay a single token under tab delimiter
+    assert pa_table.to_pydict()["tokens"] == [["tok1", "phrase with spaces"], ["tok2"]]
+
+
+def test_generate_tables_conllu_with_comment_prefix(conllu_file):
+    builder = Conll(
+        column_names=["id", "form", "lemma", "upos"],
+        delimiter="\t",
+        comment_prefix="#",
+    )
+    generator = builder._generate_tables(base_files=[conllu_file], files_iterables=[[conllu_file]])
+    pa_table = pa.concat_tables([table for _, table in generator])
+    data = pa_table.to_pydict()
+    # 2 sentences, comments stripped, # never appears in any column
+    assert data["form"] == [["The", "cat", "sat", "."], ["Hello"]]
+    assert data["upos"] == [["DET", "NOUN", "VERB", "PUNCT"], ["INTJ"]]
+    for col in data.values():
+        for sentence in col:
+            for cell in sentence:
+                assert not cell.startswith("#")
+
+
+def test_generate_tables_pads_rows_with_fewer_columns(tmp_path):
+    """A row with fewer columns than column_names should pad with empty strings."""
+    filename = tmp_path / "padded.conll"
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write("only_token\nfull token POS\n")
+    builder = Conll(column_names=["tokens", "pos_tags"])
+    generator = builder._generate_tables(base_files=[str(filename)], files_iterables=[[str(filename)]])
+    pa_table = pa.concat_tables([table for _, table in generator])
+    assert pa_table.to_pydict()["tokens"] == [["only_token", "full"]]
+    # First row has empty POS (padded), second has "token" because line had 3 parts but column_names has 2 -> truncated
+    assert pa_table.to_pydict()["pos_tags"] == [["", "token"]]
+
+
+def test_generate_tables_truncates_rows_with_extra_columns(tmp_path):
+    """A row with more columns than column_names should truncate."""
+    filename = tmp_path / "extra.conll"
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write("tok POS chunk ner extra_col\n")
+    builder = Conll(column_names=["tokens", "pos_tags"])
+    generator = builder._generate_tables(base_files=[str(filename)], files_iterables=[[str(filename)]])
+    pa_table = pa.concat_tables([table for _, table in generator])
+    assert pa_table.to_pydict()["tokens"] == [["tok"]]
+    assert pa_table.to_pydict()["pos_tags"] == [["POS"]]
+
+
+def test_generate_tables_empty_column_names_raises(tmp_path):
+    """Empty column_names should raise a clear ValueError."""
+    filename = tmp_path / "x.conll"
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write("a b\n")
+    builder = Conll(column_names=[])
+    generator = builder._generate_tables(base_files=[str(filename)], files_iterables=[[str(filename)]])
+    with pytest.raises(ValueError, match="column_names must be a non-empty list"):
+        list(generator)

tests/packaged_modules/test_csv.py

@@ -0,0 +1,153 @@
+import os
+import textwrap
+
+import pyarrow as pa
+import pytest
+
+from datasets import ClassLabel, Features, Image
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesList
+from datasets.packaged_modules.csv.csv import Csv, CsvConfig
+
+from ..utils import require_pil
+
+
+@pytest.fixture
+def csv_file(tmp_path):
+    filename = tmp_path / "file.csv"
+    data = textwrap.dedent(
+        """\
+        header1,header2
+        1,2
+        10,20
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def malformed_csv_file(tmp_path):
+    filename = tmp_path / "malformed_file.csv"
+    data = textwrap.dedent(
+        """\
+        header1,header2
+        1,2
+        10,20,
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def csv_file_with_image(tmp_path, image_file):
+    filename = tmp_path / "csv_with_image.csv"
+    data = textwrap.dedent(
+        f"""\
+        image
+        {image_file}
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def csv_file_with_label(tmp_path):
+    filename = tmp_path / "csv_with_label.csv"
+    data = textwrap.dedent(
+        """\
+        label
+        good
+        bad
+        good
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def csv_file_with_int_list(tmp_path):
+    filename = tmp_path / "csv_with_int_list.csv"
+    data = textwrap.dedent(
+        """\
+        int_list
+        1 2 3
+        4 5 6
+        7 8 9
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = CsvConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = CsvConfig(name="name", data_files=data_files)
+
+
+def test_csv_generate_tables_raises_error_with_malformed_csv(csv_file, malformed_csv_file, caplog):
+    csv = Csv()
+    base_files = [csv_file, malformed_csv_file]
+    files_iterables = [[file] for file in base_files]
+    generator = csv._generate_tables(base_files=base_files, files_iterables=files_iterables)
+    with pytest.raises(ValueError, match="Error tokenizing data"):
+        for _ in generator:
+            pass
+    assert any(
+        record.levelname == "ERROR"
+        and "Failed to read file" in record.message
+        and os.path.basename(malformed_csv_file) in record.message
+        for record in caplog.records
+    )
+
+
+@require_pil
+def test_csv_cast_image(csv_file_with_image):
+    with open(csv_file_with_image, encoding="utf-8") as f:
+        image_file = f.read().splitlines()[1]
+    csv = Csv(encoding="utf-8", features=Features({"image": Image()}))
+    base_files = [csv_file_with_image]
+    files_iterables = [[file] for file in base_files]
+    generator = csv._generate_tables(base_files=base_files, files_iterables=files_iterables)
+    pa_table = pa.concat_tables([table for _, table in generator])
+    assert pa_table.schema.field("image").type == Image()()
+    generated_content = pa_table.to_pydict()["image"]
+    assert generated_content == [{"path": image_file, "bytes": None}]
+
+
+def test_csv_cast_label(csv_file_with_label):
+    with open(csv_file_with_label, encoding="utf-8") as f:
+        labels = f.read().splitlines()[1:]
+    csv = Csv(encoding="utf-8", features=Features({"label": ClassLabel(names=["good", "bad"])}))
+    base_files = [csv_file_with_label]
+    files_iterables = [[file] for file in base_files]
+    generator = csv._generate_tables(base_files=base_files, files_iterables=files_iterables)
+    pa_table = pa.concat_tables([table for _, table in generator])
+    assert pa_table.schema.field("label").type == ClassLabel(names=["good", "bad"])()
+    generated_content = pa_table.to_pydict()["label"]
+    assert generated_content == [ClassLabel(names=["good", "bad"]).str2int(label) for label in labels]
+
+
+def test_csv_convert_int_list(csv_file_with_int_list):
+    csv = Csv(encoding="utf-8", sep=",", converters={"int_list": lambda x: [int(i) for i in x.split()]})
+    base_files = [csv_file_with_int_list]
+    files_iterables = [[file] for file in base_files]
+    generator = csv._generate_tables(base_files=base_files, files_iterables=files_iterables)
+    pa_table = pa.concat_tables([table for _, table in generator])
+    assert pa.types.is_list(pa_table.schema.field("int_list").type)
+    generated_content = pa_table.to_pydict()["int_list"]
+    assert generated_content == [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

tests/packaged_modules/test_folder_based_builder.py

@@ -0,0 +1,500 @@
+import importlib
+import shutil
+import textwrap
+
+import pytest
+
+from datasets import ClassLabel, DownloadManager
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesDict, DataFilesList, get_data_patterns
+from datasets.download.streaming_download_manager import StreamingDownloadManager
+from datasets.packaged_modules.folder_based_builder.folder_based_builder import (
+    FolderBasedBuilder,
+    FolderBasedBuilderConfig,
+)
+
+
+remote_files = [
+    "https://huggingface.co/datasets/hf-internal-testing/textfolder/resolve/main/hallo.txt",
+    "https://huggingface.co/datasets/hf-internal-testing/textfolder/resolve/main/hello.txt",
+    "https://huggingface.co/datasets/hf-internal-testing/textfolder/resolve/main/class1/bonjour.txt",
+    "https://huggingface.co/datasets/hf-internal-testing/textfolder/resolve/main/class1/bonjour2.txt",
+]
+
+
+class DummyFeature:
+    pass
+
+
+class DummyFolderBasedBuilder(FolderBasedBuilder):
+    BASE_FEATURE = DummyFeature
+    BASE_COLUMN_NAME = "base"
+    BUILDER_CONFIG_CLASS = FolderBasedBuilderConfig
+    EXTENSIONS = [".txt"]
+
+
+@pytest.fixture
+def cache_dir(tmp_path):
+    return str(tmp_path / "autofolder_cache_dir")
+
+
+@pytest.fixture
+def auto_text_file(text_file):
+    return str(text_file)
+
+
+@pytest.fixture
+def data_files_with_labels_no_metadata(tmp_path, auto_text_file):
+    data_dir = tmp_path / "data_files_with_labels_no_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    subdir_class_0 = data_dir / "class0"
+    subdir_class_0.mkdir(parents=True, exist_ok=True)
+    subdir_class_1 = data_dir / "class1"
+    subdir_class_1.mkdir(parents=True, exist_ok=True)
+
+    filename = subdir_class_0 / "file0.txt"
+    shutil.copyfile(auto_text_file, filename)
+    filename2 = subdir_class_1 / "file1.txt"
+    shutil.copyfile(auto_text_file, filename2)
+
+    data_files_with_labels_no_metadata = DataFilesDict.from_patterns(
+        get_data_patterns(str(data_dir)), data_dir.as_posix()
+    )
+
+    return data_files_with_labels_no_metadata
+
+
+@pytest.fixture
+def data_files_with_different_levels_no_metadata(tmp_path, auto_text_file):
+    data_dir = tmp_path / "data_files_with_different_levels"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    subdir_class_0 = data_dir / "class0"
+    subdir_class_0.mkdir(parents=True, exist_ok=True)
+    subdir_class_1 = data_dir / "subdir" / "class1"
+    subdir_class_1.mkdir(parents=True, exist_ok=True)
+
+    filename = subdir_class_0 / "file0.txt"
+    shutil.copyfile(auto_text_file, filename)
+    filename2 = subdir_class_1 / "file1.txt"
+    shutil.copyfile(auto_text_file, filename2)
+
+    data_files_with_different_levels = DataFilesDict.from_patterns(
+        get_data_patterns(str(data_dir)), data_dir.as_posix()
+    )
+
+    return data_files_with_different_levels
+
+
+@pytest.fixture
+def data_files_with_one_label_no_metadata(tmp_path, auto_text_file):
+    # only one label found = all files in a single dir/in a root dir
+    data_dir = tmp_path / "data_files_with_one_label"
+    data_dir.mkdir(parents=True, exist_ok=True)
+
+    filename = data_dir / "file0.txt"
+    shutil.copyfile(auto_text_file, filename)
+    filename2 = data_dir / "file1.txt"
+    shutil.copyfile(auto_text_file, filename2)
+
+    data_files_with_one_label = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+
+    return data_files_with_one_label
+
+
+@pytest.fixture
+def files_with_labels_and_duplicated_label_key_in_metadata(tmp_path, auto_text_file):
+    data_dir = tmp_path / "files_with_labels_and_label_key_in_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    subdir_class_0 = data_dir / "class0"
+    subdir_class_0.mkdir(parents=True, exist_ok=True)
+    subdir_class_1 = data_dir / "class1"
+    subdir_class_1.mkdir(parents=True, exist_ok=True)
+
+    filename = subdir_class_0 / "file_class0.txt"
+    shutil.copyfile(auto_text_file, filename)
+    filename2 = subdir_class_1 / "file_class1.txt"
+    shutil.copyfile(auto_text_file, filename2)
+
+    metadata_filename = tmp_path / data_dir / "metadata.jsonl"
+    metadata = textwrap.dedent(
+        """\
+        {"file_name": "class0/file_class0.txt", "additional_feature": "First dummy file", "label": "CLASS_0"}
+        {"file_name": "class1/file_class1.txt", "additional_feature": "Second dummy file", "label": "CLASS_1"}
+        """
+    )
+    with open(metadata_filename, "w", encoding="utf-8") as f:
+        f.write(metadata)
+
+    return str(filename), str(filename2), str(metadata_filename)
+
+
+@pytest.fixture
+def file_with_metadata(tmp_path, text_file):
+    filename = tmp_path / "file.txt"
+    shutil.copyfile(text_file, filename)
+    metadata_filename = tmp_path / "metadata.jsonl"
+    metadata = textwrap.dedent(
+        """\
+        {"file_name": "file.txt", "additional_feature": "Dummy file"}
+        """
+    )
+    with open(metadata_filename, "w", encoding="utf-8") as f:
+        f.write(metadata)
+    return str(filename), str(metadata_filename)
+
+
+@pytest.fixture
+def data_files_with_one_split_and_metadata(tmp_path, auto_text_file):
+    data_dir = tmp_path / "autofolder_data_dir_with_metadata_one_split"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    subdir = data_dir / "subdir"
+    subdir.mkdir(parents=True, exist_ok=True)
+
+    filename = data_dir / "file.txt"
+    shutil.copyfile(auto_text_file, filename)
+    filename2 = data_dir / "file2.txt"
+    shutil.copyfile(auto_text_file, filename2)
+    filename3 = subdir / "file3.txt"  # in subdir
+    shutil.copyfile(auto_text_file, filename3)
+
+    metadata_filename = data_dir / "metadata.jsonl"
+    metadata = textwrap.dedent(
+        """\
+        {"file_name": "file.txt", "additional_feature": "Dummy file"}
+        {"file_name": "file2.txt", "additional_feature": "Second dummy file"}
+        {"file_name": "./subdir/file3.txt", "additional_feature": "Third dummy file"}
+        """
+    )
+    with open(metadata_filename, "w", encoding="utf-8") as f:
+        f.write(metadata)
+    data_files_with_one_split_and_metadata = DataFilesDict.from_patterns(
+        get_data_patterns(str(data_dir)), data_dir.as_posix()
+    )
+    assert len(data_files_with_one_split_and_metadata) == 1
+    assert len(data_files_with_one_split_and_metadata["train"]) == 4
+    return data_files_with_one_split_and_metadata
+
+
+@pytest.fixture
+def data_files_with_two_splits_and_metadata(tmp_path, auto_text_file):
+    data_dir = tmp_path / "autofolder_data_dir_with_metadata_two_splits"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    train_dir = data_dir / "train"
+    train_dir.mkdir(parents=True, exist_ok=True)
+    test_dir = data_dir / "test"
+    test_dir.mkdir(parents=True, exist_ok=True)
+
+    filename = train_dir / "file.txt"  # train
+    shutil.copyfile(auto_text_file, filename)
+    filename2 = train_dir / "file2.txt"  # train
+    shutil.copyfile(auto_text_file, filename2)
+    filename3 = test_dir / "file3.txt"  # test
+    shutil.copyfile(auto_text_file, filename3)
+
+    train_metadata_filename = train_dir / "metadata.jsonl"
+    train_metadata = textwrap.dedent(
+        """\
+        {"file_name": "file.txt", "additional_feature": "Train dummy file"}
+        {"file_name": "file2.txt", "additional_feature": "Second train dummy file"}
+        """
+    )
+    with open(train_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(train_metadata)
+    test_metadata_filename = test_dir / "metadata.jsonl"
+    test_metadata = textwrap.dedent(
+        """\
+        {"file_name": "file3.txt", "additional_feature": "Test dummy file"}
+        """
+    )
+    with open(test_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(test_metadata)
+    data_files_with_two_splits_and_metadata = DataFilesDict.from_patterns(
+        get_data_patterns(str(data_dir)), data_dir.as_posix()
+    )
+    assert len(data_files_with_two_splits_and_metadata) == 2
+    assert len(data_files_with_two_splits_and_metadata["train"]) == 3
+    assert len(data_files_with_two_splits_and_metadata["test"]) == 2
+    return data_files_with_two_splits_and_metadata
+
+
+@pytest.fixture
+def data_files_with_zip_archives(tmp_path, auto_text_file):
+    data_dir = tmp_path / "autofolder_data_dir_with_zip_archives"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    archive_dir = data_dir / "archive"
+    archive_dir.mkdir(parents=True, exist_ok=True)
+    subdir = archive_dir / "subdir"
+    subdir.mkdir(parents=True, exist_ok=True)
+
+    filename = archive_dir / "file.txt"
+    shutil.copyfile(auto_text_file, filename)
+    filename2 = subdir / "file2.txt"  # in subdir
+    shutil.copyfile(auto_text_file, filename2)
+
+    metadata_filename = archive_dir / "metadata.jsonl"
+    metadata = textwrap.dedent(
+        """\
+        {"file_name": "file.txt", "additional_feature": "Dummy file"}
+        {"file_name": "subdir/file2.txt", "additional_feature": "Second dummy file"}
+        """
+    )
+    with open(metadata_filename, "w", encoding="utf-8") as f:
+        f.write(metadata)
+
+    shutil.make_archive(archive_dir, "zip", archive_dir)
+    shutil.rmtree(str(archive_dir))
+
+    data_files_with_zip_archives = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+
+    assert len(data_files_with_zip_archives) == 1
+    assert len(data_files_with_zip_archives["train"]) == 1
+    return data_files_with_zip_archives
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = FolderBasedBuilderConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = FolderBasedBuilderConfig(name="name", data_files=data_files)
+
+
+def test_inferring_labels_from_data_dirs(data_files_with_labels_no_metadata, cache_dir):
+    autofolder = DummyFolderBasedBuilder(
+        data_files=data_files_with_labels_no_metadata, cache_dir=cache_dir, drop_labels=False
+    )
+    gen_kwargs = autofolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    assert autofolder.info.features["label"] == ClassLabel(names=["class0", "class1"])
+    generator = autofolder._generate_examples(**gen_kwargs)
+    assert all(example["label"] in {"class0", "class1"} for _, example in generator)
+
+
+def test_default_folder_builder_not_usable(data_files_with_labels_no_metadata, cache_dir):
+    # builder would try to access non-existing attributes of a default `BuilderConfig` class
+    # as a custom one is not provided
+    with pytest.raises(AttributeError):
+        _ = FolderBasedBuilder(
+            data_files=data_files_with_labels_no_metadata,
+            cache_dir=cache_dir,
+        )
+
+
+# test that AutoFolder is extended for streaming when it's child class is instantiated:
+# see line 115 in src/datasets/streaming.py
+def test_streaming_patched():
+    _ = DummyFolderBasedBuilder(data_dir=".")
+    module = importlib.import_module(FolderBasedBuilder.__module__)
+    assert hasattr(module, "_patched_for_streaming")
+    assert module._patched_for_streaming
+
+
+@pytest.mark.parametrize("drop_metadata", [None, True, False])
+@pytest.mark.parametrize("drop_labels", [None, True, False])
+def test_generate_examples_drop_labels(
+    data_files_with_labels_no_metadata, auto_text_file, drop_metadata, drop_labels, cache_dir
+):
+    autofolder = DummyFolderBasedBuilder(
+        data_files=data_files_with_labels_no_metadata,
+        drop_metadata=drop_metadata,
+        drop_labels=drop_labels,
+        cache_dir=cache_dir,
+    )
+    gen_kwargs = autofolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    # removing labels explicitly requires drop_labels=True
+    assert gen_kwargs["add_labels"] is not bool(drop_labels)
+    assert gen_kwargs["add_metadata"] is False
+    generator = autofolder._generate_examples(**gen_kwargs)
+    if not drop_labels:
+        assert all(
+            example.keys() == {"base", "label"} and all(val is not None for val in example.values())
+            for _, example in generator
+        )
+    else:
+        assert all(
+            example.keys() == {"base"} and all(val is not None for val in example.values()) for _, example in generator
+        )
+
+
+@pytest.mark.parametrize("drop_metadata", [None, True, False])
+@pytest.mark.parametrize("drop_labels", [None, True, False])
+def test_generate_examples_drop_metadata(file_with_metadata, drop_metadata, drop_labels, cache_dir):
+    file, metadata_file = file_with_metadata
+    autofolder = DummyFolderBasedBuilder(
+        data_files=[file, metadata_file],
+        drop_metadata=drop_metadata,
+        drop_labels=drop_labels,
+        cache_dir=cache_dir,
+    )
+    gen_kwargs = autofolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    # since the dataset has metadata, removing the metadata explicitly requires drop_metadata=True
+    assert gen_kwargs["add_metadata"] is not bool(drop_metadata)
+    # since the dataset has metadata, adding the labels explicitly requires drop_labels=False
+    assert gen_kwargs["add_labels"] is False
+    generator = autofolder._generate_examples(**gen_kwargs)
+    expected_columns = {"base"}
+    if gen_kwargs["add_metadata"]:
+        expected_columns.add("additional_feature")
+    if gen_kwargs["add_labels"]:
+        expected_columns.add("label")
+    result = [example for _, example in generator]
+    assert len(result) == 1
+    example = result[0]
+    assert example.keys() == expected_columns
+    for column in expected_columns:
+        assert example[column] is not None
+
+
+@pytest.mark.parametrize("remote", [True, False])
+@pytest.mark.parametrize("drop_labels", [None, True, False])
+def test_data_files_with_different_levels_no_metadata(
+    data_files_with_different_levels_no_metadata, drop_labels, remote, cache_dir
+):
+    data_files = remote_files if remote else data_files_with_different_levels_no_metadata
+    autofolder = DummyFolderBasedBuilder(
+        data_files=data_files,
+        cache_dir=cache_dir,
+        drop_labels=drop_labels,
+    )
+    gen_kwargs = autofolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    generator = autofolder._generate_examples(**gen_kwargs)
+    if drop_labels is not False:
+        # with None (default) we should drop labels if files are on different levels in dir structure
+        assert "label" not in autofolder.info.features
+        assert all(example.keys() == {"base"} for _, example in generator)
+    else:
+        assert "label" in autofolder.info.features
+        assert isinstance(autofolder.info.features["label"], ClassLabel)
+        assert all(example.keys() == {"base", "label"} for _, example in generator)
+
+
+@pytest.mark.parametrize("remote", [False, True])
+@pytest.mark.parametrize("drop_labels", [None, True, False])
+def test_data_files_with_one_label_no_metadata(data_files_with_one_label_no_metadata, drop_labels, remote, cache_dir):
+    data_files = remote_files[:2] if remote else data_files_with_one_label_no_metadata
+    autofolder = DummyFolderBasedBuilder(
+        data_files=data_files,
+        cache_dir=cache_dir,
+        drop_labels=drop_labels,
+    )
+    gen_kwargs = autofolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    generator = autofolder._generate_examples(**gen_kwargs)
+    if drop_labels is not False:
+        # with None (default) we should drop labels if only one label is found (=if there is a single dir)
+        assert "label" not in autofolder.info.features
+        assert all(example.keys() == {"base"} for _, example in generator)
+    else:
+        assert "label" in autofolder.info.features
+        assert isinstance(autofolder.info.features["label"], ClassLabel)
+        assert all(example.keys() == {"base", "label"} for _, example in generator)
+
+
+@pytest.mark.parametrize("streaming", [False, True])
+@pytest.mark.parametrize("n_splits", [1, 2])
+def test_data_files_with_metadata_and_splits(
+    streaming, cache_dir, n_splits, data_files_with_one_split_and_metadata, data_files_with_two_splits_and_metadata
+):
+    data_files = data_files_with_one_split_and_metadata if n_splits == 1 else data_files_with_two_splits_and_metadata
+    autofolder = DummyFolderBasedBuilder(
+        data_files=data_files,
+        cache_dir=cache_dir,
+    )
+    download_manager = StreamingDownloadManager() if streaming else DownloadManager()
+    generated_splits = autofolder._split_generators(download_manager)
+    for (split, files), generated_split in zip(data_files.items(), generated_splits):
+        assert split == generated_split.name
+        expected_num_of_examples = len(files) - 1
+        generated_examples = list(autofolder._generate_examples(**generated_split.gen_kwargs))
+        assert len(generated_examples) == expected_num_of_examples
+        assert len({example["base"] for _, example in generated_examples}) == expected_num_of_examples
+        assert len({example["additional_feature"] for _, example in generated_examples}) == expected_num_of_examples
+        assert all(example["additional_feature"] is not None for _, example in generated_examples)
+
+
+@pytest.mark.parametrize("streaming", [False, True])
+def test_data_files_with_metadata_and_archives(streaming, cache_dir, data_files_with_zip_archives):
+    autofolder = DummyFolderBasedBuilder(data_files=data_files_with_zip_archives, cache_dir=cache_dir)
+    download_manager = StreamingDownloadManager() if streaming else DownloadManager()
+    generated_splits = autofolder._split_generators(download_manager)
+    for (split, files), generated_split in zip(data_files_with_zip_archives.items(), generated_splits):
+        assert split == generated_split.name
+        num_of_archives = len(files)
+        expected_num_of_examples = 2 * num_of_archives
+        generated_examples = list(autofolder._generate_examples(**generated_split.gen_kwargs))
+        assert len(generated_examples) == expected_num_of_examples
+        assert len({example["base"] for _, example in generated_examples}) == expected_num_of_examples
+        assert len({example["additional_feature"] for _, example in generated_examples}) == expected_num_of_examples
+        assert all(example["additional_feature"] is not None for _, example in generated_examples)
+
+
+def test_data_files_with_wrong_metadata_file_name(cache_dir, tmp_path, auto_text_file):
+    data_dir = tmp_path / "data_dir_with_bad_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(auto_text_file, data_dir / "file.txt")
+    metadata_filename = data_dir / "bad_metadata.jsonl"  # bad file
+    metadata = textwrap.dedent(
+        """\
+        {"file_name": "file.txt", "additional_feature": "Dummy file"}
+        """
+    )
+    with open(metadata_filename, "w", encoding="utf-8") as f:
+        f.write(metadata)
+
+    data_files_with_bad_metadata = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+    autofolder = DummyFolderBasedBuilder(data_files=data_files_with_bad_metadata, cache_dir=cache_dir)
+    gen_kwargs = autofolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    generator = autofolder._generate_examples(**gen_kwargs)
+    assert all("additional_feature" not in example for _, example in generator)
+
+
+def test_data_files_with_custom_file_name_column_in_metadata_file(cache_dir, tmp_path, auto_text_file):
+    data_dir = tmp_path / "data_dir_with_custom_file_name_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(auto_text_file, data_dir / "file.txt")
+    metadata_filename = data_dir / "metadata.jsonl"
+    metadata = textwrap.dedent(  # with bad column "bad_file_name" instead of "file_name"
+        """\
+        {"text_file_name": "file.txt", "additional_feature": "Dummy file"}
+        """
+    )
+    with open(metadata_filename, "w", encoding="utf-8") as f:
+        f.write(metadata)
+
+    data_files_with_bad_metadata = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+    autofolder = DummyFolderBasedBuilder(data_files=data_files_with_bad_metadata, cache_dir=cache_dir)
+    gen_kwargs = autofolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    generator = autofolder._generate_examples(**gen_kwargs)
+    assert all("text" in example and "text_file_name" not in example for _, example in generator)
+
+
+def test_data_files_with_custom_file_names_column_in_metadata_file_large_string_list(
+    cache_dir, tmp_path, auto_text_file
+):
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    data_dir = tmp_path / "data_dir_with_custom_file_names_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(auto_text_file, data_dir / "file.txt")
+    metadata_filename = data_dir / "metadata.parquet"
+    pq.write_table(
+        pa.Table.from_arrays(
+            [
+                pa.array([["file.txt"]], type=pa.list_(pa.large_string())),
+                pa.array(["Dummy file"], type=pa.large_string()),
+            ],
+            names=["text_file_names", "additional_feature"],
+        ),
+        metadata_filename,
+    )
+
+    data_files_with_metadata = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+    autofolder = DummyFolderBasedBuilder(data_files=data_files_with_metadata, cache_dir=cache_dir)
+    gen_kwargs = autofolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    generator = autofolder._generate_examples(**gen_kwargs)
+    examples = [example for _, example in generator]
+    assert len(examples) == 1
+    assert "text" in examples[0] and "text_file_names" not in examples[0]
+    assert len(examples[0]["text"]) == 1 and examples[0]["text"][0].endswith("file.txt")

tests/packaged_modules/test_hdf5.py

@@ -0,0 +1,829 @@
+import h5py
+import numpy as np
+import pytest
+
+from datasets import Array2D, Array3D, Array4D, Features, List, Value, load_dataset
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesList
+from datasets.exceptions import DatasetGenerationError
+from datasets.packaged_modules.hdf5.hdf5 import HDF5, HDF5Config
+
+
+@pytest.fixture
+def hdf5_file(tmp_path):
+    """Create a basic HDF5 file with numeric datasets."""
+    filename = tmp_path / "basic.h5"
+    n_rows = 5
+
+    with h5py.File(filename, "w") as f:
+        f.create_dataset("int32", data=np.arange(n_rows, dtype=np.int32))
+        f.create_dataset("float32", data=np.arange(n_rows, dtype=np.float32) / 10.0)
+        f.create_dataset("bool", data=np.array([True, False, True, False, True]))
+
+    return str(filename)
+
+
+@pytest.fixture
+def hdf5_file_with_groups(tmp_path):
+    """Create an HDF5 file with nested groups."""
+    filename = tmp_path / "nested.h5"
+    n_rows = 3
+
+    with h5py.File(filename, "w") as f:
+        f.create_dataset("root_data", data=np.arange(n_rows, dtype=np.int32))
+        grp = f.create_group("group1")
+        grp.create_dataset("group_data", data=np.arange(n_rows, dtype=np.float32))
+        subgrp = grp.create_group("subgroup")
+        subgrp.create_dataset("sub_data", data=np.arange(n_rows, dtype=np.int64))
+
+    return str(filename)
+
+
+@pytest.fixture
+def hdf5_file_with_arrays(tmp_path):
+    """Create an HDF5 file with multi-dimensional arrays."""
+    filename = tmp_path / "arrays.h5"
+    n_rows = 4
+
+    with h5py.File(filename, "w") as f:
+        # 2D array (should become Array2D)
+        f.create_dataset("matrix_2d", data=np.random.randn(n_rows, 3, 4).astype(np.float32))
+        # 3D array (should become Array3D)
+        f.create_dataset("tensor_3d", data=np.random.randn(n_rows, 2, 3, 4).astype(np.float64))
+        # 4D array (should become Array4D)
+        f.create_dataset("tensor_4d", data=np.random.randn(n_rows, 2, 3, 4, 5).astype(np.float32))
+        # 5D array (should become Array5D)
+        f.create_dataset("tensor_5d", data=np.random.randn(n_rows, 2, 3, 4, 5, 6).astype(np.float64))
+        # 1D array (should become Value)
+        f.create_dataset("vector_1d", data=np.random.randn(n_rows, 10).astype(np.float32))
+
+    return str(filename)
+
+
+@pytest.fixture
+def hdf5_file_with_different_dtypes(tmp_path):
+    """Create an HDF5 file with various numeric dtypes."""
+    filename = tmp_path / "dtypes.h5"
+    n_rows = 3
+
+    with h5py.File(filename, "w") as f:
+        f.create_dataset("int8", data=np.arange(n_rows, dtype=np.int8))
+        f.create_dataset("int16", data=np.arange(n_rows, dtype=np.int16))
+        f.create_dataset("int64", data=np.arange(n_rows, dtype=np.int64))
+        f.create_dataset("uint8", data=np.arange(n_rows, dtype=np.uint8))
+        f.create_dataset("uint16", data=np.arange(n_rows, dtype=np.uint16))
+        f.create_dataset("uint32", data=np.arange(n_rows, dtype=np.uint32))
+        f.create_dataset("uint64", data=np.arange(n_rows, dtype=np.uint64))
+        f.create_dataset("float16", data=np.arange(n_rows, dtype=np.float16) / 10.0)
+        f.create_dataset("float64", data=np.arange(n_rows, dtype=np.float64) / 10.0)
+        f.create_dataset("bytes", data=np.array([b"row_%d" % i for i in range(n_rows)], dtype="S10"))
+
+    return str(filename)
+
+
+@pytest.fixture
+def hdf5_file_with_vlen_arrays(tmp_path):
+    """Create an HDF5 file with variable-length arrays using HDF5's vlen_dtype."""
+    filename = tmp_path / "vlen.h5"
+    n_rows = 4
+
+    with h5py.File(filename, "w") as f:
+        # Variable-length arrays of different sizes using vlen_dtype
+        vlen_arrays = [[1, 2, 3], [4, 5], [6, 7, 8, 9], [10]]
+        # Create variable-length int dataset using vlen_dtype
+        dt = h5py.vlen_dtype(np.dtype("int32"))
+        dset = f.create_dataset("vlen_ints", (n_rows,), dtype=dt)
+        for i, arr in enumerate(vlen_arrays):
+            dset[i] = arr
+
+        # Mixed types (some empty arrays) - use variable-length with empty arrays
+        mixed_data = [
+            [1, 2, 3],
+            [],  # Empty array
+            [4, 5],
+            [6],
+        ]
+        dt_mixed = h5py.vlen_dtype(np.dtype("int32"))
+        dset_mixed = f.create_dataset("mixed_data", (n_rows,), dtype=dt_mixed)
+        for i, arr in enumerate(mixed_data):
+            dset_mixed[i] = arr
+
+    return str(filename)
+
+
+@pytest.fixture
+def hdf5_file_with_variable_length_strings(tmp_path):
+    """Create an HDF5 file with variable-length string datasets."""
+    filename = tmp_path / "var_strings.h5"
+    n_rows = 4
+
+    with h5py.File(filename, "w") as f:
+        # Variable-length string dataset
+        var_strings = ["short", "medium length string", "very long string with many characters", "tiny"]
+        # Create variable-length string dataset using vlen_dtype
+        dt = h5py.vlen_dtype(str)
+        dset = f.create_dataset("var_strings", (n_rows,), dtype=dt)
+        for i, s in enumerate(var_strings):
+            dset[i] = s
+
+        # Variable-length bytes dataset
+        var_bytes = [b"short", b"medium length bytes", b"very long bytes with many characters", b"tiny"]
+        dt_bytes = h5py.vlen_dtype(bytes)
+        dset_bytes = f.create_dataset("var_bytes", (n_rows,), dtype=dt_bytes)
+        for i, b in enumerate(var_bytes):
+            dset_bytes[i] = b
+
+    return str(filename)
+
+
+@pytest.fixture
+def hdf5_file_with_complex_data(tmp_path):
+    """Create an HDF5 file with complex number datasets."""
+    filename = tmp_path / "complex.h5"
+
+    with h5py.File(filename, "w") as f:
+        # Complex numbers
+        complex_data = np.array([1 + 2j, 3 + 4j, 5 + 6j, 7 + 8j], dtype=np.complex64)
+        f.create_dataset("complex_64", data=complex_data)
+
+        # Complex double precision
+        complex_double = np.array([1.5 + 2.5j, 3.5 + 4.5j, 5.5 + 6.5j, 7.5 + 8.5j], dtype=np.complex128)
+        f.create_dataset("complex_128", data=complex_double)
+
+        # Complex array
+        complex_array = np.array(
+            [[1 + 2j, 3 + 4j], [5 + 6j, 7 + 8j], [9 + 10j, 11 + 12j], [13 + 14j, 15 + 16j]], dtype=np.complex64
+        )
+        f.create_dataset("complex_array", data=complex_array)
+
+    return str(filename)
+
+
+@pytest.fixture
+def hdf5_file_with_compound_data(tmp_path):
+    """Create an HDF5 file with compound/structured datasets."""
+    filename = tmp_path / "compound.h5"
+
+    with h5py.File(filename, "w") as f:
+        # Simple compound type
+        dt_simple = np.dtype([("x", "i4"), ("y", "f8")])
+        compound_simple = np.array([(1, 2.5), (3, 4.5), (5, 6.5)], dtype=dt_simple)
+        f.create_dataset("simple_compound", data=compound_simple)
+
+        # Compound type with complex numbers
+        dt_complex = np.dtype([("real", "f4"), ("imag", "f4")])
+        compound_complex = np.array([(1.0, 2.0), (3.0, 4.0), (5.0, 6.0)], dtype=dt_complex)
+        f.create_dataset("complex_compound", data=compound_complex)
+
+        # Nested compound type
+        dt_nested = np.dtype([("position", [("x", "i4"), ("y", "i4")]), ("velocity", [("vx", "f4"), ("vy", "f4")])])
+        compound_nested = np.array([((1, 2), (1.5, 2.5)), ((3, 4), (3.5, 4.5)), ((5, 6), (5.5, 6.5))], dtype=dt_nested)
+        f.create_dataset("nested_compound", data=compound_nested)
+
+    return str(filename)
+
+
+@pytest.fixture
+def hdf5_file_with_compound_complex_arrays(tmp_path):
+    """Create an HDF5 file with compound datasets containing complex arrays."""
+    filename = tmp_path / "compound_complex_arrays.h5"
+
+    with h5py.File(filename, "w") as f:
+        # Compound type with complex arrays
+        dt_complex_arrays = np.dtype(
+            [
+                ("position", [("x", "i4"), ("y", "i4")]),
+                ("complex_field", "c8"),
+                ("complex_array", "c8", (2, 3)),
+                ("nested_complex", [("real", "f4"), ("imag", "f4")]),
+            ]
+        )
+
+        # Create data with complex numbers
+        compound_data = np.array(
+            [
+                (
+                    (1, 2),
+                    1.0 + 2.0j,
+                    [[1.0 + 2.0j, 3.0 + 4.0j, 5.0 + 6.0j], [7.0 + 8.0j, 9.0 + 10.0j, 11.0 + 12.0j]],
+                    (1.5, 2.5),
+                ),
+                (
+                    (3, 4),
+                    3.0 + 4.0j,
+                    [[13.0 + 14.0j, 15.0 + 16.0j, 17.0 + 18.0j], [19.0 + 20.0j, 21.0 + 22.0j, 23.0 + 24.0j]],
+                    (3.5, 4.5),
+                ),
+                (
+                    (5, 6),
+                    5.0 + 6.0j,
+                    [[25.0 + 26.0j, 27.0 + 28.0j, 29.0 + 30.0j], [31.0 + 32.0j, 33.0 + 34.0j, 35.0 + 36.0j]],
+                    (5.5, 6.5),
+                ),
+            ],
+            dtype=dt_complex_arrays,
+        )
+
+        f.create_dataset("compound_with_complex", data=compound_data)
+
+    return str(filename)
+
+
+@pytest.fixture
+def hdf5_file_with_mismatched_lengths(tmp_path):
+    """Create an HDF5 file with datasets of different lengths (should raise error)."""
+    filename = tmp_path / "mismatched.h5"
+
+    with h5py.File(filename, "w") as f:
+        f.create_dataset("data1", data=np.arange(5, dtype=np.int32))
+        # Dataset with 3 rows (mismatched)
+        f.create_dataset("data2", data=np.arange(3, dtype=np.int32))
+        f.create_dataset("data3", data=np.random.randn(5, 3, 4).astype(np.float32))
+        f.create_dataset("data4", data=np.arange(5, dtype=np.float64) / 10.0)
+        f.create_dataset("data5", data=np.array([True, False, True, False, True]))
+        var_strings = ["short", "medium length", "very long string", "tiny", "another string"]
+        dt = h5py.vlen_dtype(str)
+        dset = f.create_dataset("data6", (5,), dtype=dt)
+        for i, s in enumerate(var_strings):
+            dset[i] = s
+
+    return str(filename)
+
+
+@pytest.fixture
+def hdf5_file_with_zero_dimensions(tmp_path):
+    """Create an HDF5 file with zero dimensions (should be handled gracefully)."""
+    filename = tmp_path / "zero_dims.h5"
+
+    with h5py.File(filename, "w") as f:
+        # Create a dataset with a zero dimension
+        f.create_dataset("zero_dim", data=np.zeros((3, 0, 2), dtype=np.float32))
+        # Create a dataset with zero in the middle dimension
+        f.create_dataset("zero_middle", data=np.zeros((3, 0), dtype=np.int32))
+        # Create a dataset with zero in the last dimension
+        f.create_dataset("zero_last", data=np.zeros((3, 2, 0), dtype=np.float64))
+
+    return str(filename)
+
+
+@pytest.fixture
+def empty_hdf5_file(tmp_path):
+    """Create an HDF5 file with no datasets (should warn and skip)."""
+    filename = tmp_path / "empty.h5"
+
+    with h5py.File(filename, "w") as f:
+        # Create only groups, no datasets
+        f.create_group("empty_group")
+        grp = f.create_group("another_group")
+        grp.create_group("subgroup")
+
+    return str(filename)
+
+
+@pytest.fixture
+def hdf5_file_with_mixed_data_types(tmp_path):
+    """Create an HDF5 file with mixed data types in the same file."""
+    filename = tmp_path / "mixed.h5"
+    n_rows = 3
+
+    with h5py.File(filename, "w") as f:
+        # Regular numeric data
+        f.create_dataset("regular_int", data=np.arange(n_rows, dtype=np.int32))
+        f.create_dataset("regular_float", data=np.arange(n_rows, dtype=np.float32))
+
+        # Complex data
+        complex_data = np.array([1 + 2j, 3 + 4j, 5 + 6j], dtype=np.complex64)
+        f.create_dataset("complex_data", data=complex_data)
+
+        # Compound data
+        dt_compound = np.dtype([("x", "i4"), ("y", "f8")])
+        compound_data = np.array([(1, 2.5), (3, 4.5), (5, 6.5)], dtype=dt_compound)
+        f.create_dataset("compound_data", data=compound_data)
+
+    return str(filename)
+
+
+def test_config_raises_when_invalid_name():
+    """Test that invalid config names raise an error."""
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = HDF5Config(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files):
+    """Test that invalid data_files parameter raises an error."""
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = HDF5Config(name="name", data_files=data_files)
+
+
+def test_hdf5_basic_functionality(hdf5_file):
+    """Test basic HDF5 loading with simple numeric datasets."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file], split="train")
+
+    assert "int32" in dataset.column_names
+    assert "float32" in dataset.column_names
+    assert "bool" in dataset.column_names
+
+    assert np.asarray(dataset.data["int32"]).dtype == np.int32
+    assert np.asarray(dataset.data["float32"]).dtype == np.float32
+    assert np.asarray(dataset.data["bool"]).dtype == np.bool_
+
+    assert dataset["int32"] == [0, 1, 2, 3, 4]
+    float32_data = dataset["float32"]
+    expected_float32 = [0.0, 0.1, 0.2, 0.3, 0.4]
+    np.testing.assert_allclose(float32_data, expected_float32, rtol=1e-6)
+
+
+def test_hdf5_nested_groups(hdf5_file_with_groups):
+    """Test HDF5 loading with nested groups."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_groups], split="train")
+
+    expected_columns = {"root_data", "group1"}
+    assert set(dataset.column_names) == expected_columns
+
+    # Check data
+    root_data = dataset["root_data"]
+    group1_data = dataset["group1"]
+    assert root_data == [0, 1, 2]
+    assert group1_data == [
+        {"group_data": 0.0, "subgroup": {"sub_data": 0}},
+        {"group_data": 1.0, "subgroup": {"sub_data": 1}},
+        {"group_data": 2.0, "subgroup": {"sub_data": 2}},
+    ]
+
+
+def test_hdf5_multi_dimensional_arrays(hdf5_file_with_arrays):
+    """Test HDF5 loading with multi-dimensional arrays."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_arrays], split="train")
+
+    expected_columns = {"matrix_2d", "tensor_3d", "tensor_4d", "tensor_5d", "vector_1d"}
+    assert set(dataset.column_names) == expected_columns
+
+    # Check shapes
+    matrix_2d = dataset["matrix_2d"]
+    assert len(matrix_2d) == 4  # 4 rows
+    assert len(matrix_2d[0]) == 3  # 3 rows in each matrix
+    assert len(matrix_2d[0][0]) == 4  # 4 columns in each matrix
+
+
+def test_hdf5_vlen_arrays(hdf5_file_with_vlen_arrays):
+    """Test HDF5 loading with variable-length arrays (int32)."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_vlen_arrays], split="train")
+
+    expected_columns = {"vlen_ints", "mixed_data"}
+    assert set(dataset.column_names) == expected_columns
+
+    # Check vlen_ints data
+    vlen_ints = dataset["vlen_ints"]
+    assert len(vlen_ints) == 4
+    assert vlen_ints[0] == [1, 2, 3]
+    assert vlen_ints[1] == [4, 5]
+    assert vlen_ints[2] == [6, 7, 8, 9]
+    assert vlen_ints[3] == [10]
+
+    # Check mixed_data (with None values)
+    mixed_data = dataset["mixed_data"]
+    assert len(mixed_data) == 4
+    assert mixed_data[0] == [1, 2, 3]
+    assert mixed_data[1] == []  # Empty array instead of None
+    assert mixed_data[2] == [4, 5]
+    assert mixed_data[3] == [6]
+
+
+def test_hdf5_variable_length_strings(hdf5_file_with_variable_length_strings):
+    """Test HDF5 loading with variable-length string datasets."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_variable_length_strings], split="train")
+    expected_columns = {"var_strings", "var_bytes"}
+    assert set(dataset.column_names) == expected_columns
+
+    # Check variable-length strings (converted to strings for usability)
+    var_strings = dataset["var_strings"]
+    assert len(var_strings) == 4
+    assert var_strings[0] == "short"
+    assert var_strings[1] == "medium length string"
+    assert var_strings[2] == "very long string with many characters"
+    assert var_strings[3] == "tiny"
+
+    # Check variable-length bytes (converted to strings for usability)
+    var_bytes = dataset["var_bytes"]
+    assert len(var_bytes) == 4
+    assert var_bytes[0] == "short"
+    assert var_bytes[1] == "medium length bytes"
+    assert var_bytes[2] == "very long bytes with many characters"
+    assert var_bytes[3] == "tiny"
+
+
+def test_hdf5_different_dtypes(hdf5_file_with_different_dtypes):
+    """Test HDF5 loading with various numeric dtypes."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_different_dtypes], split="train")
+    expected_columns = {"int8", "int16", "int64", "uint8", "uint16", "uint32", "uint64", "float16", "float64", "bytes"}
+    assert set(dataset.column_names) == expected_columns
+
+    # Check specific dtypes
+    int8_data = dataset["int8"]
+    assert int8_data == [0, 1, 2]
+
+    bytes_data = dataset["bytes"]
+    assert bytes_data == [b"row_0", b"row_1", b"row_2"]
+
+
+def test_hdf5_batch_processing(hdf5_file):
+    """Test HDF5 loading with custom batch size."""
+    config = HDF5Config(batch_size=2)
+    hdf5 = HDF5()
+    hdf5.config = config
+    generator = hdf5._generate_tables([hdf5_file])
+
+    tables = list(generator)
+    # Should have 3 batches: [0,1], [2,3], [4]
+    assert len(tables) == 3
+
+    # Check first batch
+    _, first_batch = tables[0]
+    assert len(first_batch) == 2
+
+    # Check last batch
+    _, last_batch = tables[2]
+    assert len(last_batch) == 1
+
+
+def test_hdf5_column_filtering(hdf5_file_with_groups):
+    """Test HDF5 loading with column filtering."""
+    features = Features({"root_data": Value("int32"), "group1": Features({"group_data": Value("float32")})})
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_groups], split="train", features=features)
+
+    expected_columns = {"root_data", "group1"}
+    assert set(dataset.column_names) == expected_columns
+
+    # Check that subgroup is filtered out
+    group1_data = dataset["group1"]
+    assert group1_data == [
+        {"group_data": 0.0},
+        {"group_data": 1.0},
+        {"group_data": 2.0},
+    ]
+
+
+def test_hdf5_feature_specification(hdf5_file):
+    """Test HDF5 loading with explicit feature specification."""
+    features = Features({"int32": Value("int32"), "float32": Value("float64"), "bool": Value("bool")})
+    dataset = load_dataset("hdf5", data_files=[hdf5_file], split="train", features=features)
+
+    # Check that features are properly cast
+    assert np.asarray(dataset.data["float32"]).dtype == np.float64
+    assert np.asarray(dataset.data["int32"]).dtype == np.int32
+    assert np.asarray(dataset.data["bool"]).dtype == np.bool_
+
+
+def test_hdf5_mismatched_lengths_error(hdf5_file_with_mismatched_lengths):
+    """Test that mismatched dataset lengths raise an error."""
+    with pytest.raises(DatasetGenerationError) as exc_info:
+        load_dataset("hdf5", data_files=[hdf5_file_with_mismatched_lengths], split="train")
+
+    assert isinstance(exc_info.value.__cause__, ValueError)
+    assert "3 but expected 5" in str(exc_info.value.__cause__)
+
+
+def test_hdf5_zero_dimensions_handling(hdf5_file_with_zero_dimensions, caplog):
+    """Test that zero dimensions are handled gracefully."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_zero_dimensions], split="train")
+
+    expected_columns = {"zero_dim", "zero_middle", "zero_last"}
+    assert set(dataset.column_names) == expected_columns
+
+    # Check that the data is loaded (should be empty arrays)
+    zero_dim_data = dataset["zero_dim"]
+    assert len(zero_dim_data) == 3  # 3 rows
+    assert all(len(row) == 0 for row in zero_dim_data)  # Each row is empty
+
+    # Check that shape info is lost
+    assert all(isinstance(col, List) and col.length == -1 for col in dataset.features.values())
+
+    # Check for the warnings
+    assert (
+        len(
+            [
+                record.message
+                for record in caplog.records
+                if record.levelname == "WARNING" and "dimension with size 0" in record.message
+            ]
+        )
+        == 3
+    )
+
+
+def test_hdf5_empty_file_warning(empty_hdf5_file, hdf5_file_with_arrays, caplog):
+    """Test that empty files (no datasets) are skipped with a warning."""
+    load_dataset("hdf5", data_files=[hdf5_file_with_arrays, empty_hdf5_file], split="train")
+
+    # Check that warning was logged
+    assert any(
+        record.levelname == "WARNING" and "contains no data, skipping" in record.message for record in caplog.records
+    )
+
+
+def test_hdf5_feature_inference(hdf5_file_with_arrays):
+    """Test automatic feature inference from HDF5 datasets."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_arrays], split="train")
+
+    # Check that features were inferred
+    assert dataset.features is not None
+
+    # Check specific feature types
+    features = dataset.features
+    # (n_rows, 3, 4) -> Array2D with shape (3, 4)
+    assert isinstance(features["matrix_2d"], Array2D)
+    assert features["matrix_2d"].shape == (3, 4)
+    # (n_rows, 2, 3, 4) -> Array3D with shape (2, 3, 4)
+    assert isinstance(features["tensor_3d"], Array3D)
+    assert features["tensor_3d"].shape == (2, 3, 4)
+    # (n_rows, 2, 3, 4, 5) -> Array4D with shape (2, 3, 4, 5)
+    assert isinstance(features["tensor_4d"], Array4D)
+    assert features["tensor_4d"].shape == (2, 3, 4, 5)
+    # (n_rows, 10) -> List of length 10
+    assert isinstance(features["vector_1d"], List)
+    assert features["vector_1d"].length == 10
+
+
+def test_hdf5_vlen_feature_inference(hdf5_file_with_vlen_arrays):
+    """Test automatic feature inference from variable-length HDF5 datasets."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_vlen_arrays], split="train")
+
+    # Check that features were inferred
+    assert dataset.features is not None
+
+    # Check specific feature types for variable-length arrays
+    features = dataset.features
+    # Variable-length arrays should become List features by default (for small datasets)
+    assert isinstance(features["vlen_ints"], List)
+    assert isinstance(features["mixed_data"], List)
+
+    # Check that the inner feature types are correct
+    assert isinstance(features["vlen_ints"].feature, Value)
+    assert features["vlen_ints"].feature.dtype == "int32"
+    assert isinstance(features["mixed_data"].feature, Value)
+    assert features["mixed_data"].feature.dtype == "int32"
+
+
+def test_hdf5_variable_string_feature_inference(hdf5_file_with_variable_length_strings):
+    """Test automatic feature inference from variable-length string datasets."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_variable_length_strings], split="train")
+
+    # Check that features were inferred
+    assert dataset.features is not None
+
+    # Check specific feature types for variable-length strings
+    features = dataset.features
+    # Variable-length strings should become Value("string") features
+    assert isinstance(features["var_strings"], Value)
+    assert isinstance(features["var_bytes"], Value)
+
+    # Check that the feature types are correct
+    assert features["var_strings"].dtype == "string"
+    assert features["var_bytes"].dtype == "string"
+
+
+def test_hdf5_invalid_features(hdf5_file_with_arrays):
+    """Test that invalid features raise an error."""
+    features = Features({"fakefeature": Value("int32")})
+    with pytest.raises(ValueError):
+        load_dataset("hdf5", data_files=[hdf5_file_with_arrays], split="train", features=features)
+
+    # try with one valid and one invalid feature
+    features = Features({"matrix_2d": Array2D(shape=(3, 4), dtype="float32"), "fakefeature": Value("int32")})
+    with pytest.raises(DatasetGenerationError):
+        load_dataset("hdf5", data_files=[hdf5_file_with_arrays], split="train", features=features)
+
+
+def test_hdf5_no_data_files_error():
+    """Test that missing data_files raises an error."""
+    config = HDF5Config(name="test", data_files=None)
+    hdf5 = HDF5()
+    hdf5.config = config
+
+    with pytest.raises(ValueError, match="At least one data file must be specified"):
+        hdf5._split_generators(None)
+
+
+def test_hdf5_complex_numbers(hdf5_file_with_complex_data):
+    """Test HDF5 loading with complex number datasets."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_complex_data], split="train")
+
+    # Check that complex numbers are represented as nested Features
+    expected_columns = {
+        "complex_64",
+        "complex_128",
+        "complex_array",
+    }
+    assert set(dataset.column_names) == expected_columns
+
+    # Check complex_64 data
+    complex_64_data = dataset["complex_64"]
+    assert len(complex_64_data) == 4
+    assert complex_64_data[0] == {"real": 1.0, "imag": 2.0}
+    assert complex_64_data[1] == {"real": 3.0, "imag": 4.0}
+    assert complex_64_data[2] == {"real": 5.0, "imag": 6.0}
+    assert complex_64_data[3] == {"real": 7.0, "imag": 8.0}
+
+    assert np.asarray(dataset.data["complex_64"].flatten()[0]).dtype == np.float32
+    assert np.asarray(dataset.data["complex_64"].flatten()[1]).dtype == np.float32
+    assert (np.asarray(dataset.data["complex_64"].flatten()[0]) == np.array([1, 3, 5, 7], dtype=np.float32)).all()
+    assert (np.asarray(dataset.data["complex_64"].flatten()[1]) == np.array([2, 4, 6, 8], dtype=np.float32)).all()
+
+    assert np.asarray(dataset.data["complex_128"].flatten()[0]).dtype == np.float64
+    assert np.asarray(dataset.data["complex_128"].flatten()[1]).dtype == np.float64
+    assert (
+        np.asarray(dataset.data["complex_128"].flatten()[0]) == np.array([1.5, 3.5, 5.5, 7.5], dtype=np.float64)
+    ).all()
+    assert (
+        np.asarray(dataset.data["complex_128"].flatten()[1]) == np.array([2.5, 4.5, 6.5, 8.5], dtype=np.float64)
+    ).all()
+
+
+def test_hdf5_compound_types(hdf5_file_with_compound_data):
+    """Test HDF5 loading with compound/structured datasets."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_compound_data], split="train")
+
+    # Check that compound types are represented as nested structures
+    expected_columns = {
+        "simple_compound",
+        "complex_compound",
+        "nested_compound",
+    }
+    assert set(dataset.column_names) == expected_columns
+
+    # Check simple compound data
+    simple_compound_data = dataset["simple_compound"]
+    assert len(simple_compound_data) == 3
+    assert simple_compound_data[0] == {"x": 1, "y": 2.5}
+    assert simple_compound_data[1] == {"x": 3, "y": 4.5}
+    assert simple_compound_data[2] == {"x": 5, "y": 6.5}
+
+
+def test_hdf5_feature_inference_complex(hdf5_file_with_complex_data):
+    """Test automatic feature inference for complex datasets."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_complex_data], split="train")
+
+    # Check that features were inferred correctly
+    assert dataset.features is not None
+    features = dataset.features
+
+    # Check complex number features
+    assert "complex_64" in features
+    # Complex features are represented as dict, not Features object
+    assert isinstance(features["complex_64"], dict)
+    assert features["complex_64"]["real"] == Value("float32")
+    assert features["complex_64"]["imag"] == Value("float32")
+
+
+def test_hdf5_feature_inference_compound(hdf5_file_with_compound_data):
+    """Test automatic feature inference for compound datasets."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_compound_data], split="train")
+
+    # Check that features were inferred correctly
+    assert dataset.features is not None
+    features = dataset.features
+
+    # Check compound type features
+    assert "simple_compound" in features
+    # Compound features are represented as dict, not Features object
+    assert isinstance(features["simple_compound"], dict)
+    assert features["simple_compound"]["x"] == Value("int32")
+    assert features["simple_compound"]["y"] == Value("float64")
+
+
+def test_hdf5_mixed_data_types(hdf5_file_with_mixed_data_types):
+    """Test HDF5 loading with mixed data types in the same file."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_mixed_data_types], split="train")
+
+    # Check all expected columns are present
+    expected_columns = {
+        "regular_int",
+        "regular_float",
+        "complex_data",
+        "compound_data",
+    }
+    assert set(dataset.column_names) == expected_columns
+
+    # Check data types
+    assert dataset["regular_int"] == [0, 1, 2]
+    assert len(dataset["complex_data"]) == 3
+    assert len(dataset["compound_data"]) == 3
+
+
+def test_hdf5_mismatched_lengths_with_column_filtering(hdf5_file_with_mismatched_lengths):
+    """Test that mismatched dataset lengths are ignored when the mismatched dataset is excluded via columns config."""
+    # Test 1: Include only the first dataset
+    features = Features({"data1": Value("int32")})
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_mismatched_lengths], split="train", features=features)
+
+    # Should work without error since we're only including the first dataset
+    expected_columns = {"data1"}
+    assert set(dataset.column_names) == expected_columns
+    assert "data2" not in dataset.column_names
+
+    # Check the data
+    data1_values = dataset["data1"]
+    assert data1_values == [0, 1, 2, 3, 4]
+
+    # Test 2: Include multiple compatible datasets (all with 5 rows)
+    features = Features(
+        {
+            "data1": Value("int32"),
+            "data3": Array2D(shape=(3, 4), dtype="float32"),
+            "data4": Value("float64"),
+            "data5": Value("bool"),
+            "data6": Value("string"),
+        }
+    )
+    dataset2 = load_dataset("hdf5", data_files=[hdf5_file_with_mismatched_lengths], split="train", features=features)
+
+    # Should work without error since we're excluding the mismatched dataset
+    expected_columns2 = {"data1", "data3", "data4", "data5", "data6"}
+    assert set(dataset2.column_names) == expected_columns2
+    assert "data2" not in dataset2.column_names
+
+    # Check data types and values
+    assert dataset2["data1"] == [0, 1, 2, 3, 4]  # int32
+    assert len(dataset2["data3"]) == 5  # Array2D
+    assert len(dataset2["data3"][0]) == 3  # 3 rows in each 2D array
+    assert len(dataset2["data3"][0][0]) == 4  # 4 columns in each 2D array
+    np.testing.assert_allclose(dataset2["data4"], [0.0, 0.1, 0.2, 0.3, 0.4], rtol=1e-6)  # float64
+    assert dataset2["data5"] == [True, False, True, False, True]  # boolean
+    assert dataset2["data6"] == [
+        "short",
+        "medium length",
+        "very long string",
+        "tiny",
+        "another string",
+    ]  # vlen string
+
+
+def test_hdf5_compound_with_complex_arrays(hdf5_file_with_compound_complex_arrays):
+    """Test HDF5 loading with compound datasets containing complex arrays."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_compound_complex_arrays], split="train")
+
+    # Check that compound types with complex arrays are represented as nested structures
+    expected_columns = {"compound_with_complex"}
+    assert set(dataset.column_names) == expected_columns
+
+    # Check compound data with complex arrays
+    compound_data = dataset["compound_with_complex"]
+    assert len(compound_data) == 3
+
+    # Check first row
+    first_row = compound_data[0]
+    assert first_row["position"]["x"] == 1
+    assert first_row["position"]["y"] == 2
+
+    # Check complex field (should be represented as real/imag structure)
+    assert first_row["complex_field"]["real"] == 1.0
+    assert first_row["complex_field"]["imag"] == 2.0
+
+    # Check complex array (should be represented as nested real/imag structures)
+    complex_array = first_row["complex_array"]
+    assert len(complex_array["real"]) == 2  # 2 rows
+    assert len(complex_array["real"][0]) == 3  # 3 columns
+
+    # Check first element of complex array
+    assert complex_array["real"][0][0] == 1.0
+    assert complex_array["imag"][0][0] == 2.0
+
+    # Check nested complex field
+    assert first_row["nested_complex"]["real"] == 1.5
+    assert first_row["nested_complex"]["imag"] == 2.5
+
+
+def test_hdf5_feature_inference_compound_complex_arrays(hdf5_file_with_compound_complex_arrays):
+    """Test automatic feature inference for compound datasets with complex arrays."""
+    dataset = load_dataset("hdf5", data_files=[hdf5_file_with_compound_complex_arrays], split="train")
+
+    # Check that features were inferred correctly
+    assert dataset.features is not None
+    features = dataset.features
+
+    # Check compound type features with complex arrays
+    assert "compound_with_complex" in features
+
+    # Check nested structure
+    compound_features = features["compound_with_complex"]
+    assert "position" in compound_features
+    assert "complex_field" in compound_features
+    assert "complex_array" in compound_features
+    assert "nested_complex" in compound_features
+
+    # Check position field (nested compound)
+    assert compound_features["position"]["x"] == Value("int32")
+    assert compound_features["position"]["y"] == Value("int32")
+
+    # Check complex field (should be real/imag structure)
+    assert compound_features["complex_field"]["real"] == Value("float32")
+    assert compound_features["complex_field"]["imag"] == Value("float32")
+
+    # Check complex array (should be nested real/imag structures)
+    assert compound_features["complex_array"]["real"] == Array2D(shape=(2, 3), dtype="float32")
+    assert compound_features["complex_array"]["imag"] == Array2D(shape=(2, 3), dtype="float32")
+
+    # Check nested complex field
+    assert compound_features["nested_complex"]["real"] == Value("float32")
+    assert compound_features["nested_complex"]["imag"] == Value("float32")

tests/packaged_modules/test_iceberg.py

@@ -0,0 +1,190 @@
+import numpy as np
+import pyarrow as pa
+import pytest
+
+from datasets import IterableDataset, load_dataset
+
+from ..utils import require_not_windows, require_pyiceberg
+
+
+@pytest.fixture
+def catalog(tmp_path):
+    from pyiceberg.catalog.sql import SqlCatalog
+
+    cat = SqlCatalog(
+        "test_catalog",
+        **{
+            "uri": f"sqlite:///{tmp_path}/catalog.db",
+            "warehouse": str(tmp_path / "warehouse"),
+        },
+    )
+    cat.create_namespace("test_db")
+    return cat
+
+
+@pytest.fixture
+def sample_table(catalog):
+    from pyiceberg.schema import Schema
+    from pyiceberg.types import DoubleType, FloatType, ListType, LongType, NestedField, StringType
+
+    schema = Schema(
+        NestedField(1, "id", LongType()),
+        NestedField(2, "name", StringType()),
+        NestedField(3, "value", DoubleType()),
+        NestedField(4, "vector", ListType(element_id=5, element_type=FloatType(), element_required=False)),
+    )
+    table = catalog.create_table("test_db.sample", schema=schema)
+    table.append(
+        pa.table(
+            {
+                "id": pa.array([1, 2, 3], type=pa.int64()),
+                "name": pa.array(["alice", "bob", "carol"], type=pa.large_string()),
+                "value": pa.array([1.1, 2.2, 3.3], type=pa.float64()),
+                "vector": pa.FixedSizeListArray.from_arrays(pa.array([0.1] * 12, pa.float32()), list_size=4),
+            }
+        )
+    )
+    return table
+
+
+@require_not_windows
+@require_pyiceberg
+def test_load_iceberg_basic(catalog, sample_table):
+    ds = load_dataset("iceberg", catalog=catalog, table="test_db.sample")
+    assert "train" in ds
+    dataset = ds["train"]
+    assert dataset.num_rows == 3
+    assert "id" in dataset.column_names
+    assert "name" in dataset.column_names
+    assert "value" in dataset.column_names
+    assert "vector" in dataset.column_names
+    assert list(dataset["id"]) == [1, 2, 3]
+    assert list(dataset["name"]) == ["alice", "bob", "carol"]
+
+
+@require_not_windows
+@require_pyiceberg
+def test_load_vectors(catalog, sample_table):
+    ds = load_dataset("iceberg", catalog=catalog, table="test_db.sample", columns=["vector"])
+    dataset = ds["train"]
+    assert "vector" in dataset.column_names
+    vectors = dataset.data["vector"].combine_chunks().values.to_numpy(zero_copy_only=False)
+    assert np.allclose(vectors, np.full(12, 0.1), atol=1e-6)
+
+
+@require_not_windows
+@require_pyiceberg
+def test_load_iceberg_columns(catalog, sample_table):
+    ds = load_dataset("iceberg", catalog=catalog, table="test_db.sample", columns=["id", "name"])
+    dataset = ds["train"]
+    assert "id" in dataset.column_names
+    assert "name" in dataset.column_names
+    assert "value" not in dataset.column_names
+
+
+@require_not_windows
+@require_pyiceberg
+def test_load_iceberg_filters(catalog, sample_table):
+    ds = load_dataset("iceberg", catalog=catalog, table="test_db.sample", filters="value > 2.0")
+    dataset = ds["train"]
+    assert dataset.num_rows == 2
+    assert list(dataset["name"]) == ["bob", "carol"]
+
+
+@require_not_windows
+@require_pyiceberg
+def test_load_iceberg_multi_split(catalog):
+    from pyiceberg.schema import Schema
+    from pyiceberg.types import LongType, NestedField
+
+    schema = Schema(
+        NestedField(1, "x", LongType()),
+    )
+    train_table = catalog.create_table("test_db.train_split", schema=schema)
+    train_table.append(pa.table({"x": pa.array([1, 2, 3], type=pa.int64())}))
+
+    test_table = catalog.create_table("test_db.test_split", schema=schema)
+    test_table.append(pa.table({"x": pa.array([10, 20], type=pa.int64())}))
+
+    ds = load_dataset(
+        "iceberg",
+        catalog=catalog,
+        table={"train": "test_db.train_split", "test": "test_db.test_split"},
+    )
+    assert "train" in ds
+    assert "test" in ds
+    assert ds["train"].num_rows == 3
+    assert ds["test"].num_rows == 2
+
+
+@require_not_windows
+@require_pyiceberg
+@pytest.mark.parametrize("streaming", [False, True])
+def test_load_iceberg_streaming(catalog, sample_table, streaming):
+    ds = load_dataset("iceberg", catalog=catalog, table="test_db.sample", split="train", streaming=streaming)
+    if streaming:
+        assert isinstance(ds, IterableDataset)
+    items = list(ds)
+    assert len(items) == 3
+    assert all("id" in item for item in items)
+
+
+@require_not_windows
+@require_pyiceberg
+def test_load_iceberg_snapshot(catalog):
+    from pyiceberg.schema import Schema
+    from pyiceberg.types import LongType, NestedField
+
+    schema = Schema(
+        NestedField(1, "id", LongType()),
+    )
+    table = catalog.create_table("test_db.versioned", schema=schema)
+    table.append(pa.table({"id": pa.array([1, 2], type=pa.int64())}))
+
+    # Capture snapshot after first append
+    first_snapshot_id = table.current_snapshot().snapshot_id
+
+    # Append more data
+    table.append(pa.table({"id": pa.array([3, 4, 5], type=pa.int64())}))
+
+    # Load at latest: should have 5 rows
+    ds_latest = load_dataset("iceberg", catalog=catalog, table="test_db.versioned")
+    assert ds_latest["train"].num_rows == 5
+
+    # Load at first snapshot: should have 2 rows
+    ds_old = load_dataset("iceberg", catalog=catalog, table="test_db.versioned", snapshot_id=first_snapshot_id)
+    assert ds_old["train"].num_rows == 2
+
+
+@require_not_windows
+@require_pyiceberg
+def test_load_iceberg_num_proc(catalog):
+    """Test that num_proc > 1 works for parallel processing."""
+    from pyiceberg.schema import Schema
+    from pyiceberg.types import LongType, NestedField
+
+    schema = Schema(
+        NestedField(1, "id", LongType()),
+    )
+    table = catalog.create_table("test_db.parallel", schema=schema)
+    table.append(pa.table({"id": pa.array([1, 2, 3], type=pa.int64())}))
+    table.append(pa.table({"id": pa.array([4, 5, 6], type=pa.int64())}))
+
+    ds = load_dataset("iceberg", catalog=catalog, table="test_db.parallel", num_proc=2)
+    dataset = ds["train"]
+    assert dataset.num_rows == 6
+    assert sorted(dataset["id"]) == [1, 2, 3, 4, 5, 6]
+
+
+@require_not_windows
+@require_pyiceberg
+def test_load_iceberg_missing_catalog_raises():
+    with pytest.raises(ValueError, match="catalog"):
+        load_dataset("iceberg", catalog=None, table="db.table")
+
+
+@require_not_windows
+@require_pyiceberg
+def test_load_iceberg_missing_table_raises(catalog):
+    with pytest.raises(ValueError, match="table"):
+        load_dataset("iceberg", catalog=catalog, table=None)

tests/packaged_modules/test_imagefolder.py

@@ -0,0 +1,444 @@
+import shutil
+import textwrap
+
+import numpy as np
+import pytest
+
+from datasets import ClassLabel, Features, Image
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesDict, DataFilesList, get_data_patterns
+from datasets.download.streaming_download_manager import StreamingDownloadManager
+from datasets.packaged_modules.imagefolder.imagefolder import ImageFolder, ImageFolderConfig
+
+from ..utils import require_pil
+
+
+@pytest.fixture
+def cache_dir(tmp_path):
+    return str(tmp_path / "imagefolder_cache_dir")
+
+
+@pytest.fixture
+def data_files_with_labels_no_metadata(tmp_path, image_file):
+    data_dir = tmp_path / "data_files_with_labels_no_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    subdir_class_0 = data_dir / "cat"
+    subdir_class_0.mkdir(parents=True, exist_ok=True)
+    subdir_class_1 = data_dir / "dog"
+    subdir_class_1.mkdir(parents=True, exist_ok=True)
+
+    image_filename = subdir_class_0 / "image_cat.jpg"
+    shutil.copyfile(image_file, image_filename)
+    image_filename2 = subdir_class_1 / "image_dog.jpg"
+    shutil.copyfile(image_file, image_filename2)
+
+    data_files_with_labels_no_metadata = DataFilesDict.from_patterns(
+        get_data_patterns(str(data_dir)), data_dir.as_posix()
+    )
+
+    return data_files_with_labels_no_metadata
+
+
+@pytest.fixture
+def image_files_with_labels_and_duplicated_label_key_in_metadata(tmp_path, image_file):
+    data_dir = tmp_path / "image_files_with_labels_and_label_key_in_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    subdir_class_0 = data_dir / "cat"
+    subdir_class_0.mkdir(parents=True, exist_ok=True)
+    subdir_class_1 = data_dir / "dog"
+    subdir_class_1.mkdir(parents=True, exist_ok=True)
+
+    image_filename = subdir_class_0 / "image_cat.jpg"
+    shutil.copyfile(image_file, image_filename)
+    image_filename2 = subdir_class_1 / "image_dog.jpg"
+    shutil.copyfile(image_file, image_filename2)
+
+    image_metadata_filename = tmp_path / data_dir / "metadata.jsonl"
+    image_metadata = textwrap.dedent(
+        """\
+        {"file_name": "cat/image_cat.jpg", "caption": "Nice image of a cat", "label": "Cat"}
+        {"file_name": "dog/image_dog.jpg", "caption": "Nice image of a dog", "label": "Dog"}
+        """
+    )
+    with open(image_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(image_metadata)
+
+    return str(image_filename), str(image_filename2), str(image_metadata_filename)
+
+
+@pytest.fixture
+def image_file_with_metadata(tmp_path, image_file):
+    image_filename = tmp_path / "image_rgb.jpg"
+    shutil.copyfile(image_file, image_filename)
+    image_metadata_filename = tmp_path / "metadata.jsonl"
+    image_metadata = textwrap.dedent(
+        """\
+        {"file_name": "image_rgb.jpg", "caption": "Nice image"}
+        """
+    )
+    with open(image_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(image_metadata)
+    return str(image_filename), str(image_metadata_filename)
+
+
+@pytest.fixture
+def image_files_with_metadata_that_misses_one_image(tmp_path, image_file):
+    image_filename = tmp_path / "image_rgb.jpg"
+    shutil.copyfile(image_file, image_filename)
+    image_filename2 = tmp_path / "image_rgb2.jpg"
+    shutil.copyfile(image_file, image_filename2)
+    image_metadata_filename = tmp_path / "metadata.jsonl"
+    image_metadata = textwrap.dedent(
+        """\
+        {"file_name": "image_rgb.jpg", "caption": "Nice image"}
+        """
+    )
+    with open(image_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(image_metadata)
+    return str(image_filename), str(image_filename2), str(image_metadata_filename)
+
+
+@pytest.fixture(params=["jsonl", "csv"])
+def data_files_with_one_split_and_metadata(request, tmp_path, image_file):
+    data_dir = tmp_path / "imagefolder_data_dir_with_metadata_one_split"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    subdir = data_dir / "subdir"
+    subdir.mkdir(parents=True, exist_ok=True)
+
+    image_filename = data_dir / "image_rgb.jpg"
+    shutil.copyfile(image_file, image_filename)
+    image_filename2 = data_dir / "image_rgb2.jpg"
+    shutil.copyfile(image_file, image_filename2)
+    image_filename3 = subdir / "image_rgb3.jpg"  # in subdir
+    shutil.copyfile(image_file, image_filename3)
+
+    image_metadata_filename = data_dir / f"metadata.{request.param}"
+    image_metadata = (
+        textwrap.dedent(
+            """\
+        {"file_name": "image_rgb.jpg", "caption": "Nice image"}
+        {"file_name": "image_rgb2.jpg", "caption": "Nice second image"}
+        {"file_name": "subdir/image_rgb3.jpg", "caption": "Nice third image"}
+        """
+        )
+        if request.param == "jsonl"
+        else textwrap.dedent(
+            """\
+        file_name,caption
+        image_rgb.jpg,Nice image
+        image_rgb2.jpg,Nice second image
+        subdir/image_rgb3.jpg,Nice third image
+        """
+        )
+    )
+    with open(image_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(image_metadata)
+    data_files_with_one_split_and_metadata = DataFilesDict.from_patterns(
+        get_data_patterns(str(data_dir)), data_dir.as_posix()
+    )
+    assert len(data_files_with_one_split_and_metadata) == 1
+    assert len(data_files_with_one_split_and_metadata["train"]) == 4
+    return data_files_with_one_split_and_metadata
+
+
+@pytest.fixture(params=["jsonl", "csv"])
+def data_files_with_two_splits_and_metadata(request, tmp_path, image_file):
+    data_dir = tmp_path / "imagefolder_data_dir_with_metadata_two_splits"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    train_dir = data_dir / "train"
+    train_dir.mkdir(parents=True, exist_ok=True)
+    test_dir = data_dir / "test"
+    test_dir.mkdir(parents=True, exist_ok=True)
+
+    image_filename = train_dir / "image_rgb.jpg"  # train image
+    shutil.copyfile(image_file, image_filename)
+    image_filename2 = train_dir / "image_rgb2.jpg"  # train image
+    shutil.copyfile(image_file, image_filename2)
+    image_filename3 = test_dir / "image_rgb3.jpg"  # test image
+    shutil.copyfile(image_file, image_filename3)
+
+    train_image_metadata_filename = train_dir / f"metadata.{request.param}"
+    image_metadata = (
+        textwrap.dedent(
+            """\
+        {"file_name": "image_rgb.jpg", "caption": "Nice train image"}
+        {"file_name": "image_rgb2.jpg", "caption": "Nice second train image"}
+        """
+        )
+        if request.param == "jsonl"
+        else textwrap.dedent(
+            """\
+        file_name,caption
+        image_rgb.jpg,Nice train image
+        image_rgb2.jpg,Nice second train image
+        """
+        )
+    )
+    with open(train_image_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(image_metadata)
+    test_image_metadata_filename = test_dir / f"metadata.{request.param}"
+    image_metadata = (
+        textwrap.dedent(
+            """\
+        {"file_name": "image_rgb3.jpg", "caption": "Nice test image"}
+        """
+        )
+        if request.param == "jsonl"
+        else textwrap.dedent(
+            """\
+        file_name,caption
+        image_rgb3.jpg,Nice test image
+        """
+        )
+    )
+    with open(test_image_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(image_metadata)
+    data_files_with_two_splits_and_metadata = DataFilesDict.from_patterns(
+        get_data_patterns(str(data_dir)), data_dir.as_posix()
+    )
+    assert len(data_files_with_two_splits_and_metadata) == 2
+    assert len(data_files_with_two_splits_and_metadata["train"]) == 3
+    assert len(data_files_with_two_splits_and_metadata["test"]) == 2
+    return data_files_with_two_splits_and_metadata
+
+
+@pytest.fixture
+def data_files_with_zip_archives(tmp_path, image_file):
+    from PIL import Image, ImageOps
+
+    data_dir = tmp_path / "imagefolder_data_dir_with_zip_archives"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    archive_dir = data_dir / "archive"
+    archive_dir.mkdir(parents=True, exist_ok=True)
+    subdir = archive_dir / "subdir"
+    subdir.mkdir(parents=True, exist_ok=True)
+
+    image_filename = archive_dir / "image_rgb.jpg"
+    shutil.copyfile(image_file, image_filename)
+    image_filename2 = subdir / "image_rgb2.jpg"  # in subdir
+    # make sure they're two different images
+    # Indeed we won't be able to compare the image.filename, since the archive is not extracted in streaming mode
+    ImageOps.flip(Image.open(image_file)).save(image_filename2)
+
+    image_metadata_filename = archive_dir / "metadata.jsonl"
+    image_metadata = textwrap.dedent(
+        """\
+        {"file_name": "image_rgb.jpg", "caption": "Nice image"}
+        {"file_name": "subdir/image_rgb2.jpg", "caption": "Nice second image"}
+        """
+    )
+    with open(image_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(image_metadata)
+
+    shutil.make_archive(archive_dir, "zip", archive_dir)
+    shutil.rmtree(str(archive_dir))
+
+    data_files_with_zip_archives = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+
+    assert len(data_files_with_zip_archives) == 1
+    assert len(data_files_with_zip_archives["train"]) == 1
+    return data_files_with_zip_archives
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = ImageFolderConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = ImageFolderConfig(name="name", data_files=data_files)
+
+
+@require_pil
+# check that labels are inferred correctly from dir names
+def test_generate_examples_with_labels(data_files_with_labels_no_metadata, cache_dir):
+    # there are no metadata.jsonl files in this test case
+    imagefolder = ImageFolder(data_files=data_files_with_labels_no_metadata, cache_dir=cache_dir, drop_labels=False)
+    imagefolder.download_and_prepare()
+    assert imagefolder.info.features == Features({"image": Image(), "label": ClassLabel(names=["cat", "dog"])})
+    dataset = list(imagefolder.as_dataset()["train"])
+    label_feature = imagefolder.info.features["label"]
+
+    assert dataset[0]["label"] == label_feature._str2int["cat"]
+    assert dataset[1]["label"] == label_feature._str2int["dog"]
+
+
+@require_pil
+@pytest.mark.parametrize("drop_metadata", [None, True, False])
+@pytest.mark.parametrize("drop_labels", [None, True, False])
+def test_generate_examples_drop_labels(data_files_with_labels_no_metadata, drop_metadata, drop_labels):
+    imagefolder = ImageFolder(
+        drop_metadata=drop_metadata, drop_labels=drop_labels, data_files=data_files_with_labels_no_metadata
+    )
+    gen_kwargs = imagefolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    # removing the labels explicitly requires drop_labels=True
+    assert gen_kwargs["add_labels"] is not bool(drop_labels)
+    assert gen_kwargs["add_metadata"] is False
+    generator = imagefolder._generate_examples(**gen_kwargs)
+    if not drop_labels:
+        assert all(
+            example.keys() == {"image", "label"} and all(val is not None for val in example.values())
+            for _, example in generator
+        )
+    else:
+        assert all(
+            example.keys() == {"image"} and all(val is not None for val in example.values())
+            for _, example in generator
+        )
+
+
+@require_pil
+@pytest.mark.parametrize("drop_metadata", [None, True, False])
+@pytest.mark.parametrize("drop_labels", [None, True, False])
+def test_generate_examples_drop_metadata(image_file_with_metadata, drop_metadata, drop_labels):
+    image_file, image_metadata_file = image_file_with_metadata
+    imagefolder = ImageFolder(
+        drop_metadata=drop_metadata, drop_labels=drop_labels, data_files={"train": [image_file, image_metadata_file]}
+    )
+    gen_kwargs = imagefolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    # since the dataset has metadata, removing the metadata explicitly requires drop_metadata=True
+    assert gen_kwargs["add_metadata"] is not bool(drop_metadata)
+    # since the dataset has metadata, adding the labels explicitly requires drop_labels=False
+    assert gen_kwargs["add_labels"] is False
+    generator = imagefolder._generate_examples(**gen_kwargs)
+    expected_columns = {"image"}
+    if gen_kwargs["add_metadata"]:
+        expected_columns.add("caption")
+    if gen_kwargs["add_labels"]:
+        expected_columns.add("label")
+    result = [example for _, example in generator]
+    assert len(result) == 1
+    example = result[0]
+    assert example.keys() == expected_columns
+    for column in expected_columns:
+        assert example[column] is not None
+
+
+@require_pil
+@pytest.mark.parametrize("streaming", [False, True])
+def test_data_files_with_metadata_and_single_split(streaming, cache_dir, data_files_with_one_split_and_metadata):
+    data_files = data_files_with_one_split_and_metadata
+    imagefolder = ImageFolder(data_files=data_files, cache_dir=cache_dir)
+    imagefolder.download_and_prepare()
+    datasets = imagefolder.as_streaming_dataset() if streaming else imagefolder.as_dataset()
+    for split, data_files in data_files.items():
+        expected_num_of_images = len(data_files) - 1  # don't count the metadata file
+        assert split in datasets
+        dataset = list(datasets[split])
+        assert len(dataset) == expected_num_of_images
+        # make sure each sample has its own image and metadata
+        assert len({example["image"].filename for example in dataset}) == expected_num_of_images
+        assert len({example["caption"] for example in dataset}) == expected_num_of_images
+        assert all(example["caption"] is not None for example in dataset)
+
+
+@require_pil
+@pytest.mark.parametrize("streaming", [False, True])
+def test_data_files_with_metadata_and_multiple_splits(streaming, cache_dir, data_files_with_two_splits_and_metadata):
+    data_files = data_files_with_two_splits_and_metadata
+    imagefolder = ImageFolder(data_files=data_files, cache_dir=cache_dir)
+    imagefolder.download_and_prepare()
+    datasets = imagefolder.as_streaming_dataset() if streaming else imagefolder.as_dataset()
+    for split, data_files in data_files.items():
+        expected_num_of_images = len(data_files) - 1  # don't count the metadata file
+        assert split in datasets
+        dataset = list(datasets[split])
+        assert len(dataset) == expected_num_of_images
+        # make sure each sample has its own image and metadata
+        assert len({example["image"].filename for example in dataset}) == expected_num_of_images
+        assert len({example["caption"] for example in dataset}) == expected_num_of_images
+        assert all(example["caption"] is not None for example in dataset)
+
+
+@require_pil
+@pytest.mark.parametrize("streaming", [False, True])
+def test_data_files_with_metadata_and_archives(streaming, cache_dir, data_files_with_zip_archives):
+    imagefolder = ImageFolder(data_files=data_files_with_zip_archives, cache_dir=cache_dir)
+    imagefolder.download_and_prepare()
+    datasets = imagefolder.as_streaming_dataset() if streaming else imagefolder.as_dataset()
+    for split, data_files in data_files_with_zip_archives.items():
+        num_of_archives = len(data_files)  # the metadata file is inside the archive
+        expected_num_of_images = 2 * num_of_archives
+        assert split in datasets
+        dataset = list(datasets[split])
+        assert len(dataset) == expected_num_of_images
+        # make sure each sample has its own image and metadata
+        assert len({np.array(example["image"])[0, 0, 0] for example in dataset}) == expected_num_of_images
+        assert len({example["caption"] for example in dataset}) == expected_num_of_images
+        assert all(example["caption"] is not None for example in dataset)
+
+
+@require_pil
+def test_data_files_with_wrong_metadata_file_name(cache_dir, tmp_path, image_file):
+    data_dir = tmp_path / "data_dir_with_bad_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(image_file, data_dir / "image_rgb.jpg")
+    image_metadata_filename = data_dir / "bad_metadata.jsonl"  # bad file
+    image_metadata = textwrap.dedent(
+        """\
+        {"file_name": "image_rgb.jpg", "caption": "Nice image"}
+        """
+    )
+    with open(image_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(image_metadata)
+
+    data_files_with_bad_metadata = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+    imagefolder = ImageFolder(data_files=data_files_with_bad_metadata, cache_dir=cache_dir)
+    imagefolder.download_and_prepare()
+    dataset = imagefolder.as_dataset(split="train")
+    # check that there are no metadata, since the metadata file name doesn't have the right name
+    assert "caption" not in dataset.column_names
+
+
+@require_pil
+def test_data_files_with_custom_image_file_name_column_in_metadata_file(cache_dir, tmp_path, image_file):
+    data_dir = tmp_path / "data_dir_with_custom_file_name_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(image_file, data_dir / "image_rgb.jpg")
+    image_metadata_filename = data_dir / "metadata.jsonl"
+    image_metadata = textwrap.dedent(  # with bad column "bad_file_name" instead of "file_name"
+        """\
+        {"picture_file_name": "image_rgb.jpg", "caption": "Nice image"}
+        """
+    )
+    with open(image_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(image_metadata)
+
+    data_files_with_bad_metadata = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+    imagefolder = ImageFolder(data_files=data_files_with_bad_metadata, cache_dir=cache_dir)
+    imagefolder.download_and_prepare()
+    dataset = imagefolder.as_dataset(split="train")
+    assert "picture" in dataset.features
+    assert "picture_file_name" not in dataset.features
+
+
+@require_pil
+def test_data_files_with_with_metadata_in_different_formats(cache_dir, tmp_path, image_file):
+    data_dir = tmp_path / "data_dir_with_metadata_in_different_format"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(image_file, data_dir / "image_rgb.jpg")
+    image_metadata_filename_jsonl = data_dir / "metadata.jsonl"
+    image_metadata_jsonl = textwrap.dedent(
+        """\
+        {"file_name": "image_rgb.jpg", "caption": "Nice image"}
+        """
+    )
+    with open(image_metadata_filename_jsonl, "w", encoding="utf-8") as f:
+        f.write(image_metadata_jsonl)
+    image_metadata_filename_csv = data_dir / "metadata.csv"
+    image_metadata_csv = textwrap.dedent(
+        """\
+        file_name,caption
+        image_rgb.jpg,Nice image
+        """
+    )
+    with open(image_metadata_filename_csv, "w", encoding="utf-8") as f:
+        f.write(image_metadata_csv)
+
+    data_files_with_bad_metadata = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+    imagefolder = ImageFolder(data_files=data_files_with_bad_metadata, cache_dir=cache_dir)
+    with pytest.raises(ValueError) as exc_info:
+        imagefolder.download_and_prepare()
+    assert "metadata files with different extensions" in str(exc_info.value)

tests/packaged_modules/test_json.py

@@ -0,0 +1,809 @@
+import json
+import textwrap
+
+import pyarrow as pa
+import pytest
+
+from datasets import Features, Value, load_dataset
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesList
+from datasets.packaged_modules.json.json import AGENT_TRACES_FEATURES, Json, JsonConfig
+
+from ..utils import require_teich
+
+
+@pytest.fixture
+def json_file(tmp_path):
+    filename = tmp_path / "file.json"
+    data = textwrap.dedent(
+        """\
+        {
+            "col_1": 1,
+            "col_2": 2
+        }
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def jsonl_file(tmp_path):
+    filename = tmp_path / "file.jsonl"
+    data = textwrap.dedent(
+        """\
+        {"col_1": -1}
+        {"col_1": 1, "col_2": 2}
+        {"col_1": 10, "col_2": 20}
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+# ndjson format is no longer maintained (see: https://github.com/ndjson/ndjson-spec/issues/35#issuecomment-1285673417)
+@pytest.fixture
+def ndjson_file(tmp_path):
+    filename = tmp_path / "file.ndjson"
+    data = textwrap.dedent(
+        """\
+        {"col_1": -1}
+        {"col_1": 1, "col_2": 2}
+        {"col_1": 10, "col_2": 20}
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def jsonl_file_utf16_encoded(tmp_path):
+    filename = tmp_path / "file_utf16_encoded.jsonl"
+    data = textwrap.dedent(
+        """\
+        {"col_1": -1}
+        {"col_1": 1, "col_2": 2}
+        {"col_1": 10, "col_2": 20}
+        """
+    )
+    with open(filename, "w", encoding="utf-16") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def json_file_with_list_of_dicts(tmp_path):
+    filename = tmp_path / "file_with_list_of_dicts.json"
+    data = textwrap.dedent(
+        """\
+        [
+            {"col_1": -1},
+            {"col_1": 1, "col_2": 2},
+            {"col_1": 10, "col_2": 20}
+        ]
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def json_file_with_list_of_strings(tmp_path):
+    filename = tmp_path / "file_with_list_of_strings.json"
+    data = textwrap.dedent(
+        """\
+        [
+            "First text.",
+            "Second text.",
+            "Third text."
+        ]
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def json_file_with_list_of_dicts_field(tmp_path):
+    filename = tmp_path / "file_with_list_of_dicts_field.json"
+    data = textwrap.dedent(
+        """\
+        {
+            "field1": 1,
+            "field2": "aabb",
+            "field3": [
+                {"col_1": -1},
+                {"col_1": 1, "col_2": 2},
+                {"col_1": 10, "col_2": 20}
+            ]
+        }
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def json_file_with_list_of_strings_field(tmp_path):
+    path = tmp_path / "file.json"
+    data = textwrap.dedent(
+        """\
+        {
+            "field1": 1,
+            "field2": "aabb",
+            "field3": [
+                "First text.",
+                "Second text.",
+                "Third text."
+            ]
+        }
+        """
+    )
+    with open(path, "w") as f:
+        f.write(data)
+    return str(path)
+
+
+@pytest.fixture
+def json_file_with_dict_of_lists_field(tmp_path):
+    path = tmp_path / "file.json"
+    data = textwrap.dedent(
+        """\
+        {
+            "field1": 1,
+            "field2": "aabb",
+            "field3": {
+                "col_1": [-1, 1, 10],
+                "col_2": [null, 2, 20]
+            }
+        }
+        """
+    )
+    with open(path, "w") as f:
+        f.write(data)
+    return str(path)
+
+
+@pytest.fixture
+def json_file_with_list_of_dicts_with_sorted_columns(tmp_path):
+    path = tmp_path / "file.json"
+    data = textwrap.dedent(
+        """\
+        [
+            {"ID": 0, "Language": "Language-0", "Topic": "Topic-0"},
+            {"ID": 1, "Language": "Language-1", "Topic": "Topic-1"},
+            {"ID": 2, "Language": "Language-2", "Topic": "Topic-2"}
+        ]
+        """
+    )
+    with open(path, "w") as f:
+        f.write(data)
+    return str(path)
+
+
+@pytest.fixture
+def json_file_with_list_of_dicts_with_sorted_columns_field(tmp_path):
+    path = tmp_path / "file.json"
+    data = textwrap.dedent(
+        """\
+        {
+            "field1": 1,
+            "field2": "aabb",
+            "field3": [
+                {"ID": 0, "Language": "Language-0", "Topic": "Topic-0"},
+                {"ID": 1, "Language": "Language-1", "Topic": "Topic-1"},
+                {"ID": 2, "Language": "Language-2", "Topic": "Topic-2"}
+            ]
+        }
+        """
+    )
+    with open(path, "w") as f:
+        f.write(data)
+    return str(path)
+
+
+@pytest.fixture
+def jsonl_file_with_mix_of_str_and_int(tmp_path):
+    filename = tmp_path / "file.jsonl"
+    data = textwrap.dedent(
+        """\
+        {"col_1": -1}
+        {"col_1": 1}
+        {"col_1": "foo"}
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def jsonl_file_with_dicts_of_varying_keys(tmp_path):
+    filename = tmp_path / "file.jsonl"
+    data = textwrap.dedent(
+        """\
+        {"col_1": {"a": 0}}
+        {"col_1": {"b": 0}}
+        {"col_1": {"c": 0}}
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def jsonl_file_with_lists_of_dicts_of_varying_keys(tmp_path):
+    filename = tmp_path / "file.jsonl"
+    data = textwrap.dedent(
+        """\
+        {"col_1": [{"a": 0}, {"b": 0}]}
+        {"col_1": [{"c": 0}, {"d": 0}]}
+        """
+    )
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def jsonl_file_with_lists_of_dicts_of_varying_keys_and_bom(tmp_path):
+    # Same content as jsonl_file_with_lists_of_dicts_of_varying_keys, but the file
+    # starts with a UTF-8 BOM (written via the utf-8-sig codec).
+    filename = tmp_path / "file_with_bom.jsonl"
+    data = textwrap.dedent(
+        """\
+        {"col_1": [{"a": 0}, {"b": 0}]}
+        {"col_1": [{"c": 0}, {"d": 0}]}
+        """
+    )
+    with open(filename, "w", encoding="utf-8-sig") as f:
+        f.write(data)
+    return str(filename)
+
+
+_messages = [
+    {"role": "user", "content": "Turn on the living room lights and play my electronic music playlist."},
+    {
+        "role": "assistant",
+        "tool_calls": [
+            {
+                "type": "function",
+                "function": {"name": "control_light", "arguments": {"room": "living room", "state": "on"}},
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "play_music",
+                    "arguments": {
+                        "playlist": "electronic"
+                    },  # mixed-type here since keys ["playlist"] and ["room", "state"] are different
+                },
+            },
+        ],
+    },
+    {"role": "tool", "name": "control_light", "content": "The lights in the living room are now on."},
+    {"role": "tool", "name": "play_music", "content": "The music is now playing."},
+    {"role": "assistant", "content": "Done!"},
+]
+
+EXPECTED_ONE = {"col_1": [1], "col_2": [2]}
+EXPECTED_THREE = {"col_1": [-1, 1, 10], "col_2": [None, 2, 20]}
+EXPECTED_LIST_OF_STRINGS = {"text": ["First text.", "Second text.", "Third text."]}
+EXPECTED_MIX = {"col_1": [-1, 1, "foo"]}
+EXPECTED_DICTS_WITH_VARYING_KEYS = {"col_1": [{"a": 0}, {"b": 0}, {"c": 0}]}
+EXPECTED_LISTS_OF_DICTS_WITH_VARYING_KEYS = {"col_1": [[{"a": 0}, {"b": 0}], [{"c": 0}, {"d": 0}]]}
+EXPECTED_MESSAGES = {"messages": [_messages]}
+
+
+@pytest.fixture
+def jsonl_file_with_messages(tmp_path):
+    filename = tmp_path / "file.jsonl"
+    data = json.dumps({"messages": _messages})
+    with open(filename, "w") as f:
+        f.write(data)
+    return str(filename)
+
+
+def write_jsonl(path, rows):
+    with open(path, "w") as f:
+        for row in rows:
+            f.write(json.dumps(row) + "\n")
+    return str(path)
+
+
+def generate_agent_traces_output(trace_file):
+    json_builder = Json(features=AGENT_TRACES_FEATURES)
+    base_files = [trace_file]
+    files_iterables = [[trace_file]]
+    original_files = list(base_files)
+    generator = json_builder._generate_tables(
+        base_files=base_files, files_iterables=files_iterables, original_files=original_files
+    )
+    pa_table = pa.concat_tables([table for _, table in generator])
+    return Features.from_arrow_schema(pa_table.schema).decode_batch(pa_table.to_pydict())
+
+
+def assert_agent_traces_output(tmp_path, filename, rows, expected, num_sessions=1):
+    trace_file = write_jsonl(tmp_path / filename, rows)
+    out = generate_agent_traces_output(trace_file)
+    for key, value in zip(AGENT_TRACE_FIELD_NAMES_TO_CHECK, expected):
+        assert out[key] == [value] * num_sessions
+    assert out["file_path"] == [trace_file] * num_sessions
+    assert isinstance(out["messages"], list)
+    assert out["messages"]
+    assert all(
+        isinstance(message["role"], str) and isinstance(message["content"], str)
+        for messages in out["messages"]
+        for message in messages
+    )
+    assert isinstance(out["tools"], list)
+    assert isinstance(out["metadata"], (dict, list))
+    assert isinstance(out["trace"], (dict, list))
+    return trace_file, out
+
+
+AGENT_TRACE_FIELD_NAMES_TO_CHECK = (
+    "harness",
+    "session_id",
+    "prompt",
+    "sent_at",
+    "num_user_messages",
+    "num_tool_calls",
+)
+CODEX_AGENT_TRACE_ROWS = [
+    {"type": "session_meta", "payload": {"id": "codex-session"}},
+    {
+        "type": "response_item",
+        "payload": {
+            "type": "message",
+            "role": "user",
+            "content": [{"type": "input_text", "text": "context-wrapped codex prompt"}],
+        },
+    },
+    {
+        "type": "response_item",
+        "payload": {
+            "type": "message",
+            "role": "user",
+            "content": [{"type": "input_text", "text": "actual codex prompt"}],
+        },
+    },
+    {
+        "timestamp": "2026-04-01T10:01:00.000Z",
+        "type": "event_msg",
+        "payload": {"type": "user_message", "message": "actual codex prompt"},
+    },
+    {
+        "type": "response_item",
+        "payload": {"type": "function_call", "name": "exec_command", "call_id": "call_1"},
+    },
+    {
+        "type": "response_item",
+        "payload": {"type": "function_call_output", "call_id": "call_1", "output": "done"},
+    },
+]
+CODEX_EXPECTED_AGENT_TRACE_FIELDS = (
+    "codex",
+    "codex-session",
+    "actual codex prompt",
+    "2026-04-01T10:01:00.000Z",
+    1,
+    1,
+)
+
+HERMES_SESSION = {
+    "id": "20260605_092247_d018ec",
+    "source": "cli",
+    "model": "Qwen/Qwen3.5-35B-A3B",
+    "system_prompt": "You are Hermes.",
+    "started_at": 1_780_665_768.307,
+    "message_count": 4,
+    "tool_call_count": 1,
+    "messages": [
+        {
+            "session_id": "20260605_092247_d018ec",
+            "role": "user",
+            "content": "Run pwd and date.",
+            "timestamp": 1_780_665_767.307,
+        },
+        {
+            "session_id": "20260605_092247_d018ec",
+            "role": "assistant",
+            "content": "",
+            "reasoning_content": "The user asked for two shell commands.",
+            "timestamp": 1_780_665_767.308,
+            "tool_calls": [
+                {
+                    "id": "call_677f321e2b3047b4b8c7a1e1",
+                    "type": "function",
+                    "function": {
+                        "name": "terminal",
+                        "arguments": json.dumps({"command": "pwd && date -u +%Y-%m-%dT%H:%M:%SZ"}),
+                    },
+                },
+            ],
+        },
+        {
+            "session_id": "20260605_092247_d018ec",
+            "role": "tool",
+            "content": json.dumps({"output": "/tmp/work\n2026-06-05T13:22:47Z", "exit_code": 0, "error": None}),
+            "tool_call_id": "call_677f321e2b3047b4b8c7a1e1",
+            "timestamp": 1_780_665_767.309,
+        },
+        {
+            "session_id": "20260605_092247_d018ec",
+            "role": "assistant",
+            "content": "Working directory: `/tmp/work`.",
+            "reasoning": "The commands executed successfully.",
+            "timestamp": 1_780_665_767.31,
+        },
+    ],
+}
+
+
+DROID_SESSION = [
+    {
+        "type": "session_start",
+        "id": "droid-session",
+        "title": "inspect the project",
+        "sessionTitle": "Inspect project files",
+        "owner": "caleb",
+        "version": 2,
+        "cwd": "/workspace/project",
+    },
+    {
+        "type": "message",
+        "id": "context-1",
+        "timestamp": "2026-06-02T18:55:29.000Z",
+        "message": {
+            "role": "user",
+            "visibility": "llm_only",
+            "content": [{"type": "text", "text": "<system-reminder>injected context</system-reminder>"}],
+        },
+        "parentId": None,
+    },
+    {
+        "type": "message",
+        "id": "message-1",
+        "timestamp": "2026-06-02T18:55:30.274Z",
+        "message": {
+            "role": "user",
+            "content": [{"type": "text", "text": "Inspect the project"}],
+        },
+        "parentId": "context-1",
+    },
+    {
+        "type": "message",
+        "id": "message-2",
+        "timestamp": "2026-06-02T18:55:35.000Z",
+        "message": {
+            "role": "assistant",
+            "content": [
+                {
+                    "type": "thinking",
+                    "thinking": "I should list the files first.",
+                    "signature": "reasoning_content",
+                    "signatureProvider": "generic-chat-completion-api",
+                    "durationMs": 1200,
+                },
+                {"type": "text", "text": "I'll list the files."},
+                {"type": "tool_use", "id": "LS_0", "name": "LS", "input": {"directory_path": "/workspace/project"}},
+            ],
+            "chatCompletionReasoningField": "reasoning_content",
+            "chatCompletionReasoningContent": "I should list the files first.",
+        },
+        "parentId": "message-1",
+    },
+    {
+        "type": "message",
+        "id": "message-3",
+        "timestamp": "2026-06-02T18:55:36.000Z",
+        "message": {
+            "role": "user",
+            "content": [
+                {"type": "tool_result", "tool_use_id": "LS_0", "is_error": False, "content": "README.md\nsrc"},
+            ],
+        },
+        "parentId": "message-2",
+    },
+]
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = JsonConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = JsonConfig(name="name", data_files=data_files)
+
+
+@pytest.mark.parametrize(
+    "file_fixture, config_kwargs, expected",
+    [
+        ("json_file", {}, EXPECTED_ONE),
+        ("jsonl_file", {}, EXPECTED_THREE),
+        ("ndjson_file", {}, EXPECTED_THREE),
+        ("jsonl_file_utf16_encoded", {"encoding": "utf-16"}, EXPECTED_THREE),
+        ("json_file_with_list_of_dicts", {}, EXPECTED_THREE),
+        ("json_file_with_list_of_dicts_field", {"field": "field3"}, EXPECTED_THREE),
+        ("json_file_with_list_of_strings", {}, EXPECTED_LIST_OF_STRINGS),
+        ("json_file_with_list_of_strings_field", {"field": "field3"}, EXPECTED_LIST_OF_STRINGS),
+        ("json_file_with_dict_of_lists_field", {"field": "field3"}, EXPECTED_THREE),
+        ("jsonl_file_with_mix_of_str_and_int", {}, EXPECTED_MIX),
+        ("jsonl_file_with_dicts_of_varying_keys", {}, EXPECTED_DICTS_WITH_VARYING_KEYS),
+        ("jsonl_file_with_lists_of_dicts_of_varying_keys", {}, EXPECTED_LISTS_OF_DICTS_WITH_VARYING_KEYS),
+        (
+            "jsonl_file_with_lists_of_dicts_of_varying_keys_and_bom",
+            {},
+            EXPECTED_LISTS_OF_DICTS_WITH_VARYING_KEYS,
+        ),
+        ("jsonl_file_with_messages", {}, EXPECTED_MESSAGES),
+    ],
+)
+def test_json_generate_tables(file_fixture, config_kwargs, expected, request):
+    json = Json(**config_kwargs)
+    base_files = [request.getfixturevalue(file_fixture)]
+    files_iterables = [[file] for file in base_files]
+    original_files = list(base_files)
+    generator = json._generate_tables(
+        base_files=base_files, files_iterables=files_iterables, original_files=original_files
+    )
+    pa_table = pa.concat_tables([table for _, table in generator])
+    out = Features.from_arrow_schema(pa_table.schema).decode_batch(pa_table.to_pydict())
+    assert out == expected
+
+
+@pytest.mark.parametrize(
+    "file_fixture, config_kwargs",
+    [
+        (
+            "jsonl_file",
+            {"features": Features({"col_1": Value("int64"), "col_2": Value("int64"), "missing_col": Value("string")})},
+        ),
+        (
+            "json_file_with_list_of_dicts",
+            {"features": Features({"col_1": Value("int64"), "col_2": Value("int64"), "missing_col": Value("string")})},
+        ),
+        (
+            "json_file_with_list_of_dicts_field",
+            {
+                "field": "field3",
+                "features": Features(
+                    {"col_1": Value("int64"), "col_2": Value("int64"), "missing_col": Value("string")}
+                ),
+            },
+        ),
+    ],
+)
+def test_json_generate_tables_with_missing_features(file_fixture, config_kwargs, request):
+    json = Json(**config_kwargs)
+    base_files = [request.getfixturevalue(file_fixture)]
+    files_iterables = [[file] for file in base_files]
+    original_files = list(base_files)
+    generator = json._generate_tables(
+        base_files=base_files, files_iterables=files_iterables, original_files=original_files
+    )
+    pa_table = pa.concat_tables([table for _, table in generator])
+    assert pa_table.to_pydict() == {"col_1": [-1, 1, 10], "col_2": [None, 2, 20], "missing_col": [None, None, None]}
+
+
+@pytest.mark.parametrize(
+    "file_fixture, config_kwargs",
+    [
+        ("json_file_with_list_of_dicts_with_sorted_columns", {}),
+        ("json_file_with_list_of_dicts_with_sorted_columns_field", {"field": "field3"}),
+    ],
+)
+def test_json_generate_tables_with_sorted_columns(file_fixture, config_kwargs, request):
+    json = Json(**config_kwargs)
+    base_files = [request.getfixturevalue(file_fixture)]
+    files_iterables = [[file] for file in base_files]
+    original_files = list(base_files)
+    generator = json._generate_tables(
+        base_files=base_files, files_iterables=files_iterables, original_files=original_files
+    )
+    pa_table = pa.concat_tables([table for _, table in generator])
+    assert pa_table.column_names == ["ID", "Language", "Topic"]
+
+
+@pytest.mark.parametrize(
+    "filename, rows, expected",
+    [
+        pytest.param("codex.jsonl", CODEX_AGENT_TRACE_ROWS, CODEX_EXPECTED_AGENT_TRACE_FIELDS, id="codex"),
+        pytest.param(
+            "codex_response_item_only.jsonl",
+            [
+                {"type": "session_meta", "payload": {"id": "codex-session"}},
+                {
+                    "type": "response_item",
+                    "payload": {
+                        "type": "message",
+                        "role": "user",
+                        "content": [{"text": "codex response item prompt"}],
+                    },
+                },
+                {
+                    "type": "response_item",
+                    "payload": {"type": "function_call", "name": "exec_command", "call_id": "call_1"},
+                },
+            ],
+            ("codex", "codex-session", None, None, 0, 1),
+            id="codex-response-item-only",
+        ),
+        pytest.param(
+            "claude.jsonl",
+            [
+                {
+                    "timestamp": "2026-04-02T10:00:00.000Z",
+                    "type": "user",
+                    "sessionId": "claude-session",
+                    "message": {"role": "user", "content": "claude prompt"},
+                },
+                {
+                    "type": "assistant",
+                    "sessionId": "claude-session",
+                    "message": {
+                        "role": "assistant",
+                        "content": [{"type": "tool_use", "id": "toolu_1"}, {"type": "tool_use", "id": "toolu_2"}],
+                    },
+                },
+                {
+                    "type": "user",
+                    "sessionId": "claude-session",
+                    "message": {"role": "user", "content": [{"type": "tool_result", "content": "done"}]},
+                },
+            ],
+            ("claude_code", "claude-session", "claude prompt", "2026-04-02T10:00:00.000Z", 1, 2),
+            id="claude",
+        ),
+        pytest.param(
+            "pi.jsonl",
+            [
+                {"type": "session", "id": "pi-session"},
+                {
+                    "timestamp": "2026-04-03T10:01:00.000Z",
+                    "type": "message",
+                    "message": {"role": "user", "content": "pi prompt"},
+                },
+                {
+                    "type": "message",
+                    "message": {
+                        "role": "assistant",
+                        "content": [{"type": "toolCall", "id": "call_1"}, {"type": "toolCall", "id": "call_2"}],
+                    },
+                },
+                {
+                    "type": "message",
+                    "message": {"role": "toolResult", "content": [{"type": "text", "text": "done"}]},
+                },
+                {
+                    "type": "message",
+                    "message": {"role": "user", "content": [{"type": "text", "text": "second pi prompt"}]},
+                },
+            ],
+            ("pi", "pi-session", "pi prompt", "2026-04-03T10:01:00.000Z", 2, 2),
+            id="pi",
+        ),
+        pytest.param(
+            "openclaw.jsonl",
+            [
+                {
+                    "type": "session",
+                    "id": "openclaw-session",
+                    "cwd": "/Users/test/.openclaw/agents/main",
+                },
+                {
+                    "timestamp": "2026-04-03T10:01:00.000Z",
+                    "type": "message",
+                    "message": {"role": "user", "content": "openclaw prompt"},
+                },
+            ],
+            ("openclaw", "openclaw-session", "openclaw prompt", "2026-04-03T10:01:00.000Z", 1, 0),
+            id="openclaw",
+        ),
+        pytest.param(
+            "hermes.jsonl",
+            [HERMES_SESSION],
+            ("hermes", "20260605_092247_d018ec", "Run pwd and date.", "2026-06-05T13:22:48.307Z", 1, 1),
+        ),
+        pytest.param(
+            "hermes_two_sessions.jsonl",
+            [HERMES_SESSION] * 2,
+            ("hermes", "20260605_092247_d018ec", "Run pwd and date.", "2026-06-05T13:22:48.307Z", 1, 1),
+        ),
+        pytest.param(
+            "droid.jsonl",
+            DROID_SESSION,
+            ("droid", "droid-session", "Inspect the project", "2026-06-02T18:55:30.274Z", 1, 1),
+            id="droid",
+        ),
+        pytest.param(
+            "missing_prompt.jsonl",
+            [
+                {"type": "session_meta", "payload": {"id": "codex-session"}},
+                {
+                    "type": "response_item",
+                    "payload": {"type": "message", "role": "assistant", "content": [{"text": "assistant response"}]},
+                },
+            ],
+            ("codex", "codex-session", None, None, 0, 0),
+            id="missing-prompt",
+        ),
+    ],
+)
+@require_teich
+def test_json_generate_tables_with_agent_trace_metadata(tmp_path, filename, rows, expected):
+    num_sessions = 2 if filename == "hermes_two_sessions.jsonl" else 1
+    _, out = assert_agent_traces_output(tmp_path, filename, rows, expected, num_sessions=num_sessions)
+    if filename == "droid.jsonl":
+        assert out["metadata"][0]["trace_type"] == "droid"
+    assert "models" not in out
+
+
+@require_teich
+@pytest.mark.parametrize(
+    "filename, rows, expected",
+    [
+        pytest.param("codex.jsonl", CODEX_AGENT_TRACE_ROWS, CODEX_EXPECTED_AGENT_TRACE_FIELDS, id="codex"),
+        pytest.param(
+            "droid.jsonl",
+            DROID_SESSION,
+            ("droid", "droid-session", "Inspect the project", "2026-06-02T18:55:30.274Z", 1, 1),
+            id="droid",
+        ),
+    ],
+)
+def test_json_load_dataset_with_agent_trace_metadata(tmp_path, filename, rows, expected):
+    trace_file = write_jsonl(tmp_path / filename, rows)
+
+    dataset = load_dataset("json", data_files=trace_file, split="train", cache_dir=str(tmp_path / "cache"))
+    row = dataset[0]
+
+    assert dataset.column_names == [
+        "harness",
+        "session_id",
+        "prompt",
+        "messages",
+        "tools",
+        "metadata",
+        "sent_at",
+        "num_user_messages",
+        "num_tool_calls",
+        "trace",
+        "file_path",
+    ]
+    for key, value in zip(AGENT_TRACE_FIELD_NAMES_TO_CHECK, expected):
+        assert row[key] == value, key
+    if filename == "droid.jsonl":
+        assert row["metadata"]["trace_type"] == "droid"
+        assert json.loads(row["trace"].splitlines()[0])["type"] == "session_start"
+
+
+def test_json_load_dataset_without_droid_marker_stays_ordinary_json(tmp_path):
+    trace_file = write_jsonl(
+        tmp_path / "droid_missing_marker.jsonl",
+        [
+            {"type": "session_start", "id": "droid-session", "version": 2},
+            {
+                "type": "message",
+                "id": "message-1",
+                "timestamp": "2026-06-02T18:55:30.274Z",
+                "message": {"role": "user", "content": [{"type": "text", "text": "Inspect the project"}]},
+            },
+        ],
+    )
+
+    dataset = load_dataset("json", data_files=trace_file, split="train", cache_dir=str(tmp_path / "cache"))
+
+    assert dataset.column_names == ["type", "id", "version", "timestamp", "message"]
+    assert dataset[0]["type"] == "session_start"

tests/packaged_modules/test_lance.py

@@ -0,0 +1,106 @@
+import lance
+import numpy as np
+import pyarrow as pa
+import pytest
+
+from datasets import load_dataset
+
+
+@pytest.fixture
+def lance_dataset(tmp_path) -> str:
+    data = pa.table(
+        {
+            "id": pa.array([1, 2, 3, 4]),
+            "value": pa.array([10.0, 20.0, 30.0, 40.0]),
+            "text": pa.array(["a", "b", "c", "d"]),
+            "vector": pa.FixedSizeListArray.from_arrays(pa.array([0.1] * 16, pa.float32()), list_size=4),
+        }
+    )
+    dataset_path = tmp_path / "test_dataset.lance"
+    lance.write_dataset(data, dataset_path)
+    return str(dataset_path)
+
+
+@pytest.fixture
+def lance_hf_dataset(tmp_path) -> str:
+    data = pa.table(
+        {
+            "id": pa.array([1, 2, 3, 4]),
+            "value": pa.array([10.0, 20.0, 30.0, 40.0]),
+            "text": pa.array(["a", "b", "c", "d"]),
+            "vector": pa.FixedSizeListArray.from_arrays(pa.array([0.1] * 16, pa.float32()), list_size=4),
+        }
+    )
+    dataset_dir = tmp_path / "data" / "train.lance"
+    dataset_dir.parent.mkdir(parents=True, exist_ok=True)
+    lance.write_dataset(data, dataset_dir)
+    lance.write_dataset(data[:2], tmp_path / "data" / "test.lance")
+
+    with open(tmp_path / "README.md", "w") as f:
+        f.write("""---
+size_categories:
+- 1M<n<10M
+source_datasets:
+- lance_test
+---
+    # Test Lance Dataset\n\n
+    # My Markdown is fancier\n
+""")
+
+    return str(tmp_path)
+
+
+def test_load_lance_dataset(lance_dataset):
+    dataset_dict = load_dataset(lance_dataset)
+    assert "train" in dataset_dict.keys()
+
+    dataset = dataset_dict["train"]
+    assert "id" in dataset.column_names
+    assert "value" in dataset.column_names
+    assert "text" in dataset.column_names
+    assert "vector" in dataset.column_names
+    ids = dataset["id"]
+    assert ids == [1, 2, 3, 4]
+
+
+@pytest.mark.parametrize("streaming", [False, True])
+def test_load_hf_dataset(lance_hf_dataset, streaming):
+    dataset_dict = load_dataset(lance_hf_dataset, columns=["id", "text"], streaming=streaming)
+    assert "train" in dataset_dict.keys()
+    assert "test" in dataset_dict.keys()
+    dataset = dataset_dict["train"]
+
+    assert "id" in dataset.column_names
+    assert "text" in dataset.column_names
+    assert "value" not in dataset.column_names
+    assert "vector" not in dataset.column_names
+    ids = list(dataset["id"])
+    assert ids == [1, 2, 3, 4]
+    text = list(dataset["text"])
+    assert text == ["a", "b", "c", "d"]
+    assert "value" not in dataset.column_names
+
+
+def test_load_vectors(lance_hf_dataset):
+    dataset_dict = load_dataset(lance_hf_dataset, columns=["vector"])
+    assert "train" in dataset_dict.keys()
+    dataset = dataset_dict["train"]
+
+    assert "vector" in dataset.column_names
+    vectors = dataset.data["vector"].combine_chunks().values.to_numpy(zero_copy_only=False)
+    assert np.allclose(vectors, np.full(16, 0.1))
+
+
+@pytest.mark.parametrize("streaming", [False, True])
+def test_load_lance_streaming_modes(lance_hf_dataset, streaming):
+    """Test loading Lance dataset in both streaming and non-streaming modes."""
+    from datasets import IterableDataset
+
+    ds = load_dataset(lance_hf_dataset, split="train", streaming=streaming)
+    if streaming:
+        assert isinstance(ds, IterableDataset)
+        items = list(ds)
+    else:
+        items = list(ds)
+    assert len(items) == 4
+    assert all("id" in item for item in items)

tests/packaged_modules/test_meshfolder.py

@@ -0,0 +1,86 @@
+import shutil
+import textwrap
+
+import pytest
+
+from datasets import ClassLabel, Features, Mesh
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesDict, get_data_patterns
+from datasets.packaged_modules.meshfolder.meshfolder import MeshFolder, MeshFolderConfig
+
+
+@pytest.fixture
+def cache_dir(tmp_path):
+    return str(tmp_path / "meshfolder_cache_dir")
+
+
+@pytest.fixture
+def data_files_with_labels_no_metadata(tmp_path, mesh_file):
+    data_dir = tmp_path / "data_files_with_labels_no_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    subdir_class_0 = data_dir / "chair"
+    subdir_class_0.mkdir(parents=True, exist_ok=True)
+    subdir_class_1 = data_dir / "table"
+    subdir_class_1.mkdir(parents=True, exist_ok=True)
+
+    mesh_filename = subdir_class_0 / "mesh_chair.glb"
+    shutil.copyfile(mesh_file, mesh_filename)
+    mesh_filename2 = subdir_class_1 / "mesh_table.glb"
+    shutil.copyfile(mesh_file, mesh_filename2)
+
+    data_files_with_labels_no_metadata = DataFilesDict.from_patterns(
+        get_data_patterns(str(data_dir)), data_dir.as_posix()
+    )
+
+    return data_files_with_labels_no_metadata
+
+
+@pytest.fixture
+def mesh_file_with_metadata(tmp_path, mesh_file):
+    mesh_filename = tmp_path / "mesh_file.glb"
+    shutil.copyfile(mesh_file, mesh_filename)
+    mesh_metadata_filename = tmp_path / "metadata.jsonl"
+    mesh_metadata = textwrap.dedent(
+        """\
+        {"file_name": "mesh_file.glb", "text": "Mesh description"}
+        """
+    )
+    with open(mesh_metadata_filename, "w", encoding="utf-8") as f:
+        f.write(mesh_metadata)
+    return str(mesh_filename), str(mesh_metadata_filename)
+
+
+def test_meshfolder_config_and_extensions():
+    # Verify extensions
+    assert MeshFolder.EXTENSIONS == [".glb", ".ply", ".stl"]
+    assert MeshFolder.BASE_FEATURE == Mesh
+    assert MeshFolder.BASE_COLUMN_NAME == "mesh"
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = MeshFolderConfig(name="name-with-*-invalid-character")
+
+
+def test_generate_examples_with_labels(data_files_with_labels_no_metadata, cache_dir):
+    # there are no metadata.jsonl files in this test case
+    meshfolder = MeshFolder(data_files=data_files_with_labels_no_metadata, cache_dir=cache_dir, drop_labels=False)
+    meshfolder.download_and_prepare()
+    assert meshfolder.info.features == Features({"mesh": Mesh(), "label": ClassLabel(names=["chair", "table"])})
+    dataset = list(meshfolder.as_dataset()["train"])
+    label_feature = meshfolder.info.features["label"]
+
+    assert dataset[0]["label"] == label_feature._str2int["chair"]
+    assert dataset[1]["label"] == label_feature._str2int["table"]
+
+
+@pytest.mark.parametrize("streaming", [False, True])
+def test_data_files_with_metadata_and_single_split(streaming, cache_dir, mesh_file_with_metadata):
+    mesh_file, mesh_metadata_file = mesh_file_with_metadata
+    meshfolder = MeshFolder(data_files={"train": [mesh_file, mesh_metadata_file]}, cache_dir=cache_dir)
+    meshfolder.download_and_prepare()
+    dataset = meshfolder.as_streaming_dataset()["train"] if streaming else meshfolder.as_dataset()["train"]
+
+    item = next(iter(dataset)) if streaming else dataset[0]
+    assert "mesh" in item
+    assert item["text"] == "Mesh description"

tests/packaged_modules/test_pandas.py

@@ -0,0 +1,16 @@
+import pytest
+
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesList
+from datasets.packaged_modules.pandas.pandas import PandasConfig
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = PandasConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = PandasConfig(name="name", data_files=data_files)

tests/packaged_modules/test_parquet.py

@@ -0,0 +1,41 @@
+import pytest
+
+from datasets import load_dataset
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesList
+from datasets.packaged_modules.parquet.parquet import ParquetConfig
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = ParquetConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = ParquetConfig(name="name", data_files=data_files)
+
+
+def test_parquet_reshard(multi_row_groups_parquet_path):
+    ds = load_dataset("parquet", data_files=multi_row_groups_parquet_path, split="train", streaming=True)
+    assert ds.num_shards == 1
+    expected = list(ds)
+    resharded_ds = ds.reshard()
+    assert resharded_ds.num_shards == 4
+    assert list(resharded_ds) == expected
+
+
+def test_parquet_columns(parquet_path):
+    ds = load_dataset("parquet", data_files=parquet_path, split="train", streaming=True)
+    full_features = ds.features
+    assert len(ds.features) == 3
+    assert len(next(iter(ds))) == 3
+    ds = load_dataset("parquet", data_files=parquet_path, split="train", streaming=True, columns=["col_1"])
+    assert len(ds.features) == 1
+    assert len(next(iter(ds))) == 1
+    ds = load_dataset(
+        "parquet", data_files=parquet_path, split="train", streaming=True, columns=["col_1"], features=full_features
+    )
+    assert len(ds.features) == 1
+    assert len(next(iter(ds))) == 1

tests/packaged_modules/test_spark.py

@@ -0,0 +1,170 @@
+from unittest.mock import patch
+
+import numpy as np
+import pyspark
+import pytest
+
+from datasets import Features, Image, IterableDataset
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesList
+from datasets.packaged_modules.spark.spark import (
+    Spark,
+    SparkConfig,
+    SparkExamplesIterable,
+    _generate_iterable_examples,
+)
+
+from ..utils import (
+    require_dill_gt_0_3_2,
+    require_not_windows,
+)
+
+
+def _get_expected_row_ids_and_row_dicts_for_partition_order(df, partition_order):
+    expected_row_ids_and_row_dicts = []
+    for part_id in partition_order:
+        partition = df.where(f"SPARK_PARTITION_ID() = {part_id}").collect()
+        for row_idx, row in enumerate(partition):
+            expected_row_ids_and_row_dicts.append(((part_id, row_idx), row.asDict()))
+    return expected_row_ids_and_row_dicts
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = SparkConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = SparkConfig(name="name", data_files=data_files)
+
+
+@require_not_windows
+@require_dill_gt_0_3_2
+def test_repartition_df_if_needed():
+    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
+    df = spark.range(100).repartition(1)
+    spark_builder = Spark(df)
+    # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
+    # that each partition can hold 2 rows.
+    spark_builder._repartition_df_if_needed(max_shard_size=16)
+    # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
+    assert spark_builder.df.rdd.getNumPartitions() == 50
+
+
+@require_not_windows
+@require_dill_gt_0_3_2
+def test_generate_iterable_examples():
+    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
+    df = spark.range(10).repartition(2)
+    partition_order = [1, 0]
+    iterator = _generate_iterable_examples(df, partition_order)  # Reverse the partitions.
+    expected_row_ids_and_row_dicts = _get_expected_row_ids_and_row_dicts_for_partition_order(df, partition_order)
+
+    for i, (row_id, row_dict) in enumerate(iterator):
+        expected_row_id, expected_row_dict = expected_row_ids_and_row_dicts[i]
+        assert row_id == expected_row_id
+        assert row_dict == expected_row_dict
+
+
+@require_not_windows
+@require_dill_gt_0_3_2
+def test_spark_examples_iterable():
+    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
+    df = spark.range(10).repartition(1)
+    it = SparkExamplesIterable(df)
+    assert it.num_shards == 1
+    for i, (row_key, row_dict) in enumerate(it):
+        assert row_key == (0, i)
+        assert row_dict == {"id": i}
+
+
+@require_not_windows
+@require_dill_gt_0_3_2
+def test_spark_examples_iterable_shuffle():
+    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
+    df = spark.range(30).repartition(3)
+    # Mock the generator so that shuffle reverses the partition indices.
+    with patch("numpy.random.Generator") as generator_mock:
+        generator_mock.shuffle.side_effect = lambda x: x.reverse()
+        expected_row_ids_and_row_dicts = _get_expected_row_ids_and_row_dicts_for_partition_order(df, [2, 1, 0])
+
+        shuffled_it = SparkExamplesIterable(df).shuffle_data_sources(generator_mock)
+        assert shuffled_it.num_shards == 3
+        for i, (row_id, row_dict) in enumerate(shuffled_it):
+            expected_row_id, expected_row_dict = expected_row_ids_and_row_dicts[i]
+            assert row_id == expected_row_id
+            assert row_dict == expected_row_dict
+
+
+@require_not_windows
+@require_dill_gt_0_3_2
+def test_spark_examples_iterable_shard():
+    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
+    df = spark.range(20).repartition(4)
+
+    # Partitions 0 and 2
+    shard_it_1 = SparkExamplesIterable(df).shard_data_sources(index=0, num_shards=2, contiguous=False)
+    assert shard_it_1.num_shards == 2
+    expected_row_ids_and_row_dicts_1 = _get_expected_row_ids_and_row_dicts_for_partition_order(df, [0, 2])
+    for i, (row_id, row_dict) in enumerate(shard_it_1):
+        expected_row_id, expected_row_dict = expected_row_ids_and_row_dicts_1[i]
+        assert row_id == expected_row_id
+        assert row_dict == expected_row_dict
+
+    # Partitions 1 and 3
+    shard_it_2 = SparkExamplesIterable(df).shard_data_sources(index=1, num_shards=2, contiguous=False)
+    assert shard_it_2.num_shards == 2
+    expected_row_ids_and_row_dicts_2 = _get_expected_row_ids_and_row_dicts_for_partition_order(df, [1, 3])
+    for i, (row_id, row_dict) in enumerate(shard_it_2):
+        expected_row_id, expected_row_dict = expected_row_ids_and_row_dicts_2[i]
+        assert row_id == expected_row_id
+        assert row_dict == expected_row_dict
+
+
+@require_not_windows
+@require_dill_gt_0_3_2
+def test_repartition_df_if_needed_max_num_df_rows():
+    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
+    df = spark.range(100).repartition(1)
+    spark_builder = Spark(df)
+    # Choose a small max_shard_size for maximum partitioning.
+    spark_builder._repartition_df_if_needed(max_shard_size=1)
+    # The new number of partitions should not be greater than the number of rows.
+    assert spark_builder.df.rdd.getNumPartitions() == 100
+
+
+@require_not_windows
+@require_dill_gt_0_3_2
+def test_iterable_image_features():
+    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
+    img_bytes = np.zeros((10, 10, 3), dtype=np.uint8).tobytes()
+    data = [(img_bytes,)]
+    df = spark.createDataFrame(data, "image: binary")
+    features = Features({"image": Image(decode=False)})
+    dset = IterableDataset.from_spark(df, features=features)
+    item = next(iter(dset))
+    assert item.keys() == {"image"}
+    assert item == {"image": {"path": None, "bytes": img_bytes}}
+
+
+@require_not_windows
+@require_dill_gt_0_3_2
+def test_iterable_image_features_decode():
+    from io import BytesIO
+
+    import PIL.Image
+
+    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
+    img = PIL.Image.fromarray(np.zeros((10, 10, 3), dtype=np.uint8), "RGB")
+    buffer = BytesIO()
+    img.save(buffer, format="PNG")
+    img_bytes = bytes(buffer.getvalue())
+    data = [(img_bytes,)]
+    df = spark.createDataFrame(data, "image: binary")
+    features = Features({"image": Image()})
+    dset = IterableDataset.from_spark(df, features=features)
+    item = next(iter(dset))
+    assert item.keys() == {"image"}
+    assert isinstance(item["image"], PIL.Image.Image)

tests/packaged_modules/test_sql.py

@@ -0,0 +1,16 @@
+import pytest
+
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesList
+from datasets.packaged_modules.sql.sql import SqlConfig
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = SqlConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = SqlConfig(name="name", data_files=data_files)

tests/packaged_modules/test_text.py

@@ -0,0 +1,90 @@
+import textwrap
+
+import pyarrow as pa
+import pytest
+
+from datasets import Features, Image
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesList
+from datasets.packaged_modules.text.text import Text, TextConfig
+
+from ..utils import require_pil
+
+
+@pytest.fixture
+def text_file(tmp_path):
+    filename = tmp_path / "text.txt"
+    data = textwrap.dedent(
+        """\
+        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
+        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
+        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
+        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
+
+        Second paragraph:
+        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
+        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
+        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
+        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
+        """
+    )
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write(data)
+    return str(filename)
+
+
+@pytest.fixture
+def text_file_with_image(tmp_path, image_file):
+    filename = tmp_path / "text_with_image.txt"
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write(image_file)
+    return str(filename)
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = TextConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = TextConfig(name="name", data_files=data_files)
+
+
+@pytest.mark.parametrize("keep_linebreaks", [True, False])
+def test_text_linebreaks(text_file, keep_linebreaks):
+    with open(text_file, encoding="utf-8") as f:
+        expected_content = f.read().splitlines(keepends=keep_linebreaks)
+    text = Text(keep_linebreaks=keep_linebreaks, encoding="utf-8")
+    generator = text._generate_tables(base_files=[text_file], files_iterables=[[text_file]])
+    generated_content = pa.concat_tables([table for _, table in generator]).to_pydict()["text"]
+    assert generated_content == expected_content
+
+
+@require_pil
+def test_text_cast_image(text_file_with_image):
+    with open(text_file_with_image, encoding="utf-8") as f:
+        image_file = f.read().splitlines()[0]
+    text = Text(encoding="utf-8", features=Features({"image": Image()}))
+    generator = text._generate_tables(base_files=[text_file_with_image], files_iterables=[[text_file_with_image]])
+    pa_table = pa.concat_tables([table for _, table in generator])
+    assert pa_table.schema.field("image").type == Image()()
+    generated_content = pa_table.to_pydict()["image"]
+    assert generated_content == [{"path": image_file, "bytes": None}]
+
+
+@pytest.mark.parametrize("sample_by", ["line", "paragraph", "document"])
+def test_text_sample_by(sample_by, text_file):
+    with open(text_file, encoding="utf-8") as f:
+        expected_content = f.read()
+    if sample_by == "line":
+        expected_content = expected_content.splitlines()
+    elif sample_by == "paragraph":
+        expected_content = expected_content.split("\n\n")
+    elif sample_by == "document":
+        expected_content = [expected_content]
+    text = Text(sample_by=sample_by, encoding="utf-8", chunksize=100)
+    generator = text._generate_tables(base_files=[text_file], files_iterables=[[text_file]])
+    generated_content = pa.concat_tables([table for _, table in generator]).to_pydict()["text"]
+    assert generated_content == expected_content

tests/packaged_modules/test_tsfile.py

@@ -0,0 +1,748 @@
+"""Tests for the per-device wide-format TsFile builder."""
+
+from __future__ import annotations
+
+import logging
+from datetime import date, datetime, timedelta, timezone
+from typing import Any, Sequence
+
+import pyarrow as pa
+import pytest
+
+
+# `tsfile` requires pyarrow<20 for python<3.14, which conflicts with datasets'
+# pyarrow>=21.0.0. It is therefore only installed in the py3.14 CI. Skip this
+# whole module (at collection time) when tsfile is not importable.
+pytest.importorskip("tsfile")
+
+from tsfile import ColumnCategory, ColumnSchema, TableSchema, Tablet, TsFileWriter  # noqa: E402
+from tsfile.constants import TSDataType  # noqa: E402
+
+from datasets import IterableDataset, load_dataset  # noqa: E402
+from datasets.builder import InvalidConfigName  # noqa: E402
+from datasets.data_files import DataFilesList  # noqa: E402
+from datasets.packaged_modules.tsfile.tsfile import TsFileConfig, _to_epoch  # noqa: E402
+
+
+# ---------------------------------------------------------------------------
+# Time-base constants
+# ---------------------------------------------------------------------------
+#
+# Every fixture's timestamps live in a disjoint epoch-ms slice off ``T0`` so
+# that, when two files of the same device are merged, the resulting
+# time-sorted order is fully determined by writer-side timestamps. This lets
+# the assertions below check both *content* and *order* unambiguously.
+T0 = 1_700_000_000_000  # base: single_device, multi_device, all_types
+T_EVOLVED = T0 + 500_000  # evolved file (after single_device's 5 points)
+T_INT32 = T0 + 1_000_000
+T_INT64 = T0 + 2_000_000
+T_FLOAT = T0 + 3_000_000
+
+
+# ---------------------------------------------------------------------------
+# Generic writer
+# ---------------------------------------------------------------------------
+
+
+# A row maps column name -> Python value, plus a special "time" -> int (epoch).
+Row = dict[str, Any]
+ColumnSpec = tuple[str, TSDataType, ColumnCategory]
+
+
+def _write_tsfile(path: str, tables: Sequence[tuple[str, Sequence[ColumnSpec], Sequence[Sequence[Row]]]]) -> None:
+    """Write one or more tables, each as one or more tablets, to ``path``.
+
+    Each ``tables`` entry is ``(table_name, columns, tablets)`` where:
+
+    - ``columns`` is the table schema as ``[(name, TSDataType, ColumnCategory), ...]``;
+      it must include exactly one TIME column called ``"time"``.
+    - ``tablets`` is a list of tablets; each tablet is a list of row dicts. A
+      row dict must carry ``"time"`` plus every TAG/FIELD column in the table.
+    """
+    writer = TsFileWriter(path)
+    try:
+        # Register schemas first so multiple-table files validate up-front.
+        for table_name, columns, _ in tables:
+            writer.register_table(TableSchema(table_name, [ColumnSchema(*c) for c in columns]))
+
+        for table_name, columns, tablets in tables:
+            non_time = [(n, t) for (n, t, c) in columns if c != ColumnCategory.TIME]
+            col_names = [n for n, _ in non_time]
+            col_types = [t for _, t in non_time]
+            for rows in tablets:
+                tablet = Tablet(col_names, col_types, len(rows))
+                tablet.set_table_name(table_name)
+                for i, row in enumerate(rows):
+                    tablet.add_timestamp(i, row["time"])
+                    for name in col_names:
+                        tablet.add_value_by_name(name, i, row[name])
+                writer.write_table(tablet)
+    finally:
+        writer.close()
+
+
+# ---------------------------------------------------------------------------
+# Per-fixture writers (each declarative + tiny)
+# ---------------------------------------------------------------------------
+
+
+def _write_single_device(path: str) -> None:
+    """One device 'd1', two DOUBLE fields, 5 points starting at T0."""
+    cols = [
+        ("time", TSDataType.TIMESTAMP, ColumnCategory.TIME),
+        ("device", TSDataType.STRING, ColumnCategory.TAG),
+        ("temperature", TSDataType.DOUBLE, ColumnCategory.FIELD),
+        ("humidity", TSDataType.DOUBLE, ColumnCategory.FIELD),
+    ]
+    rows = [{"time": T0 + i * 1000, "device": "d1", "temperature": 20.0 + i, "humidity": 50.0 + i} for i in range(5)]
+    _write_tsfile(path, [("mytable", cols, [rows])])
+
+
+def _write_multi_device(path: str) -> None:
+    """Three devices, 3 points each, all sharing the same field schema."""
+    cols = [
+        ("time", TSDataType.TIMESTAMP, ColumnCategory.TIME),
+        ("device", TSDataType.STRING, ColumnCategory.TAG),
+        ("temperature", TSDataType.DOUBLE, ColumnCategory.FIELD),
+        ("humidity", TSDataType.DOUBLE, ColumnCategory.FIELD),
+    ]
+    tablets = [
+        [{"time": T0 + i * 1000, "device": dev, "temperature": 10.0 + i, "humidity": 50.0 + i} for i in range(3)]
+        for dev in ("d1", "d2", "d3")
+    ]
+    _write_tsfile(path, [("plant", cols, tablets)])
+
+
+def _write_evolved(path: str) -> None:
+    """Same table+device as single_device, plus a new ``voltage`` field."""
+    cols = [
+        ("time", TSDataType.TIMESTAMP, ColumnCategory.TIME),
+        ("device", TSDataType.STRING, ColumnCategory.TAG),
+        ("temperature", TSDataType.DOUBLE, ColumnCategory.FIELD),
+        ("humidity", TSDataType.DOUBLE, ColumnCategory.FIELD),
+        ("voltage", TSDataType.DOUBLE, ColumnCategory.FIELD),
+    ]
+    rows = [
+        {
+            "time": T_EVOLVED + i * 1000,
+            "device": "d1",
+            "temperature": 30.0 + i,
+            "humidity": 60.0 + i,
+            "voltage": 220.0 + i,
+        }
+        for i in range(3)
+    ]
+    _write_tsfile(path, [("mytable", cols, [rows])])
+
+
+def _write_numeric_field(path: str, ts_type: TSDataType, base_ts: int, value_fn) -> None:
+    """One device 'd1' with a single numeric ``temperature`` field."""
+    cols = [
+        ("time", TSDataType.TIMESTAMP, ColumnCategory.TIME),
+        ("device", TSDataType.STRING, ColumnCategory.TAG),
+        ("temperature", ts_type, ColumnCategory.FIELD),
+    ]
+    rows = [{"time": base_ts + i * 1000, "device": "d1", "temperature": value_fn(i)} for i in range(3)]
+    _write_tsfile(path, [("mytable", cols, [rows])])
+
+
+def _write_two_tables(path: str) -> None:
+    """Two distinct tables in one file: ``table_a`` (registered first) and ``table_b``."""
+    a_cols = [
+        ("time", TSDataType.TIMESTAMP, ColumnCategory.TIME),
+        ("device", TSDataType.STRING, ColumnCategory.TAG),
+        ("a", TSDataType.DOUBLE, ColumnCategory.FIELD),
+    ]
+    b_cols = [
+        ("time", TSDataType.TIMESTAMP, ColumnCategory.TIME),
+        ("device", TSDataType.STRING, ColumnCategory.TAG),
+        ("b", TSDataType.DOUBLE, ColumnCategory.FIELD),
+    ]
+    a_rows = [{"time": 1_000 + i, "device": "d1", "a": float(i)} for i in range(2)]
+    b_rows = [{"time": 2_000 + i, "device": "d1", "b": 100.0 + i} for i in range(2)]
+    _write_tsfile(path, [("table_a", a_cols, [a_rows]), ("table_b", b_cols, [b_rows])])
+
+
+def _write_all_types(path: str) -> None:
+    """Every supported TSDataType represented as a FIELD."""
+    cols = [
+        ("time", TSDataType.TIMESTAMP, ColumnCategory.TIME),
+        ("tag", TSDataType.STRING, ColumnCategory.TAG),
+        ("col_boolean", TSDataType.BOOLEAN, ColumnCategory.FIELD),
+        ("col_int32", TSDataType.INT32, ColumnCategory.FIELD),
+        ("col_int64", TSDataType.INT64, ColumnCategory.FIELD),
+        ("col_float", TSDataType.FLOAT, ColumnCategory.FIELD),
+        ("col_double", TSDataType.DOUBLE, ColumnCategory.FIELD),
+        ("col_text", TSDataType.TEXT, ColumnCategory.FIELD),
+        ("col_string", TSDataType.STRING, ColumnCategory.FIELD),
+        ("col_timestamp", TSDataType.TIMESTAMP, ColumnCategory.FIELD),
+        ("col_date", TSDataType.DATE, ColumnCategory.FIELD),
+        ("col_blob", TSDataType.BLOB, ColumnCategory.FIELD),
+    ]
+    rows = [
+        {
+            "time": T0 + i * 1000,
+            "tag": "d1",
+            "col_boolean": i % 2 == 0,
+            "col_int32": 100 + i,
+            "col_int64": 1_000_000 + i,
+            "col_float": 1.5 + i,
+            "col_double": 100.5 + i,
+            "col_text": f"text_{i}",
+            "col_string": f"str_{i}",
+            "col_timestamp": 1_600_000_000_000 + i * 1000,
+            "col_date": date(2024, 1, 1 + i),
+            "col_blob": f"blob{i}".encode(),
+        }
+        for i in range(3)
+    ]
+    _write_tsfile(path, [("alltypes", cols, [rows])])
+
+
+def _write_large_device(path: str, n_points: int = 200) -> None:
+    """Single device with many points, used to exercise multi-batch concat."""
+    cols = [
+        ("time", TSDataType.TIMESTAMP, ColumnCategory.TIME),
+        ("device", TSDataType.STRING, ColumnCategory.TAG),
+        ("v", TSDataType.INT64, ColumnCategory.FIELD),
+    ]
+    rows = [{"time": T0 + i, "device": "d1", "v": i} for i in range(n_points)]
+    _write_tsfile(path, [("mytable", cols, [rows])])
+
+
+def _write_two_devices_subset(path: str, devices: Sequence[str], base_ts: int) -> None:
+    """A multi-device fixture used to assemble cross-file device sets."""
+    cols = [
+        ("time", TSDataType.TIMESTAMP, ColumnCategory.TIME),
+        ("device", TSDataType.STRING, ColumnCategory.TAG),
+        ("v", TSDataType.DOUBLE, ColumnCategory.FIELD),
+    ]
+    tablets = [[{"time": base_ts + i * 1000, "device": dev, "v": float(i)} for i in range(3)] for dev in devices]
+    _write_tsfile(path, [("mytable", cols, tablets)])
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def make_tsfile(tmp_path):
+    """Factory fixture: ``make_tsfile("name", writer_fn, *args, **kwargs)``."""
+
+    def _make(name: str, writer_fn, *args, **kwargs) -> str:
+        p = str(tmp_path / f"{name}.tsfile")
+        writer_fn(p, *args, **kwargs)
+        return p
+
+    return _make
+
+
+@pytest.fixture
+def tsfile_path(make_tsfile):
+    return make_tsfile("sample", _write_single_device)
+
+
+@pytest.fixture
+def multi_device_tsfile_path(make_tsfile):
+    return make_tsfile("multi", _write_multi_device)
+
+
+@pytest.fixture
+def evolved_tsfile_path(make_tsfile):
+    return make_tsfile("evolved", _write_evolved)
+
+
+@pytest.fixture
+def two_tables_tsfile_path(make_tsfile):
+    return make_tsfile("two_tables", _write_two_tables)
+
+
+@pytest.fixture
+def all_types_tsfile_path(make_tsfile):
+    return make_tsfile("alltypes", _write_all_types)
+
+
+# ---------------------------------------------------------------------------
+# Config-level
+# ---------------------------------------------------------------------------
+
+
+def test_config_raises_when_invalid_name():
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        TsFileConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files):
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        TsFileConfig(name="name", data_files=data_files)
+
+
+@pytest.mark.parametrize(
+    "kwargs, match",
+    [
+        ({"input_batch_size": 0}, "input_batch_size"),
+        ({"output_batch_size": 0}, "output_batch_size"),
+        ({"columns": []}, "non-empty"),
+        ({"timestamp_unit": "minute"}, "timestamp_unit"),
+        ({"on_bad_files": "boom"}, "on_bad_files"),
+    ],
+)
+def test_config_rejects_invalid_values(kwargs, match):
+    with pytest.raises(ValueError, match=match):
+        TsFileConfig(name="x", **kwargs)
+
+
+def test_config_normalizes_time_bounds():
+    cfg = TsFileConfig(
+        name="x",
+        start_time=pa.scalar(1500, type=pa.timestamp("ms")),
+        end_time=2000,
+    )
+    assert cfg.start_time == 1500
+    assert cfg.end_time == 2000
+
+
+# ---------------------------------------------------------------------------
+# _to_epoch unit tests
+# ---------------------------------------------------------------------------
+
+
+def test_to_epoch_int_passthrough():
+    assert _to_epoch(1234, "ms") == 1234
+
+
+def test_to_epoch_naive_datetime():
+    assert _to_epoch(datetime(1970, 1, 1, 0, 0, 1), "ms") == 1000
+
+
+@pytest.mark.parametrize(
+    "aware",
+    [
+        datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone(timedelta(hours=8))),
+        "2024-01-01T00:00:00+08:00",
+    ],
+    ids=["datetime", "iso_string"],
+)
+def test_to_epoch_aware_inputs_normalized_to_utc(aware):
+    # 2024-01-01T00:00:00 in UTC+8 == 2023-12-31T16:00:00 UTC.
+    naive_utc = datetime(2023, 12, 31, 16, 0, 0)
+    assert _to_epoch(aware, "ms") == _to_epoch(naive_utc, "ms")
+
+
+def test_to_epoch_date():
+    assert _to_epoch(date(1970, 1, 2), "ms") == 86_400_000
+
+
+def test_to_epoch_iso_string():
+    assert _to_epoch("1970-01-01T00:00:01", "ms") == 1000
+
+
+def test_to_epoch_pa_scalar():
+    assert _to_epoch(pa.scalar(1500, type=pa.timestamp("ms")), "ms") == 1500
+
+
+def test_to_epoch_rejects_bool():
+    with pytest.raises(TypeError, match="bool"):
+        _to_epoch(True, "ms")
+
+
+@pytest.mark.parametrize("value", [object(), b"bytes", "not-a-date"])
+def test_to_epoch_rejects_garbage(value):
+    with pytest.raises(TypeError, match="must be a"):
+        _to_epoch(value, "ms")
+
+
+# ---------------------------------------------------------------------------
+# End-to-end: single device, full table
+# ---------------------------------------------------------------------------
+
+
+def test_load_full_table(tsfile_path):
+    ds = load_dataset("tsfile", data_files=tsfile_path)["train"]
+
+    # One row per device. TAG = scalar string; time + fields = lists.
+    assert ds.column_names == ["device", "time", "temperature", "humidity"]
+    assert len(ds) == 1
+    row = ds[0]
+    assert row["device"] == "d1"
+    assert len(row["time"]) == 5
+    assert row["time"][0] == datetime(2023, 11, 14, 22, 13, 20)
+    assert row["time"][-1] == datetime(2023, 11, 14, 22, 13, 24)
+    assert row["temperature"] == [20.0, 21.0, 22.0, 23.0, 24.0]
+    assert row["humidity"] == [50.0, 51.0, 52.0, 53.0, 54.0]
+
+
+def test_load_with_field_subset(tsfile_path):
+    ds = load_dataset("tsfile", data_files=tsfile_path, columns=["temperature"])["train"]
+    assert ds.column_names == ["device", "time", "temperature"]
+    assert ds[0]["temperature"] == [20.0, 21.0, 22.0, 23.0, 24.0]
+
+
+def test_columns_are_lowercased(tsfile_path):
+    ds = load_dataset("tsfile", data_files=tsfile_path, columns=["TEMPERATURE", "Humidity"])["train"]
+    assert ds.column_names == ["device", "time", "temperature", "humidity"]
+
+
+def test_columns_request_tag_is_silently_ignored(tsfile_path):
+    """Passing a TAG name in `columns` is a no-op (TAGs are always emitted)."""
+    ds = load_dataset("tsfile", data_files=tsfile_path, columns=["device", "temperature"])["train"]
+
+    assert ds.column_names == ["device", "time", "temperature"]
+    assert ds.features["device"].dtype == "string"
+    assert ds.features["temperature"].feature.dtype == "float64"
+    assert ds["device"] == ["d1"]
+    assert ds[0]["temperature"] == [20.0, 21.0, 22.0, 23.0, 24.0]
+
+
+def test_columns_request_time_is_silently_ignored(tsfile_path):
+    """Passing the TIME column name in `columns` is a no-op (TIME is always emitted)."""
+    ds = load_dataset("tsfile", data_files=tsfile_path, columns=["time", "temperature"])["train"]
+
+    # `time` should appear exactly once, and as the real timestamp list — not
+    # as a duplicate all-null float64 list column.
+    assert ds.column_names == ["device", "time", "temperature"]
+    assert ds.features["time"].feature.dtype.startswith("timestamp")
+    row = ds[0]
+    assert len(row["time"]) == 5
+    assert row["time"][0] == datetime(2023, 11, 14, 22, 13, 20)
+    assert row["time"][-1] == datetime(2023, 11, 14, 22, 13, 24)
+    assert row["temperature"] == [20.0, 21.0, 22.0, 23.0, 24.0]
+
+
+def test_columns_request_only_time(tsfile_path):
+    """`columns=["time"]` should still produce TAG + TIME, with no FIELD list columns."""
+    ds = load_dataset("tsfile", data_files=tsfile_path, columns=["time"])["train"]
+
+    assert ds.column_names == ["device", "time"]
+    assert ds.features["time"].feature.dtype.startswith("timestamp")
+    row = ds[0]
+    assert row["device"] == "d1"
+    assert len(row["time"]) == 5
+    assert row["time"][0] == datetime(2023, 11, 14, 22, 13, 20)
+    assert row["time"][-1] == datetime(2023, 11, 14, 22, 13, 24)
+
+
+def test_columns_unknown_field_filled_with_null(tsfile_path):
+    ds = load_dataset(
+        "tsfile",
+        data_files=tsfile_path,
+        columns=["temperature", "voltage"],  # voltage is absent
+    )["train"]
+
+    assert ds.column_names == ["device", "time", "temperature", "voltage"]
+    row = ds[0]
+    assert row["temperature"] == [20.0, 21.0, 22.0, 23.0, 24.0]
+    assert row["voltage"] == [None] * 5
+
+
+def test_columns_all_unknown_still_returns_time_and_tags(tsfile_path):
+    ds = load_dataset(
+        "tsfile",
+        data_files=tsfile_path,
+        columns=["nonexistent_a", "nonexistent_b"],
+    )["train"]
+
+    assert ds.column_names == ["device", "time", "nonexistent_a", "nonexistent_b"]
+    row = ds[0]
+    assert row["device"] == "d1"
+    assert len(row["time"]) == 5
+    assert row["nonexistent_a"] == [None] * 5
+    assert row["nonexistent_b"] == [None] * 5
+
+
+# ---------------------------------------------------------------------------
+# Time-range filtering
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "start, end",
+    [
+        # pa.scalar from datetime
+        (
+            pa.scalar(datetime(2023, 11, 14, 22, 13, 21), type=pa.timestamp("ms")),
+            pa.scalar(datetime(2023, 11, 14, 22, 13, 23), type=pa.timestamp("ms")),
+        ),
+        # pa.scalar from int epoch
+        (
+            pa.scalar(T0 + 1000, type=pa.timestamp("ms")),
+            pa.scalar(T0 + 3000, type=pa.timestamp("ms")),
+        ),
+        # plain int epoch
+        (T0 + 1000, T0 + 3000),
+        # datetime
+        (datetime(2023, 11, 14, 22, 13, 21), datetime(2023, 11, 14, 22, 13, 23)),
+        # ISO-8601 string
+        ("2023-11-14T22:13:21", "2023-11-14T22:13:23"),
+    ],
+)
+def test_load_with_time_range_inputs(tsfile_path, start, end):
+    ds = load_dataset("tsfile", data_files=tsfile_path, start_time=start, end_time=end)["train"]
+    assert len(ds[0]["time"]) == 3
+
+
+# ---------------------------------------------------------------------------
+# Multi-device & cross-file folding
+# ---------------------------------------------------------------------------
+
+
+def test_load_multi_device_one_row_per_device(multi_device_tsfile_path):
+    ds = load_dataset("tsfile", data_files=multi_device_tsfile_path)["train"]
+
+    assert len(ds) == 3
+    assert sorted(ds["device"]) == ["d1", "d2", "d3"]
+    for row in ds:
+        assert len(row["time"]) == 3
+        assert row["temperature"] == [10.0, 11.0, 12.0]
+        assert row["humidity"] == [50.0, 51.0, 52.0]
+
+
+def test_schema_evolution_merges_same_device(tsfile_path, evolved_tsfile_path):
+    """Same device d1 in two files → one row, lists merged in time order."""
+    ds = load_dataset("tsfile", data_files=[tsfile_path, evolved_tsfile_path])["train"]
+
+    assert "voltage" in ds.column_names
+    assert len(ds) == 1
+    row = ds[0]
+    assert row["device"] == "d1"
+    # 5 (old) + 3 (new) points, fully time-ordered.
+    assert len(row["time"]) == 8
+    # Old file lacked `voltage` → null on its 5 points; new file fills the rest.
+    assert row["voltage"] == [None] * 5 + [220.0, 221.0, 222.0]
+    # `temperature` is present in both files, contiguous in the merged order.
+    assert row["temperature"] == [20.0, 21.0, 22.0, 23.0, 24.0, 30.0, 31.0, 32.0]
+
+
+def test_multi_file_multi_device_partial_overlap(make_tsfile):
+    """Two files × two devices each, with one device shared.
+
+    File A: devices {d1, d2}; file B: devices {d2, d3}. The merged dataset
+    must have 3 rows (one per unique device), and d2 must have *6* points
+    (3 from each file) sorted by time.
+    """
+    fa = make_tsfile("a", _write_two_devices_subset, devices=["d1", "d2"], base_ts=T0)
+    fb = make_tsfile("b", _write_two_devices_subset, devices=["d2", "d3"], base_ts=T0 + 100_000)
+
+    ds = load_dataset("tsfile", data_files=[fa, fb])["train"]
+    by_dev = {row["device"]: row for row in ds}
+    assert set(by_dev) == {"d1", "d2", "d3"}
+    assert len(by_dev["d1"]["time"]) == 3
+    assert len(by_dev["d3"]["time"]) == 3
+    # Shared device gets all 6 points, time-sorted (file A first, then B).
+    assert len(by_dev["d2"]["time"]) == 6
+    assert by_dev["d2"]["v"] == [0.0, 1.0, 2.0, 0.0, 1.0, 2.0]
+
+
+# ---------------------------------------------------------------------------
+# Type promotion across files
+# ---------------------------------------------------------------------------
+
+
+def test_type_promotion_int32_to_int64(make_tsfile):
+    int32_path = make_tsfile("narrow", _write_numeric_field, TSDataType.INT32, T_INT32, lambda i: 10 + i)
+    int64_path = make_tsfile("wide", _write_numeric_field, TSDataType.INT64, T_INT64, lambda i: 1_000_000 + i)
+
+    ds = load_dataset("tsfile", data_files=[int32_path, int64_path])["train"]
+    assert len(ds) == 1
+    assert ds.features["temperature"].feature.dtype == "int64"
+    # int32 timestamps come earlier (T_INT32 < T_INT64).
+    assert ds[0]["temperature"] == [10, 11, 12, 1_000_000, 1_000_001, 1_000_002]
+
+
+def test_type_promotion_float_to_double(make_tsfile):
+    float_path = make_tsfile("narrow", _write_numeric_field, TSDataType.FLOAT, T_FLOAT, lambda i: 1.5 + i)
+    double_path = make_tsfile("wide", _write_single_device)
+
+    ds = load_dataset("tsfile", data_files=[float_path, double_path])["train"]
+    assert len(ds) == 1
+    assert ds.features["temperature"].feature.dtype == "float64"
+    # double fixture lives at T0..T0+4s; float at T_FLOAT (later).
+    assert ds[0]["temperature"] == [20.0, 21.0, 22.0, 23.0, 24.0, 1.5, 2.5, 3.5]
+
+
+def test_type_promotion_int32_to_double(make_tsfile):
+    int32_path = make_tsfile("int", _write_numeric_field, TSDataType.INT32, T_INT32, lambda i: 10 + i)
+    double_path = make_tsfile("double", _write_single_device)
+
+    ds = load_dataset("tsfile", data_files=[int32_path, double_path])["train"]
+    assert len(ds) == 1
+    # INT32 + DOUBLE → DOUBLE (two-step widening).
+    assert ds.features["temperature"].feature.dtype == "float64"
+    assert ds[0]["temperature"] == [20.0, 21.0, 22.0, 23.0, 24.0, 10.0, 11.0, 12.0]
+
+
+# ---------------------------------------------------------------------------
+# All-types
+# ---------------------------------------------------------------------------
+
+
+def test_load_all_supported_types(all_types_tsfile_path):
+    ds = load_dataset("tsfile", data_files=all_types_tsfile_path)["train"]
+
+    assert len(ds) == 1
+    assert ds.column_names == [
+        "tag",
+        "time",
+        "col_boolean",
+        "col_int32",
+        "col_int64",
+        "col_float",
+        "col_double",
+        "col_text",
+        "col_string",
+        "col_timestamp",
+        "col_date",
+        "col_blob",
+    ]
+    row = ds[0]
+    assert row["tag"] == "d1"
+    assert row["col_boolean"] == [True, False, True]
+    assert row["col_int32"] == [100, 101, 102]
+    assert row["col_int64"] == [1_000_000, 1_000_001, 1_000_002]
+    assert row["col_float"] == [1.5, 2.5, 3.5]
+    assert row["col_double"] == [100.5, 101.5, 102.5]
+    assert row["col_text"] == ["text_0", "text_1", "text_2"]
+    assert row["col_string"] == ["str_0", "str_1", "str_2"]
+    assert row["col_timestamp"][0] == datetime(2020, 9, 13, 12, 26, 40)
+    assert row["col_date"][0] == date(2024, 1, 1)
+    assert row["col_date"][2] == date(2024, 1, 3)
+    assert row["col_blob"][0] == b"blob0"
+    assert row["col_blob"][2] == b"blob2"
+
+
+# ---------------------------------------------------------------------------
+# Multi-table file: explicit `table_name` selection
+# ---------------------------------------------------------------------------
+
+
+def test_default_table_is_first(two_tables_tsfile_path):
+    ds = load_dataset("tsfile", data_files=two_tables_tsfile_path)["train"]
+    # `table_a` registered first → default pick.
+    assert "a" in ds.column_names
+    assert "b" not in ds.column_names
+
+
+def test_explicit_table_name(two_tables_tsfile_path):
+    ds = load_dataset("tsfile", data_files=two_tables_tsfile_path, table_name="table_b")["train"]
+    assert "b" in ds.column_names
+    assert "a" not in ds.column_names
+
+
+# ---------------------------------------------------------------------------
+# Streaming (IterableDataset)
+# ---------------------------------------------------------------------------
+
+
+def test_streaming_yields_same_rows(multi_device_tsfile_path):
+    ds = load_dataset("tsfile", data_files=multi_device_tsfile_path, streaming=True)["train"]
+    assert isinstance(ds, IterableDataset)
+    rows = list(ds)
+    assert len(rows) == 3
+    assert sorted(r["device"] for r in rows) == ["d1", "d2", "d3"]
+    for r in rows:
+        assert r["temperature"] == [10.0, 11.0, 12.0]
+
+
+# ---------------------------------------------------------------------------
+# Timezone
+# ---------------------------------------------------------------------------
+
+
+def test_load_with_timezone(make_tsfile):
+    """`timestamp_tz="UTC"` round-trips: list values come back tz-aware."""
+    path = make_tsfile("tz", _write_single_device)
+    ds = load_dataset("tsfile", data_files=path, timestamp_tz="UTC")["train"]
+    ts = ds[0]["time"][0]
+    assert ts.tzinfo is not None
+    # Same wall-clock as the naive case, attached to UTC.
+    assert ts == datetime(2023, 11, 14, 22, 13, 20, tzinfo=timezone.utc)
+
+
+# ---------------------------------------------------------------------------
+# Large-batch / multi-chunk concat
+# ---------------------------------------------------------------------------
+
+
+def test_large_device_with_small_batch_size(make_tsfile):
+    """Force multiple Arrow batches per device → exercise the concat path."""
+    path = make_tsfile("big", _write_large_device, n_points=200)
+    ds = load_dataset("tsfile", data_files=path, input_batch_size=64)["train"]
+    assert len(ds) == 1
+    row = ds[0]
+    assert len(row["time"]) == 200
+    assert row["v"] == list(range(200))
+
+
+# ---------------------------------------------------------------------------
+# Duplicate-timestamp detection (cross-file)
+# ---------------------------------------------------------------------------
+
+
+def test_duplicate_timestamp_across_files_raises(make_tsfile):
+    """Same device, same ts in two files → `_finalize_device` must raise."""
+    cols = [
+        ("time", TSDataType.TIMESTAMP, ColumnCategory.TIME),
+        ("device", TSDataType.STRING, ColumnCategory.TAG),
+        ("v", TSDataType.DOUBLE, ColumnCategory.FIELD),
+    ]
+    rows = [{"time": 5_000, "device": "d1", "v": 1.0}]
+
+    a = make_tsfile("dupA", lambda p: _write_tsfile(p, [("mytable", cols, [rows])]))
+    b = make_tsfile("dupB", lambda p: _write_tsfile(p, [("mytable", cols, [rows])]))
+
+    with pytest.raises(Exception) as excinfo:
+        load_dataset("tsfile", data_files=[a, b])
+    # The ValueError is wrapped by `_prepare_split_single` into a
+    # DatasetGenerationError; check the cause chain for the original message.
+    chain = [excinfo.value, *(_iter_causes(excinfo.value))]
+    assert any("Duplicate timestamp" in str(e) for e in chain)
+
+
+def _iter_causes(exc: BaseException):
+    while exc.__cause__ is not None:
+        exc = exc.__cause__
+        yield exc
+
+
+# ---------------------------------------------------------------------------
+# Error handling
+# ---------------------------------------------------------------------------
+
+
+def test_on_bad_files_skip(tmp_path, tsfile_path):
+    bad = tmp_path / "broken.tsfile"
+    bad.write_bytes(b"not a real tsfile")
+
+    ds = load_dataset(
+        "tsfile",
+        data_files=[tsfile_path, str(bad)],
+        on_bad_files="skip",
+    )["train"]
+    assert len(ds) == 1
+    assert len(ds[0]["time"]) == 5
+
+
+def test_on_bad_files_warn(tmp_path, tsfile_path, caplog):
+    bad = tmp_path / "broken.tsfile"
+    bad.write_bytes(b"not a real tsfile")
+
+    with caplog.at_level(logging.WARNING, logger="datasets.packaged_modules.tsfile.tsfile"):
+        ds = load_dataset(
+            "tsfile",
+            data_files=[tsfile_path, str(bad)],
+            on_bad_files="warn",
+        )["train"]
+    assert len(ds) == 1
+    assert any("Skipping bad file" in rec.message for rec in caplog.records)
+
+
+def test_on_bad_files_default_raises(tmp_path, tsfile_path):
+    bad = tmp_path / "broken.tsfile"
+    bad.write_bytes(b"not a real tsfile")
+
+    with pytest.raises(Exception) as excinfo:
+        load_dataset("tsfile", data_files=[tsfile_path, str(bad)])
+    chain = [excinfo.value, *(_iter_causes(excinfo.value))]
+    assert any("not a valid TsFile" in str(e) for e in chain)

tests/packaged_modules/test_videofolder.py

@@ -0,0 +1,132 @@
+import shutil
+import textwrap
+from pathlib import Path
+
+import pytest
+
+from datasets import ClassLabel, DownloadManager, Features, Video
+from datasets.builder import InvalidConfigName
+from datasets.data_files import DataFilesDict, DataFilesList, get_data_patterns
+from datasets.download.streaming_download_manager import StreamingDownloadManager
+from datasets.packaged_modules.videofolder.videofolder import VideoFolder, VideoFolderConfig
+
+
+@pytest.fixture
+def cache_dir(tmp_path):
+    return str(tmp_path / "videofolder_cache_dir")
+
+
+@pytest.fixture
+def video_file_path():
+    return Path(__file__).resolve().parents[1] / "features" / "data" / "test_video_66x50.mov"
+
+
+@pytest.fixture
+def data_files_with_labels_no_metadata(tmp_path, video_file_path):
+    data_dir = tmp_path / "data_files_with_labels_no_metadata"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    subdir_class_0 = data_dir / "cat"
+    subdir_class_0.mkdir(parents=True, exist_ok=True)
+    subdir_class_1 = data_dir / "dog"
+    subdir_class_1.mkdir(parents=True, exist_ok=True)
+
+    video_filename = subdir_class_0 / "video_cat.mov"
+    shutil.copyfile(video_file_path, video_filename)
+    video_filename2 = subdir_class_1 / "video_dog.mov"
+    shutil.copyfile(video_file_path, video_filename2)
+
+    data_files = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+    return data_files
+
+
+@pytest.fixture
+def video_file_with_metadata(tmp_path, video_file_path):
+    video_filename = tmp_path / "video.mov"
+    shutil.copyfile(video_file_path, video_filename)
+    metadata_filename = tmp_path / "metadata.jsonl"
+    metadata = textwrap.dedent(
+        """\
+        {"file_name": "video.mov", "caption": "A short video"}
+        """
+    )
+    with open(metadata_filename, "w", encoding="utf-8") as f:
+        f.write(metadata)
+    return str(video_filename), str(metadata_filename)
+
+
+@pytest.fixture
+def data_files_with_zip_archives(tmp_path, video_file_path):
+    data_dir = tmp_path / "videofolder_data_dir_with_zip_archives"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    archive_dir = data_dir / "archive"
+    archive_dir.mkdir(parents=True, exist_ok=True)
+    subdir = archive_dir / "subdir"
+    subdir.mkdir(parents=True, exist_ok=True)
+
+    video_filename = archive_dir / "video.mov"
+    shutil.copyfile(video_file_path, video_filename)
+    video_filename2 = subdir / "video2.mov"
+    shutil.copyfile(video_file_path, video_filename2)
+
+    metadata_filename = archive_dir / "metadata.jsonl"
+    metadata = textwrap.dedent(
+        """\
+        {"file_name": "video.mov", "caption": "First video"}
+        {"file_name": "subdir/video2.mov", "caption": "Second video"}
+        """
+    )
+    with open(metadata_filename, "w", encoding="utf-8") as f:
+        f.write(metadata)
+
+    shutil.make_archive(archive_dir, "zip", archive_dir)
+    shutil.rmtree(str(archive_dir))
+
+    data_files = DataFilesDict.from_patterns(get_data_patterns(str(data_dir)), data_dir.as_posix())
+    assert len(data_files) == 1
+    assert len(data_files["train"]) == 1
+    return data_files
+
+
+def test_config_raises_when_invalid_name() -> None:
+    with pytest.raises(InvalidConfigName, match="Bad characters"):
+        _ = VideoFolderConfig(name="name-with-*-invalid-character")
+
+
+@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])])
+def test_config_raises_when_invalid_data_files(data_files) -> None:
+    with pytest.raises(ValueError, match="Expected a DataFilesDict"):
+        _ = VideoFolderConfig(name="name", data_files=data_files)
+
+
+def test_generate_examples_with_labels(data_files_with_labels_no_metadata, cache_dir):
+    videofolder = VideoFolder(data_files=data_files_with_labels_no_metadata, cache_dir=cache_dir, drop_labels=False)
+    gen_kwargs = videofolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    assert videofolder.info.features == Features({"video": Video(), "label": ClassLabel(names=["cat", "dog"])})
+    generator = videofolder._generate_examples(**gen_kwargs)
+    assert all(example["label"] in {"cat", "dog"} for _, example in generator)
+
+
+def test_generate_examples_with_metadata(video_file_with_metadata, cache_dir):
+    video_file, metadata_file = video_file_with_metadata
+    videofolder = VideoFolder(data_files=[video_file, metadata_file], cache_dir=cache_dir)
+    gen_kwargs = videofolder._split_generators(StreamingDownloadManager())[0].gen_kwargs
+    generated_examples = [example for _, example in videofolder._generate_examples(**gen_kwargs)]
+    assert len(generated_examples) == 1
+    assert generated_examples[0].keys() == {"video", "caption"}
+    assert generated_examples[0]["video"].endswith("video.mov")
+    assert generated_examples[0]["caption"] == "A short video"
+
+
+@pytest.mark.parametrize("streaming", [False, True])
+def test_data_files_with_metadata_and_archives(streaming, cache_dir, data_files_with_zip_archives):
+    videofolder = VideoFolder(data_files=data_files_with_zip_archives, cache_dir=cache_dir)
+    download_manager = StreamingDownloadManager() if streaming else DownloadManager()
+    generated_splits = videofolder._split_generators(download_manager)
+    for (split, files), generated_split in zip(data_files_with_zip_archives.items(), generated_splits):
+        assert split == generated_split.name
+        num_of_archives = len(files)
+        expected_num_of_examples = 2 * num_of_archives
+        generated_examples = list(videofolder._generate_examples(**generated_split.gen_kwargs))
+        assert len(generated_examples) == expected_num_of_examples
+        assert len({example["video"] for _, example in generated_examples}) == expected_num_of_examples
+        assert len({example["caption"] for _, example in generated_examples}) == expected_num_of_examples

tests/packaged_modules/test_webdataset.py

@@ -0,0 +1,365 @@
+import json
+import tarfile
+from pathlib import Path
+
+import pytest
+
+from datasets import Audio, DownloadManager, Features, Image, List, Value, Video
+from datasets.packaged_modules.webdataset.webdataset import WebDataset
+
+from ..utils import (
+    require_numpy1_on_windows,
+    require_pil,
+    require_torch,
+    require_torchcodec,
+)
+
+
+@pytest.fixture
+def gzipped_text_wds_file(tmp_path, text_gz_path):
+    filename = tmp_path / "file.tar"
+    num_examples = 3
+    with tarfile.open(str(filename), "w") as f:
+        for example_idx in range(num_examples):
+            f.add(text_gz_path, f"{example_idx:05d}.txt.gz")
+    return str(filename)
+
+
+@pytest.fixture
+def image_wds_file(tmp_path, image_file):
+    json_file = tmp_path / "data.json"
+    filename = tmp_path / "file.tar"
+    num_examples = 3
+    with json_file.open("w", encoding="utf-8") as f:
+        f.write(json.dumps({"caption": "this is an image"}))
+    with tarfile.open(str(filename), "w") as f:
+        for example_idx in range(num_examples):
+            f.add(json_file, f"{example_idx:05d}.json")
+            f.add(image_file, f"{example_idx:05d}.jpg")
+    return str(filename)
+
+
+@pytest.fixture
+def upper_lower_case_file(tmp_path):
+    tar_path = tmp_path / "file.tar"
+    num_examples = 3
+    variants = [
+        ("INFO1", "json"),
+        ("info2", "json"),
+        ("info3", "JSON"),
+        ("info3", "json"),  # should probably remove if testing on a case insensitive filesystem
+    ]
+    with tarfile.open(tar_path, "w") as tar:
+        for example_idx in range(num_examples):
+            example_name = f"{example_idx:05d}_{'a' if example_idx % 2 else 'A'}"
+            for tag, ext in variants:
+                caption_path = tmp_path / f"{example_name}.{tag}.{ext}"
+                caption_text = {"caption": f"caption for {example_name}.{tag}.{ext}"}
+                caption_path.write_text(json.dumps(caption_text), encoding="utf-8")
+                tar.add(caption_path, arcname=f"{example_name}.{tag}.{ext}")
+    return str(tar_path)
+
+
+@pytest.fixture
+def audio_wds_file(tmp_path, audio_file):
+    json_file = tmp_path / "data.json"
+    filename = tmp_path / "file.tar"
+    num_examples = 3
+    with json_file.open("w", encoding="utf-8") as f:
+        f.write(json.dumps({"transcript": "this is a transcript"}))
+    with tarfile.open(str(filename), "w") as f:
+        for example_idx in range(num_examples):
+            f.add(json_file, f"{example_idx:05d}.json")
+            f.add(audio_file, f"{example_idx:05d}.wav")
+    return str(filename)
+
+
+@pytest.fixture
+def video_wds_file(tmp_path):
+    json_file = tmp_path / "data.json"
+    filename = tmp_path / "file.tar"
+    video_file = Path(__file__).resolve().parents[1] / "features" / "data" / "test_video_66x50.mov"
+    num_examples = 3
+    with json_file.open("w", encoding="utf-8") as f:
+        f.write(json.dumps({"caption": "this is a video"}))
+    with tarfile.open(str(filename), "w") as f:
+        for example_idx in range(num_examples):
+            f.add(json_file, f"{example_idx:05d}.json")
+            f.add(video_file, f"{example_idx:05d}.mov")
+    return str(filename)
+
+
+@pytest.fixture
+def bad_wds_file(tmp_path, image_file, text_file):
+    json_file = tmp_path / "data.json"
+    filename = tmp_path / "bad_file.tar"
+    with json_file.open("w", encoding="utf-8") as f:
+        f.write(json.dumps({"caption": "this is an image"}))
+    with tarfile.open(str(filename), "w") as f:
+        f.add(image_file)
+        f.add(json_file)
+    return str(filename)
+
+
+@pytest.fixture
+def tensor_wds_file(tmp_path, tensor_file):
+    json_file = tmp_path / "data.json"
+    filename = tmp_path / "file.tar"
+    num_examples = 3
+    with json_file.open("w", encoding="utf-8") as f:
+        f.write(json.dumps({"text": "this is a text"}))
+    with tarfile.open(str(filename), "w") as f:
+        for example_idx in range(num_examples):
+            f.add(json_file, f"{example_idx:05d}.json")
+            f.add(tensor_file, f"{example_idx:05d}.pth")
+    return str(filename)
+
+
+@require_pil
+def test_gzipped_text_webdataset(gzipped_text_wds_file, text_path):
+    data_files = {"train": [gzipped_text_wds_file]}
+    webdataset = WebDataset(data_files=data_files)
+    split_generators = webdataset._split_generators(DownloadManager())
+    assert webdataset.info.features == Features(
+        {
+            "__key__": Value("string"),
+            "__url__": Value("string"),
+            "txt.gz": Value("string"),
+        }
+    )
+    assert len(split_generators) == 1
+    split_generator = split_generators[0]
+    assert split_generator.name == "train"
+    generator = webdataset._generate_examples(**split_generator.gen_kwargs)
+    _, examples = zip(*generator)
+    assert len(examples) == 3
+    assert isinstance(examples[0]["txt.gz"], str)
+    with open(text_path, "r") as f:
+        assert examples[0]["txt.gz"].replace("\r\n", "\n") == f.read().replace("\r\n", "\n")
+
+
+@require_pil
+def test_image_webdataset(image_wds_file):
+    import PIL.Image
+
+    data_files = {"train": [image_wds_file]}
+    webdataset = WebDataset(data_files=data_files)
+    split_generators = webdataset._split_generators(DownloadManager())
+    assert webdataset.info.features == Features(
+        {
+            "__key__": Value("string"),
+            "__url__": Value("string"),
+            "json": {"caption": Value("string")},
+            "jpg": Image(),
+        }
+    )
+    assert len(split_generators) == 1
+    split_generator = split_generators[0]
+    assert split_generator.name == "train"
+    generator = webdataset._generate_examples(**split_generator.gen_kwargs)
+    _, examples = zip(*generator)
+    assert len(examples) == 3
+    assert isinstance(examples[0]["json"], dict)
+    assert isinstance(examples[0]["json"]["caption"], str)
+    assert isinstance(examples[0]["jpg"], dict)  # keep encoded to avoid unecessary copies
+    encoded = webdataset.info.features.encode_example(examples[0])
+    decoded = webdataset.info.features.decode_example(encoded)
+    assert isinstance(decoded["json"], dict)
+    assert isinstance(decoded["json"]["caption"], str)
+    assert isinstance(decoded["jpg"], PIL.Image.Image)
+
+
+def test_upper_lower_case(upper_lower_case_file):
+    variants = [
+        ("INFO1", "json"),
+        ("info2", "json"),
+        ("info3", "JSON"),
+        ("info3", "json"),
+    ]
+
+    data_files = {"train": [upper_lower_case_file]}
+    webdataset = WebDataset(data_files=data_files)
+    split_generators = webdataset._split_generators(DownloadManager())
+
+    variant_keys = [f"{tag}.{ext}" for tag, ext in variants]
+    assert webdataset.info.features == Features(
+        {
+            "__key__": Value("string"),
+            "__url__": Value("string"),
+            **{k: {"caption": Value("string")} for k in variant_keys},
+        }
+    )
+
+    assert len(split_generators) == 1
+    split_generator = split_generators[0]
+    assert split_generator.name == "train"
+    generator = webdataset._generate_examples(**split_generator.gen_kwargs)
+    _, examples = zip(*generator)
+
+    assert len(examples) == 3
+    for example_idx, example in enumerate(examples):
+        example_name = example["__key__"]
+        expected_example_name = f"{example_idx:05d}_{'a' if example_idx % 2 else 'A'}"
+
+        assert example_name == expected_example_name
+        for key in variant_keys:
+            assert isinstance(example[key], dict)
+            assert example[key]["caption"] == f"caption for {example_name}.{key}"
+
+        encoded = webdataset.info.features.encode_example(example)
+        decoded = webdataset.info.features.decode_example(encoded)
+        for key in variant_keys:
+            assert decoded[key]["caption"] == example[key]["caption"]
+
+
+@require_pil
+def test_image_webdataset_missing_keys(image_wds_file):
+    import PIL.Image
+
+    data_files = {"train": [image_wds_file]}
+    features = Features(
+        {
+            "__key__": Value("string"),
+            "__url__": Value("string"),
+            "json": {"caption": Value("string")},
+            "jpg": Image(),
+            "jpeg": Image(),  # additional field
+            "txt": Value("string"),  # additional field
+        }
+    )
+    webdataset = WebDataset(data_files=data_files, features=features)
+    split_generators = webdataset._split_generators(DownloadManager())
+    assert webdataset.info.features == features
+    split_generator = split_generators[0]
+    assert split_generator.name == "train"
+    generator = webdataset._generate_examples(**split_generator.gen_kwargs)
+    _, example = next(iter(generator))
+    encoded = webdataset.info.features.encode_example(example)
+    decoded = webdataset.info.features.decode_example(encoded)
+    assert isinstance(decoded["json"], dict)
+    assert isinstance(decoded["json"]["caption"], str)
+    assert isinstance(decoded["jpg"], PIL.Image.Image)
+    assert decoded["jpeg"] is None
+    assert decoded["txt"] is None
+
+
+@require_torchcodec
+def test_audio_webdataset(audio_wds_file):
+    from torchcodec.decoders import AudioDecoder
+
+    data_files = {"train": [audio_wds_file]}
+    webdataset = WebDataset(data_files=data_files)
+    split_generators = webdataset._split_generators(DownloadManager())
+    assert webdataset.info.features == Features(
+        {
+            "__key__": Value("string"),
+            "__url__": Value("string"),
+            "json": {"transcript": Value("string")},
+            "wav": Audio(),
+        }
+    )
+    assert len(split_generators) == 1
+    split_generator = split_generators[0]
+    assert split_generator.name == "train"
+    generator = webdataset._generate_examples(**split_generator.gen_kwargs)
+    _, examples = zip(*generator)
+    assert len(examples) == 3
+    assert isinstance(examples[0]["json"], dict)
+    assert isinstance(examples[0]["json"]["transcript"], str)
+    assert isinstance(examples[0]["wav"], dict)
+    assert isinstance(examples[0]["wav"]["bytes"], bytes)  # keep encoded to avoid unecessary copies
+    encoded = webdataset.info.features.encode_example(examples[0])
+    decoded = webdataset.info.features.decode_example(encoded)
+    assert isinstance(decoded["json"], dict)
+    assert isinstance(decoded["json"]["transcript"], str)
+    assert isinstance(decoded["wav"], AudioDecoder)
+
+
+def test_video_webdataset(video_wds_file):
+    data_files = {"train": [video_wds_file]}
+    webdataset = WebDataset(data_files=data_files)
+    split_generators = webdataset._split_generators(DownloadManager())
+    assert webdataset.info.features == Features(
+        {
+            "__key__": Value("string"),
+            "__url__": Value("string"),
+            "json": {"caption": Value("string")},
+            "mov": Video(),
+        }
+    )
+    assert len(split_generators) == 1
+    split_generator = split_generators[0]
+    assert split_generator.name == "train"
+    generator = webdataset._generate_examples(**split_generator.gen_kwargs)
+    _, examples = zip(*generator)
+    assert len(examples) == 3
+    assert isinstance(examples[0]["json"], dict)
+    assert isinstance(examples[0]["json"]["caption"], str)
+    assert isinstance(examples[0]["mov"], dict)
+
+
+def test_webdataset_errors_on_bad_file(bad_wds_file):
+    data_files = {"train": [bad_wds_file]}
+    webdataset = WebDataset(data_files=data_files)
+    with pytest.raises(ValueError):
+        webdataset._split_generators(DownloadManager())
+
+
+@require_pil
+def test_webdataset_with_features(image_wds_file):
+    import PIL.Image
+
+    data_files = {"train": [image_wds_file]}
+    features = Features(
+        {
+            "__key__": Value("string"),
+            "__url__": Value("string"),
+            "json": {"caption": Value("string"), "additional_field": Value("int64")},
+            "jpg": Image(),
+        }
+    )
+    webdataset = WebDataset(data_files=data_files, features=features)
+    split_generators = webdataset._split_generators(DownloadManager())
+    assert webdataset.info.features == features
+    split_generator = split_generators[0]
+    assert split_generator.name == "train"
+    generator = webdataset._generate_examples(**split_generator.gen_kwargs)
+    _, example = next(iter(generator))
+    encoded = webdataset.info.features.encode_example(example)
+    decoded = webdataset.info.features.decode_example(encoded)
+    assert decoded["json"]["additional_field"] is None
+    assert isinstance(decoded["json"], dict)
+    assert isinstance(decoded["json"]["caption"], str)
+    assert isinstance(decoded["jpg"], PIL.Image.Image)
+
+
+@require_numpy1_on_windows
+@require_torch
+def test_tensor_webdataset(tensor_wds_file):
+    import torch
+
+    data_files = {"train": [tensor_wds_file]}
+    webdataset = WebDataset(data_files=data_files)
+    split_generators = webdataset._split_generators(DownloadManager())
+    assert webdataset.info.features == Features(
+        {
+            "__key__": Value("string"),
+            "__url__": Value("string"),
+            "json": {"text": Value("string")},
+            "pth": List(Value("float32")),
+        }
+    )
+    assert len(split_generators) == 1
+    split_generator = split_generators[0]
+    assert split_generator.name == "train"
+    generator = webdataset._generate_examples(**split_generator.gen_kwargs)
+    _, examples = zip(*generator)
+    assert len(examples) == 3
+    assert isinstance(examples[0]["json"], dict)
+    assert isinstance(examples[0]["json"]["text"], str)
+    assert isinstance(examples[0]["pth"], torch.Tensor)  # keep encoded to avoid unecessary copies
+    encoded = webdataset.info.features.encode_example(examples[0])
+    decoded = webdataset.info.features.decode_example(encoded)
+    assert isinstance(decoded["json"], dict)
+    assert isinstance(decoded["json"]["text"], str)
+    assert isinstance(decoded["pth"], list)