Route DataFrame key columns through the native animation-key encoder - #327
Conversation
#319 moved stable animation `key=` encoding into Rust, but routed on the container rather than the values: only a bare ndarray or a Python list/tuple reached the kernel. `data=df, key="id"` resolves to a Series, so the documented idiom was the one shape that never took the fast path. Route on values instead. A `to_numpy`-carrying column unwraps to NumPy, and object storage qualifies whenever every row holds the same builtin scalar type — which is how a pandas string column arrives. The conversion guard that bounds fixed-width temporaries applies to all of them, so skewed key lengths still fall back. Also narrow the NUL rule to the case that is actually ambiguous. Only a key that *ends* in NUL is indistinguishable from fixed-width padding; interior NULs survive the round trip and now encode natively, matching what NumPy `U`/`S` records already did. Separate "declined this data" from "invalid layout" at the ABI. Both were status 1, and status 1 means "fall back to the oracle" — so a drift between the ctypes dtype gate and Rust's `valid_layout` would have cost the whole speedup with every assertion still green. Argument errors are status 4 and raise; the panic shield defaults there too. The CodSpeed row now asserts its input routes, since nothing about the result distinguishes the two paths. Interleaved A/B against 96b5a2c, 100k rows, medians of 3x9: list[str] 26.4 ms -> 21.8 ms list[int] 20.3 ms -> 14.9 ms object ndarray[str] 89.8 ms -> 22.5 ms pandas str Series 88.5 ms -> 23.5 ms pandas int64 Series 88.0 ms -> 12.3 ms The hot-path gain is from hoisting per-element `.item()`/`type()` work out of the guard into `set(map(type, ...))` and `map(len, ...)`. That costs a full scan before a mixed sequence can bail, so the early-bail fallback goes 94.2 ms -> 98.0 ms; it is dominated by the reference encode that follows either way. Docs: the F-order planes only ship copy-free while unreordered — the line-like geometry sort and finite-row selection both return C-order through advanced indexing. And the perf-audit finding is downgraded to "mostly resolved": keys are still encoded for every row before `_transition_entry` drops them above MAX_ANIMATION_MATCH_ROWS, now ~144 ms for 1M rows instead of ~800 ms. Verified: 2564 passed / 4 skipped; cargo test 130 passed; clippy clean; abi_smoke 140 checks; 200-draw randomized parity across list/tuple/ object-ndarray/Series/ndarray containers, and a 13-case pandas dtype matrix, all bit-identical to the Python oracle.
|
Warning Review limit reached
Next review available in: 21 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (10)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Merging this PR will improve performance by 26.9%
|
| Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|
| ⚡ | test_animation_encode_100k_stable_keys |
348.8 ms | 274.8 ms | +26.9% |
Tip
Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.
Comparing alek/animation-key-followups (ae21d81) with main (bb0059c)
Footnotes
-
2 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩
#319 landed `xy_transition_keys_fixed` as ABI 43 while this branch was claiming 43 for `xy_range_indices_rows`, and open PR #327 already claims 44. Both edits touch the same two lines with the same value, so such a merge is clean and nothing complains — which is the problem. One number would name two different ABIs, and `scripts/abi_smoke.py` compares the built library against the wrapper's constant, so it cannot catch a collision where both sides agree: a cdylib built from either tree would load against the other's wrapper and fail at dlsym instead of at the version check that exists to prevent exactly that. 45 is unclaimed across every open PR. It also fails loudly rather than silently if #327 lands second, because its 43 -> 44 edit then conflicts against a main at 45 instead of merging into agreement.
claiming 43 for `xy_range_indices_rows`, and open PR #327 already claims 44. Both edits touch the same two lines with the same value, so such a merge is clean and nothing complains — which is the problem. One number would name two different ABIs, and `scripts/abi_smoke.py` compares the built library against the wrapper's constant, so it cannot catch a collision where both sides agree: a cdylib built from either tree would load against the other's wrapper and fail at dlsym instead of at the version check that exists to prevent exactly that. 45 is unclaimed across every open PR. It also fails loudly rather than silently if #327 lands second, because its 43 -> 44 edit then conflicts against a main at 45 instead of merging into agreement.
* Scan selection candidates in place (ABI 42 -> 43) `select_range` prunes by zone map before it scans, and then threw the saving away: the pruned branch gathered `x[candidates]` and `y[candidates]` into two fresh f64 columns purely so it could call `range_indices`, which only knows how to walk a whole column. Sixteen bytes per candidate, which produces the backwards result that selecting *part* of a trace costs far more than selecting all of it — on a 10M-row scatter, a half-domain box peaked at 134.6 MB while a full-domain box peaked at 38.1 MB. `xy_range_indices_rows` is the rectangular twin of the lasso's `xy_polygon_select`: same `(x, y, len, rows, n_rows, ..., out)` shape, same inclusive comparison chain as `range_scan_scalar`, same order-preserving gather-down across worker segments. The pruned branch now reads through the row ids and gets canonical ids straight back, so both the gathers and the `candidates[hits]` re-index are gone. The half-domain case drops to 57.6 MB and the full-domain case is untouched. ABI_VERSION goes 42 -> 43 in `src/lib.rs` and `python/xy/_native.py` together. Parity is the gate, not inspection: the new kernel is tested against `range_indices` over the gathered form for five windows (including empty, degenerate-point and whole-domain) crossed with five row subsets and five sizes up to 100k, plus a 2M-row case that exercises the threaded path, NaN and infinities on both axes, inclusive bounds, unsorted row order, invalid windows and an out-of-range row id. `test_selection_prunes_non_overlapping_zone_chunks` now also asserts that exactly one scan happens and that it reads the whole column, which is what proves the gather is gone. 118 Rust tests, 133 ABI smoke checks and the full Python suite pass; a 97-artifact output fingerprint sweep is identical to main. * Answer an out-of-range row id instead of aborting on it `range_indices_rows` and `polygon_select` index `x[row]` and left the report to Rust's bounds check, which panics. `ffi_guard` turns a panic into the entry point's sentinel only where panics unwind, and the PyEmscripten wheel is built `-C panic=abort` precisely so they cannot: there the same id aborted the Pyodide instance (exit 134) rather than raising ValueError. `xy.kernels` is public API, so these row arrays are caller data — the documented contract has to hold on every target, not just the ones that unwind. Both kernels now read through `get` and return `Option`, with the FFI layer mapping `None` to the existing sentinel. This is free: the same two bounds checks happen either way, so only the reporting changes. A separate validating pass over `rows` was not free and is deliberately not what this does — one serial sweep of 5M u32 cost more than the entire parallel scan (1.0 ms -> 2.4 ms), where reading through `get` measures 1.01 ms against 1.02 ms, and `polygon_select` 36.3 ms against 36.4 ms. * Take ABI 45: 43 shipped with #319 and #327 claims 44 claiming 43 for `xy_range_indices_rows`, and open PR #327 already claims 44. Both edits touch the same two lines with the same value, so such a merge is clean and nothing complains — which is the problem. One number would name two different ABIs, and `scripts/abi_smoke.py` compares the built library against the wrapper's constant, so it cannot catch a collision where both sides agree: a cdylib built from either tree would load against the other's wrapper and fail at dlsym instead of at the version check that exists to prevent exactly that. 45 is unclaimed across every open PR. It also fails loudly rather than silently if #327 lands second, because its 43 -> 44 edit then conflicts against a main at 45 instead of merging into agreement. * Range-check row ids before the u32 cast, not after The out-of-range guarantee the previous commit made stops at the ctypes boundary: the wrapper reached the kernel through `ascontiguousarray(rows, dtype=np.uint32)`, an unchecked C cast. An id of `2**32 + 3` arrives as row 3 — in range, indistinguishable from a real id, so the bounds check passes and the call returns a row the caller never asked for. `kernels.range_indices_rows(x, y, [2**32 + 3], ...)` answered `[3]`; a negative id only errored by accident, because `-1` wraps to 4294967295 and that happened to exceed the column length. Both row-restricted kernels now widen and range-check before casting. The internal caller passes uint32 already (`_expand_zone_chunks`), which takes the pass-through branch, so the hot path is untouched — 5M candidates still 1.00 ms. * Refuse float row ids instead of casting them The range check widened through `astype(np.int64)`, which still let floats through silently and, worse, decided them differently per architecture: `[3.9]` selected row 3, and `[nan]` selected row 0 on arm64 because the NaN cast saturates, where x86_64 traps it to INT64_MIN and the range check would have caught it. A guard whose answer depends on the wheel is not a guard. The cast also emitted a RuntimeWarning from inside the wrapper. Integer dtypes only now, with an empty array exempted because `np.asarray([])` is float64 and selecting no rows is not a dtype error.
Rebasing onto #319 made `result` Fortran-order — the native encoder's layout, so the payload ships `values[:, 0]` copy-free — and that breaks the `reshape(-1).view(np.uint64)` the uniqueness test used. For a one-row prefix numpy can satisfy the reshape with a strided *view* rather than a copy, and re-viewing that as a wider dtype raises "the last axis must be contiguous". Every fallback key error at row 1 came back as that numpy message instead of its own: five tests on main caught it, including `test_invalid_stable_keys_fail_clearly` and the DataFrame missing-value case. Both columns are contiguous in this order, so the digests are packed from them directly. Correct for either layout, same 8 bytes a row the reshape copy cost, and no dependence on when numpy decides to copy. The packing is hi-word-first rather than the wire's little-endian order, which is fine because only injectivity matters here. Date keys are the case that reaches this — they always take the Python fallback, as do mixed types, trailing-NUL strings and oversized ints, while #319 and #327 route str/int/float/object-ndarray to Rust. At 400k date keys the fallback is 404 ms and 17.1 MB peak against main's 496 ms and 100.1 MB, and the duplicate paths are back to parity (7.6 vs 6.4 MB early, 101.1 vs 100.1 MB late) rather than double.
…#321) * Cap, share and narrow: four more peak-memory cuts in the Python paths Second-pass audit follow-ups, each byte-identical and each measured on its own workload. None of them changes what any path emits; a 97-artifact fingerprint sweep (payloads, SVG/PNG/JPEG/WebP/HTML exports, selections through both zone pruning branches, density views, append, pyplot) is identical to main. - Selection events built `sorted(int(v) for v in selected[tid])` over the whole selection and then sliced it to SELECTION_EVENT_ID_LIMIT. That is a Python int per selected row, ~40 bytes each, to keep at most 10,000 of them. Partition the array first and box only the survivors: a 3M-row box selection peaks at 37.2 MB instead of 138.1 MB. The ids are still the numerically smallest ones in ascending order, and the shared id budget still drains in trace order. - A truecolor heatmap ingested `rgba[..., 0].reshape(-1)` twice — once as the scalar `grid`, once as `rgba_grid[0]`. The source is a strided view, so each reshape materializes its own contiguous copy, and the ColumnStore kept both for the figure's lifetime. Reuse the column: a 1600x1200 RGB heatmap holds 58.6 MB of canonical columns instead of 73.2 MB. - `_encode_transition_keys` carried `seen: dict[bytes, int]` and `digests: dict[bytes, bytes]` across the whole loop to produce eight bytes per row. Equal tokens hash equally, so distinct digests prove both distinctness and no collision — one `np.unique` over the result stands in for both dicts, and a conflict re-walks with the original bookkeeping so the message and the rows it names are unchanged. 200k keys peak at 7.1 MB instead of 48.1 MB. - PNG chunks are built as join-ready parts. `tag + data` copied the compressed IDAT once, wrapping it copied it again, and each `+` in the document chain copied it once more. CRC is accumulated over the tag and then the data, which is exactly `crc32(tag + data)` without materializing the concatenation: a 1800x1400 truecolor encode peaks at 32.6 MB instead of 49.1 MB. Ruff clean. * Keep a duplicate key ahead of a later invalid one Moving the uniqueness test after the row walk also moved it behind every later bad row, so an invalid key masked an earlier duplicate entirely: `["a", "a", None]` reported `duplicate value at rows 0 and 1` before this and `animation key is missing at row 2` after it. Same for a non-finite or wrong-typed row behind a duplicate. The claim that the messages and the rows they name were unchanged only held when the input was wrong in one way at a time. On a token error the rows already walked have complete digests, so a conflict among them belongs to an earlier row than the one that failed; re-check that prefix before propagating, and let the existing re-walk produce the message. Error path only — the walk itself is untouched, and 200k unique keys still peak at 7.5 MB against the merge base's 49.9 MB. * Do not let the superseded key error surface, or narrow it to ValueError Two follow-ups on the prefix re-check. It runs inside an `except`, so the token error it supersedes was being chained onto the duplicate it raises instead. The traceback then led with `animation key is missing at row 2` and a "During handling of the above exception" banner before reaching the duplicate — printing, first, the one message the re-check exists to avoid reporting. Every raise in the re-walk is now `from None`; the second call site is outside any handler, where that costs nothing. 31 traceback lines back down to 24. And the handler caught ValueError, when the rule it implements is that the first bad row wins — which row that is does not depend on what the token raised. A key type whose `encode`, `isoformat` or `__format__` raises anything else could still slip past an earlier duplicate; catching Exception closes it, and a token failure with no earlier duplicate still propagates unchanged. * Test key uniqueness on growing prefixes, not once at the end Testing only after the walk meant a duplicate paid for the whole walk and then a full re-walk, which made peak *worse* than the dictionaries this replaced — on the one path a memory change should never regress. At 2M keys with a non-unique id column, the ordinary form of this mistake: merge base 16.2 ms 32.0 MB peak this branch 1636 ms 65.1 MB peak with prefix 18.1 ms 33.1 MB peak Uniqueness is now tested on each prefix as it completes, starting at 4096 rows and growing x8, so both the wasted walk and the `np.unique` transient are bounded by the rows needed to reach the duplicate rather than by the column. Growth is x8, not x2, to keep the intermediate tests at roughly an eighth of the final one's cost. Blocked rather than one compare per row: `arr[start:stop]` is a view of an object array, so the inner loop is unchanged, where a per-row compare cost 8% of the success path at 2M rows. The success path pays ~2% for the insurance and is still 14% faster than the merge base (1709 ms vs 1984 ms) at 6.7x lower peak (65.0 MB vs 437.7 MB). A late duplicate is still ~1.6x the merge base in time (3.0 s vs 1.9 s), since the walk must complete before the re-walk can name the rows; its peak is at parity. * Pack prefix digests from the columns, not through a reshape Rebasing onto #319 made `result` Fortran-order — the native encoder's layout, so the payload ships `values[:, 0]` copy-free — and that breaks the `reshape(-1).view(np.uint64)` the uniqueness test used. For a one-row prefix numpy can satisfy the reshape with a strided *view* rather than a copy, and re-viewing that as a wider dtype raises "the last axis must be contiguous". Every fallback key error at row 1 came back as that numpy message instead of its own: five tests on main caught it, including `test_invalid_stable_keys_fail_clearly` and the DataFrame missing-value case. Both columns are contiguous in this order, so the digests are packed from them directly. Correct for either layout, same 8 bytes a row the reshape copy cost, and no dependence on when numpy decides to copy. The packing is hi-word-first rather than the wire's little-endian order, which is fine because only injectivity matters here. Date keys are the case that reaches this — they always take the Python fallback, as do mixed types, trailing-NUL strings and oversized ints, while #319 and #327 route str/int/float/object-ndarray to Rust. At 400k date keys the fallback is 404 ms and 17.1 MB peak against main's 496 ms and 100.1 MB, and the duplicate paths are back to parity (7.6 vs 6.4 MB early, 101.1 vs 100.1 MB late) rather than double.
Follow-ups from the review of #319. Correctness there held up under a 400-case fuzz and a 29-dtype sweep — these are the coverage, contract, and doc-accuracy gaps that review turned up.
The main one: the documented idiom never took the fast path
#319 routed on the container — bare
ndarray, orlist/tuple. But_resolve(data, key)returnsdata[key], sodata=df, key="id"hands the encoder a pandas Series, and a string column arrives as object storage even afterto_numpy(). Both fell straight through to the reference encoder.Routing now follows the values: a
to_numpy-carrying column unwraps to NumPy, and object storage qualifies whenever every row holds the same builtin scalar type. The fixed-width conversion guard applies to all of them, so skewed key lengths still fall back and still bound the temporary.Narrower NUL rule
Only a key that ends in NUL is ambiguous against fixed-width padding. Interior NULs survive the round trip, so they now encode natively — which is what NumPy
U/Srecords already did, and what #319 already tested for them.Argument errors are no longer a silent slow path
Status 1 meant both "this data isn't mine, use the oracle" and "you sent a layout the ABI doesn't define". Since status 1 falls back, any future drift between the ctypes dtype gate in
_native.pyand Rust'svalid_layoutwould have silently cost the entire speedup — and no assertion would have caught it, because the F-order result the new benchmark checks is now produced by both paths.Argument errors are status 4 and raise. The
ffi_guardpanic default moves there too. The CodSpeed row asserts its input actually routes. ABI 43 → 44.Measurements
Interleaved A/B against 96b5a2c, 100k rows, medians of 3×9:
list[str]list[int]The hot-path gain comes from hoisting per-element
.item()/type()out of the guard intoset(map(type, ...))andmap(len, ...). Tradeoff: that scans the whole sequence before a mixed one can bail, so the early-bail fallback goes 94.2 ms → 98.0 ms. It's dominated by the reference encode that runs either way, and a Python-level early-exit loop would be slower on the homogeneous case we care about.Docs corrected
keys[np.argsort(...)]) and any finite-row selection (keys[sel]) both return C-order from NumPy advanced indexing, putting the per-column copy back at ship time.animation.mdand the_native.pydocstring said otherwise._transition_entrydiscards them aboveMAX_ANIMATION_MATCH_ROWS. A 1M-row keyed mark spends ~144 ms on identities nothing reads — down from ~800 ms, so much less urgent, but not gone.Verification
pytest tests/— 2564 passed, 4 skippedcargo test— 130 passed;cargo clippy --all-targets -- -D warningscleanscripts/abi_smoke.py— 140 checksruff@0.15.8 check/format --checkcleanInt64/boolean,category,datetime64, shuffled index), 8-case object-ndarray matrix, 7-case NUL-placement matrix, and 200 randomized draws replayed across list / tuple / object-ndarray / Series / typed-ndarray containers — all bit-identical to the Python oracle, withmissing at row Nandfinite at row Ndiagnostics preserved.Four
test_transition_key_object_policy_uses_python_referenceparams moved to the parity test: their routing widened on purpose.