Skip to content

Move animation key encoding to Rust - #319

Merged
Alek99 merged 1 commit into
mainfrom
agent/native-animation-key-encoding
Jul 26, 2026
Merged

Move animation key encoding to Rust#319
Alek99 merged 1 commit into
mainfrom
agent/native-animation-key-encoding

Conversation

@Alek99

@Alek99 Alek99 commented Jul 26, 2026

Copy link
Copy Markdown
Member

Summary

  • add an ABI v43 Rust kernel that canonicalizes and hashes homogeneous fixed-width animation keys in one row scan
  • route eligible NumPy arrays and homogeneous Python sequences through the native path while preserving the existing 64-bit identities and duplicate-row errors
  • keep mixed objects, dates, NUL-bearing Python sequences, non-finite diagnostics, and excessively padded string/bytes sequences on the Python oracle
  • add parity, endian, boundary, error, ABI, memory-amplification, and CodSpeed coverage; update the animation/Rust-engine specs and changelog

Why

Stable animation key= values were tokenized and hashed one row at a time in Python. At 100k rows, the Python object/token/hash loop dominated encoding even though the result is only two u32 columns. The native kernel borrows the fixed-width input, emits caller-owned lo/hi planes directly, and preserves the public wire contract.

The Python routing guard caps fixed-width sequence temporaries at 256 MiB and falls back when padding exceeds 8× source bytes, avoiding pathological N × longest-key memory amplification.

Performance

Local CodSpeed best time for test_animation_encode_100k_stable_keys:

  • main: 83.3 ms
  • this branch: 12.1 ms
  • improvement: ~6.9×

Verification

  • rustup run 1.97.0 cargo test — 129 passed
  • rustup run 1.97.0 cargo clippy --all-targets -- -D warnings
  • rustup run 1.97.0 cargo build --release
  • .venv/bin/python scripts/abi_smoke.py — 138 checks passed
  • .venv/bin/pytest tests/test_animation.py -q — 79 passed
  • uv run --with pre-commit pre-commit run --all-files
  • uv run ruff check .
  • uv run ruff format --check .

@coderabbitai

coderabbitai Bot commented Jul 26, 2026

Copy link
Copy Markdown

Warning

Review limit reached

@Alek99, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 59 seconds

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

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 configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 28201bd1-529f-47d9-b1a9-94771a108919

📥 Commits

Reviewing files that changed from the base of the PR and between d35731d and 96b5a2c.

📒 Files selected for processing (12)
  • CHANGELOG.md
  • benchmarks/test_codspeed_animation.py
  • python/xy/_native.py
  • python/xy/components.py
  • python/xy/kernels.py
  • scripts/abi_smoke.py
  • spec/design/animation.md
  • spec/design/rust-engine.md
  • spec/process/perf-audit-2026-07-22.md
  • src/lib.rs
  • src/transition.rs
  • tests/test_animation.py
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch agent/native-animation-key-encoding

Comment @coderabbitai help to get the list of available commands.

@codspeed-hq

codspeed-hq Bot commented Jul 26, 2026

Copy link
Copy Markdown

Merging this PR will improve performance by 98.77%

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 2 improved benchmarks
✅ 101 untouched benchmarks
⏩ 2 skipped benchmarks1

Performance Changes

Benchmark BASE HEAD Efficiency
test_animation_encode_100k_stable_keys 884.6 ms 349.7 ms ×2.5
test_animation_keyed_payload_100k 2.9 ms 1.9 ms +56.17%

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 agent/native-animation-key-encoding (96b5a2c) with main (d35731d)

Open in CodSpeed

Footnotes

  1. 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.

@Alek99
Alek99 marked this pull request as ready for review July 26, 2026 22:47
@Alek99
Alek99 merged commit bb0059c into main Jul 26, 2026
28 checks passed
@Alek99
Alek99 deleted the agent/native-animation-key-encoding branch July 26, 2026 23:08
Alek99 added a commit that referenced this pull request Jul 26, 2026
#319 landed `xy_transition_keys_fixed` as ABI 43 while this branch was
claiming 43 for `xy_range_indices_rows`. Both edits touch the same two
lines with the same value, so the 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 then load against the other's wrapper and
fail at dlsym instead of at the version check that exists to prevent it.
Alek99 added a commit that referenced this pull request Jul 26, 2026
#319 landed `xy_transition_keys_fixed` as ABI 43 while this branch was
claiming 43 for `xy_range_indices_rows`. Both edits touch the same two
lines with the same value, so the 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 then load against the other's wrapper and
fail at dlsym instead of at the version check that exists to prevent it.
Alek99 added a commit that referenced this pull request Jul 26, 2026
#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.
Alek99 added a commit that referenced this pull request Jul 26, 2026
…327)

#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.
Alek99 added a commit that referenced this pull request Jul 26, 2026
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.
Alek99 added a commit that referenced this pull request Jul 27, 2026
* 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.
Alek99 added a commit that referenced this pull request Jul 27, 2026
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.
Alek99 added a commit that referenced this pull request Jul 27, 2026
…#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.
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