feat(rust, python): Add by argument for Expr.top_k and Expr.bottom_k

[CanglongCl]

Closes #10054

If by is provided, columns input will be treated entirely. Just a simple wrapper of Expr.sort_by().head().

If by is not provided, each column will be sort individually (same as before modification).

Example

df = pl.DataFrame(
    {
        "a": [1, 2, 3, 4, 5, 6],
        "b": [12, 11, 10, 9, 8, 7],
        "c": ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"],
    }
)
df

shape: (6, 3)
┌─────┬─────┬────────┐
│ a   ┆ b   ┆ c      │
│ --- ┆ --- ┆ ---    │
│ i64 ┆ i64 ┆ str    │
╞═════╪═════╪════════╡
│ 1   ┆ 12  ┆ Apple  │
│ 2   ┆ 11  ┆ Orange │
│ 3   ┆ 10  ┆ Apple  │
│ 4   ┆ 9   ┆ Apple  │
│ 5   ┆ 8   ┆ Banana │
│ 6   ┆ 7   ┆ Banana │
└─────┴─────┴────────┘

Top k rows by specific column.

df2.select(
    pl.all()
    .top_k(2, by="a")
    .name.suffix("_top_by_a"),
    pl.all()
    .top_k(2, by="b")
    .name.suffix("_top_by_b"),
)

shape: (2, 6)
┌────────────┬────────────┬────────────┬────────────┬────────────┬────────────┐
│ a_top_by_a ┆ b_top_by_a ┆ c_top_by_a ┆ a_top_by_b ┆ b_top_by_b ┆ c_top_by_b │
│ ---        ┆ ---        ┆ ---        ┆ ---        ┆ ---        ┆ ---        │
│ i64        ┆ i64        ┆ str        ┆ i64        ┆ i64        ┆ str        │
╞════════════╪════════════╪════════════╪════════════╪════════════╪════════════╡
│ 6          ┆ 7          ┆ Banana     ┆ 1          ┆ 12         ┆ Apple      │
│ 5          ┆ 8          ┆ Banana     ┆ 2          ┆ 11         ┆ Orange     │
└────────────┴────────────┴────────────┴────────────┴────────────┴────────────┘

Top k by multiple columns with given order.

df2.select(
    pl.all()
    .top_k(2, by=["c", "a"], descending=[False, True])
    .name.suffix("_top_by_ca"),
    pl.all()
    .top_k(2, by=["c", "b"], descending=[False, True])
    .name.suffix("_top_by_cb"),
)

shape: (2, 6)
┌─────────────┬─────────────┬─────────────┬─────────────┬─────────────┬─────────────┐
│ a_top_by_ca ┆ b_top_by_ca ┆ c_top_by_ca ┆ a_top_by_cb ┆ b_top_by_cb ┆ c_top_by_cb │
│ ---         ┆ ---         ┆ ---         ┆ ---         ┆ ---         ┆ ---         │
│ i64         ┆ i64         ┆ str         ┆ i64         ┆ i64         ┆ str         │
╞═════════════╪═════════════╪═════════════╪═════════════╪═════════════╪═════════════╡
│ 2           ┆ 11          ┆ Orange      ┆ 2           ┆ 11          ┆ Orange      │
│ 5           ┆ 8           ┆ Banana      ┆ 6           ┆ 7           ┆ Banana      │
└─────────────┴─────────────┴─────────────┴─────────────┴─────────────┴─────────────┘

Get the top k rows by specific column in each group (#10054).

(
    df2
    .group_by('c', maintain_order=True)
    .agg(
        pl.all().bottom_k(2, by="a")
    )
    .explode(pl.all().exclude("c"))
)

shape: (5, 3)
┌────────┬─────┬─────┐
│ c      ┆ a   ┆ b   │
│ ---    ┆ --- ┆ --- │
│ str    ┆ i64 ┆ i64 │
╞════════╪═════╪═════╡
│ Apple  ┆ 1   ┆ 12  │
│ Apple  ┆ 3   ┆ 10  │
│ Orange ┆ 2   ┆ 11  │
│ Banana ┆ 5   ┆ 8   │
│ Banana ┆ 6   ┆ 7   │
└────────┴─────┴─────┘

When by is not provided, top k of each columns will be returned. (Values of a and b in a row is not in the same in the same row in original DataFrame, which is the same as before modification.)

(
    df2
    .group_by('c', maintain_order=True)
    .agg(
        pl.all().bottom_k(2)
    )
    .explode(pl.all().exclude("c"))
)

shape: (5, 3)
┌────────┬─────┬─────┐
│ c      ┆ a   ┆ b   │
│ ---    ┆ --- ┆ --- │
│ str    ┆ i64 ┆ i64 │
╞════════╪═════╪═════╡
│ Apple  ┆ 1   ┆ 9   │
│ Apple  ┆ 3   ┆ 10  │
│ Orange ┆ 2   ┆ 11  │
│ Banana ┆ 5   ┆ 7   │
│ Banana ┆ 6   ┆ 8   │
└────────┴─────┴─────┘

[codecov]

Codecov Report

Attention: Patch coverage is 82.02247% with 48 lines in your changes are missing coverage. Please review.

Project coverage is 81.36%. Comparing base (5253e46) to head (e644a5f).

❗ Current head e644a5f differs from pull request most recent head 939194d. Consider uploading reports for the commit 939194d to get more accurate results

Files	Patch %	Lines
crates/polars-ops/src/chunked_array/top_k.rs	64.16%	43 Missing ⚠️
crates/polars-plan/src/dsl/function_expr/mod.rs	28.57%	5 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main   #15468       +/-   ##
===========================================
+ Coverage   39.16%   81.36%   +42.19%     
===========================================
  Files        1363     1374       +11     
  Lines      167739   176680     +8941     
  Branches     3034     2552      -482     
===========================================
+ Hits        65703   143754    +78051     
+ Misses     101568    32443    -69125     
- Partials      468      483       +15     

Flag	Coverage Δ
python	?
rust	?

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[codspeed-hq]

CodSpeed Performance Report

Merging #15468 will not alter performance

_{Comparing CanglongCl:top_k_by (939194d) with main (5253e46)}

Summary

✅ 13 untouched benchmarks

[ritchie46]

Looks great @CanglongCl. Thanks a lot. I shall do a rebase and merge it in.