feat: add list aggregate methods

docs/api-reference/expr_list.md

@@ -7,6 +7,11 @@
         - contains
         - get
         - len
+        - max
+        - mean
+        - median
+        - min
+        - sum
         - unique
       show_source: false
       show_bases: false

docs/api-reference/series_list.md

@@ -7,6 +7,11 @@
         - contains
         - get
         - len
+        - max
+        - mean
+        - median
+        - min
+        - sum
         - unique
       show_source: false
       show_bases: false

narwhals/_arrow/series_list.py

@@ -5,7 +5,7 @@
 import pyarrow as pa
 import pyarrow.compute as pc
 
-from narwhals._arrow.utils import ArrowSeriesNamespace
+from narwhals._arrow.utils import ArrowSeriesNamespace, list_agg
 from narwhals._compliant.any_namespace import ListNamespace
 from narwhals._utils import not_implemented
 
@@ -20,5 +20,20 @@ def len(self) -> ArrowSeries:
     def get(self, index: int) -> ArrowSeries:
         return self.with_native(pc.list_element(self.native, index))
 
+    def min(self) -> ArrowSeries:
+        return self.with_native(list_agg(self.native, "min"))
+
+    def max(self) -> ArrowSeries:
+        return self.with_native(list_agg(self.native, "max"))
+
+    def mean(self) -> ArrowSeries:
+        return self.with_native(list_agg(self.native, "mean"))
+
+    def median(self) -> ArrowSeries:
+        return self.with_native(list_agg(self.native, "approximate_median"))
+
+    def sum(self) -> ArrowSeries:
+        return self.with_native(list_agg(self.native, "sum"))
+
     unique = not_implemented()
     contains = not_implemented()

narwhals/_arrow/utils.py

@@ -11,6 +11,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Iterator, Mapping
+    from typing import Literal
 
     from typing_extensions import TypeAlias, TypeIs
 
@@ -494,3 +495,40 @@ def arange(start: int, end: int, step: int) -> ArrayAny:
         return pa.array(np.arange(start, end, step))
     # NOTE: Added in https://github.com/apache/arrow/pull/46778
     return pa.arange(start, end, step)  # type: ignore[attr-defined]
+
+
+def list_agg(
+    array: ChunkedArrayAny,
+    func: Literal["min", "max", "mean", "approximate_median", "sum"],
+) -> ChunkedArrayAny:
+    lit_: Incomplete = lit
+    aggregation = (
+        ("values", func, pc.ScalarAggregateOptions(min_count=0))
+        if func == "sum"
+        else ("values", func)
+    )
+    agg = pa.array(
+        pa.Table.from_arrays(
+            [pc.list_flatten(array), pc.list_parent_indices(array)],
+            names=["values", "offsets"],
+        )
+        .group_by("offsets")
+        .aggregate([aggregation])
+        .sort_by("offsets")
+        .column(f"values_{func}")
+    )
+    non_empty_mask = pa.array(pc.not_equal(pc.list_value_length(array), lit(0)))
+    if func == "sum":
+        # Make sure sum of empty list is 0.
+        base_array = pc.if_else(non_empty_mask.is_null(), None, 0)
+    else:
+        base_array = pa.repeat(lit_(None, type=agg.type), len(array))
+    return pa.chunked_array(
+        [
+            pc.replace_with_mask(
+                base_array,
+                non_empty_mask.fill_null(False),  # type: ignore[arg-type]
+                agg,
+            )
+        ]
+    )

narwhals/_compliant/any_namespace.py

@@ -70,6 +70,11 @@ def get(self, index: int) -> CompliantT_co: ...
     def len(self) -> CompliantT_co: ...
     def unique(self) -> CompliantT_co: ...
     def contains(self, item: NonNestedLiteral) -> CompliantT_co: ...
+    def min(self) -> CompliantT_co: ...
+    def max(self) -> CompliantT_co: ...
+    def mean(self) -> CompliantT_co: ...
+    def median(self) -> CompliantT_co: ...
+    def sum(self) -> CompliantT_co: ...
 
 
 class NameNamespace(_StoresCompliant[CompliantT_co], Protocol[CompliantT_co]):

narwhals/_compliant/expr.py

