perf: avoid full broadcast in horizontal functions (#3199)

Author: MarcoGorelli

narwhals/_arrow/namespace.py

@@ -23,7 +23,12 @@
 if TYPE_CHECKING:
     from collections.abc import Iterator, Sequence
 
-    from narwhals._arrow.typing import ArrayOrScalar, ChunkedArrayAny, Incomplete
+    from narwhals._arrow.typing import (
+        ArrayOrScalar,
+        ChunkedArrayAny,
+        Incomplete,
+        ScalarAny,
+    )
     from narwhals._compliant.typing import ScalarKwargs
     from narwhals._utils import Version
     from narwhals.typing import IntoDType, NonNestedLiteral
@@ -49,6 +54,11 @@ def _series(self) -> type[ArrowSeries]:
     def __init__(self, *, version: Version) -> None:
         self._version = version
 
+    def extract_native(
+        self, *series: ArrowSeries
+    ) -> Iterator[ChunkedArrayAny | ScalarAny]:
+        return (s.native[0] if s._broadcast else s.native for s in series)
+
     def len(self) -> ArrowExpr:
         # coverage bug? this is definitely hit
         return self._expr(  # pragma: no cover
@@ -83,10 +93,9 @@ def _lit_arrow_series(_: ArrowDataFrame) -> ArrowSeries:
     def all_horizontal(self, *exprs: ArrowExpr, ignore_nulls: bool) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             series: Iterator[ArrowSeries] = chain.from_iterable(e(df) for e in exprs)
-            align = self._series._align_full_broadcast
             if ignore_nulls:
                 series = (s.fill_null(True, None, None) for s in series)
-            return [reduce(operator.and_, align(*series))]
+            return [reduce(operator.and_, series)]
 
         return self._expr._from_callable(
             func=func,
@@ -100,10 +109,9 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
     def any_horizontal(self, *exprs: ArrowExpr, ignore_nulls: bool) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             series: Iterator[ArrowSeries] = chain.from_iterable(e(df) for e in exprs)
-            align = self._series._align_full_broadcast
             if ignore_nulls:
                 series = (s.fill_null(False, None, None) for s in series)
-            return [reduce(operator.or_, align(*series))]
+            return [reduce(operator.or_, series)]
 
         return self._expr._from_callable(
             func=func,
@@ -118,8 +126,7 @@ def sum_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             it = chain.from_iterable(expr(df) for expr in exprs)
             series = (s.fill_null(0, strategy=None, limit=None) for s in it)
-            align = self._series._align_full_broadcast
-            return [reduce(operator.add, align(*series))]
+            return [reduce(operator.add, series)]
 
         return self._expr._from_callable(
             func=func,
@@ -135,11 +142,8 @@ def mean_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr:
 
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             expr_results = tuple(chain.from_iterable(expr(df) for expr in exprs))
-            align = self._series._align_full_broadcast
-            series = align(
-                *(s.fill_null(0, strategy=None, limit=None) for s in expr_results)
-            )
-            non_na = align(*(1 - s.is_null().cast(int_64) for s in expr_results))
+            series = [s.fill_null(0, strategy=None, limit=None) for s in expr_results]
+            non_na = [1 - s.is_null().cast(int_64) for s in expr_results]
             return [reduce(operator.add, series) / reduce(operator.add, non_na)]
 
         return self._expr._from_callable(
@@ -153,9 +157,7 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
 
     def min_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            align = self._series._align_full_broadcast
             init_series, *series = tuple(chain.from_iterable(expr(df) for expr in exprs))
-            init_series, *series = align(init_series, *series)
             native_series = reduce(
                 pc.min_element_wise, [s.native for s in series], init_series.native
             )
@@ -174,9 +176,7 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
 
     def max_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            align = self._series._align_full_broadcast
             init_series, *series = tuple(chain.from_iterable(expr(df) for expr in exprs))
-            init_series, *series = align(init_series, *series)
             native_series = reduce(
                 pc.max_element_wise, [s.native for s in series], init_series.native
             )
@@ -227,16 +227,13 @@ def concat_str(
         self, *exprs: ArrowExpr, separator: str, ignore_nulls: bool
     ) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            align = self._series._align_full_broadcast
-            compliant_series_list = align(
-                *(chain.from_iterable(expr(df) for expr in exprs))
-            )
-            name = compliant_series_list[0].name
+            series = list(chain.from_iterable(expr(df) for expr in exprs))
+            name = series[0].name
             null_handling: Literal["skip", "emit_null"] = (
                 "skip" if ignore_nulls else "emit_null"
             )
             it, separator_scalar = cast_to_comparable_string_types(
-                *(s.native for s in compliant_series_list), separator=separator
+                *self.extract_native(*series), separator=separator
             )
             # NOTE: stubs indicate `separator` must also be a `ChunkedArray`
             # Reality: `str` is fine

narwhals/_arrow/series.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal, cast, overload
+from typing import TYPE_CHECKING, Any, Callable, Literal, cast, overload
 
 import pyarrow as pa
 import pyarrow.compute as pc
@@ -148,6 +148,16 @@ def _with_native(
             result._broadcast = self._broadcast
         return result
 
+    def _with_binary(self, op: Callable[..., ArrayOrScalar], other: Any) -> Self:
+        ser, other_native = extract_native(self, other)
+        preserve_broadcast = self._broadcast and getattr(other, "_broadcast", True)
+        return self._with_native(
+            op(ser, other_native), preserve_broadcast=preserve_broadcast
+        ).alias(self.name)
+
+    def _with_binary_right(self, op: Callable[..., ArrayOrScalar], other: Any) -> Self:
+        return self._with_binary(lambda x, y: op(y, x), other).alias(self.name)
+
     @classmethod
     def from_iterable(
         cls,
@@ -214,106 +224,89 @@ def __narwhals_namespace__(self) -> ArrowNamespace:
         return ArrowNamespace(version=self._version)
 
     def __eq__(self, other: object) -> Self:  # type: ignore[override]
-        other = cast("PythonLiteral | ArrowSeries | None", other)
-        ser, rhs = extract_native(self, other)
-        return self._with_native(pc.equal(ser, rhs))
+        return self._with_binary(pc.equal, other)
 
     def __ne__(self, other: object) -> Self:  # type: ignore[override]
-        other = cast("PythonLiteral | ArrowSeries | None", other)
-        ser, rhs = extract_native(self, other)
-        return self._with_native(pc.not_equal(ser, rhs))
+        return self._with_binary(pc.not_equal, other)
 
     def __ge__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.greater_equal(ser, other))
+        return self._with_binary(pc.greater_equal, other)
 
     def __gt__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.greater(ser, other))
+        return self._with_binary(pc.greater, other)
 
     def __le__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.less_equal(ser, other))
+        return self._with_binary(pc.less_equal, other)
 
     def __lt__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.less(ser, other))
+        return self._with_binary(pc.less, other)
 
     def __and__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.and_kleene(ser, other))  # type: ignore[arg-type]
+        return self._with_binary(pc.and_kleene, other)
 
     def __rand__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.and_kleene(other, ser))  # type: ignore[arg-type]
+        return self._with_binary_right(pc.and_kleene, other)
 
     def __or__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.or_kleene(ser, other))  # type: ignore[arg-type]
+        return self._with_binary_right(pc.or_kleene, other)
 
     def __ror__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.or_kleene(other, ser))  # type: ignore[arg-type]
