feat(DRAFT): Porting (#3332)

Tried to keep everything as close to original as possible
Next step is simplifying everything and fixing `list.sum`

Author: dangotbanned

narwhals/_plan/arrow/expr.py

@@ -15,8 +15,14 @@
     is_seq_column,
 )
 from narwhals._plan.arrow import functions as fn
+from narwhals._plan.arrow.group_by import AggSpec
 from narwhals._plan.arrow.series import ArrowSeries as Series
-from narwhals._plan.arrow.typing import ChunkedOrScalarAny, NativeScalar, StoresNativeT_co
+from narwhals._plan.arrow.typing import (
+    ChunkedOrArrayAny,
+    ChunkedOrScalarAny,
+    NativeScalar,
+    StoresNativeT_co,
+)
 from narwhals._plan.common import temp
 from narwhals._plan.compliant.accessors import (
     ExprCatNamespace,
@@ -994,11 +1000,47 @@ def join(self, node: FExpr[lists.Join], frame: Frame, name: str) -> Expr | Scala
             )
         return self.with_native(result, name)
 
-    min = not_implemented()
-    max = not_implemented()
-    mean = not_implemented()
-    median = not_implemented()
-    sum = not_implemented()
+    def aggregate(
+        self, node: FExpr[lists.Aggregation], frame: Frame, name: str
+    ) -> Expr | Scalar:
+        previous = node.input[0].dispatch(self.compliant, frame, name)
+        func = node.function
+        if isinstance(previous, ArrowScalar):
+            msg = f"TODO: ArrowScalar.{func!r}"
+            raise NotImplementedError(msg)
+
+        native = previous.native
+        lists = native
+        # TODO @dangotbanned: Experiment with explode step
+        # These options are to mirror `main`, but setting them to `True` may simplify everything after?
+        builder = fn.ExplodeBuilder(empty_as_null=False, keep_nulls=False)
+        explode_w_idx = builder.explode_with_indices(lists)
+        idx, v = "idx", "values"
+        agg_result = (
+            AggSpec._from_agg(type(func), v)
+            .over(explode_w_idx, [idx])
+            .sort_by(idx)
+            .column(v)
+        )
+        dtype: pa.DataType = agg_result.type
+        non_empty_mask = fn.not_eq(fn.list_len(lists), fn.lit(0))
+        base_array: ChunkedOrArrayAny
+        if isinstance(func, ir.lists.Sum):
+            # Make sure sum of empty list is 0.
+            base_array = fn.when_then(fn.is_not_null(non_empty_mask), fn.lit(0, dtype))
+        else:
+            base_array = fn.repeat_unchecked(fn.lit(None, dtype), len(lists))
+        replaced = fn.replace_with_mask(
+            base_array, fn.fill_null(non_empty_mask, False), agg_result
+        )
+        result = fn.chunked_array(replaced)
+        return self.with_native(result, name)
+
+    min = aggregate
+    max = aggregate
+    mean = aggregate
+    median = aggregate
+    sum = aggregate
 
 
 class ArrowStringNamespace(

narwhals/_plan/arrow/functions.py

@@ -442,6 +442,17 @@ def explode(
         return chunked_array(_list_explode(safe))
 
     def explode_with_indices(self, native: ChunkedList | ListArray) -> pa.Table:
+        """Explode list elements, expanding one-level into a table indexing the origin.
+
+        Returns a 2-column table, with names `"idx"` and `"values"`:
+
+            >>> from narwhals._plan.arrow import functions as fn
+            >>>
+            >>> arr = fn.array([[1, 2, 3], None, [4, 5, 6], []])
+            >>> fn.ExplodeBuilder().explode_with_indices(arr).to_pydict()
+            {'idx': [0, 0, 0, 1, 2, 2, 2, 3], 'values': [1, 2, 3, None, 4, 5, 6, None]}
+            # ^ Which sublist we came from     ^ The exploded values themselves
+        """
         safe = self._fill_with_null(native) if self.options.any() else native
         arrays = [_list_parent_indices(safe), _list_explode(safe)]
         return concat_horizontal(arrays, ["idx", "values"])
@@ -1042,6 +1053,12 @@ def _str_zfill_compat(
     )
 
 
+@t.overload
+def when_then(
+    predicate: ChunkedArray[BooleanScalar], then: ScalarAny
+) -> ChunkedArrayAny: ...
+@t.overload
+def when_then(predicate: Array[BooleanScalar], then: ScalarAny) -> ArrayAny: ...
 @t.overload
 def when_then(
     predicate: Predicate, then: SameArrowT, otherwise: SameArrowT | None
@@ -1059,6 +1076,11 @@ def when_then(
 def when_then(
     predicate: Predicate, then: ArrowAny, otherwise: ArrowAny | NonNestedLiteral = None
 ) -> Incomplete:
+    """Thin wrapper around `pyarrow.compute.if_else`.
+
+    - Supports a 2-arg form, like `pl.when(...).then(...)`
+    - Accepts python literals, but only in the `otherwise` position
+    """
     if is_non_nested_literal(otherwise):
         otherwise = lit(otherwise, then.type)
     return pc.if_else(predicate, then, otherwise)

narwhals/_plan/arrow/group_by.py

@@ -13,7 +13,7 @@
 from narwhals._plan.common import temp
 from narwhals._plan.compliant.group_by import EagerDataFrameGroupBy
 from narwhals._plan.expressions import aggregation as agg
-from narwhals._utils import Implementation
+from narwhals._utils import Implementation, qualified_type_name
 from narwhals.exceptions import InvalidOperationError
 
 if TYPE_CHECKING:
@@ -51,6 +51,13 @@
     agg.Last: "hash_last",
     fn.MinMax: "hash_min_max",
 }
+SUPPORTED_LIST_AGG: Mapping[type[ir.lists.Aggregation], type[agg.AggExpr]] = {
+    ir.lists.Mean: agg.Mean,
+    ir.lists.Median: agg.Median,
+    ir.lists.Max: agg.Max,
+    ir.lists.Min: agg.Min,
+    ir.lists.Sum: agg.Sum,
+}
 SUPPORTED_IR: Mapping[type[ir.ExprIR], acero.Aggregation] = {
     ir.Len: "hash_count_all",
     ir.Column: "hash_list",  # `hash_aggregate` only
@@ -141,6 +148,14 @@ def from_expr_ir(cls, expr: ir.ExprIR, name: acero.OutputName) -> Self:
     def _from_function(cls, tp: type[ir.Function], name: str) -> Self:
         return cls(name, SUPPORTED_FUNCTION[tp], options.FUNCTION.get(tp), name)
 
+    @classmethod
+    def _from_agg(cls, tp: type[ir.lists.Aggregation | agg.AggExpr], name: str) -> Self:
+        tp_agg = SUPPORTED_LIST_AGG[tp] if issubclass(tp, ir.lists.ListFunction) else tp
+        if tp_agg in {agg.Std, agg.Var}:
+            msg = f"TODO: {qualified_type_name(agg)!r} needs access to `ddof`, so can't be passed in without an instance"
+            raise NotImplementedError(msg)
+        return cls(name, SUPPORTED_AGG[tp_agg], options.AGG.get(tp_agg), name)
+
     @classmethod
     def any(cls, name: str) -> Self:
         return cls._from_function(ir.boolean.Any, name)

narwhals/_plan/expressions/lists.py

@@ -11,7 +11,7 @@
 from narwhals.exceptions import InvalidOperationError
 
 if TYPE_CHECKING:
-    from typing_extensions import Self
+    from typing_extensions import Self, TypeAlias
 
     from narwhals._plan.expr import Expr
     from narwhals._plan.expressions import ExprIR, FunctionExpr as FExpr
@@ -45,6 +45,9 @@ def unwrap_input(self, node: FExpr[Self], /) -> tuple[ExprIR, ExprIR]:
         return expr, item
 
 
+Aggregation: TypeAlias = "Min | Max | Mean | Median | Sum"
+
+
 class IRListNamespace(IRNamespace):
     len: ClassVar = Len
     unique: ClassVar = Unique

tests/plan/list_agg_test.py

@@ -29,34 +29,32 @@ def data_median(data: Data) -> Data:
 cast_a = a.cast(nw.List(nw.Int32))
 
 
-XFAIL_NOT_IMPL = pytest.mark.xfail(
-    reason="TODO: ArrowExpr.list.<agg>", raises=NotImplementedError
-)
-
-
-@XFAIL_NOT_IMPL
 @pytest.mark.parametrize(
     ("exprs", "expected"),
     [
         (a.list.max(), {"a": [4, -1, None, None, None]}),
         (a.list.mean(), {"a": [2.75, -1, None, None, None]}),
         (a.list.min(), {"a": [2, -1, None, None, None]}),
-        (a.list.sum(), {"a": [11, -1, None, 0, 0]}),
+        pytest.param(
+            a.list.sum(),
+            {"a": [11, -1, None, 0, 0]},
+            marks=pytest.mark.xfail(
+                reason="Mismatch at index 3, key a: None != 0", raises=AssertionError
+            ),
+        ),
     ],
+    ids=["max", "mean", "min", "sum"],
 )
-def test_list_agg(
-    data: Data, exprs: OneOrIterable[nwp.Expr], expected: Data
-) -> None:  # pragma: no cover
+def test_list_agg(data: Data, exprs: OneOrIterable[nwp.Expr], expected: Data) -> None:
     df = dataframe(data).with_columns(cast_a)
     result = df.select(exprs)
     assert_equal_data(result, expected)
 
 
-@XFAIL_NOT_IMPL
 @pytest.mark.xfail(
     is_windows() and sys.version_info < (3, 10), reason="Old pyarrow windows bad?"
 )
-def test_list_median(data_median: Data) -> None:  # pragma: no cover
+def test_list_median(data_median: Data) -> None:
     df = dataframe(data_median).with_columns(cast_a)
     result = df.select(a.list.median())