refactor(perf): Side-port (#3435) {min,max}_horizontal fix

See #3435 (comment)

Also
- tweaked some typing
- added docs
- included `coalesce` in the same path

Author: dangotbanned

narwhals/_plan/arrow/functions/__init__.py

@@ -95,7 +95,11 @@
     dtype_native,
     string_type,
 )
-from narwhals._plan.arrow.functions._horizontal import max_horizontal, min_horizontal
+from narwhals._plan.arrow.functions._horizontal import (
+    coalesce,
+    max_horizontal,
+    min_horizontal,
+)
 from narwhals._plan.arrow.functions._lists import ExplodeBuilder
 from narwhals._plan.arrow.functions._multiplex import (
     fill_nan,
@@ -161,6 +165,7 @@
     "clip",
     "clip_lower",
     "clip_upper",
+    "coalesce",
     "concat_horizontal",
     "concat_tables",
     "concat_vertical",

narwhals/_plan/arrow/functions/_horizontal.py

@@ -2,14 +2,16 @@
 
 import pyarrow.compute as pc  # ignore-banned-import
 
-__all__ = ["max_horizontal", "min_horizontal"]
+__all__ = ["coalesce", "max_horizontal", "min_horizontal"]
 
 # TODO @dangotbanned: Wrap horizontal functions with correct typing
 # Should only return scalar if all elements are as well
 # NOTE: Changing typing will propagate to a lot of places (so be careful!):
 # - `_round.{clip,clip_lower,clip_upper}`
 # - `acero.join_asof_tables`
 # - `ArrowNamespace.{min,max}_horizontal`
+# - `ArrowNamespace.coalesce`
 # - `ArrowSeries.rolling_var`
 min_horizontal = pc.min_element_wise
 max_horizontal = pc.max_element_wise
+coalesce = pc.coalesce

narwhals/_plan/arrow/namespace.py

@@ -5,7 +5,6 @@
 from typing import TYPE_CHECKING, Any, Literal, cast, overload
 
 import pyarrow as pa  # ignore-banned-import
-import pyarrow.compute as pc  # ignore-banned-import
 
 from narwhals._arrow.utils import narwhals_to_native_dtype
 from narwhals._plan._guards import is_tuple_of
@@ -18,13 +17,21 @@
 from narwhals.exceptions import InvalidOperationError
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterable, Iterator, Sequence
+    from collections.abc import Iterable, Iterator, Sequence
+
+    from typing_extensions import TypeAlias
 
     from narwhals._arrow.typing import ChunkedArrayAny
+    from narwhals._plan._dispatch import BoundMethod
     from narwhals._plan.arrow.dataframe import ArrowDataFrame as Frame
     from narwhals._plan.arrow.expr import ArrowExpr as Expr, ArrowScalar as Scalar
     from narwhals._plan.arrow.series import ArrowSeries as Series
-    from narwhals._plan.arrow.typing import ChunkedArray, IntegerScalar
+    from narwhals._plan.arrow.typing import (
+        BinaryFunction,
+        ChunkedArray,
+        IntegerScalar,
+        VariadicFunction,
+    )
     from narwhals._plan.expressions import expr, functions as F
     from narwhals._plan.expressions.boolean import AllHorizontal, AnyHorizontal
     from narwhals._plan.expressions.expr import FunctionExpr as FExpr, RangeExpr
@@ -40,6 +47,8 @@
         PythonLiteral,
     )
 
+    Wrapper: TypeAlias = BoundMethod[FExpr[Any], Frame, Expr | Scalar]
+
 
 Int64 = Version.MAIN.dtypes.Int64()
 
@@ -102,53 +111,71 @@ def lit(
             nw_ser.to_native(), name or node.name, nw_ser.version
         )
 