+        return self._with_binary_right(pc.or_kleene, other)
 
     def __add__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.add(ser, other))
+        return self._with_binary(pc.add, other)
 
     def __radd__(self, other: Any) -> Self:
-        return self + other
+        return self._with_binary_right(pc.add, other)
 
     def __sub__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.subtract(ser, other))
+        return self._with_binary(pc.subtract, other)
 
     def __rsub__(self, other: Any) -> Self:
-        return (self - other) * (-1)
+        return self._with_binary_right(pc.subtract, other)
 
     def __mul__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.multiply(ser, other))
+        return self._with_binary(pc.multiply, other)
 
     def __rmul__(self, other: Any) -> Self:
-        return self * other
+        return self._with_binary_right(pc.multiply, other)
 
     def __pow__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.power(ser, other))
+        return self._with_binary(pc.power, other)
 
     def __rpow__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.power(other, ser))
+        return self._with_binary_right(pc.power, other)
 
     def __floordiv__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(floordiv_compat(ser, other))
+        return self._with_binary(floordiv_compat, other)
 
     def __rfloordiv__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(floordiv_compat(other, ser))
+        return self._with_binary_right(floordiv_compat, other)
 
     def __truediv__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.divide(*cast_for_truediv(ser, other)))  # type: ignore[type-var]
+        return self._with_binary(lambda x, y: pc.divide(*cast_for_truediv(x, y)), other)
 
     def __rtruediv__(self, other: Any) -> Self:
-        ser, other = extract_native(self, other)
-        return self._with_native(pc.divide(*cast_for_truediv(other, ser)))  # type: ignore[type-var]
+        return self._with_binary_right(
+            lambda x, y: pc.divide(*cast_for_truediv(x, y)), other
+        )
 
     def __mod__(self, other: Any) -> Self:
