test_spark_dataframe.py

from unittest import mock

import pytest

_ = pytest.importorskip("duckdb.experimental.spark")


from spark_namespace import USE_ACTUAL_SPARK
from spark_namespace.errors import PySparkTypeError, PySparkValueError
from spark_namespace.sql.column import Column
from spark_namespace.sql.functions import col, struct, when
from spark_namespace.sql.types import (
    ArrayType,
    BooleanType,
    IntegerType,
    LongType,
    MapType,
    Row,
    StringType,
    StructField,
    StructType,
)


def assert_column_objects_equal(col1: Column, col2: Column):
    assert type(col1) is type(col2)
    if not USE_ACTUAL_SPARK:
        assert col1.expr == col2.expr


class TestDataFrame:
    def test_dataframe_from_list_of_tuples(self, spark):
        # Valid
        address = [(1, "14851 Jeffrey Rd", "DE"), (2, "43421 Margarita St", "NY"), (3, "13111 Siemon Ave", "CA")]
        df = spark.createDataFrame(address, ["id", "address", "state"])
        res = df.collect()
        assert res == [
            Row(id=1, address="14851 Jeffrey Rd", state="DE"),
            Row(id=2, address="43421 Margarita St", state="NY"),
            Row(id=3, address="13111 Siemon Ave", state="CA"),
        ]

        # Tuples of different sizes
        address = [(1, "14851 Jeffrey Rd", "DE"), (2, "43421 Margarita St", "NY"), (3, "13111 Siemon Ave")]

        if USE_ACTUAL_SPARK:
            from py4j.protocol import Py4JJavaError

            with pytest.raises(Py4JJavaError):
                spark.createDataFrame(address, ["id", "address", "state"])
        else:
            with pytest.raises(PySparkTypeError, match="LENGTH_SHOULD_BE_THE_SAME"):
                spark.createDataFrame(address, ["id", "address", "state"])

        # Dataframe instead of list
        with pytest.raises(PySparkTypeError, match="SHOULD_NOT_DATAFRAME"):
            spark.createDataFrame(df, ["id", "address", "state"])

        # Not a list
        with pytest.raises(TypeError, match="not iterable"):
            spark.createDataFrame(5, ["id", "address", "test"])

        # Empty list
        if not USE_ACTUAL_SPARK:
            # TODO: Spark raises PySparkValueError [CANNOT_INFER_EMPTY_SCHEMA]  # noqa: TD002, TD003
            df = spark.createDataFrame([], ["id", "address", "test"])
            res = df.collect()
            assert res == []

        # Duplicate column names
        address = [(1, "14851 Jeffrey Rd", "DE"), (2, "43421 Margarita St", "NY"), (3, "13111 Siemon Ave", "DE")]
        df = spark.createDataFrame(address, ["id", "address", "id"])
        res = df.collect()
        exptected_res_str = (
            "[Row(id=1, address='14851 Jeffrey Rd', id='DE'), Row(id=2, address='43421 "
            "Margarita St', id='NY'), Row(id=3, address='13111 Siemon Ave', id='DE')]"
        )
        if USE_ACTUAL_SPARK:
            # Spark uses string for both ID columns. DuckDB correctly infers the types.
            exptected_res_str = (
                exptected_res_str.replace("id=1", "id='1'").replace("id=2", "id='2'").replace("id=3", "id='3'")
            )
        assert str(res) == exptected_res_str

        # Not enough column names
        if not USE_ACTUAL_SPARK:
            # TODO: Spark does not raise this error  # noqa: TD002, TD003
            with pytest.raises(PySparkValueError, match="number of columns in the DataFrame don't match"):
                spark.createDataFrame(address, ["id", "address"])

        # Empty column names list
        # Columns are filled in with default names
        # TODO: check against Spark behavior  # noqa: TD002, TD003
        df = spark.createDataFrame(address, [])
        res = df.collect()
        assert res == [
            Row(col0=1, col1="14851 Jeffrey Rd", col2="DE"),
            Row(col0=2, col1="43421 Margarita St", col2="NY"),
            Row(col0=3, col1="13111 Siemon Ave", col2="DE"),
        ]

        # Too many column names
        if not USE_ACTUAL_SPARK:
            # In Spark, this leads to an IndexError
            with pytest.raises(PySparkValueError, match="number of columns in the DataFrame don't match"):
                spark.createDataFrame(address, ["id", "address", "one", "two", "three"])

        # Column names is not a list (but is iterable)
        if not USE_ACTUAL_SPARK:
            # These things do not work in Spark or throw different errors
            df = spark.createDataFrame(address, {"a": 5, "b": 6, "c": 42})
            res = df.collect()
            assert res == [
                Row(a=1, b="14851 Jeffrey Rd", c="DE"),
                Row(a=2, b="43421 Margarita St", c="NY"),
                Row(a=3, b="13111 Siemon Ave", c="DE"),
            ]

            # Column names is not a list (string, becomes a single column name)
            with pytest.raises(PySparkValueError, match="number of columns in the DataFrame don't match"):
                spark.createDataFrame(address, "a")

            with pytest.raises(TypeError, match="must be an iterable, not int"):
                spark.createDataFrame(address, 5)

    def test_dataframe(self, spark):
        # Create DataFrame
        df = spark.createDataFrame([("Scala", 25000), ("Spark", 35000), ("PHP", 21000)])
        res = df.collect()
        assert res == [Row(col0="Scala", col1=25000), Row(col0="Spark", col1=35000), Row(col0="PHP", col1=21000)]

    @pytest.mark.skipif(USE_ACTUAL_SPARK, reason="We can't create tables with our Spark test setup")
    def test_writing_to_table(self, spark):
        # Create Hive table & query it.
        spark.sql(
            """
            create table sample_table("_1" bool, "_2" integer)
        """
        )
        spark.sql("insert into sample_table VALUES (True, 42)")
        spark.table("sample_table").write.saveAsTable("sample_hive_table")
        df3 = spark.sql("SELECT _1,_2 FROM sample_hive_table")
        res = df3.collect()
        assert res == [Row(_1=True, _2=42)]
        schema = df3.schema
        assert schema == StructType([StructField("_1", BooleanType(), True), StructField("_2", IntegerType(), True)])

    def test_dataframe_collect(self, spark):
        df = spark.createDataFrame([(42,), (21,)]).toDF("a")
        res = df.collect()
        assert str(res) == "[Row(a=42), Row(a=21)]"

    def test_dataframe_from_rows(self, spark):
        columns = ["language", "users_count"]
        data = [("Java", "20000"), ("Python", "100000"), ("Scala", "3000")]

        rowData = (Row(*x) for x in data)
        df = spark.createDataFrame(rowData, columns)
        res = df.collect()
        assert res == [
            Row(language="Java", users_count="20000"),
            Row(language="Python", users_count="100000"),
            Row(language="Scala", users_count="3000"),
        ]

    def test_empty_df(self, spark):
        schema = StructType(
            [
                StructField("firstname", StringType(), True),
                StructField("middlename", StringType(), True),
                StructField("lastname", StringType(), True),
            ]
        )
        df = spark.createDataFrame([], schema=schema)
        res = df.collect()
        # TODO: assert that the types and column names are correct  # noqa: TD002, TD003
        assert res == []

    def test_df_from_pandas(self, spark):
        import pandas as pd

        df = spark.createDataFrame(pd.DataFrame({"a": [42, 21], "b": [True, False]}))
        res = df.collect()
        assert res == [Row(a=42, b=True), Row(a=21, b=False)]

    def test_df_from_struct_type(self, spark):
        schema = StructType([StructField("a", LongType()), StructField("b", BooleanType())])
        df = spark.createDataFrame([(42, True), (21, False)], schema)
        res = df.collect()
        assert res == [Row(a=42, b=True), Row(a=21, b=False)]

    def test_df_from_name_list(self, spark):
        df = spark.createDataFrame([(42, True), (21, False)], ["a", "b"])
        res = df.collect()
        assert res == [Row(a=42, b=True), Row(a=21, b=False)]

    def test_df_creation_coverage(self, spark):
        from spark_namespace.sql.types import IntegerType, StringType, StructField, StructType

        data2 = [
            ("James", "", "Smith", "36636", "M", 3000),
            ("Michael", "Rose", "", "40288", "M", 4000),
            ("Robert", "", "Williams", "42114", "M", 4000),
            ("Maria", "Anne", "Jones", "39192", "F", 4000),
            ("Jen", "Mary", "Brown", "", "F", -1),
        ]

        schema = StructType(
            [
                StructField("firstname", StringType(), True),
                StructField("middlename", StringType(), True),
                StructField("lastname", StringType(), True),
                StructField("id", StringType(), True),
                StructField("gender", StringType(), True),
                StructField("salary", IntegerType(), True),
            ]
        )

        df = spark.createDataFrame(data=data2, schema=schema)
        res = df.collect()
        assert res == [
            Row(firstname="James", middlename="", lastname="Smith", id="36636", gender="M", salary=3000),
            Row(firstname="Michael", middlename="Rose", lastname="", id="40288", gender="M", salary=4000),
            Row(firstname="Robert", middlename="", lastname="Williams", id="42114", gender="M", salary=4000),
            Row(firstname="Maria", middlename="Anne", lastname="Jones", id="39192", gender="F", salary=4000),
            Row(firstname="Jen", middlename="Mary", lastname="Brown", id="", gender="F", salary=-1),
        ]

    def test_df_nested_struct(self, spark):
        structureData = [
            (("James", "", "Smith"), "36636", "M", 3100),
            (("Michael", "Rose", ""), "40288", "M", 4300),
            (("Robert", "", "Williams"), "42114", "M", 1400),
            (("Maria", "Anne", "Jones"), "39192", "F", 5500),
            (("Jen", "Mary", "Brown"), "", "F", -1),
        ]
        structureSchema = StructType(
            [
                StructField(
                    "name",
                    StructType(
                        [
                            StructField("firstname", StringType(), True),
                            StructField("middlename", StringType(), True),
                            StructField("lastname", StringType(), True),
                        ]
                    ),
                ),
                StructField("id", StringType(), True),
                StructField("gender", StringType(), True),
                StructField("salary", IntegerType(), True),
            ]
        )

        df2 = spark.createDataFrame(data=structureData, schema=structureSchema)
        res = df2.collect()
        expected_res = [
            Row(
                name={"firstname": "James", "middlename": "", "lastname": "Smith"}, id="36636", gender="M", salary=3100
            ),
            Row(
                name={"firstname": "Michael", "middlename": "Rose", "lastname": ""}, id="40288", gender="M", salary=4300
            ),
            Row(
                name={"firstname": "Robert", "middlename": "", "lastname": "Williams"},
                id="42114",
                gender="M",
                salary=1400,
            ),
            Row(
                name={"firstname": "Maria", "middlename": "Anne", "lastname": "Jones"},
                id="39192",
                gender="F",
                salary=5500,
            ),
            Row(name={"firstname": "Jen", "middlename": "Mary", "lastname": "Brown"}, id="", gender="F", salary=-1),
        ]
        if USE_ACTUAL_SPARK:
            expected_res = [Row(name=Row(**r.name), id=r.id, gender=r.gender, salary=r.salary) for r in expected_res]
        assert res == expected_res
        schema = df2.schema
        assert schema == StructType(
            [
                StructField(
                    "name",
                    StructType(
                        [
                            StructField("firstname", StringType(), True),
                            StructField("middlename", StringType(), True),
                            StructField("lastname", StringType(), True),
                        ]
                    ),
                    True,
                ),
                StructField("id", StringType(), True),
                StructField("gender", StringType(), True),
                StructField("salary", IntegerType(), True),
            ]
        )

    def test_df_columns(self, spark):
        from spark_namespace.sql.functions import col

        structureData = [
            (("James", "", "Smith"), "36636", "M", 3100),
            (("Michael", "Rose", ""), "40288", "M", 4300),
            (("Robert", "", "Williams"), "42114", "M", 1400),
            (("Maria", "Anne", "Jones"), "39192", "F", 5500),
            (("Jen", "Mary", "Brown"), "", "F", -1),
        ]
        structureSchema = StructType(
            [
                StructField(
                    "name",
                    StructType(
                        [
                            StructField("firstname", StringType(), True),
                            StructField("middlename", StringType(), True),
                            StructField("lastname", StringType(), True),
                        ]
                    ),
                ),
                StructField("id", StringType(), True),
                StructField("gender", StringType(), True),
                StructField("salary", IntegerType(), True),
            ]
        )

        df2 = spark.createDataFrame(data=structureData, schema=structureSchema)
        updatedDF = df2.withColumn(
            "OtherInfo",
            struct(
                col("id").alias("identifier"),
                col("gender").alias("gender"),
                col("salary").alias("salary"),
                when(col("salary").cast(IntegerType()) < 2000, "Low")
                .when(col("salary").cast(IntegerType()) < 4000, "Medium")
                .otherwise("High")
                .alias("Salary_Grade"),
            ),
        ).drop("id", "gender", "salary")

        assert "OtherInfo" in updatedDF.columns

    def test_array_and_map_type(self, spark):
        """Array & Map."""
        StructType(
            [
                StructField(
                    "name",
                    StructType(
                        [
                            StructField("firstname", StringType(), True),
                            StructField("middlename", StringType(), True),
                            StructField("lastname", StringType(), True),
                        ]
                    ),
                ),
                StructField("hobbies", ArrayType(StringType()), True),
                StructField("properties", MapType(StringType(), StringType()), True),
            ]
        )

    def test_getitem_filter(self, spark):
        data = [(56, "Carol"), (20, "Alice"), (3, "Dave"), (3, "Anna"), (1, "Ben")]
        df = spark.createDataFrame(data, ["age", "name"])
        expected = [
            Row(age=56, name="Carol"),
            Row(age=20, name="Alice"),
        ]

        df = df[col("age") > 18]
        assert df.collect() == expected

    def test_getitem_column(self, spark):
        data = [(56, "Carol"), (20, "Alice")]
        df = spark.createDataFrame(data, ["age", "name"])

        assert_column_objects_equal(df["age"], col("age"))

    def test_getitem_dataframe(self, spark):
        data = [(56, "Ben", "Street1"), (20, "Tom", "Street2")]
        df = spark.createDataFrame(data, ["age", "name", "address"])

        res = df[["name", "address"]].collect()
        expected = df.select("name", "address").collect()

        assert res == expected

    def test_getattr_dataframe(self, spark):
        data = [(56, "Ben", "Street1"), (20, "Tom", "Street2")]
        df = spark.createDataFrame(data, ["age", "name", "address"])

        assert_column_objects_equal(df.age, col("age"))

    def test_head_first(self, spark):
        data = [(56, "Carol"), (20, "Alice"), (3, "Dave"), (3, "Anna"), (1, "Ben")]
        df = spark.createDataFrame(data, ["age", "name"])
        expected = Row(age=1, name="Ben")

        head = df.orderBy(df.age).head()
        first = df.orderBy(df.age).first()
        assert head == first == expected

    def test_head_take_n(self, spark):
        data = [(56, "Carol"), (20, "Alice"), (3, "Dave"), (3, "Anna"), (1, "Ben")]
        df = spark.createDataFrame(data, ["age", "name"])
        expected = [
            Row(age=1, name="Ben"),
            Row(age=3, name="Anna"),
        ]
        df = df.orderBy(df.age, df.name)
        rows = df.head(2)
        take = df.take(2)
        assert rows == take == expected

    def test_drop(self, spark):
        data = [(1, 2, 3, 4)]
        df = spark.createDataFrame(data, ["one", "two", "three", "four"])
        expected = ["one", "four"]
        assert df.drop("two", "three").columns == expected
        assert df.drop("two", col("three")).columns == expected
        assert df.drop("two", col("three"), col("missing")).columns == expected

    def test_cache(self, spark):
        data = [(1, 2, 3, 4)]
        df = spark.createDataFrame(data, ["one", "two", "three", "four"])
        cached = df.cache()
        assert df is not cached
        assert cached.collect() == df.collect()
        assert cached.collect() == [Row(one=1, two=2, three=3, four=4)]

    def test_dtypes(self, spark):
        data = [("Alice", 25, 5000.0), ("Bob", 30, 6000.0)]
        df = spark.createDataFrame(data, ["name", "age", "salary"])
        dtypes = df.dtypes

        assert isinstance(dtypes, list)
        assert len(dtypes) == 3
        for col_name, col_type in dtypes:
            assert isinstance(col_name, str)
            assert isinstance(col_type, str)

        col_names = [name for name, _ in dtypes]
        assert col_names == ["name", "age", "salary"]
        for _, col_type in dtypes:
            assert len(col_type) > 0

    def test_dtypes_complex_types(self, spark):
        from spark_namespace.sql.types import ArrayType, IntegerType, StringType, StructField, StructType

        schema = StructType(
            [
                StructField("name", StringType(), True),
                StructField("scores", ArrayType(IntegerType()), True),
                StructField(
                    "address",
                    StructType([StructField("city", StringType(), True), StructField("zip", StringType(), True)]),
                    True,
                ),
            ]
        )
        data = [
            ("Alice", [90, 85, 88], {"city": "NYC", "zip": "10001"}),
            ("Bob", [75, 80, 82], {"city": "LA", "zip": "90001"}),
        ]
        df = spark.createDataFrame(data, schema)
        dtypes = df.dtypes

        assert len(dtypes) == 3
        col_names = [name for name, _ in dtypes]
        assert col_names == ["name", "scores", "address"]

    def test_printSchema(self, spark, capsys):
        data = [("Alice", 25, 5000), ("Bob", 30, 6000)]
        df = spark.createDataFrame(data, ["name", "age", "salary"])
        df.printSchema()
        captured = capsys.readouterr()
        output = captured.out

        assert "root" in output
        assert "name" in output
        assert "age" in output
        assert "salary" in output
        assert "string" in output or "varchar" in output.lower()
        assert "int" in output.lower() or "bigint" in output.lower()

    def test_printSchema_nested(self, spark, capsys):
        from spark_namespace.sql.types import ArrayType, IntegerType, StringType, StructField, StructType

        schema = StructType(
            [
                StructField("id", IntegerType(), True),
                StructField(
                    "person",
                    StructType([StructField("name", StringType(), True), StructField("age", IntegerType(), True)]),
                    True,
                ),
                StructField("hobbies", ArrayType(StringType()), True),
            ]
        )
        data = [
            (1, {"name": "Alice", "age": 25}, ["reading", "coding"]),
            (2, {"name": "Bob", "age": 30}, ["gaming", "music"]),
        ]
        df = spark.createDataFrame(data, schema)
        df.printSchema()
        captured = capsys.readouterr()
        output = captured.out

        assert "root" in output
        assert "person" in output
        assert "hobbies" in output

    def test_printSchema_negative_level(self, spark):
        data = [("Alice", 25)]
        df = spark.createDataFrame(data, ["name", "age"])

        with pytest.raises(PySparkValueError):
            df.printSchema(level=-1)

    def test_treeString_basic(self, spark):
        data = [("Alice", 25, 5000)]
        df = spark.createDataFrame(data, ["name", "age", "salary"])
        tree = df.schema.treeString()

        assert tree.startswith("root\n")
        assert " |-- name:" in tree
        assert " |-- age:" in tree
        assert " |-- salary:" in tree
        assert "(nullable = true)" in tree
        assert tree.count(" |-- ") == 3

    def test_treeString_nested_struct(self, spark):
        from spark_namespace.sql.types import IntegerType, StringType, StructField, StructType

        schema = StructType(
            [
                StructField("id", IntegerType(), True),
                StructField(
                    "person",
                    StructType([StructField("name", StringType(), True), StructField("age", IntegerType(), True)]),
                    True,
                ),
            ]
        )
        data = [(1, {"name": "Alice", "age": 25})]
        df = spark.createDataFrame(data, schema)
        tree = df.schema.treeString()

        assert "root\n" in tree
        assert " |-- id:" in tree
        assert " |-- person: struct (nullable = true)" in tree
        assert "name:" in tree
        assert "age:" in tree

    def test_treeString_with_level(self, spark):
        from spark_namespace.sql.types import IntegerType, StringType, StructField, StructType

        schema = StructType(
            [
                StructField("id", IntegerType(), True),
                StructField(
                    "person",
                    StructType(
                        [
                            StructField("name", StringType(), True),
                            StructField("details", StructType([StructField("address", StringType(), True)]), True),
                        ]
                    ),
                    True,
                ),
            ]
        )

        data = [(1, {"name": "Alice", "details": {"address": "123 Main St"}})]
        df = spark.createDataFrame(data, schema)

        # Level 1 should only show top-level fields
        tree_level_1 = df.schema.treeString(level=1)
        assert " |-- id:" in tree_level_1
        assert " |-- person: struct" in tree_level_1
        # Should not show nested field names at level 1
        lines = tree_level_1.split("\n")
        assert len([line for line in lines if line.strip()]) <= 3

    def test_treeString_array_type(self, spark):
        from spark_namespace.sql.types import ArrayType, StringType, StructField, StructType

        schema = StructType(
            [StructField("name", StringType(), True), StructField("hobbies", ArrayType(StringType()), True)]
        )

        data = [("Alice", ["reading", "coding"])]
        df = spark.createDataFrame(data, schema)
        tree = df.schema.treeString()

        assert "root\n" in tree
        assert " |-- name:" in tree
        assert " |-- hobbies: array<" in tree
        assert "(nullable = true)" in tree

    def test_method_is_empty(self, spark):
        data = [(1, "Alice"), (2, "Bob")]
        df = spark.createDataFrame(data, ["id", "name"])
        empty_df = spark.createDataFrame([], schema=df.schema)

        assert not df.isEmpty()
        assert empty_df.isEmpty()

    def test_dataframe_foreach(self, spark):
        data = [(56, "Carol"), (20, "Alice"), (3, "Dave")]
        df = spark.createDataFrame(data, ["age", "name"])
        expected = [Row(age=56, name="Carol"), Row(age=20, name="Alice"), Row(age=3, name="Dave")]

        mock_callable = mock.MagicMock()
        df.foreach(mock_callable)
        mock_callable.assert_has_calls(
            [mock.call(expected[0]), mock.call(expected[1]), mock.call(expected[2])],
            any_order=True,
        )

    def test_dataframe_foreach_partition(self, spark):
        data = [(56, "Carol"), (20, "Alice"), (3, "Dave")]
        df = spark.createDataFrame(data, ["age", "name"])
        expected = [Row(age=56, name="Carol"), Row(age=20, name="Alice"), Row(age=3, name="Dave")]

        mock_callable = mock.MagicMock()
        df.foreachPartition(mock_callable)
        mock_callable.assert_called_once_with(expected)

    def test_to_local_iterator(self, spark):
        data = [(56, "Carol"), (20, "Alice"), (3, "Dave")]
        df = spark.createDataFrame(data, ["age", "name"])
        expected = [Row(age=56, name="Carol"), Row(age=20, name="Alice"), Row(age=3, name="Dave")]

        res = list(df.toLocalIterator())
        assert res == expected