feat: add list.sort

[raisadz]

Description

I tried to have some workarounds for the edge cases for pyarrow that we discussed #3332 (comment) (None, empty lists and lists with only None elements) with pc.sort_indices and pc.replace_with_mask but the latter doesn't seem to work for the list types. There is an open issue apache/arrow#48060 that can make it work for pyarrow, pandas and modin when solved.

Also, I opened a couple of issues in sqlframe and ibis related to this PR:

eakmanrq/sqlframe#559
eakmanrq/sqlframe#560
ibis-project/ibis#11735

What type of PR is this? (check all applicable)

• 💾 Refactor
• ✨ Feature
• 🐛 Bug Fix
• 🔧 Optimization
• 📝 Documentation
• ✅ Test
• 🐳 Other

Related issues

• Related issue #<issue number>
• Closes #<issue number>

Checklist

• Code follows style guide (ruff)
• Tests added
• Documented the changes

[dangotbanned]

@raisadz I still need to loop back to it and clean things up, but I think list.sort will be possible for pyarrow 🙂

• feat(expr-ir): Add list.sort #3359

Here's a rough equivalent in polars:

Outdated

Note: I've hidden this because of (#3356 (comment))

import polars as pl


def list_sort(
    native: pl.Series, *, descending: bool = False, nulls_last: bool = False
) -> pl.Series:
    idx, name = "index", native.name
    indexed = native.to_frame().with_row_index(name=idx)

    len_gt_1 = pl.col(name).list.len() > 1
    valid = indexed.filter(len_gt_1)
    invalid = indexed.filter(pl.col(name).is_null() | ~len_gt_1)

    exploded = valid.explode(name, empty_as_null=False, keep_nulls=False)
    valid_finished = (
        exploded.sort(idx, name, descending=[False, descending], nulls_last=nulls_last)
        .group_by(idx, maintain_order=True)
        .agg(name)
    )
    return pl.concat([valid_finished, invalid], how="vertical").sort(idx).get_column(name)

And then against the test suite (https://github.com/narwhals-dev/narwhals/pull/3356/files#diff-24b32e940c35026188c24063b47ae536444a402c28879bdf7fa4853c7c5f5a0b)

data = {"a": [[3, 2, 2, 4, -10, None, None], [-1], None, [None, None, None], []]}
ser = pl.DataFrame(data).to_series()

print(list_sort(ser, descending=True, nulls_last=True).to_list())
print(list_sort(ser, descending=True, nulls_last=False).to_list())
print(list_sort(ser, descending=False, nulls_last=True).to_list())
print(list_sort(ser, descending=False, nulls_last=False).to_list())

[[4, 3, 2, 2, -10, None, None], [-1], None, [None, None, None], []]
[[None, None, 4, 3, 2, 2, -10], [-1], None, [None, None, None], []]
[[-10, 2, 2, 3, 4, None, None], [-1], None, [None, None, None], []]
[[None, None, -10, 2, 2, 3, 4], [-1], None, [None, None, None], []]

[MarcoGorelli]

thanks for working on this! just some comments

[MarcoGorelli]

maybe we can skip this test if the below already tests all possibilities?

[MarcoGorelli]

@FBruzzesi factored some repeated logic here out using _apply_pyarrow_compute_func, is it possible to use that here?

[MarcoGorelli]

same

[raisadz]

Thanks for the review, @MarcoGorelli ! I addressed your comments

[MarcoGorelli]

thanks @raisadz and @dangotbanned !