feat: Add support for spark connect

.github/workflows/pytest-pyspark.yml

@@ -35,3 +35,48 @@ jobs:
         run: uv pip freeze
       - name: Run pytest
         run: pytest tests --cov=narwhals/_spark_like --cov-fail-under=95 --runslow --constructors pyspark
+
+
+  pytest-pyspark-connect-constructor:
+    if: ${{ contains(github.event.pull_request.labels.*.name, 'pyspark-connect') || contains(github.event.pull_request.labels.*.name, 'release') }}
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11"]
+        os: [ubuntu-latest]
+    env:
+      SPARK_VERSION: 3.5.1
+      SPARK_PORT: 15002
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: "true"
+          cache-suffix: ${{ matrix.python-version }}
+          cache-dependency-glob: "pyproject.toml"
+      - name: install-reqs
+        run: uv pip install -e . --group core-tests --group extra --system
+      - name: install pyspark
+        run: echo "setuptools<78" | uv pip install -b - -e . "pyspark[connect]==3.5.1" --system
+      - name: show-deps
+        run: uv pip freeze
+
+      - name: Download Spark
+        run: |
+          curl -sL https://downloads.apache.org/spark/spark-${SPARK_VERSION}/spark-${SPARK_VERSION}-bin-hadoop3.tgz | tar xz
+          echo "SPARK_HOME=$PWD/spark-${SPARK_VERSION}-bin-hadoop3" >> $GITHUB_ENV
+          echo "$PWD/spark-${SPARK_VERSION}-bin-hadoop3/bin" >> $GITHUB_PATH
+
+      - name: Start Spark Connect Server
+        run: |
+          nohup spark-connect-server \
+            --conf spark.connect.grpc.binding=localhost:$SPARK_PORT > connect.log 2>&1 &
+          # wait for port to open
+          for i in {1..10}; do nc -z localhost $SPARK_PORT && break || sleep 1; done
+
+      - name: Run pytest
+        run: pytest tests --cov=narwhals/_spark_like --cov-fail-under=95 --runslow --constructors "pyspark[connect]""

narwhals/_namespace.py

@@ -20,6 +20,7 @@
 from narwhals.dependencies import get_pyarrow
 from narwhals.dependencies import is_dask_dataframe
 from narwhals.dependencies import is_duckdb_relation
+from narwhals.dependencies import is_pyspark_connect_dataframe
 from narwhals.dependencies import is_pyspark_dataframe
 from narwhals.dependencies import is_sqlframe_dataframe
 from narwhals.utils import Implementation
@@ -34,6 +35,7 @@
     import polars as pl
     import pyarrow as pa
     import pyspark.sql as pyspark_sql
+    from pyspark.sql.connect.dataframe import DataFrame as PySparkConnectDataFrame
     from typing_extensions import TypeAlias
     from typing_extensions import TypeIs
 
@@ -72,7 +74,10 @@
         _PandasLike, Implementation.PANDAS, Implementation.CUDF, Implementation.MODIN
     ]
     SparkLike: TypeAlias = Literal[
-        _SparkLike, Implementation.PYSPARK, Implementation.SQLFRAME
+        _SparkLike,
+        Implementation.PYSPARK,
+        Implementation.SQLFRAME,
+        Implementation.PYSPARK_CONNECT,
     ]
     EagerOnly: TypeAlias = "PandasLike | Arrow"
     EagerAllowed: TypeAlias = "EagerOnly | Polars"
@@ -111,7 +116,10 @@ class _ModinSeries(Protocol):
     _NativePandasLike: TypeAlias = "_NativePandas | _NativeCuDF | _NativeModin"
     _NativeSQLFrame: TypeAlias = "SQLFrameDataFrame"
     _NativePySpark: TypeAlias = "pyspark_sql.DataFrame"