-    def _horizontal_function(
-        self, fn_native: Callable[[Any, Any], Any], /, fill: NonNestedLiteral = None
-    ) -> Callable[[FExpr[Any], Frame, str], Expr | Scalar]:
+    @overload
+    def _horizontal(
+        self, function: BinaryFunction, /, fill: NonNestedLiteral = None
+    ) -> Wrapper: ...
+    @overload
+    def _horizontal(
+        self, function: VariadicFunction, /, *, variadic: Literal[True]
+    ) -> Wrapper: ...
+    def _horizontal(
+        self,
+        function: BinaryFunction | VariadicFunction,
+        /,
+        fill: NonNestedLiteral = None,
+        *,
+        variadic: bool = False,
+    ) -> Wrapper:
+        """Generate a horizontal wrapper function.
+
+        Arguments:
+            function: Native binary or variadic function.
+            fill: Fill value to use when nulls should *not* be ignored.
+            variadic: If False (default), perform a binary reduction.
+                Otherwise, assume we can unpack directly into `function`.
+        """
+
         def func(node: FExpr[Any], frame: Frame, name: str) -> Expr | Scalar:
             it = (self._expr.from_ir(e, frame, name).native for e in node.input)
             if fill is not None:
                 it = (fn.fill_null(native, fill) for native in it)
-            result = reduce(fn_native, it)
+            result = function(*it) if variadic else reduce(function, it)
             if isinstance(result, pa.Scalar):
                 return self._scalar.from_native(result, name, self.version)
             return self._expr.from_native(result, name, self.version)
 
         return func
 
     def coalesce(self, node: FExpr[F.Coalesce], frame: Frame, name: str) -> Expr | Scalar:
-        it = (self._expr.from_ir(e, frame, name).native for e in node.input)
-        result = pc.coalesce(*it)
-        if isinstance(result, pa.Scalar):
-            return self._scalar.from_native(result, name, self.version)
-        return self._expr.from_native(result, name, self.version)
+        return self._horizontal(fn.coalesce, variadic=True)(node, frame, name)
 
     def any_horizontal(
         self, node: FExpr[AnyHorizontal], frame: Frame, name: str
     ) -> Expr | Scalar:
         fill = False if node.function.ignore_nulls else None
-        return self._horizontal_function(fn.or_, fill)(node, frame, name)
+        return self._horizontal(fn.or_, fill)(node, frame, name)
 
     def all_horizontal(
         self, node: FExpr[AllHorizontal], frame: Frame, name: str
     ) -> Expr | Scalar:
         fill = True if node.function.ignore_nulls else None
-        return self._horizontal_function(fn.and_, fill)(node, frame, name)
+        return self._horizontal(fn.and_, fill)(node, frame, name)
 
     def sum_horizontal(
         self, node: FExpr[F.SumHorizontal], frame: Frame, name: str
     ) -> Expr | Scalar:
-        return self._horizontal_function(fn.add, fill=0)(node, frame, name)
+        return self._horizontal(fn.add, fill=0)(node, frame, name)
 
     def min_horizontal(
         self, node: FExpr[F.MinHorizontal], frame: Frame, name: str
     ) -> Expr | Scalar:
-        return self._horizontal_function(fn.min_horizontal)(node, frame, name)
+        return self._horizontal(fn.min_horizontal, variadic=True)(node, frame, name)
 
     def max_horizontal(
         self, node: FExpr[F.MaxHorizontal], frame: Frame, name: str
     ) -> Expr | Scalar:
-        return self._horizontal_function(fn.max_horizontal)(node, frame, name)
+        return self._horizontal(fn.max_horizontal, variadic=True)(node, frame, name)
 
     def mean_horizontal(
         self, node: FExpr[F.MeanHorizontal], frame: Frame, name: str

narwhals/_plan/arrow/typing.py

@@ -206,6 +206,12 @@ class BinaryComp(
 class BinaryLogical(BinaryFunction["BooleanScalar", "pa.BooleanScalar"], Protocol): ...
 
 
+# TODO @dangotbanned: Use stricter typing & revisit with
+# https://github.com/narwhals-dev/narwhals/blob/0d81b8b8cd1d24a68d360e6f8cf742b45cc2bdec/narwhals/_plan/arrow/functions/_horizontal.py#L7-L15
+class VariadicFunction(Protocol):
+    def __call__(self, *args: Arrow) -> Any: ...
+
+
 BinaryNumericTemporal: TypeAlias = BinaryFunction[
     NumericOrTemporalScalarT, NumericOrTemporalScalarT
 ]