Skip to content

Migrate WindowFeatures to narwhals, add polars support - #1077

Open
solegalli wants to merge 3 commits into
narwhals-migrationfrom
narwhals-window-features
Open

solegalli wants to merge 3 commits into
narwhals-migrationfrom
narwhals-window-features

Conversation

@solegalli

Copy link
Copy Markdown
Collaborator

Stacked on #1065 (time series base). Until #1065 is merged, this diff also shows its commit; the WindowFeatures change is the last commit.

Summary

WindowFeatures takes pandas and polars dataframes.

  • It implements the base hook _add_features(nw_X). transform() keeps its docstring and calls super().transform(X), and the temporary .to_native() from Migrate BaseForecastTransformer to narwhals, add polars support #1065 is gone. drop_original, drop_na and transform_x_y come from the base.
  • pandas keeps a native path on the index: time-span windows like "3D", freq, min_periods (passed to rolling(), as in Pass min_periods to rolling() in WindowFeatures #1043) and every rolling function. It now calls the rolling methods (rolling.mean(), ...) instead of rolling.agg(functions), which is faster with many variables (numbers below), and the output is identical.
  • polars: the rows are used in the order given, window is a number of rows and the features are shifted by periods rows, as with pandas. mean, sum, std, var use narwhals' rolling expressions with min_samples; count is a narwhals rolling sum of the non-null mask. min, max, median, skew, kurt use polars expressions (rolling_min, ... rolling_skew(bias=False), rolling_kurtosis(bias=False)), because narwhals has no such expressions and polars beat numpy by a wide margin. std/var use ddof=1, and min_periods=None means the full window, both as in pandas.
  • _check_index(X) calls super()._check_index(X). Then, when the input isn't pandas, it raises NotImplementedError in the same style as the base's freq message for:
    • a window that isn't an integer (time spans);
    • functions polars can't do (sem, first, last, nunique, rank, ...), with the list of the supported ones;
    • min_periods=0, which narwhals rejects and which would give different sums on all-null windows.
  • The docstring and user guide say what works with polars. The user guide has a new "Window features with polars" section. I ran it and pasted the real output.

Benchmarks

Median times in ms. The machine was heavily loaded by other jobs (load average about 40), so the ratios are noisy. Where they differ, the direction was the same in every run.

pandas, whole transform(): base ref (rolling().agg()) vs this PR (rolling methods). The base ref's transform() ran in the same process, alternating with the new one:

rows x vars params base new new/base
10k x 1 w=3 mean 3.06 2.83 0.92
10k x 10 w=[3,7,14] mean,std,max 104.8 41.8 0.40
10k x 10 w="1h" mean,max freq="15min" 37.5 15.0 0.40
100k x 1 w=[3,7,14] mean,std,max 30.9 23.4 0.76
100k x 10 w=3 mean 24.9 19.9 0.80
100k x 10 w=[3,7,14] mean,std,max 280.8 217.6 0.77
100k x 10 w="1h" mean,max freq="15min" 121.2 73.7 0.61
500k x 1 w=3 mean 11.9 18.2 1.53 (noise; 1.05 and 1.37 in other runs)
500k x 10 w=[3,7,14] mean,std,max 2048 1731 0.85
500k x 10 w=[2,5] mean drop_na 1236 795 0.64

pandas rolling().agg() vs the rolling methods alone, on the rolling step only, with 3 functions: 10k x 10: 23.3 vs 10.5; 100k x 10: 98.0 vs 74.0; 500k x 10: 470 vs 411. The same step through narwhals' rolling expressions on pandas was about as fast as the methods (100k x 10 mean+std: pandas agg 92.6, narwhals 82.6). Narwhals can't do time spans, freq or min/max/median, so pandas keeps its own path.

polars, functions narwhals lacks (min/max/median), including shift:

rows x vars window function polars expr numpy sliding window narwhals min/max_horizontal of shifts
500k x 1 30 min 6.5 16.3 21.2
500k x 10 3 max 12.5 103.9 14.7
500k x 10 30 median 55.7 3857 n/a
2M x 1 3 min 19.4 34.6 11.2
2M x 10 30 max 52.8 565.6 1068

Polars won or tied everywhere except window=3 with 1 variable. There the shifts version was a little faster, but it scales with the window size.

polars, narwhals + polars expressions (this PR) vs all polars expressions in one with_columns:

rows x vars params narwhals + polars all polars
500k x 1 w=[3,7,14] mean,std,max 22.3 14.6
500k x 1 w=[3,7] mean,min,max,median 32.5 23.6
500k x 10 w=[3,7,14] mean,std,max 102.6 99.6
2M x 1 w=[3,7,14] mean,std,max 74.6 49.5
2M x 10 w=[3,7] mean,min,max,median 546.6 469.1
2M x 10 w=3 mean 53.4 52.4

See "Needs decision" 2.

polars, whole transform() (this PR): 500k x 10, w=[3,7,14] mean,std,max: 98 ms; 2M x 10: 415 ms; 500k x 10, w=[3,7] median,min: 81.5 ms.

Behaviour

  • pandas: identical to the base ref (assert_frame_equal, exact, plus column index dtype) in 375 of 377 recorded cases. The cases cover windows (int, list, time spans), functions (mean; mean+std; all 15 rolling functions), min_periods None/1/2, periods, freq, drop_original, drop_na, sort_index on shuffled rows, transform_x_y, integer column names, NaN in the data, variable subsets and return_empty. The 2 differences:
    • return_empty=True with no numerical variables and a list of windows or functions used to raise ValueError: No objects to concatenate. It now returns the dataframe unchanged, as return_empty documents. This is a bug fix; the test fails on the base ref.
    • An invalid function name like "pizza" still raises AttributeError, but the message is now pandas' 'Rolling' object has no attribute 'pizza' instead of 'pizza' is not a valid function for 'Rolling' object.
  • polars vs pandas: same values in 408 of 576 cases. These cover windows 1, 3 and [2, 5], "mean" and all 10 supported functions, min_periods None/0/1/2, periods 1/3, drop_original, drop_na, missing data, and shuffled rows. The rest:
    • 144 cases are min_periods=0, which raises on polars on purpose.
    • 24 cases are skew with missing data. After a window with fewer than 3 values, pandas 3.0.3's rolling skew returns NaN for the rest of the series. polars matches scipy.stats.skew(bias=False), so this is a pandas bug.
  • polars keeps its natural dtypes: sum/min/max of an integer column stay integers, while pandas returns floats. The values are the same.
  • polars float NaN (not null) is a value in polars, so it propagates through the windows. pandas treats NaN as missing. This only matters with missing_values="ignore", because the NaN check catches it otherwise.

Tests

tests/test_time_series/test_forecasting/test_window_features.py is rewritten to the conventions:

  • # init parameters: one test per error message, with wrong values and types, and test_init_param_assignment.
  • # fit and transform: make_df tests on both backends, with plain-dict data and explicit expected values. They cover every polars-supported function, periods, min_periods (including the Pass min_periods to rolling() in WindowFeatures #1043 cases), missing data, drop_original, drop_na, transform_x_y with make_series, return_empty, and get_feature_names_out and its errors.
  • pandas-only tests: freq, time-span windows, pandas-only functions, sort_index, y aligned on the index, and integer column names.
  • polars-only tests: row order and the three new errors, each also for a transformer fitted on pandas and transforming polars.
  • The drop_original/drop_na init-error tests were removed because test_base_forecast_transformer.py covers them.

The old test file still passes against the new code (36/36).

tests/test_time_series: base ref 189 passed, 0 failed; this PR 230 passed, 0 failed. No other test folder imports WindowFeatures. flake8 feature_engine tests is clean, and mypy feature_engine shows the same 2 errors as the base ref.

Needs decision

  1. Minimum polars version. The polars expressions are called with min_samples=, which polars 1.21 introduced (before that it was min_periods). rolling_skew and rolling_kurtosis are marked unstable in polars, and rolling_kurtosis is recent. The tests extra says polars>=1.0.0. narwhals handles this for mean/sum/std/var but not for these five. The options are:
    • raise the minimum polars version;
    • pass min_periods or min_samples depending on the polars version;
    • drop min/max/median/skew/kurt on polars.
  2. narwhals + polars expressions vs all polars expressions. I followed the brief: narwhals where it has the function, and polars expressions only for the rest. When functions are mixed and there are few variables, one with_columns of polars expressions was up to about 1.5x faster, and it was equal with 10 variables. It would also remove the reorder select. Should I switch?
  3. _check_index now validates window, functions and min_periods for non-pandas input, because it is the one hook the base calls in both fit and transform with the native X. The name no longer fits. A base change could rename it (for example to _check_input).

Pre-existing issues, not fixed

  • window and min_periods aren't validated at init. For example, window=0 or min_periods > window raise pandas' or polars' own errors at transform.
  • functions="corr" or "cov" return pairwise rolling results, which duplicates the rows (same as before).
  • The init messages for duplicated windows or functions don't end with "Got ... instead."
  • _get_new_features_name's docstring says "lag features".

solegalli and others added 3 commits September 19, 2026 11:22
The base of the forecasting transformers takes pandas and polars dataframes.
With pandas the index keeps ordering the rows in time; with polars the rows
are used in the order given, sort_index does not apply and freq raises an
error. The base now provides transform() and transform_x_y() on top of an
_add_features() hook for the subclasses.

Also sets AGENTS.md's priorities to speed, readability and simplicity.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
WindowFeatures takes pandas and polars dataframes. With pandas it keeps
rolling on the index, including time-span windows, freq and all rolling
functions; it now calls the rolling methods instead of agg(), which is
faster with many variables. With polars the rows are used in the order
given, windows are numbers of rows, and mean, sum, std, var and count use
narwhals, while min, max, median, skew and kurt use polars expressions.
Time-span windows, other functions and min_periods=0 raise a clear error
with polars.

Also stops return_empty=True from failing with lists of windows or
functions when there are no numerical variables.

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