-    _NativeSparkLike: TypeAlias = "_NativeSQLFrame | _NativePySpark"
+    _NativePySparkConnect: TypeAlias = "PySparkConnectDataFrame"
+    _NativeSparkLike: TypeAlias = (
+        "_NativeSQLFrame | _NativePySpark | _NativePySparkConnect"
+    )
 
     NativeKnown: TypeAlias = "_NativePolars | _NativeArrow | _NativePandasLike | _NativeSparkLike | _NativeDuckDB | _NativeDask"
     NativeUnknown: TypeAlias = (
@@ -292,6 +300,8 @@ def from_native_object(
             return cls.from_backend(
                 Implementation.SQLFRAME
                 if is_native_sqlframe(native)
+                else Implementation.PYSPARK_CONNECT
+                if is_native_pyspark_connect(native)
                 else Implementation.PYSPARK
             )
         elif is_native_dask(native):
@@ -326,6 +336,7 @@ def is_native_dask(obj: Any) -> TypeIs[_NativeDask]:
 is_native_duckdb: _Guard[_NativeDuckDB] = is_duckdb_relation
 is_native_sqlframe: _Guard[_NativeSQLFrame] = is_sqlframe_dataframe
 is_native_pyspark: _Guard[_NativePySpark] = is_pyspark_dataframe
+is_native_pyspark_connect: _Guard[_NativePySparkConnect] = is_pyspark_connect_dataframe
 
 
 def is_native_pandas(obj: Any) -> TypeIs[_NativePandas]:
@@ -351,4 +362,8 @@ def is_native_pandas_like(obj: Any) -> TypeIs[_NativePandasLike]:
 
 
 def is_native_spark_like(obj: Any) -> TypeIs[_NativeSparkLike]:
-    return is_native_pyspark(obj) or is_native_sqlframe(obj)
+    return (
+        is_native_sqlframe(obj)
+        or is_native_pyspark(obj)
+        or is_native_pyspark_connect(obj)
+    )

narwhals/_spark_like/dataframe.py

@@ -179,6 +179,13 @@ def _collect_to_arrow(self) -> pa.Table:
                     return pa.Table.from_pydict(data, schema=pa.schema(schema))
                 else:  # pragma: no cover
                     raise
+        elif (
+            self._implementation is Implementation.PYSPARK_CONNECT
+            and self._backend_version < (4,)
+        ):
+            import pyarrow as pa  # ignore-banned-import
+
+            return pa.Table.from_pandas(self.native.toPandas())
         else:
             return self.native.toArrow()
 
@@ -294,7 +301,7 @@ def drop(self, columns: Sequence[str], *, strict: bool) -> Self:
         return self._with_native(self.native.drop(*columns_to_drop))
 
     def head(self, n: int) -> Self:
-        return self._with_native(self.native.limit(num=n))
+        return self._with_native(self.native.limit(n))
 
     def group_by(self, *keys: str, drop_null_keys: bool) -> SparkLikeLazyGroupBy:
         from narwhals._spark_like.group_by import SparkLikeLazyGroupBy
@@ -444,7 +451,10 @@ def explode(self, columns: Sequence[str]) -> Self:
             )
             raise NotImplementedError(msg)
 
-        if self._implementation.is_pyspark():
+        if self._implementation in {
+            Implementation.PYSPARK,
+            Implementation.PYSPARK_CONNECT,
+        }:
             return self._with_native(
                 self.native.select(
                     *[

narwhals/_spark_like/expr.py

@@ -438,7 +438,8 @@ def mean(self) -> Self:
     def median(self) -> Self:
         def _median(_input: Column) -> Column:
             if (
-                self._implementation.is_pyspark()
+                self._implementation
+                in {Implementation.PYSPARK, Implementation.PYSPARK_CONNECT}
                 and (pyspark := get_pyspark()) is not None
                 and parse_version(pyspark) < (3, 4)
             ):  # pragma: no cover

narwhals/_spark_like/utils.py

@@ -245,6 +245,10 @@ def import_functions(implementation: Implementation, /) -> ModuleType:
     if implementation is Implementation.PYSPARK:
         from pyspark.sql import functions
 
+        return functions
+    if implementation is Implementation.PYSPARK_CONNECT:
+        from pyspark.sql.connect import functions
+
         return functions
     from sqlframe.base.session import _BaseSession
 
@@ -255,6 +259,10 @@ def import_native_dtypes(implementation: Implementation, /) -> ModuleType:
     if implementation is Implementation.PYSPARK:
         from pyspark.sql import types
 
+        return types
+    if implementation is Implementation.PYSPARK_CONNECT:
+        from pyspark.sql.connect import types
+
         return types
     from sqlframe.base.session import _BaseSession
 
@@ -265,6 +273,11 @@ def import_window(implementation: Implementation, /) -> type[Any]:
     if implementation is Implementation.PYSPARK:
         from pyspark.sql import Window
 
+        return Window
+
+    if implementation is Implementation.PYSPARK_CONNECT:
+        from pyspark.sql.connect.window import Window
+
         return Window
     from sqlframe.base.session import _BaseSession
 

narwhals/dependencies.py

@@ -17,6 +17,7 @@
     import polars as pl
     import pyarrow as pa
     import pyspark.sql as pyspark_sql
+    from pyspark.sql.connect.dataframe import DataFrame as PySparkConnectDataFrame
     from typing_extensions import TypeGuard
     from typing_extensions import TypeIs
 
@@ -112,6 +113,11 @@ def get_pyspark_sql() -> Any:
     return sys.modules.get("pyspark.sql", None)
 
 
+def get_pyspark_connect() -> Any:
+    """Get pyspark.sql.connect module (if already imported - else return None)."""
+    return sys.modules.get("pyspark.sql", None)
+
+
 def get_sqlframe() -> Any:
     """Get sqlframe module (if already imported - else return None)."""
     return sys.modules.get("sqlframe", None)
@@ -230,6 +236,14 @@ def is_pyspark_dataframe(df: Any) -> TypeIs[pyspark_sql.DataFrame]:
     )
 
 
+def is_pyspark_connect_dataframe(df: Any) -> TypeIs[PySparkConnectDataFrame]:
+    """Check whether `df` is a PySpark Connect DataFrame without importing PySpark."""
+    return bool(
+        (pyspark_sql := get_pyspark_connect()) is not None
+        and isinstance(df, pyspark_sql.connect.dataframe.DataFrame)
+    )
+
+
 def is_sqlframe_dataframe(df: Any) -> TypeIs[SQLFrameDataFrame]:
     """Check whether `df` is a SQLFrame DataFrame without importing SQLFrame."""
     if get_sqlframe() is not None:

narwhals/utils.py

@@ -33,6 +33,7 @@
 from narwhals.dependencies import get_polars
 from narwhals.dependencies import get_pyarrow
 from narwhals.dependencies import get_pyspark
+from narwhals.dependencies import get_pyspark_connect
 from narwhals.dependencies import get_pyspark_sql
 from narwhals.dependencies import get_sqlframe
 from narwhals.dependencies import is_cudf_series
@@ -199,6 +200,8 @@ class Implementation(Enum):
     """Ibis implementation."""
     SQLFRAME = auto()
     """SQLFrame implementation."""
+    PYSPARK_CONNECT = auto()
+    """PySpark Connect implementation."""
 
     UNKNOWN = auto()
     """Unknown implementation."""
@@ -226,6 +229,7 @@ def from_native_namespace(
             get_duckdb(): Implementation.DUCKDB,
             get_ibis(): Implementation.IBIS,
             get_sqlframe(): Implementation.SQLFRAME,
+            get_pyspark_connect(): Implementation.PYSPARK_CONNECT,
         }
         return mapping.get(native_namespace, Implementation.UNKNOWN)
 
@@ -252,6 +256,7 @@ def from_string(
             "duckdb": Implementation.DUCKDB,
             "ibis": Implementation.IBIS,
             "sqlframe": Implementation.SQLFRAME,
+            "pyspark_connect": Implementation.PYSPARK_CONNECT,
         }
         return mapping.get(backend_name, Implementation.UNKNOWN)
 
@@ -320,6 +325,11 @@ def to_native_namespace(self) -> ModuleType:
 
             return sqlframe
 
+        if self is Implementation.PYSPARK_CONNECT:
+            import pyspark.sql  # ignore-banned-import
+
+            return pyspark.sql
+
         msg = "Not supported Implementation"  # pragma: no cover
         raise AssertionError(msg)
 
@@ -373,7 +383,11 @@ def is_spark_like(self) -> bool:
             >>> df.implementation.is_spark_like()
             False
         """
-        return self in {Implementation.PYSPARK, Implementation.SQLFRAME}
+        return self in {
+            Implementation.PYSPARK,
+            Implementation.SQLFRAME,
+            Implementation.PYSPARK_CONNECT,
+        }
 
     def is_polars(self) -> bool:
         """Return whether implementation is Polars.
@@ -545,11 +559,12 @@ def _backend_version(self) -> tuple[int, ...]:
         into_version: Any
         if self not in {
             Implementation.PYSPARK,
+            Implementation.PYSPARK_CONNECT,
             Implementation.DASK,
             Implementation.SQLFRAME,
         }:
             into_version = native
-        elif self is Implementation.PYSPARK:
+        elif self in {Implementation.PYSPARK, Implementation.PYSPARK_CONNECT}:
             into_version = get_pyspark()  # pragma: no cover
         elif self is Implementation.DASK:
             into_version = get_dask()
@@ -566,6 +581,7 @@ def _backend_version(self) -> tuple[int, ...]:
     Implementation.CUDF: (24, 10),
     Implementation.PYARROW: (11,),
     Implementation.PYSPARK: (3, 5),
+    Implementation.PYSPARK_CONNECT: (3, 5),
     Implementation.POLARS: (0, 20, 3),
     Implementation.DASK: (2024, 8),
     Implementation.DUCKDB: (1,),

pyproject.toml

@@ -39,6 +39,7 @@ modin = ["modin"]
 cudf = ["cudf>=24.10.0"]
 pyarrow = ["pyarrow>=11.0.0"]
 pyspark = ["pyspark>=3.5.0"]
+pyspark-connect = ["pyspark[connect]>=3.5.0"]
 polars = ["polars>=0.20.3"]
 dask = ["dask[dataframe]>=2024.8"]
 duckdb = ["duckdb>=1.0"]
@@ -217,6 +218,9 @@ filterwarnings = [
   'ignore:.*The distutils package is deprecated and slated for removal in Python 3.12:DeprecationWarning:pyspark',
   'ignore:.*distutils Version classes are deprecated. Use packaging.version instead.*:DeprecationWarning:pyspark',
   'ignore:.*is_datetime64tz_dtype is deprecated and will be removed in a future version.*:DeprecationWarning:pyspark',
+  'ignore:.*No Partition Defined for Window operation! Moving all data to a single partition.*:UserWarning:pyspark',
+  # TODO(FBruzzesi): This is raised in `hist_test.py` from pandas??? How does spark-connect affect/changes pandas behavior?
+  'ignore:.*invalid value encountered in cast.*:RuntimeWarning',
   # Warning raised by PyArrow nightly just by importing pandas
   'ignore:.*Python binding for RankQuantileOptions not exposed:RuntimeWarning:pyarrow',
   'ignore:.*pandas only supports SQLAlchemy:UserWarning:sqlframe',
@@ -231,6 +235,7 @@ env = [
   "MODIN_ENGINE=python",
   "PYARROW_IGNORE_TIMEZONE=1",
   "TZ=UTC",
+  "SPARK_REMOTE='sc://localhost'",
 ]
 
 [tool.coverage.run]