feat: add list.sort

docs/api-reference/expr_list.md

@@ -11,6 +11,7 @@
         - mean
         - median
         - min
+        - sort
         - sum
         - unique
       show_source: false

docs/api-reference/series_list.md

@@ -11,6 +11,7 @@
         - mean
         - median
         - min
+        - sort
         - sum
         - unique
       show_source: 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, list_agg
+from narwhals._arrow.utils import ArrowSeriesNamespace, list_agg, list_sort
 from narwhals._compliant.any_namespace import ListNamespace
 from narwhals._utils import not_implemented
 
@@ -35,5 +35,10 @@ def median(self) -> ArrowSeries:
     def sum(self) -> ArrowSeries:
         return self.with_native(list_agg(self.native, "sum"))
 
+    def sort(self, *, descending: bool, nulls_last: bool) -> ArrowSeries:
+        return self.with_native(
+            list_sort(self.native, descending=descending, nulls_last=nulls_last)
+        )
+
     unique = not_implemented()
     contains = not_implemented()

narwhals/_arrow/utils.py

@@ -532,3 +532,34 @@ def list_agg(
             )
         ]
     )
+
+
+def list_sort(
+    array: ChunkedArrayAny, *, descending: bool, nulls_last: bool
+) -> ChunkedArrayAny:
+    sort_direction: Literal["ascending", "descending"] = (
+        "descending" if descending else "ascending"
+    )
+    nulls_position: Literal["at_start", "at_end"] = "at_end" if nulls_last else "at_start"
+    idx, v = "idx", "values"
+    len_gt_0 = pc.greater(pc.list_value_length(array), lit(0))
+    arange = pa.arange(0, len(array))  # type: ignore[attr-defined]
+    indexed = pa.Table.from_arrays([arange, array], names=[idx, v])
+    valid = indexed.filter(len_gt_0)
+    invalid = indexed.filter(pc.or_kleene(array.is_null(), pc.invert(len_gt_0)))
+    agg = pa.Table.from_arrays(
+        [pc.list_flatten(array), pc.list_parent_indices(array)], names=[v, idx]
+    )
+    sorted_indices = pc.sort_indices(
+        agg,
+        sort_keys=[(idx, "ascending"), (v, sort_direction)],
+        null_placement=nulls_position,
+    )
+    offsets = valid.column(v).combine_chunks().offsets  # type: ignore[attr-defined]
+    sorted_imploded = pa.ListArray.from_arrays(
+        offsets, pa.array(agg.take(sorted_indices).column(v))
+    )
+    valid_finished = pa.Table.from_arrays(
+        [valid.column(idx), sorted_imploded], names=[idx, v]
+    )
+    return pa.concat_tables([valid_finished, invalid]).sort_by(idx).column(v)

narwhals/_compliant/any_namespace.py

@@ -75,6 +75,7 @@ def max(self) -> CompliantT_co: ...
     def mean(self) -> CompliantT_co: ...
     def median(self) -> CompliantT_co: ...
     def sum(self) -> CompliantT_co: ...
+    def sort(self, *, descending: bool, nulls_last: bool) -> CompliantT_co: ...
 
 
 class NameNamespace(_StoresCompliant[CompliantT_co], Protocol[CompliantT_co]):

narwhals/_compliant/expr.py

@@ -1014,6 +1014,11 @@ def median(self) -> EagerExprT:
     def sum(self) -> EagerExprT:
         return self.compliant._reuse_series_namespace("list", "sum")
 
+    def sort(self, *, descending: bool, nulls_last: bool) -> EagerExprT:
+        return self.compliant._reuse_series_namespace(
+            "list", "sort", descending=descending, nulls_last=nulls_last
+        )
+
 
 class CompliantExprNameNamespace(  # type: ignore[misc]
     _ExprNamespace[CompliantExprT_co],

narwhals/_duckdb/expr_list.py

@@ -64,3 +64,10 @@ def func(expr: Expression) -> Expression:
             )
 
         return self.compliant._with_callable(func)
+
+    def sort(self, *, descending: bool, nulls_last: bool) -> DuckDBExpr:
+        sort_direction = "DESC" if descending else "ASC"
+        nulls_position = "NULLS LAST" if nulls_last else "NULLS FIRST"
+        return self.compliant._with_elementwise(
+            lambda expr: F("list_sort", expr, lit(sort_direction), lit(nulls_position))
+        )

narwhals/_ibis/expr_list.py

@@ -52,4 +52,18 @@ def func(expr: ir.ArrayColumn) -> ir.Value:
 
         return self.compliant._with_callable(func)
 
