feat: Allow spark-like backends in .lazy(backend=...)

docs/api-reference/expr.md

@@ -9,6 +9,7 @@
         - all
         - any
         - cast
+        - clip
         - count
         - cum_count
         - cum_max
@@ -21,7 +22,6 @@
         - exp
         - fill_null
         - filter
-        - clip
         - is_between
         - is_close
         - is_duplicated
@@ -41,8 +41,8 @@
         - median
         - min
         - mode
-        - null_count
         - n_unique
+        - null_count
         - over
         - pipe
         - quantile

narwhals/_arrow/dataframe.py

@@ -45,6 +45,7 @@
         Order,
     )
     from narwhals._compliant.typing import CompliantDataFrameAny, CompliantLazyFrameAny
+    from narwhals._spark_like.utils import SparkSession
     from narwhals._translate import IntoArrowTable
     from narwhals._typing import _EagerAllowedImpl, _LazyAllowedImpl
     from narwhals._utils import Version, _LimitedContext
@@ -527,17 +528,22 @@ def tail(self, n: int) -> Self:
             )
         return self._with_native(df.slice(abs(n)), validate_column_names=False)
 
-    def lazy(self, backend: _LazyAllowedImpl | None = None) -> CompliantLazyFrameAny:
+    def lazy(
+        self,
+        backend: _LazyAllowedImpl | None = None,
+        *,
+        session: SparkSession | None = None,
+    ) -> CompliantLazyFrameAny:
         if backend is None:
             return self
         if backend is Implementation.DUCKDB:
             import duckdb  # ignore-banned-import
 
             from narwhals._duckdb.dataframe import DuckDBLazyFrame
 
-            df = self.native  # noqa: F841
+            _df = self.native
             return DuckDBLazyFrame(
-                duckdb.table("df"), validate_backend_version=True, version=self._version
+                duckdb.table("_df"), validate_backend_version=True, version=self._version
             )
         if backend is Implementation.POLARS:
             import polars as pl  # ignore-banned-import
@@ -559,7 +565,7 @@ def lazy(self, backend: _LazyAllowedImpl | None = None) -> CompliantLazyFrameAny
                 validate_backend_version=True,
                 version=self._version,
             )
-        if backend.is_ibis():
+        if backend is Implementation.IBIS:
             import ibis  # ignore-banned-import
 
             from narwhals._ibis.dataframe import IbisLazyFrame
@@ -570,6 +576,35 @@ def lazy(self, backend: _LazyAllowedImpl | None = None) -> CompliantLazyFrameAny
                 version=self._version,
             )
 
