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

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


def _check_sql_dataset(dataset, expected_features):
    assert isinstance(dataset, Dataset)
    assert dataset.num_rows == 4
    assert dataset.num_columns == 3
    assert dataset.column_names == ["col_1", "col_2", "col_3"]
    for feature, expected_dtype in expected_features.items():
        assert dataset.features[feature].dtype == expected_dtype


@require_sqlalchemy
@pytest.mark.parametrize("con_type", ["string", "engine"])
def test_dataset_from_sql_con_type(con_type, sqlite_path, tmp_path):
    cache_dir = tmp_path / "cache"

Review comment (suggestion):
```suggestion
    df = spark.createDataFrame(data, "idx: int, image: array<array<array<int>>>")
    features = Features({"idx": Value("int64"), "image": Image()})
```

Code after:
    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"]
    assert isinstance(dataset[0]["image"], PIL.Image.Image)
    assert dataset.features == features
    assert_arrow_metadata_are_synced_with_dataset_features(dataset)


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

    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
    df = spark.createDataFrame([("0", 0)], "col_1: string, col_2: int")
    dataset = Dataset.from_spark(df)
    assert isinstance(dataset, Dataset)