+    def sort(self, *, descending: bool, nulls_last: bool) -> IbisExpr:
+        if descending:
+            msg = "Descending sort is not currently supported for Ibis."
+            raise NotImplementedError(msg)
+
+        def func(expr: ir.ArrayColumn) -> ir.ArrayValue:
+            if nulls_last:
+                return expr.sort()
+            expr_no_nulls = expr.filter(lambda x: x.notnull())
+            expr_nulls = expr.filter(lambda x: x.isnull())
+            return expr_nulls.concat(expr_no_nulls.sort())
+
+        return self.compliant._with_callable(func)
+
     median = not_implemented()

narwhals/_pandas_like/series_list.py

@@ -34,10 +34,6 @@ def len(self) -> PandasLikeSeries:
         )
         return self.with_native(result.astype(dtype)).alias(self.native.name)
 
-    unique = not_implemented()
-
-    contains = not_implemented()
-
     def get(self, index: int) -> PandasLikeSeries:
         result = self.native.list[index]
         result.name = self.native.name
@@ -80,3 +76,27 @@ def median(self) -> PandasLikeSeries:
 
     def sum(self) -> PandasLikeSeries:
         return self._agg("sum")
+
+    def sort(self, *, descending: bool, nulls_last: bool) -> 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_sort, native_to_narwhals_dtype
+
+        ca = self.native.array._pa_array
+        result_arr = list_sort(ca, descending=descending, nulls_last=nulls_last)
+        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)
+
+    unique = not_implemented()
+    contains = not_implemented()

narwhals/_polars/utils.py

@@ -370,6 +370,8 @@ def len(self) -> CompliantT: ...
 
     min: Method[CompliantT]
 
+    sort: Method[CompliantT]
+
     sum: Method[CompliantT]
 
 

narwhals/_spark_like/expr_list.py

@@ -82,3 +82,15 @@ def func(expr: Column) -> Column:  # pragma: no cover
             )
 
         return self.compliant._with_elementwise(func)
+
+    def sort(self, *, descending: bool, nulls_last: bool) -> SparkLikeExpr:
+        def func(expr: Column) -> Column:
+            F = self.compliant._F
+            if not descending and nulls_last:
+                return F.array_sort(expr)
+            if descending and not nulls_last:  # pragma: no cover
+                # https://github.com/eakmanrq/sqlframe/issues/559
+                return F.reverse(F.array_sort(expr))
+            return F.sort_array(expr, asc=not descending)
+
+        return self.compliant._with_elementwise(func)

narwhals/expr_list.py