+        if backend.is_spark_like():
+            from narwhals._spark_like.dataframe import SparkLikeLazyFrame
+            from narwhals.schema import Schema
+
+            if session is None:
+                msg = "Spark like backends require `session` to be not None."
+                raise ValueError(msg)
+            spark_like_schema = Schema(self.schema)._to_spark_like(
+                backend=backend, session=session
+            )
+
+            can_create_from_arrow = backend in {
+                Implementation.PYSPARK,
+                Implementation.PYSPARK_CONNECT,
+            } and backend._backend_version() >= (4, 0, 0)
+
+            data = (
+                self.native
+                if can_create_from_arrow
+                else tuple(self.iter_rows(named=True, buffer_size=512))
+            )
+
+            return SparkLikeLazyFrame(
+                session.createDataFrame(data, schema=spark_like_schema),  # type: ignore[arg-type]
+                version=self._version,
+                implementation=backend,
+                validate_backend_version=True,
+            )
+
         raise AssertionError  # pragma: no cover
 
     def collect(

narwhals/_compliant/dataframe.py

@@ -50,6 +50,7 @@
 
     from narwhals._compliant.group_by import CompliantGroupBy, DataFrameGroupBy
     from narwhals._compliant.namespace import EagerNamespace
+    from narwhals._spark_like.utils import SparkSession
     from narwhals._translate import IntoArrowTable
     from narwhals._typing import _EagerAllowedImpl, _LazyAllowedImpl
     from narwhals._utils import Implementation, _LimitedContext
@@ -199,7 +200,9 @@ def join_asof(
         strategy: AsofJoinStrategy,
         suffix: str,
     ) -> Self: ...
-    def lazy(self, backend: _LazyAllowedImpl | None) -> CompliantLazyFrameAny: ...
+    def lazy(
+        self, backend: _LazyAllowedImpl | None, *, session: SparkSession | None
+    ) -> CompliantLazyFrameAny: ...
     def pivot(
         self,
         on: Sequence[str],

narwhals/_duckdb/dataframe.py

@@ -48,6 +48,7 @@
     from narwhals._duckdb.group_by import DuckDBGroupBy
     from narwhals._duckdb.namespace import DuckDBNamespace
     from narwhals._duckdb.series import DuckDBInterchangeSeries
+    from narwhals._spark_like.utils import SparkSession
     from narwhals._typing import _EagerAllowedImpl, _LazyAllowedImpl
     from narwhals._utils import _LimitedContext
     from narwhals.dataframe import LazyFrame
@@ -189,7 +190,12 @@ def drop(self, columns: Sequence[str], *, strict: bool) -> Self:
         selection = (name for name in self.columns if name not in columns_to_drop)
         return self._with_native(self.native.select(*selection))
 
-    def lazy(self, backend: _LazyAllowedImpl | None = None) -> Self:
+    def lazy(
+        self,
+        backend: _LazyAllowedImpl | None = None,
+        *,
+        session: SparkSession | None = None,
+    ) -> Self:
         # The `backend`` argument has no effect but we keep it here for
         # backwards compatibility because in `narwhals.stable.v1`
         # function `.from_native()` will return a DataFrame for DuckDB.

narwhals/_ibis/dataframe.py

@@ -34,6 +34,7 @@
     from narwhals._ibis.group_by import IbisGroupBy
     from narwhals._ibis.namespace import IbisNamespace
     from narwhals._ibis.series import IbisInterchangeSeries
+    from narwhals._spark_like.utils import SparkSession
     from narwhals._typing import _EagerAllowedImpl, _LazyAllowedImpl
     from narwhals._utils import _LimitedContext
     from narwhals.dataframe import LazyFrame
@@ -165,7 +166,12 @@ def drop(self, columns: Sequence[str], *, strict: bool) -> Self:
         selection = (col for col in self.columns if col not in columns_to_drop)
         return self._with_native(self.native.select(*selection))
 
-    def lazy(self, backend: _LazyAllowedImpl | None = None) -> Self:
+    def lazy(
+        self,
+        backend: _LazyAllowedImpl | None = None,
+        *,
+        session: SparkSession | None = None,
+    ) -> Self:
         # The `backend`` argument has no effect but we keep it here for
         # backwards compatibility because in `narwhals.stable.v1`
         # function `.from_native()` will return a DataFrame for Ibis.

narwhals/_pandas_like/dataframe.py

@@ -47,6 +47,7 @@
     from narwhals._pandas_like.expr import PandasLikeExpr
     from narwhals._pandas_like.group_by import PandasLikeGroupBy
     from narwhals._pandas_like.namespace import PandasLikeNamespace
+    from narwhals._spark_like.utils import SparkSession
     from narwhals._translate import IntoArrowTable
     from narwhals._typing import _EagerAllowedImpl, _LazyAllowedImpl
     from narwhals._utils import Version, _LimitedContext
@@ -772,7 +773,12 @@ def unique(
         )
 
     # --- lazy-only ---
-    def lazy(self, backend: _LazyAllowedImpl | None = None) -> CompliantLazyFrameAny:
+    def lazy(
+        self,
+        backend: _LazyAllowedImpl | None = None,
+        *,
+        session: SparkSession | None = None,
+    ) -> CompliantLazyFrameAny:
         pandas_df = self.to_pandas()
         if backend is None:
             return self
@@ -806,7 +812,7 @@ def lazy(self, backend: _LazyAllowedImpl | None = None) -> CompliantLazyFrameAny
                 validate_backend_version=True,
                 version=self._version,
             )
-        if backend.is_ibis():
+        if backend is Implementation.IBIS:
             import ibis  # ignore-banned-import
 
             from narwhals._ibis.dataframe import IbisLazyFrame
@@ -816,6 +822,25 @@ def lazy(self, backend: _LazyAllowedImpl | None = None) -> CompliantLazyFrameAny
                 validate_backend_version=True,
                 version=self._version,
             )
+
+        if backend.is_spark_like():
+            from narwhals._spark_like.dataframe import SparkLikeLazyFrame
+            from narwhals.schema import Schema
+
+            if session is None:
+                msg = "Spark like backends require `session` to be not None."
+                raise ValueError(msg)
+
+            spark_like_schema = Schema(self.schema)._to_spark_like(
+                backend=backend, session=session
+            )
+            return SparkLikeLazyFrame(
+                session.createDataFrame(pandas_df, schema=spark_like_schema),
+                version=self._version,
+                implementation=backend,
+                validate_backend_version=True,
+            )
+
         raise AssertionError  # pragma: no cover
 
     @property

narwhals/_polars/dataframe.py

@@ -40,6 +40,7 @@
     from narwhals._compliant.typing import CompliantDataFrameAny, CompliantLazyFrameAny
     from narwhals._polars.expr import PolarsExpr
     from narwhals._polars.group_by import PolarsGroupBy, PolarsLazyGroupBy
+    from narwhals._spark_like.utils import SparkSession
     from narwhals._translate import IntoArrowTable
     from narwhals._typing import _EagerAllowedImpl, _LazyAllowedImpl
     from narwhals._utils import Version, _LimitedContext
