[WIP] Spark Connect support

requirements-dev.txt

@@ -1,5 +1,7 @@
 # pyspark
 pyspark==3.5.3
+# pyspark-connect
+pyspark[connect]==3.5.3
 # linters
 flake8==7.1.1
 pylint==3.3.2

tests/_connect/test_connect.py

@@ -0,0 +1,22 @@
+from pyspark.sql import SparkSession
+from pyspark.sql.types import LongType
+
+from typedspark import Column, Schema
+
+
+class A(Schema):
+    a: Column[LongType]
+    b: Column[LongType]
+    c: Column[LongType]
+
+
+def test_regular_spark_works(spark: SparkSession):
+    df = spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"])
+    assert df.count() == 3
+    assert type(df).__module__ == "pyspark.sql.dataframe"
+
+
+def test_spark_connect_works(sparkConnect: SparkSession):
+    df = sparkConnect.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"])
+    assert df.count() == 3
+    assert type(df).__module__ == "pyspark.sql.connect.dataframe"

tests/_connect/test_dataset.py

@@ -0,0 +1,140 @@
+import functools
+
+import pandas as pd
+import pytest
+from pyspark import StorageLevel
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql.types import LongType, StringType
+
+from typedspark import Column, DataSet, Schema
+from typedspark._core.dataset import DataSetImplements
+from typedspark._utils.create_dataset import create_empty_dataset
+
+
+class A(Schema):
+    a: Column[LongType]
+    b: Column[StringType]
+
+
+class B(Schema):
+    a: Column[LongType]
+    b: Column[StringType]
+
+
+def create_dataframe(sparkConnect: SparkSession, d):
+    return sparkConnect.createDataFrame(pd.DataFrame(d))
+
+
+def test_dataset(sparkConnect: SparkSession):
+    d = dict(
+        a=[1, 2, 3],
+        b=["a", "b", "c"],
+    )
+    df = create_dataframe(sparkConnect, d)
+    DataSet[A](df)
+
+
+def test_dataset_allow_underscored_columns_not_in_schema(sparkConnect: SparkSession):
+    d = {"a": [1, 2, 3], "b": ["a", "b", "c"], "__c": [1, 2, 3]}
+    df = create_dataframe(sparkConnect, d)
+    DataSet[A](df)
+
+
+def test_dataset_single_underscored_column_should_raise(sparkConnect: SparkSession):
+    d = {"a": [1, 2, 3], "b": ["a", "b", "c"], "_c": [1, 2, 3]}
+    df = create_dataframe(sparkConnect, d)
+    with pytest.raises(TypeError):
+        DataSet[A](df)
+
+
+def test_dataset_missing_colnames(sparkConnect: SparkSession):
+    d = dict(
+        a=[1, 2, 3],
+    )
+    df = create_dataframe(sparkConnect, d)
+    with pytest.raises(TypeError):
+        DataSet[A](df)
+
+
+def test_dataset_too_many_colnames(sparkConnect: SparkSession):
+    d = dict(
+        a=[1, 2, 3],
+        b=["a", "b", "c"],
+        c=[1, 2, 3],
+    )
+    df = create_dataframe(sparkConnect, d)
+    with pytest.raises(TypeError):
+        DataSet[A](df)
+
+
+def test_wrong_type(sparkConnect: SparkSession):
+    d = dict(
+        a=[1, 2, 3],
+        b=[1, 2, 3],
+    )
+    df = create_dataframe(sparkConnect, d)
+    with pytest.raises(TypeError):
+        DataSet[A](df)
+
+
+def test_inherrited_functions(sparkConnect: SparkSession):
+    df = create_empty_dataset(sparkConnect, A)
+
+    df.distinct()
+    cached1: DataSet[A] = df.cache()
+    cached2: DataSet[A] = df.persist(StorageLevel.MEMORY_AND_DISK)
+    df.filter(A.a == 1)
+    df.orderBy(A.a)
+    df.transform(lambda df: df)
+
+    cached1.unpersist(True)
+    cached2.unpersist(True)
+
+
+def test_inherrited_functions_with_other_dataset(sparkConnect: SparkSession):
+    df_a = create_empty_dataset(sparkConnect, A)
+    df_b = create_empty_dataset(sparkConnect, A)
+
+    df_a.join(df_b, A.a.str)
+    df_a.unionByName(df_b)
+
+
+def test_schema_property_of_dataset(sparkConnect: SparkSession):
+    df = create_empty_dataset(sparkConnect, A)
+    assert df.typedspark_schema == A
+
+
+def test_initialize_dataset_implements(sparkConnect: SparkSession):
+    with pytest.raises(NotImplementedError):
+        DataSetImplements()
+
+
+def test_reduce(sparkConnect: SparkSession):
+    functools.reduce(
+        DataSet.unionByName,
+        [create_empty_dataset(sparkConnect, A), create_empty_dataset(sparkConnect, A)],
+    )
+
+
+def test_resetting_of_schema_annotations(sparkConnect: SparkSession):
+    df = create_empty_dataset(sparkConnect, A)
+
+    a: DataFrame
+
+    # if no schema is specified, the annotation should be None
+    a = DataSet(df)
+    assert a._schema_annotations is None
+
+    # when we specify a schema, the class variable will be set to A, but afterwards it should be
+    # reset to None again when we initialize a new object without specifying a schema
+    DataSet[A]
+    a = DataSet(df)
+    assert a._schema_annotations is None
+
+    # and then to B
+    a = DataSet[B](df)
+    assert a._schema_annotations == B
+
+    # and then to None again
+    a = DataSet(df)
+    assert a._schema_annotations is None

tests/conftest.py

@@ -14,3 +14,17 @@ def spark():
     spark = SparkSession.Builder().getOrCreate()
     yield spark
     spark.stop()
+
+
+@pytest.fixture(scope="session")
+def sparkConnect():
+    """Fixture for creating a spark session."""
+
+    spark = (
+        SparkSession.Builder()
+        .config("spark.jars.packages", "org.apache.spark:spark-connect_2.12:3.5.3")
+        .remote('local')
+        .getOrCreate()
+    )
+    yield spark
+    spark.stop()