Migrate DropHighPSIFeatures to narwhals, add polars support

[solegalli]

Stacked on #1070 (selection base). Its commit shows in the diff until #1070 merges.

Summary

Migrates DropHighPSIFeatures (feature_engine/selection/drop_psi_features.py) to narwhals. fit() accepts pandas, polars or any other narwhals-supported dataframe, and transform() (inherited from BaseSelector) returns the same library it receives. The module no longer imports pandas.

How fit() works now:

• Split. Instead of copying X into a basis and a test dataframe, _basis_mask() returns a numpy boolean array, and each feature is sliced with it. The pandas branch keeps the old pandas code for the reference column / index (quantile, value_counts().sort_index().cumsum(), isin, <=). The narwhals branch does the same with narwhals series (value_counts(normalize=True), cum_sum, is_in, <=).
• Numerical features. The intervals are computed with numpy from the basis values, with the same code as EqualFrequencyDiscretiser / EqualWidthDiscretiser (np.quantile with the qcut nudge, or linspace with the 0.1% widening for constant values). Then np.searchsorted + np.bincount count the basis and test values per interval. The discretisers are no longer instantiated (see benchmarks and "Needs decision").
• Categorical features. pandas: Series.factorize(sort=True) once, then np.bincount on the basis and test codes. Other backends: one narwhals group_by(feature).agg(nw.col("__basis__").sum(), nw.len()) (simple aggregations only), sorted by category.
• PSI. Computed with numpy from the two count arrays. Intervals or categories that are empty in both sets are left out, as before (they were not in the value_counts merge), so the sum is done over the same values in the same order: the pandas PSI values are bit-identical to main.

Splitting by the index (split_col=None)

• pandas: unchanged. The index values are the reference, so a shuffled or datetime index works as before.
• polars (and other backends without an index): the position of the rows, 0 to n-1, is the reference. This follows your decision for the time series transformers in Migrate BaseForecastTransformer to narwhals, add polars support #1065: the row order is the time order. With the default split_frac=0.5, the first half of the rows goes to the basis set; cut_off=250 sends rows 0 to 250; a list cut-off selects those positions. cut_off_ holds the position (e.g. 499.5 for 1000 rows), which is exactly what pandas returns for a default RangeIndex, so both backends give the same PSI values for the same data.

The docstring (split_col, cut_off_) and the user guide (split_col section and a new "With polars" section) explain this for users.

Benchmarks

Median of repeats with the versions alternated (whole fit(): each version in its own subprocess, 5 repeats of best-of-3; steps: 7-21 repeats). Data: normal numerical columns plus 2 string columns (5 and 100 categories), variables="all", split by the index / row order.

