Repository: huggingface/datasets
File: tests/test_arrow_dataset.py (Python)

Code before:
    cache_dir = tmp_path / "cache"
    default_expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    split = NamedSplit("validation")
    expected_features = features.copy() if features else default_expected_features
    features = (
        Features({feature: Value(dtype) for feature, dtype in features.items()}) if features is not None else None
    )
    dataset = Dataset.from_generator(data_generator, features=features, cache_dir=cache_dir, split=split)
    _check_generator_dataset(dataset, expected_features, split)


@require_not_windows
@require_dill_gt_0_3_2
@require_pyspark
def test_from_spark():
    import pyspark

    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
    data = [
        ("0", 0, 0.0),
        ("1", 1, 1.0),
        ("2", 2, 2.0),
        ("3", 3, 3.0),
    ]
    df = spark.createDataFrame(data, "col_1: string, col_2: int, col_3: float")
    dataset = Dataset.from_spark(df)
    assert isinstance(dataset, Dataset)
    assert dataset.num_rows == 4
    assert dataset.num_columns == 3
    assert dataset.column_names == ["col_1", "col_2", "col_3"]


@require_not_windows
@require_dill_gt_0_3_2
@require_pyspark
def test_from_spark_features():
    import PIL.Image
    import pyspark

    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
    data = [(0, np.arange(4 * 4 * 3).reshape(4, 4, 3).tolist())]
    df = spark.createDataFrame(data, "idx: int, image: array<array<array<int>>>")
    features = Features({"idx": Value("int64"), "image": Image()})
    dataset = Dataset.from_spark(
        df,
        features=features,
    )
    assert isinstance(dataset, Dataset)
    assert dataset.num_rows == 1
    assert dataset.num_columns == 2
    assert dataset.column_names == ["idx", "image"]

Review comment (suggestion):
I would add a specific `test_dataset_from_generator_split` with a parametrized `split` values, such as not passing any value, passing `NamedSplit("train")`, passing literal `"train"`, passing other NamedSplit, etc.

Code after:
def test_dataset_from_generator_features(features, data_generator, tmp_path):
    cache_dir = tmp_path / "cache"
    default_expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    expected_features = features.copy() if features else default_expected_features
    features = (
        Features({feature: Value(dtype) for feature, dtype in features.items()}) if features is not None else None
    )
    dataset = Dataset.from_generator(data_generator, features=features, cache_dir=cache_dir)
    _check_generator_dataset(dataset, expected_features, NamedSplit("train"))


@pytest.mark.parametrize(
    "split",
    [None, NamedSplit("train"), "train", NamedSplit("foo"), "foo"],
)
def test_dataset_from_generator_split(split, data_generator, tmp_path):
    cache_dir = tmp_path / "cache"
    default_expected_split = "train"
    expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    expected_split = split if split else default_expected_split
    if split:
        dataset = Dataset.from_generator(data_generator, cache_dir=cache_dir, split=split)
    else:
        dataset = Dataset.from_generator(data_generator, cache_dir=cache_dir)
    _check_generator_dataset(dataset, expected_features, expected_split)


@require_not_windows
@require_dill_gt_0_3_2
@require_pyspark
def test_from_spark():
    import pyspark

    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
    data = [
        ("0", 0, 0.0),
        ("1", 1, 1.0),
        ("2", 2, 2.0),
        ("3", 3, 3.0),
    ]
    df = spark.createDataFrame(data, "col_1: string, col_2: int, col_3: float")
    dataset = Dataset.from_spark(df)
    assert isinstance(dataset, Dataset)
    assert dataset.num_rows == 4
    assert dataset.num_columns == 3
    assert dataset.column_names == ["col_1", "col_2", "col_3"]


@require_not_windows
@require_dill_gt_0_3_2
@require_pyspark