@@ -455,7 +456,12 @@ def iter_columns(self) -> Iterator[PolarsSeries]:
         for series in self.native.iter_columns():
             yield PolarsSeries.from_native(series, context=self)
 
-    def lazy(self, backend: _LazyAllowedImpl | None = None) -> CompliantLazyFrameAny:
+    def lazy(
+        self,
+        backend: _LazyAllowedImpl | None = None,
+        *,
+        session: SparkSession | None = None,
+    ) -> CompliantLazyFrameAny:
         if backend is None or backend is Implementation.POLARS:
             return PolarsLazyFrame.from_native(self.native.lazy(), context=self)
         if backend is Implementation.DUCKDB:
@@ -478,7 +484,7 @@ def lazy(self, backend: _LazyAllowedImpl | None = None) -> CompliantLazyFrameAny
                 validate_backend_version=True,
                 version=self._version,
             )
-        if backend.is_ibis():
+        if backend is Implementation.IBIS:
             import ibis  # ignore-banned-import
 
             from narwhals._ibis.dataframe import IbisLazyFrame
@@ -489,6 +495,40 @@ def lazy(self, backend: _LazyAllowedImpl | None = None) -> CompliantLazyFrameAny
                 version=self._version,
             )
 
+        if backend.is_spark_like():
+            from importlib.util import find_spec
+
+            from narwhals._spark_like.dataframe import SparkLikeLazyFrame
+            from narwhals.schema import Schema
+
+            if session is None:
+                msg = "Spark like backends require `session` to be not None."
+                raise ValueError(msg)
+
+            # pyspark.sql requires pyarrow to be installed from v4.0.0
+            can_create_from_arrow = backend in {
+                Implementation.PYSPARK,
+                Implementation.PYSPARK_CONNECT,
+            } and backend._backend_version() >= (4, 0, 0)
+
+            data: Any
+            if can_create_from_arrow:
+                data = self.to_arrow()
+            elif find_spec("pandas") is not None:
+                data = self.to_pandas()
+            else:
+                data = self.iter_rows(named=True)
+
+            spark_like_schema = Schema(self.schema)._to_spark_like(
+                backend=backend, session=session
+            )
+            return SparkLikeLazyFrame(
+                session.createDataFrame(data, schema=spark_like_schema),
+                version=self._version,
+                implementation=backend,
+                validate_backend_version=True,
+            )
+
         raise AssertionError  # pragma: no cover
 
     @overload

narwhals/_typing.py

@@ -51,7 +51,6 @@
 # - https://github.com/narwhals-dev/narwhals/pull/2971#discussion_r2277137003
 # - https://github.com/narwhals-dev/narwhals/pull/3002#issuecomment-3194267667
 _LazyFrameCollectImpl: TypeAlias = Literal[_PandasImpl, _PolarsImpl, _ArrowImpl]  # noqa: PYI047
-_DataFrameLazyImpl: TypeAlias = Literal[_PolarsImpl, _DaskImpl, _DuckDBImpl, _IbisImpl]  # noqa: PYI047
 
 # `str | Implementation` aliases
 Pandas: TypeAlias = Literal[_Pandas, _PandasImpl]

narwhals/_utils.py

@@ -79,8 +79,8 @@
     from narwhals._typing import (
         Backend,
         IntoBackend,
-        _DataFrameLazyImpl,
         _EagerAllowedImpl,
+        _LazyAllowedImpl,
         _LazyFrameCollectImpl,
     )
     from narwhals.dataframe import DataFrame, LazyFrame
@@ -1648,9 +1648,9 @@ def is_compliant_expr(
 def is_eager_allowed(impl: Implementation, /) -> TypeIs[_EagerAllowedImpl]:
     """Return True if `impl` allows eager operations."""
     return impl in {
-        Implementation.PANDAS,
-        Implementation.MODIN,
         Implementation.CUDF,
+        Implementation.MODIN,
+        Implementation.PANDAS,
         Implementation.POLARS,
         Implementation.PYARROW,
     }
@@ -1661,13 +1661,16 @@ def can_lazyframe_collect(impl: Implementation, /) -> TypeIs[_LazyFrameCollectIm
     return impl in {Implementation.PANDAS, Implementation.POLARS, Implementation.PYARROW}
 
 
-def can_dataframe_lazy(impl: Implementation, /) -> TypeIs[_DataFrameLazyImpl]:
+def can_dataframe_lazy(impl: Implementation, /) -> TypeIs[_LazyAllowedImpl]:
     """Return True if `DataFrame.lazy(impl)` is allowed."""
     return impl in {
         Implementation.DASK,
         Implementation.DUCKDB,
-        Implementation.POLARS,
         Implementation.IBIS,
+        Implementation.POLARS,
+        Implementation.PYSPARK,
+        Implementation.PYSPARK_CONNECT,
+        Implementation.SQLFRAME,
     }