Skip to content

Migrate SelectByShuffling to narwhals, add polars support - #1084

Open
solegalli wants to merge 3 commits into
narwhals-migrationfrom
narwhals-select-by-shuffling
Open

solegalli wants to merge 3 commits into
narwhals-migrationfrom
narwhals-select-by-shuffling

Conversation

@solegalli

Copy link
Copy Markdown
Collaborator

Stacked on #1070 (selection base classes). Its commit shows in the diff until #1070 is merged; this PR's own changes are in the last commit.

Summary

  • SelectByShuffling accepts pandas, polars and other narwhals-supported dataframes and returns the same library it receives (through the inherited BaseSelector.transform). pandas is no longer imported.
  • Same results as before, and the same on every backend, for a given random_state. pandas Series.sample(frac=1, random_state=rs) calls rs.choice(n, n, replace=False), which is rs.permutation(n). The selector now calls random_state.permutation(n) directly, so every column gets the same shuffle as before, whatever the backend.
  • The hot loop was rewritten. The old loop copied X[variables_] for every feature, shuffled one column, and then sliced every validation fold again. Now the validation-fold dataframes (and y per fold) are built once. For each feature, only the fold's rows of the shuffled column are gathered (column[permutation[idx]]). pandas replaces that column in place in the fold copy and restores it after scoring. narwhals uses with_columns on the fold frame.
  • Validation folds come from cross_validate(..., return_indices=True) instead of a second check_cv(...).split(...) (this is a bug fix, see below). The X/y reset_index calls are gone because all indexing is positional (iloc / _safe_indexing).
  • The threshold init check now follows the conventions: threshold is not None and not isinstance(threshold, (int, float)), with the message threshold must be an integer, a float or None. Got {threshold} instead.
  • Docstring: added a "With polars" example. User guide: added a "Using polars" section, refreshed outputs that had drifted in the last digits (initial_model_performance_ was 0.488702767247119, the code returns 0.48870212980353145), and fixed two sentences that said features are dropped when their drift is greater than the mean (it's smaller). I ran every example and pasted the real output.

Benchmarks

The machine was heavily loaded during the runs (load average 40-65 on 10 cores), so the timings are noisy. BLAS was single-threaded, each number is the median of 9 repeats (whole fit: 5), and the order of the versions alternated. Data: normal floats, LinearRegression, r2, 3-fold CV. I picked this cheap estimator so the selector's own work is visible. With real models, predict dominates.

pandas: shuffle + re-score loop, all features (seconds)

rows x cols current (copy + sample + iloc) copy fold, setitem setitem + restore (chosen) iloc[:, j] write + restore shallow copy + setitem
10k x 10 0.110 0.084 0.103 0.092 0.070
10k x 50 1.297 1.056 0.875 0.760 0.822
100k x 10 0.388 0.398 0.336 0.240 0.255
100k x 50 4.097 2.418 2.002 1.883 2.097
500k x 10 0.569 0.495 0.415 0.366 0.506
500k x 50 7.656 4.771 3.822 3.383 3.831

Same loop without scoring, which isolates the selector's own work: at 500k x 50 it was 3.25 s for the current code and 1.08 s for setitem + restore (100k x 50: 0.63 vs 0.11). All strategies gave identical drifts (asserted in the benchmark).

pandas: whole fit() (seconds, two alternating runs)

rows x cols main this PR
10k x 10 0.033 / 0.029 0.033 / 0.038
10k x 50 0.215 / 0.214 0.177 / 0.187
100k x 10 0.096 / 0.097 0.088 / 0.090
100k x 50 0.857 / 1.931 1.180 / 0.774
500k x 10 0.396 / 0.324 0.439 / 0.308
500k x 50 4.226 / 5.657 2.968 / 4.513

Why setitem + restore over the iloc[:, j] in-place write: the two were within noise in the whole fit. The iloc write keeps the fold frame as one block, so it was slightly faster in the scored loop. But it needs the column position, and it relies on in-place writes that behave differently between pandas 2.2 and pandas 3 (copy-on-write) and with extension dtypes. Replacing the column works the same everywhere and keeps nullable dtypes (.array.take).

polars: shuffle + re-score loop (seconds)

rows x cols narwhals (chosen) polars gather numpy + pl.Series polars shuffle(seed), different results
100k x 10 0.153 0.167 0.158 0.129
500k x 10 0.534 0.505 0.503 0.411
500k x 50 7.643 9.228 8.232 9.503
2M x 10 1.531 1.582 1.660 1.075
2M x 30 10.774 8.964 10.572 5.900
no scoring, 500k x 50 0.386 0.416 0.371 0.203
no scoring, 2M x 30 1.515 1.612 1.726 0.899

narwhals, native polars gather and numpy were within noise of each other, so the code keeps the narwhals path with no native polars branch. polars' own shuffle is faster (see "Needs decision"). Whole fit() on polars: 500k x 10: 0.84 s, 500k x 50: 4.57 s, 2M x 10: 2.15 s, 2M x 30: 6.90 s.

Behaviour

Compared with the pandas code on main over 22 scenarios: RF / logistic / HistGradientBoosting classifiers, linear and tree regressors, roc_auc / accuracy / r2 / neg_mean_squared_error, threshold set / None / 0, int / splitter / generator / shuffled KFold cv, a variables subset, confirm_variables, categorical columns, reordered columns, NaN, nullable Int64, integer column names, a non-default index, y as a Series / one-column DataFrame / list / array, sample_weight, and the docstring example.

  • pandas: identical. initial_model_performance_, every performance_drifts_ and performance_drifts_std_ value (compared exactly, not approximately), features_to_drop_, variables_, feature_names_in_ and the transform output all match.
  • polars: identical to pandas, bit for bit, in the 19 scenarios that apply to polars.
  • performance_drifts_ / performance_drifts_std_ are dicts for every backend, as before.

Bug fixed

  • Folds could differ between training and scoring. The old code trained the models in cross_validate and then called check_cv(cv).split(X, y) again to get the validation folds. With a splitter that gives different folds on each call, the models were scored on rows they had been trained on. Examples are KFold(shuffle=True) without random_state, ShuffleSplit(), or a splitter seeded with a RandomState instance. The folds now come from cross_validate(return_indices=True). Test: test_performance_is_evaluated_on_the_folds_of_each_model. It uses a splitter that alternates its folds, and the same comparison fails on main. Deterministic splitters (the default) give the same results as before.
  • threshold validation. Falsy invalid values ("", []) slipped through the old if threshold and ... check and then silently meant "use the mean". They now raise, with the conventional message.

Tests

tests/test_selection/test_shuffle_features.py is rewritten to the conventions:

  • Init errors, parametrized (including confirm_variables), and test_init_param_assignment.
  • Both backends via make_df and the data_classification fixture: fit attributes, transform, the mean-drift threshold, regression with r2, and neg_mse with a DataFrame target. Also cv as int / splitter / generator, the fold bug, non-numerical columns ignored, confirm_variables, missing values, sample_weight, list and array targets, and a check that the input is not modified.
  • pandas-only: integer column names, a non-default index, and nullable Int64.

All expected values were checked against the pandas code on main.

tests/test_selection failing tests, base (narwhals-selection-base) vs this branch: 137 → 126, no new failures. The fixed ones are 6 old test_shuffle_features tests plus 5 test_check_estimator_selectors checks for SelectByShuffling. flake8 feature_engine tests is clean. mypy feature_engine gives the same 2 errors as the base (datetime_subtraction, log).

Needs decision

  1. Faster polars shuffle vs identical results. polars' Series.shuffle(seed) cuts the selector's own shuffle work by about 40-45% at 500k-2M rows (for example 1.52 s → 0.90 s at 2M x 30 without scoring). Most of the remaining cost is numpy's legacy RandomState.permutation (about 14 ms per column at 500k rows, vs about 8.6 ms for polars shuffle). But polars' generator gives different shuffles, so polars results would no longer match pandas for the same random_state, and it would need an int seed drawn from random_state. A numpy Generator (about 6 ms) would also be faster, but it changes the pandas results too. With any non-trivial estimator, predict dominates the run time, so I kept identical results. Switch?
  2. threshold=0 means "use the mean drift". The check is if not self.threshold:, so threshold=0 is treated like None. That looks like a bug (someone passing 0 probably wants "drop features whose drift is negative"), but SelectBySingleFeaturePerformance and SelectByTargetMeanPerformance use the same pattern, so I kept the behaviour for consistency. Fix it in all three (is None) in a follow-up?

Pre-existing issues, not fixed

  • sample_weight is used to train the models but not to score them (neither the initial score nor the shuffled ones). Same as before.

solegalli and others added 3 commits September 19, 2026 11:44
…s support

BaseSelector.transform() returns the retained features in the train set
order, in the same library as the input (pandas X[features], narwhals
select otherwise). BaseRecursiveSelector.fit() trains the estimators on
native frames and returns (nw_X, y). The helpers in
base_selection_functions no longer import pandas: correlations are
computed with numpy (np.corrcoef, or matrix products for pairwise
complete observations when there are missing values), and feature
importances are pandas Series for pandas input and dicts otherwise.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
SelectByShuffling works with pandas, polars and other narwhals-supported
dataframes. For a given random_state the shuffles, drifts and selected
features are the same as before and the same with every library.

The validation folds come from cross_validate(return_indices=True), so
each model is scored on the fold it was not trained on, also when the
splitter gives different folds on every call. The fold dataframes are
built once and one column is shuffled at a time, instead of copying the
whole dataframe for every feature.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant