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

Code before:
    "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)


def test_from_spark():
    spark = pyspark.sql.SparkSession.builder.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"]


def test_from_spark_features():
    spark = pyspark.sql.SparkSession.builder.appName("pyspark").getOrCreate()
    data = [(0, np.arange(4 * 4 * 3).reshape(4, 4, 3).tolist())]
    df = spark.createDataFrame(data, "idx: int, img_bytes: array<array<array<int>>>")
    features = Features({"idx": Value("int64"), "img_bytes": 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", "img_bytes"]
    assert dataset.features == features
    assert_arrow_metadata_are_synced_with_dataset_features(dataset)

Review comment (performance):
Could you add a test to make sure that calling `from_spark` twice in a row with two different DataFrames results in two different Dataset objects ?

The `DatasetBuilder` uses a `cache_dir` with subdirectories defined after the name of the builder ("spark") and the arguments of the dataset to generate. Here since the `**kwargs` passed to `DatasetBuilder.__init__` won't change between two `from_spark`, it may reuse the same cache and return the same dataset.

If this is the case we can add a `df_id` parameter to the `SparkConfig`, which can be set to to be the python `id` of the Spark DataFrame for example (or any identifier of the Spark DataFrame).

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)


@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"]