Narwhals drop missing data - #1016
Open
solegalli wants to merge 3 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>
Split transform()/return_na_data() by backend: benchmarked (10k-100k rows x 1-10 cols) a numpy-backed pandas mask (X[vars].notna().to_numpy().sum (axis=1) / .isnull().to_numpy().any(axis=1)) against both pandas' own axis=1 isnull()/notna().sum() and a narwhals-generic any_horizontal/ sum_horizontal path on pandas input. The numpy mask won consistently - e.g. the threshold check at 100k rows x 10 cols: 1.04ms numpy vs 4.56ms pandas-native vs 2.64ms narwhals-on-pandas (up to ~9x over the naive narwhals path, since pandas' axis=1 reductions are a known-slow case) - so pandas keeps this dedicated fast path; polars/other backends use narwhals' any_horizontal/sum_horizontal, which is fastest of all on native polars input. fit()'s missing_only variable-detection loop keeps the same pandas-loop/narwhals-null_count() split already established by MissingIndicator's migration. Found and fixed a real, pre-existing complementary-logic bug in return_na_data(): its threshold branch computed `isnull_frac >= threshold` as "dropped", when the true complement of transform()'s dropna (kept if non-null count >= n_vars*threshold) is `non_null_count < n_vars*threshold`. These aren't algebraic complements except by coincidence at threshold=0.5, and even there the boundary row was double- counted: kept by transform() AND returned by return_na_data(). Verified against the old code (predates this migration, present on origin/main): with threshold=0.5, transform() kept row 2 (2/4 non-null, meets the threshold) while return_na_data() also returned it; at threshold=1 the bug was worse - return_na_data() silently dropped 2 of 3 truly-missing rows from its output entirely. Fixed by deriving transform() and return_na_data() from one "keep" mask/expression, negated for the drop side (_select_rows(X, keep)), so the two outputs are an exact partition by construction - added test_transform_and_return_na_data_partition_input to verify this explicitly across every threshold value, plus corrected test_return_na_data_method's threshold=0.5 expectation, which had baked the bug's wrong output into the assertion. Also fixed find_all_variables(X, self.return_empty) - a positional-arg bug (return_empty was landing in the exclude_datetime slot) present on origin/main; the same bug pattern is repeated in random_sample.py, categorical.py and missing_indicator.py but those are out of scope here. Guarded the narwhals row-filter path against variables_ == [] (a real case: missing_only=True on a clean training set finds nothing to check) since narwhals' any_horizontal/sum_horizontal raise on an empty expression list, unlike pandas' dropna(subset=[]) which silently keeps every row - added a test for it. Fixed a latent bug in TransformXyMixin.transform_x_y's narwhals branch: it injects a temporary row-index column before calling self.transform(), but BaseImputer._transform() validates X's column count/names against feature_names_in_/n_features_in_ first and rejected the extra column - this combination (TransformXyMixin + a strict-validating transform()) was never exercised before since no prior narwhals migration combined both on a row-dropping transformer. Fixed by widening feature_names_in_/ n_features_in_ just for that call and restoring them after. Rewrote tests as one parametrized test per behavior over pd.DataFrame/pl.DataFrame with a shared DATA dict, replacing pandas .index-based assertions (meaningless for polars) with value-based checks via a backend-agnostic _cols() helper. Verified: tests/test_imputation full suite unchanged except for the new cases (106 passed, same 7 pre-existing test_check_estimator_imputers.py failures that predate this change). flake8 clean; mypy clean on this file, and introduces zero new errors in mixins.py (8 pre-existing attr-defined errors, inherent to the mixin pattern, unchanged). Module's own import chain verified pandas-free with pandas blocked, run successfully against polars input. Every doc example in DropMissingData.rst re-verified against actual output; added a "With polars" section. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…n't set widened self.feature_names_in_/n_features_in_ unconditionally to smuggle a row-index marker column through transform()'s column-count validation. DropMissingData's own tests exercise transform_x_y() before fit() has run in some paths, where feature_names_in_ doesn't exist yet, raising AttributeError. Guard with hasattr() so the widening only happens when there's something to widen - identical behavior for every caller that already had feature_names_in_ set (OutlierTrimmer, forecasting base), verified via the existing mixin/imputation test suites. 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
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.