Skip to content

Migrate EqualFrequencyDiscretiser.fit() to narwhals, add polars support - #1039

Open
solegalli wants to merge 2 commits into
narwhals-migrationfrom
narwhals-equal-frequency-discretiser
Open

Migrate EqualFrequencyDiscretiser.fit() to narwhals, add polars support#1039
solegalli wants to merge 2 commits into
narwhals-migrationfrom
narwhals-equal-frequency-discretiser

Conversation

@solegalli

Copy link
Copy Markdown
Collaborator

Migrates EqualFrequencyDiscretiser.fit() to narwhals with polars support.

fit()'s only pandas dependency was pd.qcut(duplicates="drop") for quantile-based bin edges per variable. Replaced with np.quantile() on each column's narwhals-extracted numpy array + np.unique() to sort and drop duplicate edges — reproducing qcut's duplicates="drop" without any per-backend branch.

Getting a bit-exact match (not just close) took two fixes verified against pandas 3.0's qcut source:

  • pandas masks out NaN before calling np.quantile(values, qs, method="linear") itself rather than using np.nanquantile (not always bit-identical). Moot here — _fit_setup() already rejects NaN in variables_.
  • qcut nudges each quantile not exactly representable in base 2 up via np.nextafter (rounding up, not to nearest). Skipping this shifted edges by ~1e-13 and broke an existing exact-equality test.

With both applied, verified np.array_equal against real pd.qcut(retbins=True) across large random floats, many-duplicate data, all-identical data, negative floats, and n<q data.

Merge vs split: benchmarked old pd.qcut vs the new numpy+narwhals path at 10k/50k/100k rows × 1/2/10 cols — the new path is consistently faster than the old pandas-native code on both backends (narwhals-on-pandas 0.19x–0.47x of old pd.qcut, narwhals-on-polars 0.12x–0.46x). A narwhals-native quantile-expression alternative was 2–3x slower than old pd.qcut on pandas. No case for a split.

Verified: tests/test_discretisation unchanged (114 passed, 5 pre-existing check_estimator failures, reproduced on the unmodified branch tip). flake8 / mypy clean, sphinx -W clean. test_equal_frequency_discretiser.py rewritten to one parametrized test per behaviour over [pd.DataFrame, pl.DataFrame]. EqualFrequencyDiscretiser.rst examples verified against real output (two stale float digits in the binner_dict_ printout reproduce with the old pd.qcut fit — doc staleness, not a regression); "uses pandas.qcut() under the hood" line corrected, "With polars" section added.


Stacked on narwhals-discretisation-base (its own PR). Until that merges this PR's diff also contains the shared BaseDiscretiser commit; review that one first.

solegalli and others added 2 commits August 25, 2026 17:08
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>
fit()'s only pandas dependency was pd.qcut(duplicates="drop"), used to
compute quantile-based bin edges per variable. Replaced it with
np.quantile() on each column's narwhals-extracted numpy array, plus
np.unique() to sort and drop duplicate edges - reproducing qcut's
duplicates="drop" behaviour without any per-backend branch, since
values come from nw_X.get_column(var).to_numpy() regardless of
backend.

Getting a bit-exact match (not just numerically close) took two fixes
verified against pandas 3.0's pandas.core.reshape.tile.qcut source:
- pandas masks out NaN before calling np.quantile(values, qs,
  method="linear") itself, rather than using np.nanquantile - the two
  are not always bit-identical. Here this distinction is moot in
  practice: _fit_setup() already rejects NaN in variables_, so no
  masking is needed - values reaching the loop are already NaN-free.
- qcut nudges each quantile that isn't exactly representable in base 2
  up via np.nextafter (np.linspace(0, 1, q+1) then
  np.putmask(quantiles, q*quantiles != np.arange(q+1),
  nextafter(quantiles, 1))), rounding up rather than to nearest.
  Skipping this shifted bin edges by ~1e-13 versus real pd.qcut
  output and broke an existing exact-equality test.
With both applied, verified bit-exact (np.array_equal) against real
pd.qcut(retbins=True) across large random floats, many-duplicate-value
data, all-identical-value data, negative floats, and n<q data.

Benchmarked old pd.qcut vs the new numpy+narwhals path at 10k/50k/100k
rows x 1/2/10 columns: the new path is consistently faster than the
old pandas-native code on BOTH backends (narwhals-on-pandas lands at
0.19x-0.47x of old pd.qcut's time, narwhals-on-polars at 0.12x-0.46x,
both converging to roughly 2x faster at realistic 50k-100k row sizes).
A narwhals-native quantile-expression alternative was also benchmarked
(one nw.col(var).quantile(qi) expr per quantile point, batched into a
single select()) - fast on polars but 2-3x *slower* than old pd.qcut
on pandas, since narwhals translates each expr to a separate
Series.quantile call there. Given the numpy path beats old pandas on
both backends, there was no case for a pandas fast-path split.

Verified: tests/test_discretisation full suite unchanged (114 passed,
5 pre-existing failures in test_check_estimator_discretisers.py,
reproduced identically on the unmodified branch tip - sklearn's
check_estimator feeds raw numpy arrays, rejected since the narwhals
migration's dataframe-only contract). flake8 and mypy clean. Module
imports with pandas blocked (loaded standalone, since sibling
discretiser files in this package aren't migrated yet). sphinx -W
build clean (only the pre-existing unrelated linkcode_resolve
warning).

test_equal_frequency_discretiser.py rewritten per AGENTS.md: each
behaviour is now one test parametrized over
@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame])
rather than pandas-only.

docs/user_guide/discretisation/EqualFrequencyDiscretiser.rst: verified
every code example against real current output. The `disc.binner_dict_`
printout had two stale float digits (8099.200000000003 ->
...004, 1601.6000000000001 -> ...004, 1717.6999999999998 ->
1717.7000000000003) - reproduced identically with the OLD pd.qcut-based
fit() on the same dataset/pandas version, so this predates the
migration and is a doc-staleness issue, not a regression. Also
corrected the "uses pandas.qcut() under the hood" line and added a
"With polars" section with a verified worked example.

Co-Authored-By: Claude Sonnet 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