Narwhals random sample imputer - #1014
Open
solegalli wants to merge 2 commits into
Open
Conversation
Shared base for the imputation module: _transform() (fit-state checks + column reorder) and transform() (fillna via imputer_dict_) are now dataframe-agnostic, with _get_feature_names_in() reading columns through narwhals on non-pandas input. Benchmarked the fillna step (select + fill from a per-column value dict) at 10k/100k/1M rows x 1/2/10 columns: pandas-native fillna runs ~1.3-1.6x faster than the narwhals-generic fill_null equivalent at the 10k-100k row sizes imputers are normally used at (the gap narrows to ~1.0x only past ~1M rows) - a real, not minimal, loss, so pandas keeps its own fast path (is_pandas = nwd.is_pandas_dataframe(X); if is_pandas is True: ... else narwhals fill_null per column). Also benchmarked a numpy rewrite (to_numpy + np.where per column, mirroring RelativeFeatures) but it did not beat pandas-native and was consistently slower than narwhals fill_null on polars, so it wasn't adopted here - unlike RelativeFeatures' arithmetic, a plain value fill is already close to a no-op for both pandas and narwhals/polars, leaving no room for a numpy win. The pandas<3 fillna-downcasting workaround (option_context + infer_objects) is preserved on the pandas branch but no longer imports pandas at module level - the module is fetched via nw.from_native(X).__native_namespace__() only once X is already confirmed to be a pandas dataframe, so no import is attempted on a polars-only install. Verified: tests/test_imputation full suite unchanged (95 passed, 7 pre-existing failures in test_check_estimator_imputers.py - sklearn's check_estimator feeds raw numpy arrays, which check_X() has always rejected per the narwhals migration's dataframe-only contract, predates this change). flake8 and mypy clean on the file. Module imports with pandas import blocked. sphinx -W build clean (only the pre-existing unrelated linkcode_resolve warning). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
fit()/transform() now accept pandas, polars, or any narwhals-supported dataframe. Split (not merged) into a pandas branch and a narwhals branch, mirroring BaseImputer's pattern, because this transformer stores a copy of the training data and draws random values from it - a correctness concern, not just a performance one. RNG/reproducibility decision: pandas' .sample() and polars'/narwhals' .sample() are backed by different random number generators, so they never draw the same values for the same seed even on identical data - this was already true within pure pandas usage across pandas versions in some cases, but is guaranteed different across backends. The contract adopted and documented (class docstring + new "With polars" user guide section) is "same seed, same backend -> same result", not cross-backend value parity. The pandas branch is the pre-migration code verbatim (still X.loc/.sample(random_state=...)/index reassignment, called directly on the pandas object already in hand - no pandas import needed per AGENTS.md), so existing pandas users see bit-identical sampled values after upgrading, seed-for-seed. The narwhals branch is a positional reimplementation for polars and other backends: null positions come from Series.is_null().arg_true(), replacement values come from Series.sample(n, with_replacement=True, seed=...) drawn from the stored training-data pool, and values are written back with Series.scatter() (mirrors the exact usage in narwhals' own scatter() docstring example). For seed="observation", pandas' per-row .loc-based seed lookup (_define_seed, kept pandas-only and untouched) is replaced for the narwhals branch by a single vectorized numpy pass over the seed columns (X.select(seed_vars).to_numpy() + sum/prod per row), since narwhals dataframes have no row-label-based access to loop against. Benchmarked fit()+transform() at 10k/50k/100k rows x 1/2/10 cols: the narwhals-generic (scatter-based) implementation running on pandas input is actually close to or faster than the pandas-native .loc-based implementation at most sizes (0.7-1.3x), so throughput alone would have allowed merging into one code path. The split is driven entirely by the backward-compatibility requirement above (existing users' random_state values must keep drawing the exact same pandas samples they did before this migration) rather than by a performance loss. Rewrote tests/test_imputation/test_random_sample_imputer.py: behavioral tests (general seed, per-observation seed with add/multiply/single variable, categorical dtype preservation, the input-validation error paths that touch a dataframe) are now single tests parametrized over pd.DataFrame/pl.DataFrame, asserting the backend-agnostic invariants that actually hold for this transformer (no nulls remain, every filled value came from the training pool, same seed + same backend reproduces the same result) rather than literal values, since literal sampled values are inherently backend-specific here. _define_seed's own test stays pandas-only (it exercises .loc label access directly, which has no narwhals equivalent). Added one dedicated pandas-only regression test asserting the exact historic literal values are unchanged post-migration, protecting the backward-compatibility guarantee above. Verified: full tests/test_imputation suite goes from 95 passed/7 pre-existing failures (baseline, via git stash) to 102 passed/same 7 pre-existing failures (MeanImputer et al. failing because sklearn's check_estimator feeds raw numpy arrays, which check_X() has rejected since the narwhals migration began - confirmed unrelated to this file). flake8 and mypy clean. sphinx -W build clean (only the pre-existing unrelated linkcode_resolve warning). random_sample.py itself contains no `import pandas` and loads standalone with pandas blocked; the feature_engine.imputation package as a whole still fails to import with pandas blocked, but only because arbitrary_imputer.py (untouched by this change, pre-existing on narwhals-imputation-base) still has a module-level `import pandas as pd` - out of scope here, flagged separately. 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.
No description provided.