Skip to content

Narwhals drop missing data - #1016

Open
solegalli wants to merge 3 commits into
narwhals-migrationfrom
narwhals-drop-missing-data
Open

Narwhals drop missing data#1016
solegalli wants to merge 3 commits into
narwhals-migrationfrom
narwhals-drop-missing-data

Conversation

@solegalli

Copy link
Copy Markdown
Collaborator

No description provided.

solegalli and others added 3 commits August 24, 2026 21:53
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>
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