@@ -999,6 +999,21 @@ def contains(self, item: NonNestedLiteral) -> EagerExprT:
     def get(self, index: int) -> EagerExprT:
         return self.compliant._reuse_series_namespace("list", "get", index=index)
 
+    def min(self) -> EagerExprT:
+        return self.compliant._reuse_series_namespace("list", "min")
+
+    def max(self) -> EagerExprT:
+        return self.compliant._reuse_series_namespace("list", "max")
+
+    def mean(self) -> EagerExprT:
+        return self.compliant._reuse_series_namespace("list", "mean")
+
+    def median(self) -> EagerExprT:
+        return self.compliant._reuse_series_namespace("list", "median")
+
+    def sum(self) -> EagerExprT:
+        return self.compliant._reuse_series_namespace("list", "sum")
+
 
 class CompliantExprNameNamespace(  # type: ignore[misc]
     _ExprNamespace[CompliantExprT_co],

narwhals/_duckdb/expr_list.py

@@ -4,7 +4,7 @@
 
 from narwhals._compliant import LazyExprNamespace
 from narwhals._compliant.any_namespace import ListNamespace
-from narwhals._duckdb.utils import F, lit, when
+from narwhals._duckdb.utils import F, col, lambda_expr, lit, when
 from narwhals._utils import requires
 
 if TYPE_CHECKING:
@@ -40,3 +40,27 @@ def get(self, index: int) -> DuckDBExpr:
         return self.compliant._with_elementwise(
             lambda expr: F("list_extract", expr, lit(index + 1))
         )
+
+    def min(self) -> DuckDBExpr:
+        return self.compliant._with_elementwise(lambda expr: F("list_min", expr))
+
+    def max(self) -> DuckDBExpr:
+        return self.compliant._with_elementwise(lambda expr: F("list_max", expr))
+
+    def mean(self) -> DuckDBExpr:
+        return self.compliant._with_elementwise(lambda expr: F("list_avg", expr))
+
+    def median(self) -> DuckDBExpr:
+        return self.compliant._with_elementwise(lambda expr: F("list_median", expr))
+
+    @requires.backend_version((1, 2))
+    def sum(self) -> DuckDBExpr:
+        def func(expr: Expression) -> Expression:
+            elem = col("_")
+            expr_no_nulls = F("list_filter", expr, lambda_expr(elem, elem.isnotnull()))
+            expr_sum = F("list_sum", expr_no_nulls)
+            return when(F("array_length", expr_no_nulls) == lit(0), lit(0)).otherwise(
+                expr_sum
+            )
+
+        return self.compliant._with_callable(func)

narwhals/_ibis/expr_list.py

@@ -2,8 +2,11 @@
 
 from typing import TYPE_CHECKING
 
+from ibis import cases, literal
+
 from narwhals._compliant import LazyExprNamespace
 from narwhals._compliant.any_namespace import ListNamespace
+from narwhals._utils import not_implemented
 
 if TYPE_CHECKING:
     import ibis.expr.types as ir
@@ -27,3 +30,26 @@ def _get(expr: ir.ArrayColumn) -> ir.Column:
             return expr[index]
 
         return self.compliant._with_callable(_get)
+
+    def min(self) -> IbisExpr:
+        return self.compliant._with_callable(lambda expr: expr.mins())
+
+    def max(self) -> IbisExpr:
+        return self.compliant._with_callable(lambda expr: expr.maxs())
+
+    def mean(self) -> IbisExpr:
+        return self.compliant._with_callable(lambda expr: expr.means())
+
+    def sum(self) -> IbisExpr:
+        def func(expr: ir.ArrayColumn) -> ir.Value:
+            expr_no_nulls = expr.filter(lambda x: x.notnull())
+            len = expr_no_nulls.length()
+            return cases(
+                (len.isnull(), literal(None)),
+                (len == literal(0), literal(0)),
+                else_=expr.sums(),
+            )
+
+        return self.compliant._with_callable(func)
+
+    median = not_implemented()

narwhals/_pandas_like/series_list.py

@@ -11,6 +11,8 @@
 from narwhals._utils import not_implemented
 
 if TYPE_CHECKING:
+    from typing import Literal
+
     from narwhals._pandas_like.series import PandasLikeSeries
 
 
@@ -40,3 +42,41 @@ def get(self, index: int) -> PandasLikeSeries:
         result = self.native.list[index]
         result.name = self.native.name
         return self.with_native(result)
+
+    def _agg(
+        self, func: Literal["min", "max", "mean", "approximate_median", "sum"]
+    ) -> PandasLikeSeries:
+        dtype_backend = get_dtype_backend(
+            self.native.dtype, self.compliant._implementation
+        )
+        if dtype_backend != "pyarrow":  # pragma: no cover
+            msg = "Only pyarrow backend is currently supported."
+            raise NotImplementedError(msg)
+
+        from narwhals._arrow.utils import list_agg, native_to_narwhals_dtype
+
+        ca = self.native.array._pa_array
+        result_arr = list_agg(ca, func)
+        nw_dtype = native_to_narwhals_dtype(result_arr.type, self.version)
+        out_dtype = narwhals_to_native_dtype(
+            nw_dtype, "pyarrow", self.implementation, self.version
+        )
+        result_native = type(self.native)(
+            result_arr, dtype=out_dtype, index=self.native.index, name=self.native.name
+        )
+        return self.with_native(result_native)
+
+    def min(self) -> PandasLikeSeries:
+        return self._agg("min")
+
+    def max(self) -> PandasLikeSeries:
+        return self._agg("max")
+
+    def mean(self) -> PandasLikeSeries:
+        return self._agg("mean")
+
+    def median(self) -> PandasLikeSeries:
+        return self._agg("approximate_median")
+
+    def sum(self) -> PandasLikeSeries:
+        return self._agg("sum")

narwhals/_polars/utils.py

@@ -362,6 +362,16 @@ def len(self) -> CompliantT: ...
 
     unique: Method[CompliantT]
 
+    max: Method[CompliantT]
+
+    mean: Method[CompliantT]
+
+    median: Method[CompliantT]
+
+    min: Method[CompliantT]
+
+    sum: Method[CompliantT]
+
 
 class PolarsStructNamespace(PolarsAnyNamespace[CompliantT, NativeT_co]):
     _accessor: ClassVar[Accessor] = "struct"

narwhals/_spark_like/expr_list.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import operator
 from typing import TYPE_CHECKING
 
 from narwhals._compliant import LazyExprNamespace
@@ -33,3 +34,51 @@ def _get(expr: Column) -> Column:
             return expr.getItem(index)
 
         return self.compliant._with_elementwise(_get)
+
+    def min(self) -> SparkLikeExpr:
+        def func(expr: Column) -> Column:
+            F = self.compliant._F
+            return F.array_min(expr)
+
+        return self.compliant._with_elementwise(func)
+
+    def max(self) -> SparkLikeExpr:
+        def func(expr: Column) -> Column:
+            F = self.compliant._F
+            return F.array_max(F.array_compact(expr))
+
+        return self.compliant._with_elementwise(func)
+
+    def sum(self) -> SparkLikeExpr:
+        def func(expr: Column) -> Column:
+            F = self.compliant._F
+            return F.aggregate(F.array_compact(expr), F.lit(0.0), operator.add)
+
+        return self.compliant._with_elementwise(func)
+
+    def mean(self) -> SparkLikeExpr:
+        def func(expr: Column) -> Column:
+            F = self.compliant._F
+            return F.try_divide(
+                F.aggregate(F.array_compact(expr), F.lit(0.0), operator.add),
+                F.array_size(F.array_compact(expr)),
+            )
+
+        return self.compliant._with_elementwise(func)
+
+    def median(self) -> SparkLikeExpr:
+        def func(expr: Column) -> Column:  # pragma: no cover
+            # sqlframe issue: https://github.com/eakmanrq/sqlframe/issues/548
+            F = self.compliant._F
+            sorted_expr = F.array_compact(F.sort_array(expr))
+            size = F.array_size(sorted_expr)
+            mid_index = (size / 2).cast("int")
+            odd_case = sorted_expr[mid_index]
+            even_case = (sorted_expr[mid_index - 1] + sorted_expr[mid_index]) / 2
+            return (
+                F.when((size.isNull()) | (size == 0), F.lit(None))
+                .when(size % 2 == 1, odd_case)
+                .otherwise(even_case)
+            )
+
+        return self.compliant._with_elementwise(func)