Migrate BaseDiscretiser to narwhals, add polars support - #1037
Open
solegalli wants to merge 1 commit into
Open
Conversation
Shared base for ArbitraryDiscretiser, EqualFrequencyDiscretiser,
EqualWidthDiscretiser and GeometricWidthDiscretiser (not
DecisionTreeDiscretiser, which extends a different base). Only
transform() needed migrating - _fit_setup(), _get_feature_names_in()
and _check_transform_input_and_state() are inherited unchanged from
BaseNumericalTransformer, already fully narwhals-migrated.
transform()'s only pandas dependency was pd.cut, applied per column to
sort values into the bins already fixed by fit() (binner_dict_).
Replaced it with a plain numpy implementation: pandas.cut is itself
built on bins.searchsorted() internally (verified against pandas 3.0's
_bins_to_cuts source), so np.searchsorted + the same include_lowest
index-1 special case reproduces its bin-index logic exactly, with no
per-backend branch needed - values come from
nw_X.get_column(feature).to_numpy() regardless of backend, and results
are re-attached via nw.new_series()/with_columns(), so the same code
path runs for pandas and polars.
Benchmarked old pd.cut vs the new numpy+narwhals path at 10k/50k/100k
rows x 1/2/10 columns:
- return_boundaries=False (bin codes): narwhals-on-pandas lands at
~1.0-1.2x of pandas-native at realistic sizes (50k-100k rows, the
~1.9x seen only at the smallest 10k-row/1-col case is fixed
per-call overhead, sub-millisecond either way) - minimal loss,
merged into a single path, no is_pandas split. narwhals-on-polars is
~1.0-1.3x *faster* than pandas-native at every size tested.
- return_boundaries=True (interval-label strings): the numpy path is
12-20x faster than pd.cut on pandas itself (e.g. 100k rows x 10
cols: 647ms old vs 40ms new) - pd.cut's Categorical/IntervalIndex
machinery has heavy per-call overhead that np.searchsorted plus
plain string formatting avoids entirely. polars is ~1.2x faster
still than the new pandas path.
Given both branches favour or are at parity with a single numpy-driven
path, there was no case for a pandas fast-path split here.
return_boundaries=True's interval-label formatting
("(lower, upper]" text, e.g. "(-0.001, 20.0]") replicates pandas.cut's
_round_frac/_infer_precision/lowest-edge-adjustment algorithm in pure
numpy so it works identically on both backends - verified against real
pd.cut(...).astype(str) output across positive/negative/duplicate-
inducing/inf-edge bins, and against the California housing dataset
used in the existing test. return_object=True now builds a nw.Object
column (narwhals' cross-backend equivalent of pandas' "O" dtype,
already used by variable_handling for categorical-column detection)
instead of a pandas-only astype("O") call.
Verified: tests/test_discretisation full suite unchanged (109 passed,
5 pre-existing failures in test_check_estimator_discretisers.py -
sklearn's check_estimator feeds raw numpy arrays, which check_X() has
rejected since the narwhals migration's dataframe-only contract;
reproduced identically on the unmodified file). Manually diffed
transform() output against real pd.cut() across ~10 edge cases (NaN,
out-of-range values on both ends, negative bins, exact-edge values,
precision auto-widening, single bin) plus the three sibling
discretisers' documented doctest examples (EqualWidthDiscretiser,
ArbitraryDiscretiser, EqualFrequencyDiscretiser value_counts()) -
all numerically identical to old pd.cut output; the "Name: x" vs
"Name: count" and bare-fit()-repr mismatches those doctests already
show are a pre-existing pandas-3.0 doc-staleness issue unrelated to
this migration (reproduced on the unmodified files too). flake8 and
mypy clean. Module imports with pandas blocked (loaded standalone,
since sibling discretiser files in this package are not yet migrated
and still import pandas at their own module level). sphinx -W build
clean (only the pre-existing unrelated linkcode_resolve warning).
test_base_discretizer.py's test_transform is now parametrized over
pd.DataFrame/pl.DataFrame per AGENTS.md - its MockClassFit hard-codes
binner_dict_ rather than actually fitting, so it needed no pandas-only
logic to begin with. The other four discretisers' own test files stay
pandas-only for now: their fit() methods still call pd.cut/pd.qcut
directly and aren't migrated by this branch.
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.
Shared base for
ArbitraryDiscretiser,EqualFrequencyDiscretiser,EqualWidthDiscretiserandGeometricWidthDiscretiser(notDecisionTreeDiscretiser, which extends a different base). Onlytransform()needed migrating —_fit_setup(),_get_feature_names_in()and_check_transform_input_and_state()are inherited unchanged from the already-migratedBaseNumericalTransformer.transform()'s only pandas dependency waspd.cutper column. Replaced with a plain numpy implementation:pandas.cutis itself built onbins.searchsorted()internally (verified against pandas 3.0's_bins_to_cutssource), sonp.searchsorted+ the sameinclude_lowestindex-1 special case reproduces its bin-index logic exactly, no per-backend branch — values come fromnw_X.get_column(feature).to_numpy()regardless of backend, results re-attached vianw.new_series()/with_columns().Merge vs split: benchmarked old
pd.cutvs the new numpy+narwhals path at 10k/50k/100k rows × 1/2/10 cols:return_boundaries=False(bin codes): narwhals-on-pandas ~1.0–1.2x at realistic sizes (the ~1.9x at 10k/1-col is sub-ms fixed overhead) — minimal loss, single path. narwhals-on-polars ~1.0–1.3x faster.return_boundaries=True(interval-label strings): the numpy path is 12–20x faster thanpd.cuton pandas itself (100k×10: 647ms → 40ms) —pd.cut's Categorical/IntervalIndex machinery has heavy per-call overhead. polars ~1.2x faster still.return_boundaries=True's"(lower, upper]"label formatting replicatespd.cut's_round_frac/_infer_precision/lowest-edge-adjustment in pure numpy (verified against realpd.cut(...).astype(str)across positive/negative/duplicate-inducing/inf-edge bins).return_object=Truebuilds anw.Objectcolumn instead of a pandas-onlyastype("O").Verified:
tests/test_discretisationunchanged (109 passed, 5 pre-existingcheck_estimatorfailures — raw numpy-array input, rejected bycheck_X()since the narwhals dataframe-only contract; reproduced on the unmodified file). Manually diffedtransform()output against realpd.cutacross ~10 edge cases plus the three sibling discretisers' doctest examples — all numerically identical. flake8 / mypy clean, sphinx -W clean.The other four discretisers' own test files stay pandas-only for now — their
fit()methods still callpd.cut/pd.qcutdirectly and aren't migrated by this branch.