Whole fit(), pandas. main = pandas code before this PR (the version on #1070 does not run). "narwhals only" = this branch with the pandas branches disabled, i.e. what the else path costs on pandas.

rows	columns	main	this PR	narwhals only
10k	5 + 2 cat	24 ms	6 ms	8 ms
10k	20 + 2 cat	93 ms	13 ms	16 ms
10k	50 + 2 cat	179 ms	25 ms	28 ms
100k	5 + 2 cat	52 ms	36 ms	39 ms
100k	20 + 2 cat	158 ms	81 ms	84 ms
100k	50 + 2 cat	372 ms	170 ms	173 ms
500k	5 + 2 cat	177 ms	170 ms	175 ms
500k	20 + 2 cat	526 ms	381 ms	389 ms
500k	50 + 2 cat	1185 ms	793 ms	801 ms

Splitting by a date column instead gives the same picture (e.g. 10k × 20: 148 → 20 ms; 500k × 20: 603 → 416 ms, measured under load).

Whole fit(), polars (there is no "before": the old code fails on polars): 500k × 5+2: 97 ms, 500k × 50+2: 718 ms, 1M × 20+2: 598 ms, 5M × 20+2: 2954 ms.

Step: binning and counting the numerical features (per-feature loop, 10 equal-frequency bins; ms for all columns):

backend	rows × cols	old (filter X, discretiser per feature, value_counts + merge)	discretiser per feature + bincount	one discretiser for all columns + bincount	numpy per column: np.quantile + searchsorted + bincount (chosen)	numpy, one 2-D array	numpy codes + pandas value_counts	numpy codes + narwhals group_by	polars value_counts
pandas	10k × 20	87.1	61.7	19.1	7.0	6.5	19.2	25.5	–
pandas	100k × 20	143.7	111.7	73.5	54.6	58.8	69.3	78.9	–
pandas	500k × 5	97.0	83.2	79.2	64.9	69.8	70.5	77.2	–
pandas	500k × 50	966.8	828.2	755.5	652.3	927.1	722.6	768.0	–
polars	500k × 20	–	278.7	276.3	258.2	341.2	–	272.3	260.0
polars	1M × 50	–	1354.1	1366.0	1284.0	1892.5	–	1307.5	1294.6
polars	5M × 20	–	2664.6	2706.7	2589.6	3469.0	–	2503.3	2573.0

numpy per column is the fastest or within 3% everywhere (at 5M × 20 on polars the narwhals group_by counting is 3% faster; np.quantile dominates all variants). The one 2-D array only wins at 10k and loses up to 1.5x at 500k. I also timed the pieces of one column at 1M rows: np.quantile 29 ms vs polars quantile ×11 74 ms; np.searchsorted 32 ms vs polars search_sorted 37 ms and polars cut + value_counts 91 ms. Only np.quantile also reproduces the discretisers' intervals bit for bit.

Step: counting the categories of one feature (ms):

backend	rows	categories	value_counts ×2 + align	Series.factorize + bincount (chosen, pandas)	narwhals group_by + len (chosen, others)	polars native group_by
pandas	10k	5 / 1000	0.68 / 1.13	0.37 / 0.57	1.21 / 1.41	–
pandas	100k	5 / 1000	3.89 / 4.05	3.08 / 3.22	4.13 / 4.07	–
pandas	500k	5 / 1000	17.77 / 17.36	15.08 / 14.98	17.03 / 16.71	–
polars	500k	5 / 1000	–	–	1.48 / 3.08	1.34 / 2.95
polars	1M	5 / 1000	–	–	2.20 / 5.21	2.01 / 5.11
polars	5M	5 / 1000	–	–	6.70 / 28.27	6.59 / 32.92

The old code (filter X twice, DataFrame.value_counts, merge) took 5.8 / 16.2 / 63.6 ms at 10k / 100k / 500k × 5 categories. On polars, the narwhals group_by and the native one are within 0.2 ms, and narwhals is faster at 5M × 1000, so the else branch stays in narwhals (it also serves other backends). The aggregation uses sum and len only, so it stays a simple group-by on pandas-like backends. Splitting and the per-feature loop are shared by both backends, so the pandas-native branches (X[feature].to_numpy(), factorize, the index handling) are what the "narwhals only" column above measures: 2-33% faster at 10k-100k, 1-3% at 500k.

Behaviour

• pandas: identical. I recorded 189 cases on origin/main (the last version where the pandas code runs: on Migrate the selection base classes and helpers to narwhals, add polars support #1070 check_X returns a narwhals frame, so 31 of the 62 old tests fail there) and compared them with this branch: psi_values_ (value and type, compared with ==, not approximately), features_to_drop_, variables_, cut_off_ (value and type), the transformed frame, get_support(), get_feature_names_out(), errors (type and message), and that the input is not modified. The cases cover every parameter (variables None/"all"/list/str, threshold number/"auto", both strategies, bins 2-50, split_frac, split_distinct, switch, min_pct_empty_bins incl. 0, missing_values, confirm_variables, p_value), split columns of every type (float, int, string, category, datetime, date objects, index with RangeIndex, shuffled, datetime and string index), cut-offs (number, 0, string, date/datetime/Timestamp, lists of numbers/strings/dates/datetimes), NaN in features and in the split column, inf, integer column names, nullable Int64/Float64, string dtype, categories with unused levels, many categories, reordered columns at transform and the error paths. Differences: 9 cases, all listed below as fixes or in "Needs decision".
• polars gives the same values as pandas: in the 165 cases that run on polars (the 24 index and pandas-dtype cases are pandas-only), 142 succeed on both backends and their PSI values are bit-identical to pandas, and features_to_drop_, variables_, cut_off_ and the transformed data are equal. Differences are only in cases that raise on both backends with a different exception (e.g. comparing a datetime column with a date cut-off: pandas TypeError, polars NotImplementedError), plus two noted below.

Bugs fixed (each has a test that fails without the fix)

1. cut_off=0 was ignored: if not self.cut_off treated 0 as "not given" and split by split_frac instead. Now cut_off is None is checked (in __init__ too). test_cut_off_zero_is_used.
2. Error messages with missing spaces: "in the referencevariable" and "cut_off orsplit_frac"; also "cannot be both set to None The current values" now has its period.
3. bins message: said "integer >= 1" but bins=1 is rejected. Now "integer >= 2".
4. Init validation: threshold=None or a list passed __init__ and failed in fit() with a TypeError; a non-numeric split_frac (e.g. "0.5") raised a TypeError from the comparison. Both now raise the ValueError of the parameter. strategy and missing_values check the type before the membership test (AGENTS.md). threshold accepts any real number (numbers.Real), so numpy numbers keep working.
5. split_frac=0 with cut_off=None raised "cut_off and split_frac cannot be both set to None" (0 is falsy); it now raises the split_frac range error.

Other differences

• A feature with only missing values in the basis or the test set (possible with missing_values="ignore") raised ValueError: 0 feature(s) (shape=(0, 1)) while a minimum of 1 is required. from the discretiser. It still raises a ValueError, now with a message that names the variable: "The variable var_1 has only missing values in the basis or in the test set, so its PSI can't be computed. Got 0 values in the basis set and 500 values in the test set."
• With missing_values="ignore", inf values were rejected by the discretiser after the split; now the inf check runs before the split (same message). Only the order of errors changes when a dataframe has both inf and a split problem.
• cut_off=[] used to fall back to split_frac (empty list is falsy); it now selects no rows and raises the "number of rows" error.
• Polars: cut_off_ computed from a datetime column is a Python datetime (pandas keeps np.datetime64); both compare equal to datetime(...).
• Polars: a Categorical split column with split_frac is sorted and compared alphabetically and works. In pandas, an unordered category split column raises TypeError: Unordered Categoricals can only compare equality or not (pre-existing, not changed).
• A list of dates as cut_off for a polars datetime column is cast to the column type, as pandas does with pd.to_datetime.

Tests

tests/test_selection/test_drop_high_psi_features.py is rewritten to the conventions: # init parameters (one test per error message, wrong values and types, match=re.escape(msg) with the full message) and test_init_param_assignment, then # fit and transform, with the data as plain dicts in the file and every behaviour on both backends through make_df (isinstance + frame_to_dict). pandas-only tests: shuffled index, datetime index, category dtype, integer column names.

Two problems in the old tests are fixed:

• pytest.approx(expected, 12) sets rel=12 (1200% tolerance), so the PSI asserts accepted almost any value. test_missing_value_ignored expected the PSI without NaN for var_3, which is not what the code returns (0.0340 vs 0.0377); it passed only because of the tolerance. The new tests use the default tolerance, with values checked against main.
• The three test_split_shuffled_df_* tests had no assert (set(base.index) == {...} alone).

test_transform_feature_to_drop_not_present is removed: BaseSelector.transform() (#1070) selects the retained features, so a missing dropped feature no longer raises; the base tests cover transform.

tests/test_selection (compared with origin/narwhals-selection-base, same command):

• before: 137 failed, 408 passed
• after: 102 failed, 588 passed. 0 new failures; the 35 fixed are the 31 old PSI tests and the 4 DropHighPSIFeatures entries in test_check_estimator_selectors.py. The remaining failures belong to selectors migrated in other PRs.

flake8 feature_engine tests is clean. mypy feature_engine: the same 2 errors as the base ref (datetime_subtraction.py, log.py).

Docs

• Docstring: split_col, cut_off_ and the data-split paragraph explain the polars behaviour; the example output was truncated at row 10 and now shows the 20 rows; a short polars example was added.
• User guide: intro mentions polars; split_col section explains row positions for polars; new "With polars" section (run, real output). All examples were run on this branch and on main and give the same output. One was stale: Case 1 used X['var_3'][250:] = ..., a chained assignment that doesn't modify X in pandas 3 (it only warns), so var_3 was not shifted. It now uses X.loc[250:, 'var_3'], and the PSI of var_3 shown is the real one, 0.8324552340063687 (the page said 1.342485289730313). The printed Index([...], dtype='object') shows dtype='str' with pandas 3; I left it as in the other user guides.

Needs decision

1. Row positions for polars when split_col=None (implemented, following Migrate BaseForecastTransformer to narwhals, add polars support #1065). The alternative is to raise an error asking for split_col with dataframes without an index. Note that pandas keeps using the index labels, so a pandas dataframe with a non-default index and the same data in polars are split differently; that is inherent to pandas having an index.
2. The discretisers are no longer instantiated. Fitting EqualFrequencyDiscretiser / EqualWidthDiscretiser on one-column frames per feature was slower than numpy: 9x at 10k rows and 1.3x at 500k with pandas, 3-8% with polars (step table above), because of the checks, frame building and with_columns on each call, so _interval_limits() repeats their ~15 lines of interval code. The results are bit-identical, but the two copies must be kept in sync. Option: move the interval computation of the discretisers into shared functions (related to issue Decide where helper functions that belong to one class should live #1054) and call them from both; that touches the discretisation module, so I did not do it here.
3. _check_variable_number() is not called. The addendum says selectors call it after selecting the variables, but DropHighPSIFeatures has always worked with one variable (e.g. variables="var_0", tested before and now). Calling it would make one-variable selections raise. I kept the old behaviour.

Pre-existing issues, not fixed

• pandas: an unordered category column in split_col with split_frac raises TypeError: Unordered Categoricals can only compare equality or not (polars works).
• A date cut-off with a datetime split column raises on both backends (the user guide already says to pass a datetime); with a list of dates it works.
• variables given as a single string equal to split_col is not rejected in __init__ (only lists are checked).
• polars: find_categorical_and_numerical_variables triggers a polars DeprecationWarning (cat.get_categories() in variable_handling/_variable_type_checks.py::_is_categories_num) for Categorical columns.
• split_col error message reads "must be a string an integer or None" (missing comma); left unchanged.