from __future__ import annotations from typing import TYPE_CHECKING from typing import Any from typing import Callable from typing import Sequence from narwhals._expression_parsing import infer_new_root_output_names from narwhals._spark_like.utils import get_column_name from narwhals._spark_like.utils import maybe_evaluate from narwhals.typing import CompliantExpr from narwhals.utils import Implementation from narwhals.utils import parse_version if TYPE_CHECKING: from pyspark.sql import Column from typing_extensions import Self from narwhals._spark_like.dataframe import SparkLikeLazyFrame from narwhals._spark_like.namespace import SparkLikeNamespace from narwhals.utils import Version class SparkLikeExpr(CompliantExpr["Column"]): _implementation = Implementation.PYSPARK def __init__( self, call: Callable[[SparkLikeLazyFrame], list[Column]], *, depth: int, function_name: str, root_names: list[str] | None, output_names: list[str] | None, # Whether the expression is a length-1 Column resulting from # a reduction, such as `nw.col('a').sum()` returns_scalar: bool, backend_version: tuple[int, ...], version: Version, kwargs: dict[str, Any], ) -> None: self._call = call self._depth = depth self._function_name = function_name self._root_names = root_names self._output_names = output_names self._returns_scalar = returns_scalar self._backend_version = backend_version self._version = version self._kwargs = kwargs def __call__(self, df: SparkLikeLazyFrame) -> Sequence[Column]: return self._call(df) def __narwhals_expr__(self) -> None: ... def __narwhals_namespace__(self) -> SparkLikeNamespace: # pragma: no cover # Unused, just for compatibility with PandasLikeExpr from narwhals._spark_like.namespace import SparkLikeNamespace return SparkLikeNamespace( # type: ignore[abstract] backend_version=self._backend_version, version=self._version ) @classmethod def from_column_names( cls: type[Self], *column_names: str, backend_version: tuple[int, ...], version: Version, ) -> Self: def func(_: SparkLikeLazyFrame) -> list[Column]: from pyspark.sql import functions as F # noqa: N812 return [F.col(col_name) for col_name in column_names] return cls( func, depth=0, function_name="col", root_names=list(column_names), output_names=list(column_names), returns_scalar=False, backend_version=backend_version, version=version, kwargs={}, ) def _from_call( self, call: Callable[..., Column], expr_name: str, *, returns_scalar: bool, **kwargs: Any, ) -> Self: def func(df: SparkLikeLazyFrame) -> list[Column]: results = [] inputs = self._call(df) _kwargs = {key: maybe_evaluate(df, value) for key, value in kwargs.items()} for _input in inputs: input_col_name = get_column_name(df, _input) column_result = call(_input, **_kwargs) if not returns_scalar: column_result = column_result.alias(input_col_name) results.append(column_result) return results root_names, output_names = infer_new_root_output_names(self, **kwargs) return self.__class__( func, depth=self._depth + 1, function_name=f"{self._function_name}->{expr_name}", root_names=root_names, output_names=output_names, returns_scalar=self._returns_scalar or returns_scalar, backend_version=self._backend_version, version=self._version, kwargs=kwargs, ) def __add__(self, other: SparkLikeExpr) -> Self: return self._from_call( lambda _input, other: _input + other, "__add__", other=other, returns_scalar=False, ) def __sub__(self, other: SparkLikeExpr) -> Self: return self._from_call( lambda _input, other: _input - other, "__sub__", other=other, returns_scalar=False, ) def __mul__(self, other: SparkLikeExpr) -> Self: return self._from_call( lambda _input, other: _input * other, "__mul__", other=other, returns_scalar=False, ) def __lt__(self, other: SparkLikeExpr) -> Self: return self._from_call( lambda _input, other: _input < other, "__lt__", other=other, returns_scalar=False, ) def __gt__(self, other: SparkLikeExpr) -> Self: return self._from_call( lambda _input, other: _input > other, "__gt__", other=other, returns_scalar=False, ) def alias(self, name: str) -> Self: def _alias(df: SparkLikeLazyFrame) -> list[Column]: return [col.alias(name) for col in self._call(df)] # Define this one manually, so that we can # override `output_names` and not increase depth return self.__class__( _alias, depth=self._depth, function_name=self._function_name, root_names=self._root_names, output_names=[name], returns_scalar=self._returns_scalar, backend_version=self._backend_version, version=self._version, kwargs={**self._kwargs, "name": name}, ) def count(self) -> Self: def _count(_input: Column) -> Column: from pyspark.sql import functions as F # noqa: N812 return F.count(_input) return self._from_call(_count, "count", returns_scalar=True) def max(self) -> Self: def _max(_input: Column) -> Column: from pyspark.sql import functions as F # noqa: N812 return F.max(_input) return self._from_call(_max, "max", returns_scalar=True) def mean(self) -> Self: def _mean(_input: Column) -> Column: from pyspark.sql import functions as F # noqa: N812 return F.mean(_input) return self._from_call(_mean, "mean", returns_scalar=True) def min(self) -> Self: def _min(_input: Column) -> Column: from pyspark.sql import functions as F # noqa: N812 return F.min(_input) return self._from_call(_min, "min", returns_scalar=True) def sum(self) -> Self: def _sum(_input: Column) -> Column: from pyspark.sql import functions as F # noqa: N812 return F.sum(_input) return self._from_call(_sum, "sum", returns_scalar=True) def std(self: Self, ddof: int) -> Self: from functools import partial import numpy as np # ignore-banned-import from narwhals._spark_like.utils import _std func = partial(_std, ddof=ddof, np_version=parse_version(np.__version__)) return self._from_call(func, "std", returns_scalar=True, ddof=ddof) def var(self: Self, ddof: int) -> Self: from functools import partial import numpy as np # ignore-banned-import from narwhals._spark_like.utils import _var func = partial(_var, ddof=ddof, np_version=parse_version(np.__version__)) return self._from_call(func, "var", returns_scalar=True, ddof=ddof)