+        preserve_broadcast = self._broadcast and getattr(other, "_broadcast", True)
         floor_div = (self // other).native
         ser, other = extract_native(self, other)
         res = pc.subtract(ser, pc.multiply(floor_div, other))
-        return self._with_native(res)
+        return self._with_native(res, preserve_broadcast=preserve_broadcast)
 
     def __rmod__(self, other: Any) -> Self:
+        preserve_broadcast = self._broadcast and getattr(other, "_broadcast", True)
         floor_div = (other // self).native
         ser, other = extract_native(self, other)
         res = pc.subtract(other, pc.multiply(floor_div, ser))
-        return self._with_native(res)
+        return self._with_native(res, preserve_broadcast=preserve_broadcast)
 
     def __invert__(self) -> Self:
-        return self._with_native(pc.invert(self.native))
+        return self._with_native(pc.invert(self.native), preserve_broadcast=True)
 
     @property
     def _type(self) -> pa.DataType:

narwhals/_arrow/utils.py

@@ -445,8 +445,8 @@ def pad_series(
 
 
 def cast_to_comparable_string_types(
-    *chunked_arrays: ChunkedArrayAny, separator: str
-) -> tuple[Iterator[ChunkedArrayAny], ScalarAny]:
+    *chunked_arrays: ChunkedArrayAny | ScalarAny, separator: str
+) -> tuple[Iterator[ChunkedArrayAny | ScalarAny], ScalarAny]:
     # Ensure `chunked_arrays` are either all `string` or all `large_string`.
     dtype = (
         pa.string()  # (PyArrow default)

narwhals/_pandas_like/namespace.py

@@ -72,8 +72,7 @@ def __init__(self, implementation: Implementation, version: Version) -> None:
 
     def coalesce(self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
-            align = self._series._align_full_broadcast
-            series = align(*(s for _expr in exprs for s in _expr(df)))
+            series = (s for _expr in exprs for s in _expr(df))
             return [
                 reduce(lambda x, y: x.fill_null(y, strategy=None, limit=None), series)
             ]
@@ -127,10 +126,8 @@ def len(self) -> PandasLikeExpr:
     # --- horizontal ---
     def sum_horizontal(self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
-            align = self._series._align_full_broadcast
             it = chain.from_iterable(expr(df) for expr in exprs)
-            series = align(*it)
-            native_series = (s.fill_null(0, None, None) for s in series)
+            native_series = (s.fill_null(0, None, None) for s in it)
             return [reduce(operator.add, native_series)]
 
         return self._expr._from_callable(
@@ -146,7 +143,6 @@ def all_horizontal(
         self, *exprs: PandasLikeExpr, ignore_nulls: bool
     ) -> PandasLikeExpr:
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
-            align = self._series._align_full_broadcast
             series = [s for _expr in exprs for s in _expr(df)]
             if not ignore_nulls and any(
                 s.native.dtype == "object" and s.is_null().any() for s in series
@@ -164,7 +160,7 @@ def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
                 if ignore_nulls
                 else iter(series)
             )
-            return [reduce(operator.and_, align(*it))]
+            return [reduce(operator.and_, it)]
 
         return self._expr._from_callable(
             func=func,
@@ -179,7 +175,6 @@ def any_horizontal(
         self, *exprs: PandasLikeExpr, ignore_nulls: bool
     ) -> PandasLikeExpr:
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
-            align = self._series._align_full_broadcast
             series = [s for _expr in exprs for s in _expr(df)]
             if not ignore_nulls and any(
                 s.native.dtype == "object" and s.is_null().any() for s in series
@@ -197,7 +192,7 @@ def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
                 if ignore_nulls
                 else iter(series)
             )
-            return [reduce(operator.or_, align(*it))]
+            return [reduce(operator.or_, it)]
 
         return self._expr._from_callable(
             func=func,
@@ -211,11 +206,8 @@ def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
     def mean_horizontal(self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
             expr_results = [s for _expr in exprs for s in _expr(df)]
-            align = self._series._align_full_broadcast
-            series = align(
-                *(s.fill_null(0, strategy=None, limit=None) for s in expr_results)
-            )
-            non_na = align(*(1 - s.is_null() for s in expr_results))
+            series = (s.fill_null(0, strategy=None, limit=None) for s in expr_results)
+            non_na = (1 - s.is_null() for s in expr_results)
             return [reduce(operator.add, series) / reduce(operator.add, non_na)]
 
         return self._expr._from_callable(
@@ -229,10 +221,7 @@ def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
 
     def min_horizontal(self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
-            it = chain.from_iterable(expr(df) for expr in exprs)
-            align = self._series._align_full_broadcast
-            series = align(*it)
-
+            series = list(chain.from_iterable(expr(df) for expr in exprs))
             return [
                 PandasLikeSeries(
                     self.concat(
@@ -254,10 +243,7 @@ def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
 
     def max_horizontal(self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
-            it = chain.from_iterable(expr(df) for expr in exprs)
-            align = self._series._align_full_broadcast
-            series = align(*it)
-
+            series = list(chain.from_iterable(expr(df) for expr in exprs))
             return [
                 PandasLikeSeries(
                     self.concat(