Migrate DecisionTreeFeatures to narwhals, add polars support - #996
Merged
Conversation
Follows the same pandas-native / narwhals-generic split established for GeoDistanceFeatures (this transformer also reimplements fit()/transform() directly, not via BaseCreation): .columns extraction, column reorder, prediction-column assignment, and drop_original all split by backend, consistent with every other operation in this module that's been benchmarked as a real (not minimal) loss when routed through narwhals on pandas. Confirmed empirically before designing: sklearn's DecisionTreeRegressor/ Classifier and GridSearchCV accept a polars DataFrame directly for both fit() and predict()/predict_proba(), so the actual tree training/inference calls are unchanged - only the surrounding column selection, extraction, and reassembly needed migrating. Fixed a pre-existing bug found while rewriting the exact code path it lived in: single-feature combos with an integer column name (e.g. DecisionTreeFeatures(features_to_combine=1) on a dataframe with columns 0, 1, ...) crashed, since the original `isinstance(features, str)` check missed the int case and fell through to plain X[features] indexing, which returns a 1D Series rather than the 2D input sklearn requires. Widened to isinstance(features, (str, int)); verified the same single-feature narwhals path (get_column().to_frame()) already handles both cleanly. Regression, binary classification, and multiclass classification paths all verified to produce identical predictions between pandas and polars input. return_empty=True + polars remains untestable here too (same nw.col([]) bug in dataframe_checks.py found during CyclicalFeatures, still tabled) - this is the second transformer it blocks. docs/user_guide/creation/DecisionTreeFeatures.rst is large (511 lines) and built around actual cross-validated tree fitting on the real California housing dataset across many sections - re-verified the cheap, deterministic parts (the raw data table) but did not re-run every tree-fitting example given the cost of repeated grid-search CV fits; unlike the other three creation-module docs this pass touched, the rest of this file's numbers are unverified. Added a self-contained "With polars" section using simple synthetic data instead, fully verified.
solegalli
commented
Aug 24, 2026
solegalli
commented
Aug 24, 2026
…AGENTS.md Two rules made explicit based on recent work: the is True/is False comparison is for flow control only, not variable assignment (per Sole's own simplification of is_pandas = nwd.is_pandas_dataframe(X) is True to just nwd.is_pandas_dataframe(X) in decision_tree_features.py); and dataframe-agnostic transformers get one parametrized test per behavior covering both pandas and polars, never separate per-backend tests. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…oss-backend suite Adds an n_jobs parameter to DecisionTreeFeatures that parallelizes tree training across feature combinations via joblib, using threads rather than processes since fitting a decision tree releases the GIL for the bulk of its computation - threads avoid the overhead of copying the whole dataframe to worker processes. Defaults to None (sequential), preserving current behavior. Benchmarked on the committed transformer (5000 rows, 10 vars, features_to_combine=3, 8-point param_grid, 175 trees): 12.17s sequential vs 5.15s at n_jobs=-1, ~2.4x. On small workloads (a handful of feature combinations, the shape of the existing unit tests) parallelizing is a net loss - thread-dispatch overhead outweighs the gain - which is why the default stays sequential. Parallelizing transform()'s predict loop the same way was also benchmarked and found to have no benefit (predict is too cheap per call), so only fit()'s tree training is parallelized. Correctness verified: identical trees/predictions regardless of n_jobs. Also rewrites test_decision_tree_features.py to the single cross-backend-parametrized-test convention used elsewhere in this migration: one test per behavior over make_df=[pd.DataFrame, pl.DataFrame], deleting the separately-added polars-only section that duplicated coverage already present once the original tests are parametrized. Adds n_jobs correctness coverage (parallel vs sequential training gives identical output, both backends). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…form transform() assigned one new tree-prediction column at a time (X[col_name] = preds), which triggers pandas' "DataFrame is highly fragmented" PerformanceWarning once there are enough feature combinations - confirmed with 10 vars/features_to_combine=3 (175 new columns). .assign(**kwargs) does NOT fix this: it inserts columns one at a time internally too, same warning. The actual fix is building all new columns into one DataFrame and joining once (single insertion). Verified: output is byte-identical to the old behavior (pd.testing.assert_frame_equal on a 3000-row/9-var/129-tree case), drop_original still works, and a new regression test confirms the warning is gone (and fails against the old code, confirming it actually catches the regression). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follows the same pandas-native / narwhals-generic split established for GeoDistanceFeatures (this transformer also reimplements fit()/transform() directly, not via BaseCreation): .columns extraction, column reorder, prediction-column assignment, and drop_original all split by backend, consistent with every other operation in this module that's been benchmarked as a real (not minimal) loss when routed through narwhals on pandas.
Confirmed empirically before designing: sklearn's DecisionTreeRegressor/ Classifier and GridSearchCV accept a polars DataFrame directly for both fit() and predict()/predict_proba(), so the actual tree training/inference calls are unchanged - only the surrounding column selection, extraction, and reassembly needed migrating.
Fixed a pre-existing bug found while rewriting the exact code path it lived in: single-feature combos with an integer column name (e.g. DecisionTreeFeatures(features_to_combine=1) on a dataframe with columns 0, 1, ...) crashed, since the original
isinstance(features, str)check missed the int case and fell through to plain X[features] indexing, which returns a 1D Series rather than the 2D input sklearn requires. Widened to isinstance(features, (str, int)); verified the same single-feature narwhals path (get_column().to_frame()) already handles both cleanly.Regression, binary classification, and multiclass classification paths all verified to produce identical predictions between pandas and polars input. return_empty=True + polars remains untestable here too (same nw.col([]) bug in dataframe_checks.py found during CyclicalFeatures, still tabled) - this is the second transformer it blocks.
docs/user_guide/creation/DecisionTreeFeatures.rst is large (511 lines) and built around actual cross-validated tree fitting on the real California housing dataset across many sections - re-verified the cheap, deterministic parts (the raw data table) but did not re-run every tree-fitting example given the cost of repeated grid-search CV fits; unlike the other three creation-module docs this pass touched, the rest of this file's numbers are unverified. Added a self-contained "With polars" section using simple synthetic data instead, fully verified.