@@ -262,3 +262,39 @@ def sum(self) -> ExprT:
             └────────────────────────┘
         """
         return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "list.sum"))
+
+    def sort(self, *, descending: bool = False, nulls_last: bool = False) -> ExprT:
+        """Sort the lists of the expression.
+
+        Arguments:
+            descending: Sort in descending order.
+            nulls_last: Place null values last.
+
+        Examples:
+            >>> import duckdb
+            >>> import narwhals as nw
+            >>> df_native = duckdb.sql(
+            ...     "SELECT * FROM VALUES ([2, -1, 1]), ([3, -4, NULL]) df(a)"
+            ... )
+            >>> df = nw.from_native(df_native)
+            >>> df.with_columns(a_sorted=nw.col("a").list.sort())
+            ┌─────────────────────────────────┐
+            |       Narwhals LazyFrame        |
+            |---------------------------------|
+            |┌───────────────┬───────────────┐|
+            |│       a       │   a_sorted    │|
+            |│    int32[]    │    int32[]    │|
+            |├───────────────┼───────────────┤|
+            |│ [2, -1, 1]    │ [-1, 1, 2]    │|
+            |│ [3, -4, NULL] │ [NULL, -4, 3] │|
+            |└───────────────┴───────────────┘|
+            └─────────────────────────────────┘
+        """
+        return self._expr._append_node(
+            ExprNode(
+                ExprKind.ELEMENTWISE,
+                "list.sort",
+                descending=descending,
+                nulls_last=nulls_last,
+            )
+        )

narwhals/series_list.py

@@ -219,3 +219,29 @@ def sum(self) -> SeriesT:
         return self._narwhals_series._with_compliant(
             self._narwhals_series._compliant_series.list.sum()
         )
+
+    def sort(self, *, descending: bool = False, nulls_last: bool = False) -> SeriesT:
+        """Sort the lists of the series.
+
+        Arguments:
+            descending: Sort in descending order.
+            nulls_last: Place null values last.
+
+        Examples:
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> s_native = pl.Series([[2, -1, 1], [3, -4, None]])
+            >>> s = nw.from_native(s_native, series_only=True)
+            >>> s.list.sort().to_native()  # doctest: +NORMALIZE_WHITESPACE
+            shape: (2,)
+            Series: '' [list[i64]]
+            [
+                    [-1, 1, 2]
+                    [null, -4, 3]
+            ]
+        """
+        return self._narwhals_series._with_compliant(
+            self._narwhals_series._compliant_series.list.sort(
+                descending=descending, nulls_last=nulls_last
+            )
+        )

tests/expr_and_series/list/sort_test.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+import narwhals as nw
+from tests.utils import POLARS_VERSION, assert_equal_data
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from tests.utils import Constructor, ConstructorEager
+
+
+data = {"a": [[3, 2, 2, 4, -10, None, None], [-1], None, [None, None, None], []]}
+expected_desc_nulls_last = [
+    [4, 3, 2, 2, -10, None, None],
+    [-1],
+    None,
+    [None, None, None],
+    [],
+]
+expected_desc_nulls_first = [
+    [None, None, 4, 3, 2, 2, -10],
+    [-1],
+    None,
+    [None, None, None],
+    [],
+]
+expected_asc_nulls_last = [
+    [-10, 2, 2, 3, 4, None, None],
+    [-1],
+    None,
+    [None, None, None],
+    [],
+]
+expected_asc_nulls_first = [
+    [None, None, -10, 2, 2, 3, 4],
+    [-1],
+    None,
+    [None, None, None],
+    [],
+]
+
+
+def test_sort_expr(request: pytest.FixtureRequest, constructor: Constructor) -> None:
+    if any(backend in str(constructor) for backend in ("dask", "cudf")):
+        request.applymarker(pytest.mark.xfail)
+    if "sqlframe" in str(constructor):
+        # https://github.com/eakmanrq/sqlframe/issues/559
+        # https://github.com/eakmanrq/sqlframe/issues/560
+        request.applymarker(pytest.mark.xfail)
+    if "polars" in str(constructor) and POLARS_VERSION < (0, 20, 5):
+        pytest.skip()
+    result = nw.from_native(constructor(data)).select(
+        nw.col("a").cast(nw.List(nw.Int32())).list.sort()
+    )
+    assert_equal_data(result, {"a": expected_asc_nulls_first})
+
+
+@pytest.mark.parametrize(
+    ("descending", "nulls_last", "expected"),
+    [
+        (True, True, expected_desc_nulls_last),
+        (True, False, expected_desc_nulls_first),
+        (False, True, expected_asc_nulls_last),
+        (False, False, expected_asc_nulls_first),
+    ],
+)
+def test_sort_expr_args(
+    request: pytest.FixtureRequest,
+    constructor: Constructor,
+    descending: bool,  # noqa: FBT001
+    nulls_last: bool,  # noqa: FBT001
+    expected: list[Any],
+) -> None:
+    if any(backend in str(constructor) for backend in ("dask", "cudf")):
+        request.applymarker(pytest.mark.xfail)
+    if "ibis" in str(constructor) and descending:
+        # https://github.com/ibis-project/ibis/issues/11735
+        request.applymarker(pytest.mark.xfail)
+    if "sqlframe" in str(constructor) and not nulls_last:
+        # https://github.com/eakmanrq/sqlframe/issues/559
+        # https://github.com/eakmanrq/sqlframe/issues/560
+        request.applymarker(pytest.mark.xfail)
+    if "polars" in str(constructor) and POLARS_VERSION < (0, 20, 5):
+        pytest.skip()
+    result = nw.from_native(constructor(data)).select(
+        nw.col("a")
+        .cast(nw.List(nw.Int32()))
+        .list.sort(descending=descending, nulls_last=nulls_last)
+    )
+    assert_equal_data(result, {"a": expected})
+
+
+def test_sort_series(
+    request: pytest.FixtureRequest, constructor_eager: ConstructorEager
+) -> None:
+    if any(backend in str(constructor_eager) for backend in ("dask", "cudf")):
+        request.applymarker(pytest.mark.xfail)
+    if "polars" in str(constructor_eager) and POLARS_VERSION < (0, 20, 5):
+        pytest.skip()
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+    result = df["a"].cast(nw.List(nw.Int32())).list.sort()
+    assert_equal_data({"a": result}, {"a": expected_asc_nulls_first})
+
+
+@pytest.mark.parametrize(
+    ("descending", "nulls_last", "expected"),
+    [
+        (True, True, expected_desc_nulls_last),
+        (True, False, expected_desc_nulls_first),
+        (False, True, expected_asc_nulls_last),
+        (False, False, expected_asc_nulls_first),
+    ],
+)
+def test_sort_series_args(
+    request: pytest.FixtureRequest,
+    constructor_eager: ConstructorEager,
+    descending: bool,  # noqa: FBT001
+    nulls_last: bool,  # noqa: FBT001
+    expected: list[Any],
+) -> None:
+    if any(backend in str(constructor_eager) for backend in ("dask", "cudf")):
+        request.applymarker(pytest.mark.xfail)
+    if "polars" in str(constructor_eager) and POLARS_VERSION < (0, 20, 5):
+        pytest.skip()
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+    result = (
+        df["a"]
+        .cast(nw.List(nw.Int32()))
+        .list.sort(descending=descending, nulls_last=nulls_last)
+    )
+    assert_equal_data({"a": result}, {"a": expected})