feat: Add new index/range based selector cs.by_index, allow multiple indices for nth

[alexander-beedie]

Enables some interesting new column selection capabilities using indices (including negative indices) and range objects, as well as their union, intersection, etc (via selector combinatorics).

• Adds a new by_index selector that works with indices and ranges.
• Allows nth to take multiple indices.

Examples

import polars.selectors as cs
import polars as pl

df = pl.DataFrame({
    "key": ["abc"],
    **{f"c{i:02}": [0.5 * i] for i in range(100)},
})
# shape: (1, 101)
# ┌─────┬─────┬─────┬─────┬───┬──────┬──────┬──────┬──────┐
# │ key ┆ c00 ┆ c01 ┆ c02 ┆ … ┆ c96  ┆ c97  ┆ c98  ┆ c99  │
# │ --- ┆ --- ┆ --- ┆ --- ┆   ┆ ---  ┆ ---  ┆ ---  ┆ ---  │
# │ str ┆ f64 ┆ f64 ┆ f64 ┆   ┆ f64  ┆ f64  ┆ f64  ┆ f64  │
# ╞═════╪═════╪═════╪═════╪═══╪══════╪══════╪══════╪══════╡
# │ abc ┆ 0.0 ┆ 0.5 ┆ 1.0 ┆ … ┆ 48.0 ┆ 48.5 ┆ 49.0 ┆ 49.5 │
# └─────┴─────┴─────┴─────┴───┴──────┴──────┴──────┴──────┘

Select "key" col and the three first/last numeric columns:
(can freely mix/match indexes and range objects)

df.select(pl.nth([0, 1, 2, 3, -3, -2, -1]))
df.select(cs.by_index(0, 1, 2, 3, -3, -2, -1))
df.select(cs.by_index(range(0,4), range(-3,0)))
# shape: (1, 7)
# ┌─────┬─────┬─────┬─────┬──────┬──────┬──────┐
# │ key ┆ c00 ┆ c01 ┆ c02 ┆ c97  ┆ c98  ┆ c99  │
# │ --- ┆ --- ┆ --- ┆ --- ┆ ---  ┆ ---  ┆ ---  │
# │ str ┆ f64 ┆ f64 ┆ f64 ┆ f64  ┆ f64  ┆ f64  │
# ╞═════╪═════╪═════╪═════╪══════╪══════╪══════╡
# │ abc ┆ 0.0 ┆ 0.5 ┆ 1.0 ┆ 48.5 ┆ 49.0 ┆ 49.5 │
# └─────┴─────┴─────┴─────┴──────┴──────┴──────┘

Select every 10th column:

df.select(cs.by_index(range(0,101,10)))
# shape: (1, 11)
# ┌─────┬─────┬─────┬──────┬───┬──────┬──────┬──────┬──────┐
# │ key ┆ c09 ┆ c19 ┆ c29  ┆ … ┆ c69  ┆ c79  ┆ c89  ┆ c99  │
# │ --- ┆ --- ┆ --- ┆ ---  ┆   ┆ ---  ┆ ---  ┆ ---  ┆ ---  │
# │ str ┆ f64 ┆ f64 ┆ f64  ┆   ┆ f64  ┆ f64  ┆ f64  ┆ f64  │
# ╞═════╪═════╪═════╪══════╪═══╪══════╪══════╪══════╪══════╡
# │ abc ┆ 4.5 ┆ 9.5 ┆ 14.5 ┆ … ┆ 34.5 ┆ 39.5 ┆ 44.5 ┆ 49.5 │
# └─────┴─────┴─────┴──────┴───┴──────┴──────┴──────┴──────┘

Select every 25th numeric column in reverse order:

