diff --git a/narwhals/_polars/expr.py b/narwhals/_polars/expr.py index f8c2fe41d2..af97d6b89a 100644 --- a/narwhals/_polars/expr.py +++ b/narwhals/_polars/expr.py @@ -39,18 +39,12 @@ def native(self) -> pl.Expr: def __repr__(self: Self) -> str: # pragma: no cover return "PolarsExpr" - def _from_native_expr(self: Self, expr: pl.Expr) -> Self: - return self.__class__( - expr, version=self._version, backend_version=self._backend_version - ) + def _with_native(self: Self, expr: pl.Expr) -> Self: + return self.__class__(expr, self._version, self._backend_version) @classmethod def _from_series(cls, series: Any) -> Self: - return cls( - series.native, - version=series._version, - backend_version=series._backend_version, - ) + return cls(series.native, series._version, series._backend_version) def broadcast(self, kind: Literal[ExprKind.AGGREGATION, ExprKind.LITERAL]) -> Self: # Let Polars do its thing. @@ -59,16 +53,17 @@ def broadcast(self, kind: Literal[ExprKind.AGGREGATION, ExprKind.LITERAL]) -> Se def __getattr__(self: Self, attr: str) -> Any: def func(*args: Any, **kwargs: Any) -> Any: args, kwargs = extract_args_kwargs(args, kwargs) # type: ignore[assignment] - return self._from_native_expr( - getattr(self._native_expr, attr)(*args, **kwargs) - ) + return self._with_native(getattr(self.native, attr)(*args, **kwargs)) return func + def _renamed_min_periods(self, min_samples: int, /) -> dict[str, Any]: + name = "min_periods" if self._backend_version < (1, 21, 0) else "min_samples" + return {name: min_samples} + def cast(self: Self, dtype: DType) -> Self: - expr = self._native_expr dtype_pl = narwhals_to_native_dtype(dtype, self._version, self._backend_version) - return self._from_native_expr(expr.cast(dtype_pl)) + return self._with_native(self.native.cast(dtype_pl)) def ewm_mean( self: Self, @@ -81,233 +76,155 @@ def ewm_mean( min_samples: int, ignore_nulls: bool, ) -> Self: - expr = self._native_expr - - extra_kwargs = ( - {"min_periods": min_samples} - if self._backend_version < (1, 21, 0) - else {"min_samples": min_samples} - ) - - native_expr = expr.ewm_mean( + native = self.native.ewm_mean( com=com, span=span, half_life=half_life, alpha=alpha, adjust=adjust, ignore_nulls=ignore_nulls, - **extra_kwargs, + **self._renamed_min_periods(min_samples), ) if self._backend_version < (1,): # pragma: no cover - return self._from_native_expr( - pl.when(~expr.is_null()).then(native_expr).otherwise(None) - ) - return self._from_native_expr(native_expr) + native = pl.when(~self.native.is_null()).then(native).otherwise(None) + return self._with_native(native) def is_nan(self: Self) -> Self: - if self._backend_version < (1, 18): # pragma: no cover - return self._from_native_expr( - pl.when(self._native_expr.is_not_null()).then(self._native_expr.is_nan()) - ) - return self._from_native_expr(self._native_expr.is_nan()) + if self._backend_version >= (1, 18): + native = self.native.is_nan() + else: # pragma: no cover + native = pl.when(self.native.is_not_null()).then(self.native.is_nan()) + return self._with_native(native) def over( - self: Self, - partition_by: Sequence[str], - order_by: Sequence[str] | None, + self: Self, partition_by: Sequence[str], order_by: Sequence[str] | None ) -> Self: if self._backend_version < (1, 9): if order_by: msg = "`order_by` in Polars requires version 1.10 or greater" raise NotImplementedError(msg) - return self._from_native_expr( - self._native_expr.over(partition_by or pl.lit(1)) - ) - return self._from_native_expr( - self._native_expr.over(partition_by or pl.lit(1), order_by=order_by) - ) + native = self.native.over(partition_by or pl.lit(1)) + else: + native = self.native.over(partition_by or pl.lit(1), order_by=order_by) + return self._with_native(native) def rolling_var( - self: Self, - window_size: int, - *, - min_samples: int, - center: bool, - ddof: int, + self: Self, window_size: int, *, min_samples: int, center: bool, ddof: int ) -> Self: if self._backend_version < (1,): # pragma: no cover msg = "`rolling_var` not implemented for polars older than 1.0" raise NotImplementedError(msg) - - extra_kwargs = ( - {"min_periods": min_samples} - if self._backend_version < (1, 21, 0) - else {"min_samples": min_samples} - ) - return self._from_native_expr( - self._native_expr.rolling_var( - window_size=window_size, - center=center, - ddof=ddof, - **extra_kwargs, # type: ignore[arg-type] - ) + kwds = self._renamed_min_periods(min_samples) + native = self.native.rolling_var( + window_size=window_size, center=center, ddof=ddof, **kwds ) + return self._with_native(native) def rolling_std( - self: Self, - window_size: int, - *, - min_samples: int, - center: bool, - ddof: int, + self: Self, window_size: int, *, min_samples: int, center: bool, ddof: int ) -> Self: if self._backend_version < (1,): # pragma: no cover msg = "`rolling_std` not implemented for polars older than 1.0" raise NotImplementedError(msg) - extra_kwargs = ( - {"min_periods": min_samples} - if self._backend_version < (1, 21, 0) - else {"min_samples": min_samples} - ) - - return self._from_native_expr( - self._native_expr.rolling_std( - window_size=window_size, - center=center, - ddof=ddof, - **extra_kwargs, # type: ignore[arg-type] - ) + kwds = self._renamed_min_periods(min_samples) + native = self.native.rolling_std( + window_size=window_size, center=center, ddof=ddof, **kwds ) + return self._with_native(native) def rolling_sum( self: Self, window_size: int, *, min_samples: int, center: bool ) -> Self: - extra_kwargs = ( - {"min_periods": min_samples} - if self._backend_version < (1, 21, 0) - else {"min_samples": min_samples} - ) - - return self._from_native_expr( - self._native_expr.rolling_sum( - window_size=window_size, - center=center, - **extra_kwargs, # type: ignore[arg-type] - ) - ) + kwds = self._renamed_min_periods(min_samples) + native = self.native.rolling_sum(window_size=window_size, center=center, **kwds) + return self._with_native(native) def rolling_mean( - self: Self, - window_size: int, - *, - min_samples: int, - center: bool, + self: Self, window_size: int, *, min_samples: int, center: bool ) -> Self: - extra_kwargs = ( - {"min_periods": min_samples} - if self._backend_version < (1, 21, 0) - else {"min_samples": min_samples} - ) - - return self._from_native_expr( - self._native_expr.rolling_mean( - window_size=window_size, - center=center, - **extra_kwargs, # type: ignore[arg-type] - ) - ) + kwds = self._renamed_min_periods(min_samples) + native = self.native.rolling_mean(window_size=window_size, center=center, **kwds) + return self._with_native(native) def map_batches( - self: Self, - function: Callable[..., Self], - return_dtype: DType | None, - ) -> Self: - if return_dtype is not None: - return_dtype_pl = narwhals_to_native_dtype( - return_dtype, self._version, self._backend_version - ) - return self._from_native_expr( - self._native_expr.map_batches(function, return_dtype_pl) - ) - else: - return self._from_native_expr(self._native_expr.map_batches(function)) - - def replace_strict( - self: Self, old: Sequence[Any], new: Sequence[Any], *, return_dtype: DType | None + self: Self, function: Callable[..., Self], return_dtype: DType | None ) -> Self: - expr = self._native_expr return_dtype_pl = ( narwhals_to_native_dtype(return_dtype, self._version, self._backend_version) if return_dtype else None ) + native = self.native.map_batches(function, return_dtype_pl) + return self._with_native(native) + + def replace_strict( + self: Self, old: Sequence[Any], new: Sequence[Any], *, return_dtype: DType | None + ) -> Self: if self._backend_version < (1,): msg = f"`replace_strict` is only available in Polars>=1.0, found version {self._backend_version}" raise NotImplementedError(msg) - return self._from_native_expr( - expr.replace_strict(old, new, return_dtype=return_dtype_pl) + return_dtype_pl = ( + narwhals_to_native_dtype(return_dtype, self._version, self._backend_version) + if return_dtype + else None ) + native = self.native.replace_strict(old, new, return_dtype=return_dtype_pl) + return self._with_native(native) def __eq__(self: Self, other: object) -> Self: # type: ignore[override] - return self._from_native_expr(self._native_expr.__eq__(extract_native(other))) # type: ignore[operator] + return self._with_native(self.native.__eq__(extract_native(other))) # type: ignore[operator] def __ne__(self: Self, other: object) -> Self: # type: ignore[override] - return self._from_native_expr(self._native_expr.__ne__(extract_native(other))) # type: ignore[operator] + return self._with_native(self.native.__ne__(extract_native(other))) # type: ignore[operator] def __ge__(self: Self, other: Any) -> Self: - return self._from_native_expr(self._native_expr.__ge__(extract_native(other))) + return self._with_native(self.native.__ge__(extract_native(other))) def __gt__(self: Self, other: Any) -> Self: - return self._from_native_expr(self._native_expr.__gt__(extract_native(other))) + return self._with_native(self.native.__gt__(extract_native(other))) def __le__(self: Self, other: Any) -> Self: - return self._from_native_expr(self._native_expr.__le__(extract_native(other))) + return self._with_native(self.native.__le__(extract_native(other))) def __lt__(self: Self, other: Any) -> Self: - return self._from_native_expr(self._native_expr.__lt__(extract_native(other))) + return self._with_native(self.native.__lt__(extract_native(other))) def __and__(self: Self, other: PolarsExpr | bool | Any) -> Self: - return self._from_native_expr(self._native_expr.__and__(extract_native(other))) # type: ignore[operator] + return self._with_native(self.native.__and__(extract_native(other))) # type: ignore[operator] def __or__(self: Self, other: PolarsExpr | bool | Any) -> Self: - return self._from_native_expr(self._native_expr.__or__(extract_native(other))) # type: ignore[operator] + return self._with_native(self.native.__or__(extract_native(other))) # type: ignore[operator] def __add__(self: Self, other: Any) -> Self: - return self._from_native_expr(self._native_expr.__add__(extract_native(other))) + return self._with_native(self.native.__add__(extract_native(other))) def __sub__(self: Self, other: Any) -> Self: - return self._from_native_expr(self._native_expr.__sub__(extract_native(other))) + return self._with_native(self.native.__sub__(extract_native(other))) def __mul__(self: Self, other: Any) -> Self: - return self._from_native_expr(self._native_expr.__mul__(extract_native(other))) + return self._with_native(self.native.__mul__(extract_native(other))) def __pow__(self: Self, other: Any) -> Self: - return self._from_native_expr(self._native_expr.__pow__(extract_native(other))) + return self._with_native(self.native.__pow__(extract_native(other))) def __truediv__(self: Self, other: Any) -> Self: - return self._from_native_expr( - self._native_expr.__truediv__(extract_native(other)) - ) + return self._with_native(self.native.__truediv__(extract_native(other))) def __floordiv__(self: Self, other: Any) -> Self: - return self._from_native_expr( - self._native_expr.__floordiv__(extract_native(other)) - ) + return self._with_native(self.native.__floordiv__(extract_native(other))) def __mod__(self: Self, other: Any) -> Self: - return self._from_native_expr(self._native_expr.__mod__(extract_native(other))) + return self._with_native(self.native.__mod__(extract_native(other))) def __invert__(self: Self) -> Self: - return self._from_native_expr(self._native_expr.__invert__()) + return self._with_native(self.native.__invert__()) def cum_count(self: Self, *, reverse: bool) -> Self: if self._backend_version < (0, 20, 4): - not_null = ~self._native_expr.is_null() - result = not_null.cum_sum(reverse=reverse) + result = (~self.native.is_null()).cum_sum(reverse=reverse) else: - result = self._native_expr.cum_count(reverse=reverse) - - return self._from_native_expr(result) + result = self.native.cum_count(reverse=reverse) + return self._with_native(result) @property def dt(self: Self) -> PolarsExprDateTimeNamespace: @@ -341,8 +258,8 @@ def __init__(self: Self, expr: PolarsExpr) -> None: def __getattr__(self: Self, attr: str) -> Callable[[Any], PolarsExpr]: def func(*args: Any, **kwargs: Any) -> PolarsExpr: args, kwargs = extract_args_kwargs(args, kwargs) # type: ignore[assignment] - return self._compliant_expr._from_native_expr( - getattr(self._compliant_expr._native_expr.dt, attr)(*args, **kwargs) + return self._compliant_expr._with_native( + getattr(self._compliant_expr.native.dt, attr)(*args, **kwargs) ) return func @@ -355,8 +272,8 @@ def __init__(self: Self, expr: PolarsExpr) -> None: def __getattr__(self: Self, attr: str) -> Callable[[Any], PolarsExpr]: def func(*args: Any, **kwargs: Any) -> PolarsExpr: args, kwargs = extract_args_kwargs(args, kwargs) # type: ignore[assignment] - return self._compliant_expr._from_native_expr( - getattr(self._compliant_expr._native_expr.str, attr)(*args, **kwargs) + return self._compliant_expr._with_native( + getattr(self._compliant_expr.native.str, attr)(*args, **kwargs) ) return func @@ -369,8 +286,8 @@ def __init__(self: Self, expr: PolarsExpr) -> None: def __getattr__(self: Self, attr: str) -> Callable[[Any], PolarsExpr]: def func(*args: Any, **kwargs: Any) -> PolarsExpr: args, kwargs = extract_args_kwargs(args, kwargs) # type: ignore[assignment] - return self._compliant_expr._from_native_expr( - getattr(self._compliant_expr._native_expr.cat, attr)(*args, **kwargs) + return self._compliant_expr._with_native( + getattr(self._compliant_expr.native.cat, attr)(*args, **kwargs) ) return func @@ -383,8 +300,8 @@ def __init__(self: Self, expr: PolarsExpr) -> None: def __getattr__(self: Self, attr: str) -> Callable[[Any], PolarsExpr]: def func(*args: Any, **kwargs: Any) -> PolarsExpr: args, kwargs = extract_args_kwargs(args, kwargs) # type: ignore[assignment] - return self._compliant_expr._from_native_expr( - getattr(self._compliant_expr._native_expr.name, attr)(*args, **kwargs) + return self._compliant_expr._with_native( + getattr(self._compliant_expr.native.name, attr)(*args, **kwargs) ) return func @@ -399,13 +316,13 @@ def len(self: Self) -> PolarsExpr: native_result = native_expr.list.len() if self._expr._backend_version < (1, 16): # pragma: no cover - native_result: pl.Expr = ( # type: ignore[no-redef] + native_result = ( pl.when(~native_expr.is_null()).then(native_result).cast(pl.UInt32()) ) elif self._expr._backend_version < (1, 17): # pragma: no cover native_result = native_result.cast(pl.UInt32()) - return self._expr._from_native_expr(native_result) + return self._expr._with_native(native_result) # TODO(FBruzzesi): Remove `pragma: no cover` once other namespace methods are added def __getattr__( @@ -413,8 +330,8 @@ def __getattr__( ) -> Callable[[Any], PolarsExpr]: # pragma: no cover def func(*args: Any, **kwargs: Any) -> PolarsExpr: args, kwargs = extract_args_kwargs(args, kwargs) # type: ignore[assignment] - return self._expr._from_native_expr( - getattr(self._expr._native_expr.list, attr)(*args, **kwargs) + return self._expr._with_native( + getattr(self._expr.native.list, attr)(*args, **kwargs) ) return func @@ -429,8 +346,8 @@ def __getattr__( ) -> Callable[[Any], PolarsExpr]: # pragma: no cover def func(*args: Any, **kwargs: Any) -> PolarsExpr: args, kwargs = extract_args_kwargs(args, kwargs) # type: ignore[assignment] - return self._expr._from_native_expr( - getattr(self._expr._native_expr.struct, attr)(*args, **kwargs) + return self._expr._with_native( + getattr(self._expr.native.struct, attr)(*args, **kwargs) ) return func