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

Code before:
        {"col_1": "string", "col_2": "string", "col_3": "string"},
        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
    ],
)
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
    dataset = Dataset.from_generator(data_generator, cache_dir=cache_dir, split=split)
    _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

Review comment (bug):
The CI is failing because in Python, "not passing any value" is different from "passing None value".

In this test, when the value of `split` is None, you should not pass it to `Dataset.from_generator`, so that you check its default value.

Code after:
    "features",
    [
        None,
        {"col_1": "string", "col_2": "int64", "col_3": "float64"},
        {"col_1": "string", "col_2": "string", "col_3": "string"},
        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
    ],
)
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)