Function test_builder_as_dataset

tests/test_builder.py:536–567

(split, expected_dataset_class, expected_dataset_length, in_memory, tmp_path)

)
@pytest.mark.parametrize("in_memory", [False, True])
def test_builder_as_dataset(split, expected_dataset_class, expected_dataset_length, in_memory, tmp_path):
    cache_dir = str(tmp_path)
    builder = DummyBuilder(cache_dir=cache_dir)
    os.makedirs(builder.cache_dir)

    builder.info.splits = SplitDict()
    builder.info.splits.add(SplitInfo("train", num_examples=10))
    builder.info.splits.add(SplitInfo("test", num_examples=10))

    for info_split in builder.info.splits:
        with ArrowWriter(
            path=os.path.join(builder.cache_dir, f"{builder.dataset_name}-{info_split}.arrow"),
            features=Features({"text": Value("string")}),
        ) as writer:
            writer.write_batch({"text": ["foo"] * 10})
            writer.finalize()

    with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
        dataset = builder.as_dataset(split=split, in_memory=in_memory)
    assert isinstance(dataset, expected_dataset_class)
    if isinstance(dataset, DatasetDict):
        assert list(dataset.keys()) == ["train", "test"]
        datasets = dataset.values()
        expected_splits = ["train", "test"]
    elif isinstance(dataset, Dataset):
        datasets = [dataset]
        expected_splits = [split]
    for dataset, expected_split in zip(datasets, expected_splits):
        assert dataset.split == expected_split
        assert len(dataset) == expected_dataset_length
        assert dataset.features == Features({"text": Value("string")})
        dataset.column_names == ["text"]


@pytest.mark.parametrize("in_memory", [False, True])