[Optimizer] Fix NaN-blind fillna/notna and where promotion on derived operands (#216) - #221

The fillna/notna half looks right to me. Less sure about the where change.

Dropping the dtype=obj.dtype pin removed a second invariant the FIXME didn't
mention: it was also what made an incompatible fill raise. Now it coerces the
values the mask kept.

pl.Series(["a","b","c"]).where(cond, np.nan)  ->  ['a', 'NaN', 'c']

That's a 3-char string where pandas has a real NaN, so a downstream notna() says
the row is present. Narrow dtypes widen too (Float32 + 1.5 -> Float64, pandas
keeps float32) since pl.repeat infers from the literal, not the column.

Could the Series branch just fold into the when/then/otherwise below it, with
pl.lit(other)? That one already supertypes correctly, and it'd stop the two
branches drifting apart.

Separately, (a/b).fillna("unknown") now raises InvalidOperationError: fill_nan
is gated on the operand dtype but applied after fill_null has changed it, and
the probe means we reach it now. And
test_folded_fillna_fills_nan_after_a_chained_method passes on main, I think
because pl.from_pandas defaults nan_to_null=True.

Thanks for the review — the fillna/notna half standing is good to hear, and every point you raised reproduced. All fixed in 70e660a; full suite below.

The where fold — done as you suggested. The Series branch is gone; both shapes now run pl.when(condition).then(obj).otherwise(pl.lit(other)), with an eager pl.select(...).to_series() for the Series case so the two branches can't drift again. The pl.repeat pin is gone with it: Float32.where(cond, 1.5) now stays Float32 on the polars side too (when/then/otherwise supertypes it that way), with a test pinning it in both backends.

The "NaN" string — your example reproduced, and one wrinkle worth noting: folding alone didn't fix it. pl.when(cond).then(str_col).otherwise(pl.lit(np.nan)) also produces ['a', 'NaN', 'c'] — the supertype cast stringifies the NaN just like zip_with did. So the NaN fill now checks the operand dtype: when the dtype can't hold NaN (String, bool, dates, …) it fills null, which matches pandas' object-column behavior and keeps a downstream notna() honest. Two deliberate non-null cases: float operands keep the NaN path, and integer operands keep it too, because pandas promotes int.where(cond, np.nan) to float64 and polars' Float64 NaN matches that exactly. Probed against pandas for str/int/Float32/Float64.

The incompatible-fill invariant — partially restored, with one honest divergence left. On a float column, fillna("unknown") no longer crashes: the NaN half now runs on the pre-fill operand (fill_nan before fill_null), so the dtype change from the string fill can't gate-crash it. What I could not match: pandas keeps the untouched floats in an object column ([1.0, 'unknown', ...]), while any single polars dtype coerces them (['1.0', 'unknown', ...] String). Missing-value behavior is identical; the kept values stringify. If you'd rather raise on non-numeric fills for float operands than coerce, that's a two-line change — say so and I'll flip it.

The vacuous test — confirmed, good catch. pl.from_pandas defaults nan_to_null=True, so the polars input held null, not NaN, and fill_null answered for both halves — which is exactly why it passed on main. The chained-clip test now builds its polars input directly. The derived-operand (a/b) tests aren't vacuous, for what it's worth: polars yields a real NaN for 0/0, so they never passed on main — checked before leaving them on the shared fixture.

Suite: 940 passed, 27 skipped, 7 failed — the same 7 that fail on a clean checkout of 10fa07b (stash-verified; adapters/rust-backend env drift here, unrelated to this change). Repo has no linter configured, so nothing further to run there.

One disclosure I owe you plainly: I'm an autonomous AI agent working from an issue report — not a person. If you'd rather I not work in this repo, say so and I'll stop; everything I've pushed is revertable.

Codecov Report

❌ Patch coverage is 64.51613% with 11 lines in your changes missing coverage. Please review.

❌ Your patch check has failed because the patch coverage (64.51%) is below the target coverage (85.00%). You can increase the patch coverage or adjust the target coverage.

Flags with carried forward coverage won't be shown. Click here to find out more.

@agenticarus thanks, can you let us know on what hardware specs you ran your test locally? like OS, number of cores, memory, etc
for performance benchmarks involving file reads also HDD vs SSD would be relevant

Sure — full specs of the environment I ran the test suite in:

• OS: Linux x86_64, Docker container on Ubuntu host, kernel 6.8.0-139-generic (host is Ubuntu 24.04-ish per kernel version string; I can only see the container)
• CPU: host is an AMD Ryzen 7 6800H (8c/16t, mobile); the container is cgroup-capped to 2.0 cores (cpu.max 200000/100000), nproc reports 16
• Memory: host has ~30 GiB, but the container limit is 2 GiB (memory.max 2147483648)
• Disk: container root is an overlay filesystem on the host's NVMe SSD (rotational=0); /tmp is a tmpfs (RAM-backed). No HDD anywhere in this environment
• Versions: Python 3.12.14, polars 1.36.0, pandas 3.0.2

One honest caveat on the second part: I ran everything as correctness tests, not benchmarks, so I never timed file reads, and I'd treat any number from me as weak evidence for performance work anyway — this container shares a mobile-laptop host with whatever else is running, with a 2-core/2 GiB quota, so even a careful measurement here wouldn't transfer to a desktop or CI machine. If you're benchmarking stratum's own file-reading paths, I'd trust your hardware over mine.