From 09661a4e5d998ea4bc5dd61a7178ce5e84b6ecac Mon Sep 17 00:00:00 2001 From: Cangyuan Li Date: Sun, 27 Jul 2025 22:34:37 -0700 Subject: [PATCH 1/6] add method version of binary ops --- narwhals/expr.py | 160 +++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 149 insertions(+), 11 deletions(-) diff --git a/narwhals/expr.py b/narwhals/expr.py index 53a4c09c63..0366182336 100644 --- a/narwhals/expr.py +++ b/narwhals/expr.py @@ -1,6 +1,8 @@ from __future__ import annotations +import functools import math +import operator from collections.abc import Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Any, Callable @@ -85,7 +87,9 @@ def _with_window(self, to_compliant_expr: Callable[[Any], Any]) -> Self: def _with_filtration(self, to_compliant_expr: Callable[[Any], Any]) -> Self: return self.__class__(to_compliant_expr, self._metadata.with_filtration()) - def _with_orderable_filtration(self, to_compliant_expr: Callable[[Any], Any]) -> Self: + def _with_orderable_filtration( + self, to_compliant_expr: Callable[[Any], Any] + ) -> Self: return self.__class__( to_compliant_expr, self._metadata.with_orderable_filtration() ) @@ -201,6 +205,14 @@ def __eq__(self, other: Self | Any) -> Self: # type: ignore[override] ExprMetadata.from_binary_op(self, other), ) + def eq(self, other: Self | Any) -> Self: + """Method equivalent of equality operator `expr == other`. + + Arguments: + other: A literal or expression value to compare against. + """ + return self.__eq__(other) + def __ne__(self, other: Self | Any) -> Self: # type: ignore[override] return self.__class__( lambda plx: apply_n_ary_operation( @@ -209,6 +221,14 @@ def __ne__(self, other: Self | Any) -> Self: # type: ignore[override] ExprMetadata.from_binary_op(self, other), ) + def ne(self, other: Self | Any) -> Self: + """Method equivalent of equality operator `expr != other`. + + Arguments: + other: A literal or expression value to compare against. + """ + return self.__ne__(other) + def __and__(self, other: Any) -> Self: return self.__class__( lambda plx: apply_n_ary_operation( @@ -217,6 +237,14 @@ def __and__(self, other: Any) -> Self: ExprMetadata.from_binary_op(self, other), ) + def and_(self, *others: Any) -> Self: + """Method equivalent of bitwise "and" operator `expr & other & ...`. + + Arguments: + *others: One or more integer or boolean expressions to evaluate/combine. + """ + return functools.reduce(operator.and_, (self, *others)) # type: ignore[no-any-return] + def __rand__(self, other: Any) -> Self: return (self & other).alias("literal") # type: ignore[no-any-return] @@ -228,6 +256,14 @@ def __or__(self, other: Any) -> Self: ExprMetadata.from_binary_op(self, other), ) + def or_(self, *others: Any) -> Self: + """Method equivalent of bitwise "or" operator `expr | other | ...`. + + Arguments: + *others: One or more integer or boolean expressions to evaluate/combine. + """ + return functools.reduce(operator.or_, (self, *others)) # type: ignore[no-any-return] + def __ror__(self, other: Any) -> Self: return (self | other).alias("literal") # type: ignore[no-any-return] @@ -239,6 +275,14 @@ def __add__(self, other: Any) -> Self: ExprMetadata.from_binary_op(self, other), ) + def add(self, other: Any) -> Self: + """Method equivalent of addition operator `expr + other`. + + Arguments: + other: numeric or string value; accepts expression input. + """ + return self.__add__(other) + def __radd__(self, other: Any) -> Self: return (self + other).alias("literal") # type: ignore[no-any-return] @@ -250,6 +294,14 @@ def __sub__(self, other: Any) -> Self: ExprMetadata.from_binary_op(self, other), ) + def sub(self, other: Any) -> Self: + """Method equivalent of subtraction operator `expr - other`. + + Arguments: + other: Numeric literal or expression value. + """ + return self.__sub__(other) + def __rsub__(self, other: Any) -> Self: return self.__class__( lambda plx: apply_n_ary_operation( @@ -266,6 +318,14 @@ def __truediv__(self, other: Any) -> Self: ExprMetadata.from_binary_op(self, other), ) + def truediv(self, other: Any) -> Self: + """Method equivalent of division operator `expr / other`. + + Arguments: + other: Numeric literal or expression value. + """ + return self.__truediv__(other) + def __rtruediv__(self, other: Any) -> Self: return self.__class__( lambda plx: apply_n_ary_operation( @@ -282,6 +342,14 @@ def __mul__(self, other: Any) -> Self: ExprMetadata.from_binary_op(self, other), ) + def mul(self, other: Any) -> Self: + """Method equivalent of multiplication operator `expr * other`. + + Arguments: + other: Numeric literal or expression value. + """ + return self.__mul__(other) + def __rmul__(self, other: Any) -> Self: return (self * other).alias("literal") # type: ignore[no-any-return] @@ -293,6 +361,14 @@ def __le__(self, other: Any) -> Self: ExprMetadata.from_binary_op(self, other), ) + def le(self, other: Any) -> Self: + """Method equivalent of "less than or equal" operator `expr <= other`. + + Arguments: + other: A literal or expression value to compare with. + """ + return self.__le__(other) + def __lt__(self, other: Any) -> Self: return self.__class__( lambda plx: apply_n_ary_operation( @@ -301,6 +377,14 @@ def __lt__(self, other: Any) -> Self: ExprMetadata.from_binary_op(self, other), ) + def lt(self, other: Any) -> Self: + """Method equivalent of "less than" operator `expr < other`. + + Arguments: + other: A literal or expression value to compare with. + """ + return self.__lt__(other) + def __gt__(self, other: Any) -> Self: return self.__class__( lambda plx: apply_n_ary_operation( @@ -309,6 +393,14 @@ def __gt__(self, other: Any) -> Self: ExprMetadata.from_binary_op(self, other), ) + def gt(self, other: Any) -> Self: + """Method equivalent of "greater than" operator `expr > other`. + + Arguments: + other: A literal or expression value to compare with. + """ + return self.__gt__(other) + def __ge__(self, other: Any) -> Self: return self.__class__( lambda plx: apply_n_ary_operation( @@ -317,6 +409,14 @@ def __ge__(self, other: Any) -> Self: ExprMetadata.from_binary_op(self, other), ) + def ge(self, other: Any) -> Self: + """Method equivalent of "greater than or equal" operator `expr >= other`. + + Arguments: + other: A literal or expression value to compare with. + """ + return self.__ge__(other) + def __pow__(self, other: Any) -> Self: return self.__class__( lambda plx: apply_n_ary_operation( @@ -325,6 +425,14 @@ def __pow__(self, other: Any) -> Self: ExprMetadata.from_binary_op(self, other), ) + def pow(self, other: Any) -> Self: + """Method equivalent of exponentiation operator `expr ** exponent`. + + Arguments: + other: Numeric literal or expression exponent value. + """ + return self.__pow__(other) + def __rpow__(self, other: Any) -> Self: return self.__class__( lambda plx: apply_n_ary_operation( @@ -341,6 +449,14 @@ def __floordiv__(self, other: Any) -> Self: ExprMetadata.from_binary_op(self, other), ) + def floordiv(self, other: Any) -> Self: + """Method equivalent of integer division operator `expr // other`. + + Arguments: + other: Numeric literal or expression value. + """ + return self.__floordiv__(other) + def __rfloordiv__(self, other: Any) -> Self: return self.__class__( lambda plx: apply_n_ary_operation( @@ -371,6 +487,10 @@ def __invert__(self) -> Self: lambda plx: self._to_compliant_expr(plx).__invert__() ) + def not_(self) -> Self: + """Method equivalent of inversion operator `~expr`.""" + return self.__invert__() + def any(self) -> Self: """Return whether any of the values in the column are `True`. @@ -703,7 +823,9 @@ def kurtosis(self) -> Self: | 0 -1.3 0.210657 | └──────────────────┘ """ - return self._with_aggregation(lambda plx: self._to_compliant_expr(plx).kurtosis()) + return self._with_aggregation( + lambda plx: self._to_compliant_expr(plx).kurtosis() + ) def sum(self) -> Expr: """Return the sum value. @@ -856,7 +978,9 @@ def n_unique(self) -> Self: | 0 5 3 | └──────────────────┘ """ - return self._with_aggregation(lambda plx: self._to_compliant_expr(plx).n_unique()) + return self._with_aggregation( + lambda plx: self._to_compliant_expr(plx).n_unique() + ) def unique(self) -> Self: """Return unique values of this expression. @@ -1079,7 +1203,9 @@ def replace_strict( """ if new is None: if not isinstance(old, Mapping): - msg = "`new` argument is required if `old` argument is not a Mapping type" + msg = ( + "`new` argument is required if `old` argument is not a Mapping type" + ) raise TypeError(msg) new = list(old.values()) @@ -1299,7 +1425,9 @@ def is_null(self) -> Self: |└───────┴────────┴───────────┴───────────┘| └──────────────────────────────────────────┘ """ - return self._with_elementwise(lambda plx: self._to_compliant_expr(plx).is_null()) + return self._with_elementwise( + lambda plx: self._to_compliant_expr(plx).is_null() + ) def is_nan(self) -> Self: """Indicate which values are NaN. @@ -1353,7 +1481,9 @@ def arg_true(self) -> Self: "See https://narwhals-dev.github.io/narwhals/backcompat/ for more information.\n" ) issue_deprecation_warning(msg, _version="1.23.0") - return self._with_filtration(lambda plx: self._to_compliant_expr(plx).arg_true()) + return self._with_filtration( + lambda plx: self._to_compliant_expr(plx).arg_true() + ) def fill_null( self, @@ -1449,9 +1579,11 @@ def fill_null( strategy=strategy, limit=limit, ), - self._metadata.with_orderable_window() - if strategy is not None - else self._metadata, + ( + self._metadata.with_orderable_window() + if strategy is not None + else self._metadata + ), ) # --- partial reduction --- @@ -2349,7 +2481,10 @@ def rolling_var( return self._with_orderable_window( lambda plx: self._to_compliant_expr(plx).rolling_var( - window_size=window_size, min_samples=min_samples, center=center, ddof=ddof + window_size=window_size, + min_samples=min_samples, + center=center, + ddof=ddof, ) ) @@ -2410,7 +2545,10 @@ def rolling_std( return self._with_orderable_window( lambda plx: self._to_compliant_expr(plx).rolling_std( - window_size=window_size, min_samples=min_samples, center=center, ddof=ddof + window_size=window_size, + min_samples=min_samples, + center=center, + ddof=ddof, ) ) From e00bd506372dee8872e0c9a621e5002bf73ef371 Mon Sep 17 00:00:00 2001 From: Cangyuan Li Date: Sun, 27 Jul 2025 22:34:52 -0700 Subject: [PATCH 2/6] replicate test_expr_binary for the method versions --- tests/expr_and_series/binary_test.py | 45 ++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/tests/expr_and_series/binary_test.py b/tests/expr_and_series/binary_test.py index 6140ead120..b6ba2729e6 100644 --- a/tests/expr_and_series/binary_test.py +++ b/tests/expr_and_series/binary_test.py @@ -47,3 +47,48 @@ def test_expr_binary(constructor: Constructor) -> None: "m": [1, 9, 4], } assert_equal_data(result, expected) + + +def test_expr_binary_method(constructor: Constructor) -> None: + if "dask" in str(constructor) and DASK_VERSION < (2024, 10): + pytest.skip() + data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8.0, 9.0]} + df_raw = constructor(data) + result = nw.from_native(df_raw).with_columns( + a=nw.lit(1).add(nw.lit(3).mul(nw.col("a"))).mul(nw.lit(1).truediv(nw.col("a"))), + b=nw.col("z").truediv(nw.lit(2).sub(nw.col("b"))), + c=nw.col("a").add(nw.col("b").truediv(2)), + d=nw.col("a").sub(nw.col("b")), + e=((nw.col("a").gt(nw.col("b"))).and_(nw.col("a").ge(nw.col("z")))).cast( + nw.Int64 + ), + f=( + (nw.col("a").lt(nw.col("b"))) + .or_(nw.col("a").le(nw.col("z"))) + .or_(nw.col("a").eq(1)) + ).cast(nw.Int64), + g=nw.col("a").ne(1), + h=(nw.lit(value=False).and_(nw.col("a").ne(1))), + i=(nw.lit(value=False).or_(nw.col("a").ne(1))), + j=nw.lit(2).pow(nw.col("a")), + k=nw.lit(2).floordiv(nw.col("a")), + l=nw.col("a").floordiv(2), + m=nw.col("a").pow(2), + ) + expected = { + "a": [4, 3.333333, 3.5], + "b": [-3.5, -4.0, -2.25], + "z": [7.0, 8.0, 9.0], + "c": [3, 5, 5], + "d": [-3, -1, -4], + "e": [0, 0, 0], + "f": [1, 1, 1], + "g": [False, True, True], + "h": [False, False, False], + "i": [False, True, True], + "j": [2, 8, 4], + "k": [2, 0, 1], + "l": [0, 1, 1], + "m": [1, 9, 4], + } + assert_equal_data(result, expected) From dcc87921e791b09bd57db00c4a74c5fc9058e842 Mon Sep 17 00:00:00 2001 From: Cangyuan Li Date: Sun, 27 Jul 2025 22:42:48 -0700 Subject: [PATCH 3/6] fix ruff errors --- narwhals/expr.py | 82 ++++++++++++++++++++++++++++++++---------------- 1 file changed, 55 insertions(+), 27 deletions(-) diff --git a/narwhals/expr.py b/narwhals/expr.py index 0366182336..509251da52 100644 --- a/narwhals/expr.py +++ b/narwhals/expr.py @@ -87,9 +87,7 @@ def _with_window(self, to_compliant_expr: Callable[[Any], Any]) -> Self: def _with_filtration(self, to_compliant_expr: Callable[[Any], Any]) -> Self: return self.__class__(to_compliant_expr, self._metadata.with_filtration()) - def _with_orderable_filtration( - self, to_compliant_expr: Callable[[Any], Any] - ) -> Self: + def _with_orderable_filtration(self, to_compliant_expr: Callable[[Any], Any]) -> Self: return self.__class__( to_compliant_expr, self._metadata.with_orderable_filtration() ) @@ -210,6 +208,9 @@ def eq(self, other: Self | Any) -> Self: Arguments: other: A literal or expression value to compare against. + + Returns: + A new expression. """ return self.__eq__(other) @@ -226,6 +227,9 @@ def ne(self, other: Self | Any) -> Self: Arguments: other: A literal or expression value to compare against. + + Returns: + A new expression. """ return self.__ne__(other) @@ -242,6 +246,9 @@ def and_(self, *others: Any) -> Self: Arguments: *others: One or more integer or boolean expressions to evaluate/combine. + + Returns: + A new expression. """ return functools.reduce(operator.and_, (self, *others)) # type: ignore[no-any-return] @@ -261,6 +268,9 @@ def or_(self, *others: Any) -> Self: Arguments: *others: One or more integer or boolean expressions to evaluate/combine. + + Returns: + A new expression. """ return functools.reduce(operator.or_, (self, *others)) # type: ignore[no-any-return] @@ -280,6 +290,9 @@ def add(self, other: Any) -> Self: Arguments: other: numeric or string value; accepts expression input. + + Returns: + A new expression. """ return self.__add__(other) @@ -299,6 +312,9 @@ def sub(self, other: Any) -> Self: Arguments: other: Numeric literal or expression value. + + Returns: + A new expression. """ return self.__sub__(other) @@ -323,6 +339,9 @@ def truediv(self, other: Any) -> Self: Arguments: other: Numeric literal or expression value. + + Returns: + A new expression. """ return self.__truediv__(other) @@ -347,6 +366,9 @@ def mul(self, other: Any) -> Self: Arguments: other: Numeric literal or expression value. + + Returns: + A new expression. """ return self.__mul__(other) @@ -366,6 +388,9 @@ def le(self, other: Any) -> Self: Arguments: other: A literal or expression value to compare with. + + Returns: + A new expression. """ return self.__le__(other) @@ -382,6 +407,9 @@ def lt(self, other: Any) -> Self: Arguments: other: A literal or expression value to compare with. + + Returns: + A new expression. """ return self.__lt__(other) @@ -398,6 +426,9 @@ def gt(self, other: Any) -> Self: Arguments: other: A literal or expression value to compare with. + + Returns: + A new expression. """ return self.__gt__(other) @@ -414,6 +445,9 @@ def ge(self, other: Any) -> Self: Arguments: other: A literal or expression value to compare with. + + Returns: + A new expression. """ return self.__ge__(other) @@ -430,6 +464,9 @@ def pow(self, other: Any) -> Self: Arguments: other: Numeric literal or expression exponent value. + + Returns: + A new expression. """ return self.__pow__(other) @@ -454,6 +491,9 @@ def floordiv(self, other: Any) -> Self: Arguments: other: Numeric literal or expression value. + + Returns: + A new expression. """ return self.__floordiv__(other) @@ -488,7 +528,11 @@ def __invert__(self) -> Self: ) def not_(self) -> Self: - """Method equivalent of inversion operator `~expr`.""" + """Method equivalent of inversion operator `~expr`. + + Returns: + A new expression. + """ return self.__invert__() def any(self) -> Self: @@ -823,9 +867,7 @@ def kurtosis(self) -> Self: | 0 -1.3 0.210657 | └──────────────────┘ """ - return self._with_aggregation( - lambda plx: self._to_compliant_expr(plx).kurtosis() - ) + return self._with_aggregation(lambda plx: self._to_compliant_expr(plx).kurtosis()) def sum(self) -> Expr: """Return the sum value. @@ -978,9 +1020,7 @@ def n_unique(self) -> Self: | 0 5 3 | └──────────────────┘ """ - return self._with_aggregation( - lambda plx: self._to_compliant_expr(plx).n_unique() - ) + return self._with_aggregation(lambda plx: self._to_compliant_expr(plx).n_unique()) def unique(self) -> Self: """Return unique values of this expression. @@ -1203,9 +1243,7 @@ def replace_strict( """ if new is None: if not isinstance(old, Mapping): - msg = ( - "`new` argument is required if `old` argument is not a Mapping type" - ) + msg = "`new` argument is required if `old` argument is not a Mapping type" raise TypeError(msg) new = list(old.values()) @@ -1425,9 +1463,7 @@ def is_null(self) -> Self: |└───────┴────────┴───────────┴───────────┘| └──────────────────────────────────────────┘ """ - return self._with_elementwise( - lambda plx: self._to_compliant_expr(plx).is_null() - ) + return self._with_elementwise(lambda plx: self._to_compliant_expr(plx).is_null()) def is_nan(self) -> Self: """Indicate which values are NaN. @@ -1481,9 +1517,7 @@ def arg_true(self) -> Self: "See https://narwhals-dev.github.io/narwhals/backcompat/ for more information.\n" ) issue_deprecation_warning(msg, _version="1.23.0") - return self._with_filtration( - lambda plx: self._to_compliant_expr(plx).arg_true() - ) + return self._with_filtration(lambda plx: self._to_compliant_expr(plx).arg_true()) def fill_null( self, @@ -2481,10 +2515,7 @@ def rolling_var( return self._with_orderable_window( lambda plx: self._to_compliant_expr(plx).rolling_var( - window_size=window_size, - min_samples=min_samples, - center=center, - ddof=ddof, + window_size=window_size, min_samples=min_samples, center=center, ddof=ddof ) ) @@ -2545,10 +2576,7 @@ def rolling_std( return self._with_orderable_window( lambda plx: self._to_compliant_expr(plx).rolling_std( - window_size=window_size, - min_samples=min_samples, - center=center, - ddof=ddof, + window_size=window_size, min_samples=min_samples, center=center, ddof=ddof ) ) From cffc94706d87688a92b7cf50b5655698589f0040 Mon Sep 17 00:00:00 2001 From: Cangyuan Li Date: Sun, 27 Jul 2025 22:48:07 -0700 Subject: [PATCH 4/6] add new methods to api reference --- docs/api-reference/expr.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/api-reference/expr.md b/docs/api-reference/expr.md index f278bc3b2c..9e20efb610 100644 --- a/docs/api-reference/expr.md +++ b/docs/api-reference/expr.md @@ -5,8 +5,10 @@ options: members: - abs + - add - alias - all + - and_ - any - arg_max - arg_min @@ -20,11 +22,15 @@ - cum_sum - diff - drop_nulls + - eq - ewm_mean - exp - fill_null - filter + - floordiv - gather_every + - ge + - gt - head - clip - is_between @@ -37,18 +43,25 @@ - is_null - is_unique - kurtosis + - le - len - log + - lt - map_batches - max - mean - median - min - mode + - mul + - ne + - not_ - null_count - n_unique + - or_ - over - pipe + - pow - quantile - rank - replace_strict @@ -63,8 +76,10 @@ - skew - sqrt - std + - sub - sum - tail + - truediv - unique - var show_source: false From 87f51bc8d2e8ba3f03db1ada503dbe58f540e35c Mon Sep 17 00:00:00 2001 From: Cangyuan Li Date: Mon, 28 Jul 2025 16:27:07 -0700 Subject: [PATCH 5/6] mark methods as expr only --- utils/check_api_reference.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/utils/check_api_reference.py b/utils/check_api_reference.py index 34aa838a89..4b6de154da 100644 --- a/utils/check_api_reference.py +++ b/utils/check_api_reference.py @@ -26,6 +26,7 @@ def _is_public_method_or_property(obj: Any) -> bool: return ( isfunction(obj) or isinstance(obj, (MethodType, property)) ) and obj.__name__.startswith(LOWERCASE) + else: def _is_public_method_or_property(obj: Any) -> bool: @@ -60,7 +61,25 @@ def read_documented_members(source: str | Path) -> list[str]: ret = 0 NAMESPACES = {"dt", "str", "cat", "name", "list", "struct"} -EXPR_ONLY_METHODS = {"over", "map_batches"} +EXPR_ONLY_METHODS = { + "add", + "and_", + "eq", + "floordiv", + "ge", + "gt", + "le", + "lt", + "mul", + "ne", + "not_", + "or_", + "pow", + "sub", + "truediv", + "over", + "map_batches", +} SERIES_ONLY_METHODS = { "dtype", "implementation", From e2b075fe014877d9203c314b4028dab2dbdeeaba Mon Sep 17 00:00:00 2001 From: Cangyuan Li Date: Mon, 28 Jul 2025 16:35:48 -0700 Subject: [PATCH 6/6] fix bad merge that resulted in outdated api reference --- docs/api-reference/expr.md | 3 --- 1 file changed, 3 deletions(-) diff --git a/docs/api-reference/expr.md b/docs/api-reference/expr.md index 6aec5544bc..fc463f96dd 100644 --- a/docs/api-reference/expr.md +++ b/docs/api-reference/expr.md @@ -25,10 +25,8 @@ - fill_null - filter - floordiv - - gather_every - ge - gt - - head - clip - is_between - is_duplicated @@ -73,7 +71,6 @@ - std - sub - sum - - tail - truediv - unique - var