diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py index 5fc84a854f..f539785407 100644 --- a/narwhals/_arrow/expr.py +++ b/narwhals/_arrow/expr.py @@ -10,7 +10,6 @@ from narwhals._arrow.series import ArrowSeries from narwhals._compliant import EagerExpr -from narwhals._expression_parsing import ExprKind from narwhals._expression_parsing import evaluate_output_names_and_aliases from narwhals._expression_parsing import is_scalar_like from narwhals.exceptions import ColumnNotFoundError @@ -23,6 +22,7 @@ from narwhals._arrow.dataframe import ArrowDataFrame from narwhals._arrow.namespace import ArrowNamespace + from narwhals._expression_parsing import ExprMetadata from narwhals.utils import Version from narwhals.utils import _FullContext @@ -52,6 +52,7 @@ def __init__( self._backend_version = backend_version self._version = version self._call_kwargs = call_kwargs or {} + self._metadata: ExprMetadata | None = None @classmethod def from_column_names( @@ -141,10 +142,10 @@ def shift(self: Self, n: int) -> Self: def over( self: Self, partition_by: Sequence[str], - kind: ExprKind, order_by: Sequence[str] | None, ) -> Self: - if partition_by and not is_scalar_like(kind): + assert self._metadata is not None # noqa: S101 + if partition_by and not is_scalar_like(self._metadata.kind): msg = "Only aggregation or literal operations are supported in grouped `over` context for PyArrow." raise NotImplementedError(msg) diff --git a/narwhals/_compliant/expr.py b/narwhals/_compliant/expr.py index ec251a6ac8..7fe41baf6f 100644 --- a/narwhals/_compliant/expr.py +++ b/narwhals/_compliant/expr.py @@ -57,6 +57,7 @@ from narwhals._compliant.namespace import EagerNamespace from narwhals._compliant.series import CompliantSeries from narwhals._expression_parsing import ExprKind + from narwhals._expression_parsing import ExprMetadata from narwhals.dtypes import DType from narwhals.typing import TimeUnit from narwhals.utils import Implementation @@ -85,6 +86,7 @@ class CompliantExpr(Protocol38[CompliantFrameT, CompliantSeriesOrNativeExprT_co] _alias_output_names: Callable[[Sequence[str]], Sequence[str]] | None _depth: int _function_name: str + _metadata: ExprMetadata | None def __call__( self, df: CompliantFrameT @@ -105,6 +107,8 @@ def from_column_names( @classmethod def from_column_indices(cls, *column_indices: int, context: _FullContext) -> Self: ... + def _with_metadata(self, metadata: ExprMetadata) -> Self: ... + def is_null(self) -> Self: ... def abs(self) -> Self: ... def all(self) -> Self: ... @@ -165,9 +169,7 @@ def replace_strict( *, return_dtype: DType | type[DType] | None, ) -> Self: ... - def over( - self: Self, keys: Sequence[str], kind: ExprKind, order_by: Sequence[str] | None - ) -> Self: ... + def over(self: Self, keys: Sequence[str], order_by: Sequence[str] | None) -> Self: ... def sample( self, n: int | None, @@ -316,6 +318,21 @@ def __narwhals_namespace__( ) -> EagerNamespace[EagerDataFrameT, EagerSeriesT, Self]: ... def __narwhals_expr__(self) -> None: ... + def _with_metadata(self, metadata: ExprMetadata) -> Self: + expr = self.__class__( + self._call, + function_name=self._function_name, + evaluate_output_names=self._evaluate_output_names, + alias_output_names=self._alias_output_names, + backend_version=self._backend_version, + version=self._version, + depth=self._depth, + implementation=self._implementation, + call_kwargs=self._call_kwargs, + ) + expr._metadata = metadata + return expr + @classmethod def _from_callable( cls, diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py index ab99a5e15e..466edd9044 100644 --- a/narwhals/_dask/expr.py +++ b/narwhals/_dask/expr.py @@ -36,6 +36,7 @@ from narwhals._dask.dataframe import DaskLazyFrame from narwhals._dask.namespace import DaskNamespace + from narwhals._expression_parsing import ExprMetadata from narwhals.dtypes import DType from narwhals.utils import Version from narwhals.utils import _FullContext @@ -66,6 +67,7 @@ def __init__( self._backend_version = backend_version self._version = version self._call_kwargs = call_kwargs or {} + self._metadata: ExprMetadata | None = None def __call__(self: Self, df: DaskLazyFrame) -> Sequence[dx.Series]: return self._call(df) @@ -78,6 +80,20 @@ def __narwhals_namespace__(self) -> DaskNamespace: # pragma: no cover return DaskNamespace(backend_version=self._backend_version, version=self._version) + def _with_metadata(self, metadata: ExprMetadata) -> Self: + expr = self.__class__( + self._call, + function_name=self._function_name, + evaluate_output_names=self._evaluate_output_names, + alias_output_names=self._alias_output_names, + backend_version=self._backend_version, + version=self._version, + depth=self._depth, + call_kwargs=self._call_kwargs, + ) + expr._metadata = metadata + return expr + def broadcast(self, kind: Literal[ExprKind.AGGREGATION, ExprKind.LITERAL]) -> Self: def func(df: DaskLazyFrame) -> list[dx.Series]: return [result[0] for result in self(df)] @@ -545,7 +561,6 @@ def null_count(self: Self) -> Self: def over( self: Self, partition_by: Sequence[str], - kind: ExprKind, order_by: Sequence[str] | None, ) -> Self: # pandas is a required dependency of dask so it's safe to import this diff --git a/narwhals/_duckdb/expr.py b/narwhals/_duckdb/expr.py index f450143afb..62bc942652 100644 --- a/narwhals/_duckdb/expr.py +++ b/narwhals/_duckdb/expr.py @@ -33,6 +33,7 @@ from narwhals._duckdb.dataframe import DuckDBLazyFrame from narwhals._duckdb.namespace import DuckDBNamespace + from narwhals._expression_parsing import ExprMetadata from narwhals.dtypes import DType from narwhals.utils import Version from narwhals.utils import _FullContext @@ -58,6 +59,7 @@ def __init__( self._alias_output_names = alias_output_names self._backend_version = backend_version self._version = version + self._metadata: ExprMetadata | None = None def __call__(self: Self, df: DuckDBLazyFrame) -> Sequence[duckdb.Expression]: return self._call(df) @@ -72,6 +74,18 @@ def __narwhals_namespace__(self) -> DuckDBNamespace: # pragma: no cover backend_version=self._backend_version, version=self._version ) + def _with_metadata(self, metadata: ExprMetadata) -> Self: + expr = self.__class__( + self._call, + function_name=self._function_name, + evaluate_output_names=self._evaluate_output_names, + alias_output_names=self._alias_output_names, + backend_version=self._backend_version, + version=self._version, + ) + expr._metadata = metadata + return expr + def broadcast(self, kind: Literal[ExprKind.AGGREGATION, ExprKind.LITERAL]) -> Self: if kind is ExprKind.AGGREGATION: msg = "Broadcasting aggregations is not yet supported for DuckDB." diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py index be93b9c0a3..b8766fec31 100644 --- a/narwhals/_pandas_like/expr.py +++ b/narwhals/_pandas_like/expr.py @@ -7,7 +7,6 @@ from typing import Sequence from narwhals._compliant import EagerExpr -from narwhals._expression_parsing import ExprKind from narwhals._expression_parsing import evaluate_output_names_and_aliases from narwhals._expression_parsing import is_elementary_expression from narwhals._pandas_like.group_by import PandasLikeGroupBy @@ -18,6 +17,7 @@ if TYPE_CHECKING: from typing_extensions import Self + from narwhals._expression_parsing import ExprMetadata from narwhals._pandas_like.dataframe import PandasLikeDataFrame from narwhals._pandas_like.namespace import PandasLikeNamespace from narwhals.utils import Implementation @@ -89,6 +89,7 @@ def __init__( self._backend_version = backend_version self._version = version self._call_kwargs = call_kwargs or {} + self._metadata: ExprMetadata | None = None def __narwhals_namespace__(self: Self) -> PandasLikeNamespace: from narwhals._pandas_like.namespace import PandasLikeNamespace @@ -196,7 +197,6 @@ def shift(self: Self, n: int) -> Self: def over( self: Self, partition_by: Sequence[str], - kind: ExprKind, order_by: Sequence[str] | None, ) -> Self: if not partition_by: diff --git a/narwhals/_polars/expr.py b/narwhals/_polars/expr.py index 8d57e31920..e3a57be6fa 100644 --- a/narwhals/_polars/expr.py +++ b/narwhals/_polars/expr.py @@ -17,6 +17,7 @@ from typing_extensions import Self from narwhals._expression_parsing import ExprKind + from narwhals._expression_parsing import ExprMetadata from narwhals.dtypes import DType from narwhals.utils import Version @@ -29,6 +30,7 @@ def __init__( self._implementation = Implementation.POLARS self._version = version self._backend_version = backend_version + self._metadata: ExprMetadata | None = None def __repr__(self: Self) -> str: # pragma: no cover return "PolarsExpr" @@ -38,6 +40,15 @@ def _from_native_expr(self: Self, expr: pl.Expr) -> Self: expr, version=self._version, backend_version=self._backend_version ) + def _with_metadata(self, metadata: ExprMetadata) -> Self: + expr = self.__class__( + self._native_expr, + backend_version=self._backend_version, + version=self._version, + ) + expr._metadata = metadata + return expr + @classmethod def _from_series(cls, series: Any) -> Self: return cls( @@ -108,7 +119,6 @@ def is_nan(self: Self) -> Self: def over( self: Self, partition_by: Sequence[str], - kind: ExprKind, order_by: Sequence[str] | None, ) -> Self: if self._backend_version < (1, 9): diff --git a/narwhals/_spark_like/expr.py b/narwhals/_spark_like/expr.py index 06d40c86b6..e9314cae20 100644 --- a/narwhals/_spark_like/expr.py +++ b/narwhals/_spark_like/expr.py @@ -30,6 +30,7 @@ from sqlframe.base.window import Window from typing_extensions import Self + from narwhals._expression_parsing import ExprMetadata from narwhals._spark_like.dataframe import SparkLikeLazyFrame from narwhals._spark_like.namespace import SparkLikeNamespace from narwhals._spark_like.typing import WindowFunction @@ -60,6 +61,7 @@ def __init__( self._version = version self._implementation = implementation self._window_function: WindowFunction | None = None + self._metadata: ExprMetadata | None = None def __call__(self: Self, df: SparkLikeLazyFrame) -> Sequence[Column]: return self._call(df) @@ -123,6 +125,37 @@ def __narwhals_namespace__(self: Self) -> SparkLikeNamespace: # pragma: no cove implementation=self._implementation, ) + def _with_metadata(self, metadata: ExprMetadata) -> Self: + expr = self.__class__( + self._call, + function_name=self._function_name, + evaluate_output_names=self._evaluate_output_names, + alias_output_names=self._alias_output_names, + backend_version=self._backend_version, + version=self._version, + implementation=self._implementation, + ) + if func := self._window_function: + expr = expr._with_window_function(func) + expr._metadata = metadata + return expr + + def _with_window_function( + self: Self, + window_function: WindowFunction, + ) -> Self: + result = self.__class__( + self._call, + function_name=self._function_name, + evaluate_output_names=self._evaluate_output_names, + alias_output_names=self._alias_output_names, + backend_version=self._backend_version, + version=self._version, + implementation=self._implementation, + ) + result._window_function = window_function + return result + @classmethod def from_column_names( cls: type[Self], @@ -190,22 +223,6 @@ def func(df: SparkLikeLazyFrame) -> list[Column]: implementation=self._implementation, ) - def _with_window_function( - self: Self, - window_function: WindowFunction, - ) -> Self: - result = self.__class__( - self._call, - function_name=self._function_name, - evaluate_output_names=self._evaluate_output_names, - alias_output_names=self._alias_output_names, - backend_version=self._backend_version, - version=self._version, - implementation=self._implementation, - ) - result._window_function = window_function - return result - def __eq__(self: Self, other: SparkLikeExpr) -> Self: # type: ignore[override] return self._from_call( lambda _input, other: _input.__eq__(other), "__eq__", other=other @@ -499,7 +516,6 @@ def _n_unique(_input: Column) -> Column: def over( self: Self, partition_by: Sequence[str], - kind: ExprKind, order_by: Sequence[str] | None, ) -> Self: if (window_function := self._window_function) is not None: diff --git a/narwhals/expr.py b/narwhals/expr.py index 42fdaf2db0..2a402f1957 100644 --- a/narwhals/expr.py +++ b/narwhals/expr.py @@ -1576,17 +1576,18 @@ def over( n_open_windows = self._metadata.n_open_windows if _order_by is not None and self._metadata.kind.is_window(): n_open_windows -= 1 - metadata = ExprMetadata( + current_meta = self._metadata + next_meta = ExprMetadata( kind, n_open_windows=n_open_windows, - is_multi_output=self._metadata.is_multi_output, + is_multi_output=current_meta.is_multi_output, ) return self.__class__( - lambda plx: self._to_compliant_expr(plx).over( - flat_partition_by, order_by=order_by, kind=self._metadata.kind - ), - metadata, + lambda plx: self._to_compliant_expr(plx) + ._with_metadata(current_meta) + .over(flat_partition_by, order_by), + next_meta, ) def is_duplicated(self: Self) -> Self: