Cap, share and narrow: four more peak-memory cuts in the Python paths - #321
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (8)
🚧 Files skipped from review as they are similar to previous changes (6)
📝 WalkthroughWalkthroughThe changes optimize PNG byte assembly, selection row-ID truncation, transition-key digest validation, and truecolor heatmap memory reuse. New tests cover serialized PNG chunks, selection limits, transition-key behavior, and heatmap memory sharing. ChangesPNG chunk assembly
Selection row ID truncation
Transition key encoding
Truecolor heatmap plane sharing
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Merging this PR will not alter performance
Comparing 
Footnotes
|
48c479b to
b73042f
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/xy/channel.py`:
- Around line 157-164: The canonical selection IDs must be transferred as raw
f32 buffers instead of JSON numbers. In python/xy/channel.py lines 157-164,
preserve the existing cap and sort behavior but encode the selected IDs into the
established raw f32 reply-buffer representation rather than assigning
kept.tolist() to canonical_row_ids. Update tests/test_selection_rows.py lines
122-187 to decode the buffer response and assert the expected ordering and
truncation cap instead of checking JSON lists.
In `@tests/test_density_mean_color.py`:
- Around line 563-571: Update the RGBA plane assertions after constructing
planes from spec["columns"] to require exactly four planes, then compare each
plane by index with the corresponding channel in rgba in serialized order.
Replace the unordered any(...) matching loop while preserving the existing
float32 conversion and allclose tolerance.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 20d795e8-d835-4f89-939d-49c501beaf41
📒 Files selected for processing (8)
python/xy/_png.pypython/xy/channel.pypython/xy/components.pypython/xy/marks.pytests/test_animation.pytests/test_density_mean_color.pytests/test_png_export.pytests/test_selection_rows.py
b73042f to
24e8b24
Compare
|
Thanks — took the second one, declining the first. Applied: RGBA plane ordering ( One correction to the framing, so the test comment isn't misleading: the Declining: canonical selection IDs as raw f32 buffers. Three reasons:
Happy to open a follow-up issue if a u32 buffer for selection ids is wanted on |
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.
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.
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.
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.
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.
bef1d7b to
152e1d1
Compare
1 participant
Four more peak-memory cuts from the second audit pass (follow-up to #309), all
Python-side and all byte-identical. Sibling PRs carry the encode paths (#320),
the Rust core, the selection kernel, and the legend cache.
Each is measured on its own workload rather than asserted.
Selection events boxed every selected row, then threw them away
canonical_row_idsships at mostSELECTION_EVENT_ID_LIMIT(10,000) ids, butthe code was:
That materializes a Python int per selected row — about 40 bytes each with
pymalloc rounding — and sorts the whole list, to keep at most 10,000 of them. A
5M-row lasso built ~200 MB of
PyLongto deliver ~360 KB.Cap on the array first with
argpartition, then box only the survivors. The idsare still the numerically smallest ones in ascending order, and the shared id
budget still drains in trace order, so the reply is unchanged.
A 3M-row box selection peaks at 37.2 MB instead of 138.1 MB.
A truecolor heatmap stored its red plane twice
For a truecolor heatmap the scalar
gridis the red plane —z_flatisrgba[..., 0].reshape(-1). Butrgba_grid[0]re-ingested the same expression,and because the source is a strided view every reshape materializes its own
contiguous copy, so the
ColumnStoreheld two identical f64 planes for thefigure's lifetime. Reuse the column.
A 1600x1200 RGB heatmap holds 58.6 MB of canonical columns instead of
73.2 MB — 8 bytes per pixel of retained duplicate, not a transient.
_encode_transition_keyscarried two dictionaries to produce 8 bytes a rowIt kept
seen: dict[bytes, int]anddigests: dict[bytes, bytes]over everyrow, holding a token and a digest bytes object per row alive to the end of the
loop, in order to emit two uint32 per row.
Equal tokens hash equally, so distinct digests prove both distinctness and no
collision: one vectorized
np.uniqueover the result stands in for both dicts.A conflict re-walks with the original bookkeeping, so the message and the rows it
names are exactly what they were — pinned by tests over str/int/float/bool/bytes
duplicates and the missing/non-finite/wrong-type row errors.
200k animation keys peak at 7.1 MB instead of 48.1 MB.
PNG chunks copied the compressed image four times
_chunkdidbody = tag + data(one copy of the whole IDAT) and then built thechunk around it (a second), and each
+in the caller's document chain copied itagain. Chunks are now assembled as join-ready parts so one
b"".joinproducesthe file with a single copy, and the CRC is accumulated over the tag and then the
data — exactly
crc32(tag + data)without materializing the concatenation.A 1800x1400 truecolor encode peaks at 32.6 MB instead of 49.1 MB.
Measured
_encode_transition_keys, 200k keys_png.encodetruecolor 1800x1400Equivalence
A 97-artifact fingerprint sweep — payloads, SVG/PNG/JPEG/WebP/HTML exports,
selections through both zone-pruning branches, density views, streaming append,
transition keys, raw encoders, pyplot savefig — is identical to main.
Full suite: 2426 passed, 4 skipped. Ruff clean. New tests cover the id cap
(smallest-N, ascending, shared budget, untruncated case), the heatmap column
sharing and its payload, the transition-key digests and every error message, and
_chunk_partsagainst the naive construction.Perf
Interleaved A/B against main over the full 103-region suite, plus a focused
drift-cancelling re-check of the rows this touches: neutral-or-better, no region
more than 5% worse.
Summary by CodeRabbit