df.select(cs.by_index(range(101,0,-25)))
# shape: (1, 4)
# ┌──────┬──────┬──────┬─────┐
# │ c75  ┆ c50  ┆ c25  ┆ c00 │
# │ ---  ┆ ---  ┆ ---  ┆ --- │
# │ f64  ┆ f64  ┆ f64  ┆ f64 │
# ╞══════╪══════╪══════╪═════╡
# │ 37.5 ┆ 25.0 ┆ 12.5 ┆ 0.0 │
# └──────┴──────┴──────┴─────┘

[codspeed-hq]

CodSpeed Performance Report

Merging #16217 will not alter performance

_{Comparing alexander-beedie:index-column-expr (8fb7aa4) with main (253fd24)}

Summary

✅ 34 untouched benchmarks

[codecov]

Codecov Report

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

Project coverage is 80.83%. Comparing base (993cc99) to head (8fb7aa4).
Report is 10 commits behind head on main.

Files	Patch %	Lines
...tes/polars-plan/src/logical_plan/expr_expansion.rs	91.07%	5 Missing ⚠️
crates/polars-plan/src/dsl/expr.rs	0.00%	1 Missing ⚠️
crates/polars-plan/src/dsl/meta.rs	0.00%	1 Missing ⚠️
crates/polars-plan/src/logical_plan/format.rs	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #16217      +/-   ##
==========================================
- Coverage   80.99%   80.83%   -0.16%     
==========================================
  Files        1393     1394       +1     
  Lines      179445   180095     +650     
  Branches     2907     2913       +6     
==========================================
+ Hits       145335   145574     +239     
- Misses      33604    34018     +414     
+ Partials      506      503       -3     

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

[stinodego]

We have the nth function for this, right? I'd prefer not to overload the col function if there is no need for it.

[alexander-beedie]

We have the nth function for this, right? I'd prefer not to overload the col function if there is no need for it.

It's not the same -nth can only take a single index, which leads to a profusion of expression objects if you want to select more than a handful:

df.select(pl.col(*range(10)))

df.select(
    pl.nth(0), 
    pl.nth(1), 
    pl.nth(2), 
    pl.nth(3), 
    pl.nth(4), 
    pl.nth(5), 
    pl.nth(6), 
    pl.nth(7), 
    pl.nth(8), 
    pl.nth(9),
)

Could probably move the multi-index support into nth, though given that col already selects by one or more names/dtypes I don't see adding index as a stretch (there's no ambiguity, and it provides some consistency - you don't have to look in different places to select columns, you just use col with a suitable argument for each selection type).

Still, nth wouldn't be an unreasonable home for multi-index support (though arguably it's less discoverable), so... Strong opinion? Want me to switch it over?

[ritchie46]

Cool one! A few comments.

[stinodego]

Could probably move the multi-index support into nth, though given that col already selects by one or more names/dtypes I don't see adding index as a stretch (there's no ambiguity, and it provides some consistency - you don't have to look in different places to select columns, you just use col with a suitable argument for each selection type).

Still, nth wouldn't be an unreasonable home for multi-index support (though arguably it's less discoverable), so... Strong opinion? Want me to switch it over?

I'm already not super happy with col taking data type inputs - I'd prefer if that were a different function. I don't think nth is a very descriptive name, but at least it has a dedicated purpose. Extending it with multi-column support makes a lot of sense.

We can add a See Also to col linking to nth if it doesn't do so already to help with discovery.

[alexander-beedie]

We can add a See Also to col linking to nth if it doesn't do so already to help with discovery.

Ok, will switch it over to nth in the next few hours✌️

(Update: @stinodego - done).

[stinodego]

This is nice functionality to have! I left a few comments for the Python side.

[stinodego]

Pity we have this *columns API, because passing nth(1,2) would feel really good. I think we should deprecate this API (in another PR) and make this a keyword argument, so we can pass nth(1,2, from=['x', 'y', 'z'])

[alexander-beedie]

I was thinking the same thing 👌 Can follow-up with a separate improvement / deprecation; I think there are a few other functions with a similar set of params, so we should check and do all at once, if so.

[stinodego]

All right, good to go!