From 317f4c1d2e2513be2857c14597e4da186cb9846c Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Sun, 26 Jul 2026 14:10:47 -0700 Subject: [PATCH 1/5] Cap, share and narrow: four more peak-memory cuts in the Python paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- python/xy/_png.py | 43 ++++++++++++++------ python/xy/channel.py | 19 +++++++-- python/xy/components.py | 22 +++++++++-- python/xy/marks.py | 8 +++- tests/test_animation.py | 64 ++++++++++++++++++++++++++++++ tests/test_density_mean_color.py | 53 +++++++++++++++++++++++++ tests/test_png_export.py | 25 ++++++++++++ tests/test_selection_rows.py | 68 ++++++++++++++++++++++++++++++++ 8 files changed, 283 insertions(+), 19 deletions(-) diff --git a/python/xy/_png.py b/python/xy/_png.py index 4ed8f48b..0bd1f5e8 100644 --- a/python/xy/_png.py +++ b/python/xy/_png.py @@ -28,9 +28,22 @@ import numpy as np +def _chunk_parts(tag: bytes, data: bytes) -> tuple[bytes, bytes, bytes, bytes]: + """A PNG chunk as join-ready parts, without ever copying `data`. + + The IDAT payload is the whole compressed image, and the obvious spelling + copies it repeatedly: `tag + data` is one copy, building the chunk around + it is a second, and each `+` in the caller's document chain is another. + Returning parts lets one `b"".join` at the top produce the file with a + single copy. CRC is accumulated over the tag and then the data, which is + exactly `crc32(tag + data)` without materializing the concatenation. + """ + crc = zlib.crc32(data, zlib.crc32(tag)) + return struct.pack(">I", len(data)), tag, data, struct.pack(">I", crc) + + def _chunk(tag: bytes, data: bytes) -> bytes: - body = tag + data - return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body)) + return b"".join(_chunk_parts(tag, data)) _SIG = b"\x89PNG\r\n\x1a\n" @@ -69,11 +82,13 @@ def png_truecolor( stride = w * 4 rows = _filtered_rows(h, stride) rows[:, 1:] = np.frombuffer(rgba, dtype=np.uint8, count=h * stride).reshape(h, stride) - return ( - _SIG - + _chunk(b"IHDR", ihdr) - + _chunk(b"IDAT", zlib.compress(rows, compression_level)) - + _chunk(b"IEND", b"") + return b"".join( + ( + _SIG, + *_chunk_parts(b"IHDR", ihdr), + *_chunk_parts(b"IDAT", zlib.compress(rows, compression_level)), + *_chunk_parts(b"IEND", b""), + ) ) @@ -84,11 +99,17 @@ def _png_indexed(w: int, h: int, rows: np.ndarray, palette: np.ndarray) -> bytes ihdr = struct.pack(">IIBBBBB", w, h, 8, 3, 0, 0, 0) plte = palette[:, :3].astype(np.uint8).tobytes() trns = palette[:, 3].astype(np.uint8).tobytes() - out = _SIG + _chunk(b"IHDR", ihdr) + _chunk(b"PLTE", plte) # tRNS may omit trailing opaque (255) entries; keep it simple and always emit. - out += _chunk(b"tRNS", trns) - out += _chunk(b"IDAT", zlib.compress(rows, _COMPRESSION_LEVEL)) + _chunk(b"IEND", b"") - return out + return b"".join( + ( + _SIG, + *_chunk_parts(b"IHDR", ihdr), + *_chunk_parts(b"PLTE", plte), + *_chunk_parts(b"tRNS", trns), + *_chunk_parts(b"IDAT", zlib.compress(rows, _COMPRESSION_LEVEL)), + *_chunk_parts(b"IEND", b""), + ) + ) def encode(img: np.ndarray) -> bytes: diff --git a/python/xy/channel.py b/python/xy/channel.py index 3368a3c3..12639d07 100644 --- a/python/xy/channel.py +++ b/python/xy/channel.py @@ -36,6 +36,8 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Optional +import numpy as np + from ._figure import Selection from ._framing import ( DEFAULT_FRAME_LIMITS, @@ -146,12 +148,21 @@ def _selection_reply( ids_remaining = SELECTION_EVENT_ID_LIMIT ids_truncated = False for tid in sorted(selected): - canonical = sorted(int(value) for value in selected[tid]) - kept = canonical[:ids_remaining] + # Cap before boxing. The event ships at most SELECTION_EVENT_ID_LIMIT + # ids, but sorting a generator of `int(value)` first materialized a + # Python int per *selected* row — ~40 bytes each, so a 5M-row lasso + # built ~200 MB of PyLongs to keep 10k of them. `argpartition` puts + # the smallest `ids_remaining` at the front for the price of one + # index array, and only the survivors are boxed. + arr = np.asarray(selected[tid]) + if arr.size > ids_remaining: + head = arr[np.argpartition(arr, ids_remaining)[:ids_remaining]] + ids_truncated = True + else: + head = arr + kept = np.sort(head).tolist() canonical_row_ids.append({"trace": int(tid), "ids": kept}) ids_remaining -= len(kept) - if len(kept) < len(canonical): - ids_truncated = True if ids_remaining == 0: ids_truncated = ids_truncated or any( len(selected[other_tid]) for other_tid in sorted(selected) if other_tid > tid diff --git a/python/xy/components.py b/python/xy/components.py index e2e899ca..cd15d85d 100644 --- a/python/xy/components.py +++ b/python/xy/components.py @@ -3998,6 +3998,24 @@ def _encode_transition_keys(value: Any, expected: int, label: str) -> np.ndarray if len(arr) != expected: raise ValueError(f"{label} must have length {expected}, got {len(arr)}") result = np.empty((expected, 2), dtype=np.uint32, order="F") + for index, raw in enumerate(arr): + token = _transition_key_token(raw, index) + digest = hashlib.blake2s(token, digest_size=8, person=b"xykeyv1").digest() + result[index, 0] = int.from_bytes(digest[:4], "little") + result[index, 1] = int.from_bytes(digest[4:], "little") + # Equal tokens hash equally, so distinct digests prove distinct keys *and* + # no collision: one vectorized uniqueness test stands in for the two + # per-row dictionaries this carried, which held a token and a digest bytes + # object for every row — hundreds of bytes of peak per 8-byte result row. + # A conflict is re-walked below so the message, and the rows it names, are + # exactly what the per-row bookkeeping produced. + if np.unique(result.reshape(-1).view(np.uint64)).size != expected: + _raise_transition_key_conflict(arr, label) + return result + + +def _raise_transition_key_conflict(arr: np.ndarray, label: str) -> None: + """Name the conflict the vectorized digest check found (error path only).""" seen: dict[bytes, int] = {} digests: dict[bytes, bytes] = {} for index, raw in enumerate(arr): @@ -4011,9 +4029,7 @@ def _encode_transition_keys(value: Any, expected: int, label: str) -> np.ndarray if collision is not None and collision != token: raise ValueError(f"{label} produced an identity digest collision") digests[digest] = token - result[index, 0] = int.from_bytes(digest[:4], "little") - result[index, 1] = int.from_bytes(digest[4:], "little") - return result + raise AssertionError(f"{label} digest conflict did not reproduce") def _original_mark_positions( diff --git a/python/xy/marks.py b/python/xy/marks.py index 83998161..0fe00039 100644 --- a/python/xy/marks.py +++ b/python/xy/marks.py @@ -2508,7 +2508,13 @@ def heatmap( grid=grid, rgba_grid=( ( - self.store.ingest(rgba[..., 0].reshape(-1)), + # `grid` already holds this plane: `z_flat` is + # `rgba[..., 0].reshape(-1)`, and re-ingesting the + # expression built a second contiguous copy (the source + # is a strided view, so each reshape materializes one) + # that the store kept for the figure's lifetime — 8 + # bytes per pixel of pure duplicate. + grid, self.store.ingest(rgba[..., 1].reshape(-1)), self.store.ingest(rgba[..., 2].reshape(-1)), self.store.ingest(rgba[..., 3].reshape(-1)), diff --git a/tests/test_animation.py b/tests/test_animation.py index 12b4c8dc..88263929 100644 --- a/tests/test_animation.py +++ b/tests/test_animation.py @@ -770,3 +770,67 @@ def test_animation_validation(kwargs: dict) -> None: def test_exit_is_not_a_supported_animation_option() -> None: with pytest.raises(TypeError, match="unexpected keyword argument 'exit'"): xy.animation(**{"exit": "fade"}) + + +# --- stable-key encoding: same answers, without the per-row dictionaries ----- + + +def test_transition_key_digests_are_stable_and_row_aligned() -> None: + """Encoding is a pure per-row hash: order-independent and reproducible.""" + from xy.components import _encode_transition_keys + + keys = ["a", "b", "c", "d"] + first = _encode_transition_keys(keys, 4, "keys") + assert first.dtype == np.uint32 and first.shape == (4, 2) + np.testing.assert_array_equal(first, _encode_transition_keys(keys, 4, "keys")) + # Row i's digest depends only on key i, so permuting the input permutes the + # output rows and nothing else. + shuffled = _encode_transition_keys(["c", "a", "d", "b"], 4, "keys") + np.testing.assert_array_equal(shuffled, first[[2, 0, 3, 1]]) + # Distinct keys must give distinct digests, which is what the vectorized + # uniqueness check the encoder now relies on is asserting. + assert len({tuple(row) for row in first}) == 4 + + +@pytest.mark.parametrize( + "keys", + [ + ["a", "b", "a"], + [1, 2, 1], + [1.5, 2.5, 1.5], + [True, False, True], + [b"x", b"y", b"x"], + ], + ids=["str", "int", "float", "bool", "bytes"], +) +def test_duplicate_transition_keys_name_both_rows(keys: list) -> None: + """The duplicate message still names the first and second offending rows. + + The encoder no longer keeps a token dictionary while hashing; a conflict is + re-walked to produce this message, so it must be unchanged. + """ + from xy.components import _encode_transition_keys + + with pytest.raises(ValueError, match=r"keys contains duplicate value at rows 0 and 2"): + _encode_transition_keys(keys, 3, "keys") + + +def test_transition_key_type_errors_still_report_their_row() -> None: + from xy.components import _encode_transition_keys + + with pytest.raises(ValueError, match="animation key is missing at row 1"): + _encode_transition_keys(["a", None, "c"], 3, "keys") + with pytest.raises(ValueError, match="animation key must be finite at row 2"): + _encode_transition_keys([1.0, 2.0, float("nan")], 3, "keys") + with pytest.raises(ValueError, match="row 1 has list"): + _encode_transition_keys(["a", [], "c"], 3, "keys") + + +def test_transition_keys_length_and_shape_are_validated() -> None: + from xy.components import _encode_transition_keys + + with pytest.raises(ValueError, match="must have length 4"): + _encode_transition_keys(["a", "b"], 4, "keys") + with pytest.raises(ValueError, match="must be one-dimensional"): + _encode_transition_keys([["a"], ["b"]], 2, "keys") + assert _encode_transition_keys([], 0, "keys").shape == (0, 2) diff --git a/tests/test_density_mean_color.py b/tests/test_density_mean_color.py index e6f52fae..e67d9419 100644 --- a/tests/test_density_mean_color.py +++ b/tests/test_density_mean_color.py @@ -521,3 +521,56 @@ def test_memory_report_itemizes_bin_color_cache_bytes(): ) plain = Figure().scatter(np.arange(1000.0), np.arange(1000.0)) assert plain.memory_report()["bin_color_bytes"] == 0 + + +def test_truecolor_heatmap_shares_its_red_plane_with_the_grid() -> None: + """`grid` and `rgba_grid[0]` are the same column, not two copies of it. + + A truecolor heatmap's scalar grid *is* the red plane. Ingesting + `rgba[..., 0].reshape(-1)` a second time built a second contiguous copy + (the source is a strided view, so every reshape materializes one) and the + store kept it for the figure's lifetime — 8 bytes per pixel of duplicate. + """ + import numpy as np + + from xy._figure import Figure + + rng = np.random.default_rng(2) + rgb = rng.random((40, 60, 3)) + fig = Figure() + fig.heatmap(rgb) + trace = fig.traces[0] + assert trace.rgba_grid is not None + assert trace.rgba_grid[0] is trace.grid + assert np.shares_memory(trace.rgba_grid[0].values, trace.grid.values) + # The other three planes stay distinct columns. + ids = {c.id for c in trace.rgba_grid} + assert len(ids) == 4 and trace.grid.id in ids + np.testing.assert_array_equal(trace.grid.values, rgb[..., 0].reshape(-1)) + + +def test_truecolor_heatmap_payload_is_unchanged_by_the_sharing() -> None: + """Sharing the column must not change what the wire carries.""" + import numpy as np + + from xy._figure import Figure + + rng = np.random.default_rng(3) + rgba = rng.random((24, 32, 4)) + fig = Figure() + fig.heatmap(rgba) + spec, blob = fig.build_payload() + planes = [ + np.frombuffer(blob, dtype=np.float32, count=c["len"], offset=c["byte_offset"]) + for c in spec["columns"] + if c["len"] == rgba.shape[0] * rgba.shape[1] + ] + # Exactly four, each matched to its own channel in serialized order. The + # duplicate this change removes was never shipped — it sat in the + # ColumnStore — so the wire is four planes before and after; what matters + # here is that sharing the column did not permute or corrupt them. An + # unordered `any(...)` match would have passed with two channels swapped. + assert len(planes) == 4 + for channel, plane in enumerate(planes): + want = rgba[..., channel].reshape(-1).astype(np.float32) + np.testing.assert_allclose(plane, want, atol=1e-6) diff --git a/tests/test_png_export.py b/tests/test_png_export.py index 8c63f45d..fd4baef3 100644 --- a/tests/test_png_export.py +++ b/tests/test_png_export.py @@ -1023,3 +1023,28 @@ def test_zero_width_characters_still_draw_nothing() -> None: # source, and an editor or a copy-paste that strips it would silently # reduce this to `_title_png("X") == _title_png("X")`. assert _title_png("X\u200b") == _title_png("X") + + +def test_chunk_parts_join_to_the_canonical_chunk(): + """Parts assembly must equal the naive `tag + data` construction. + + PNG chunks are built as join-ready parts so the compressed IDAT — the whole + image — is copied once into the file instead of once per `+`. The CRC is + accumulated over the tag and then the data, which is exactly + `crc32(tag + data)` without materializing the concatenation. + """ + import struct + import zlib + + from xy import _png + + for tag, data in ( + (b"IHDR", b"\x00" * 13), + (b"IEND", b""), + (b"IDAT", bytes(range(256)) * 40), + (b"tRNS", b"\xff\x00\x7f"), + ): + body = tag + data + canonical = struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body)) + assert b"".join(_png._chunk_parts(tag, data)) == canonical + assert _png._chunk(tag, data) == canonical diff --git a/tests/test_selection_rows.py b/tests/test_selection_rows.py index 257e7567..0ffa1a62 100644 --- a/tests/test_selection_rows.py +++ b/tests/test_selection_rows.py @@ -117,3 +117,71 @@ def test_polygon_clear_and_legacy_selection_reply_shapes(): assert legacy_reply is not None legacy, _ = legacy_reply assert set(legacy) == {"type", "traces", "total"} + + +def test_capped_ids_are_the_smallest_and_sorted(): + """The cap is applied before boxing, and picks the same ids as before. + + The reply keeps at most `SELECTION_EVENT_ID_LIMIT` canonical ids, and they + are the numerically smallest ones in ascending order — the shape a + `sorted(...)[:limit]` over the whole selection produced. That form + materialized a Python int per *selected* row (~40 bytes each) just to throw + almost all of them away, so the cap now runs on the array first; this pins + that the answer did not move. + """ + n = 5 * SELECTION_EVENT_ID_LIMIT + values = np.arange(n, dtype=np.float64) + fig = Figure().scatter(values, values) + + reply = handle_message( + fig, + {"type": "select", "x0": -1, "x1": n, "y0": -1, "y1": n, "include_rows": True}, + ) + assert reply is not None + message, _ = reply + (group,) = message["canonical_row_ids"] + ids = group["ids"] + assert ids == list(range(SELECTION_EVENT_ID_LIMIT)) + assert all(isinstance(v, int) for v in ids) + assert message["truncated"] is True + json.dumps(message, allow_nan=False) + + +def test_id_budget_is_shared_across_traces_in_trace_order(): + """A trace that exhausts the budget leaves later traces with none.""" + half = SELECTION_EVENT_ID_LIMIT + a = np.arange(half, dtype=np.float64) + fig = Figure() + fig.scatter(a, a) + fig.scatter(a, a) + + reply = handle_message( + fig, + { + "type": "select", + "x0": -1, + "x1": half, + "y0": -1, + "y1": half, + "include_rows": True, + }, + ) + assert reply is not None + message, _ = reply + counts = [len(group["ids"]) for group in message["canonical_row_ids"]] + assert sum(counts) == SELECTION_EVENT_ID_LIMIT + assert counts[0] == SELECTION_EVENT_ID_LIMIT + assert message["truncated"] is True + + +def test_uncapped_selection_reports_every_id_untruncated(): + values = np.arange(8, dtype=np.float64) + fig = Figure().scatter(values, values) + reply = handle_message( + fig, {"type": "select", "x0": -1, "x1": 9, "y0": -1, "y1": 9, "include_rows": True} + ) + assert reply is not None + message, _ = reply + (group,) = message["canonical_row_ids"] + assert group["ids"] == list(range(8)) + assert message["truncated"] is False From 469b42634e10537172264e77b5e5b9a1ce80c1db Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Sun, 26 Jul 2026 16:07:39 -0700 Subject: [PATCH 2/5] Keep a duplicate key ahead of a later invalid one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- python/xy/components.py | 15 ++++++++++++++- tests/test_animation.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/python/xy/components.py b/python/xy/components.py index cd15d85d..705ce882 100644 --- a/python/xy/components.py +++ b/python/xy/components.py @@ -3999,7 +3999,20 @@ def _encode_transition_keys(value: Any, expected: int, label: str) -> np.ndarray raise ValueError(f"{label} must have length {expected}, got {len(arr)}") result = np.empty((expected, 2), dtype=np.uint32, order="F") for index, raw in enumerate(arr): - token = _transition_key_token(raw, index) + try: + token = _transition_key_token(raw, index) + except ValueError: + # Deferring the uniqueness test to after the walk also defers every + # duplicate behind every *later* bad row, which reorders the errors: + # `["a", "a", None]` reported the row-0/1 duplicate before this and + # must keep doing so. The rows already walked tokenized cleanly, so + # their digests are complete and a conflict among them belongs to an + # earlier row than this one. Error path only, and the prefix is + # re-walked below, so the message and the rows it names are the + # per-row bookkeeping's. + if np.unique(result[:index].reshape(-1).view(np.uint64)).size != index: + _raise_transition_key_conflict(arr[:index], label) + raise digest = hashlib.blake2s(token, digest_size=8, person=b"xykeyv1").digest() result[index, 0] = int.from_bytes(digest[:4], "little") result[index, 1] = int.from_bytes(digest[4:], "little") diff --git a/tests/test_animation.py b/tests/test_animation.py index 88263929..97fdd05c 100644 --- a/tests/test_animation.py +++ b/tests/test_animation.py @@ -5,6 +5,7 @@ import datetime as dt import hashlib import json +import re import numpy as np import pytest @@ -826,6 +827,39 @@ def test_transition_key_type_errors_still_report_their_row() -> None: _encode_transition_keys(["a", [], "c"], 3, "keys") +@pytest.mark.parametrize( + ("keys", "rows"), + [ + pytest.param(["a", "a", None], (0, 1), id="missing"), + pytest.param([1.0, 1.0, float("nan")], (0, 1), id="non-finite"), + pytest.param(["a", "a", {}], (0, 1), id="wrong-type"), + pytest.param(["a", "b", "a", None], (0, 2), id="later-duplicate"), + ], +) +def test_a_duplicate_outranks_a_later_invalid_row(keys: list, rows: tuple[int, int]) -> None: + """First bad row wins, whichever kind of bad it is. + + The uniqueness test runs after the walk, so a duplicate would otherwise be + reported only once every later row had tokenized — and any invalid row + behind it would mask the duplicate entirely. Both inputs are wrong twice + over; which error surfaces is the contract. + """ + from xy.components import _encode_transition_keys + + expected = f"keys contains duplicate value at rows {rows[0]} and {rows[1]}" + with pytest.raises(ValueError, match=re.escape(expected)): + _encode_transition_keys(keys, len(keys), "keys") + + +def test_an_invalid_row_before_a_duplicate_still_wins() -> None: + """The converse: nothing about the prefix re-check promotes a later + duplicate over an earlier bad row.""" + from xy.components import _encode_transition_keys + + with pytest.raises(ValueError, match="animation key is missing at row 0"): + _encode_transition_keys([None, "a", "a"], 3, "keys") + + def test_transition_keys_length_and_shape_are_validated() -> None: from xy.components import _encode_transition_keys From fa7274cd8e33766de27b73ecbb0afad530674ac0 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Sun, 26 Jul 2026 16:22:33 -0700 Subject: [PATCH 3/5] Do not let the superseded key error surface, or narrow it to ValueError MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- python/xy/components.py | 22 ++++++++++++++++++---- tests/test_animation.py | 40 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 4 deletions(-) diff --git a/python/xy/components.py b/python/xy/components.py index 705ce882..732d26ac 100644 --- a/python/xy/components.py +++ b/python/xy/components.py @@ -4001,7 +4001,7 @@ def _encode_transition_keys(value: Any, expected: int, label: str) -> np.ndarray for index, raw in enumerate(arr): try: token = _transition_key_token(raw, index) - except ValueError: + except Exception: # Deferring the uniqueness test to after the walk also defers every # duplicate behind every *later* bad row, which reorders the errors: # `["a", "a", None]` reported the row-0/1 duplicate before this and @@ -4010,6 +4010,11 @@ def _encode_transition_keys(value: Any, expected: int, label: str) -> np.ndarray # earlier row than this one. Error path only, and the prefix is # re-walked below, so the message and the rows it names are the # per-row bookkeeping's. + # + # Catching Exception, not ValueError: the rule is "the first bad row + # wins", and which exception the token raised does not change which + # row was first. A key type whose __format__/isoformat/encode raises + # something else would otherwise still mask the earlier duplicate. if np.unique(result[:index].reshape(-1).view(np.uint64)).size != index: _raise_transition_key_conflict(arr[:index], label) raise @@ -4028,19 +4033,28 @@ def _encode_transition_keys(value: Any, expected: int, label: str) -> np.ndarray def _raise_transition_key_conflict(arr: np.ndarray, label: str) -> None: - """Name the conflict the vectorized digest check found (error path only).""" + """Name the conflict the vectorized digest check found (error path only). + + Every raise here is `from None`. One caller is the prefix re-check inside an + `except`, where the token error it is deliberately superseding would + otherwise be chained in front of this one — so the first thing printed would + be the very error this function exists to *not* report. The other caller is + outside any handler, where `from None` costs nothing. + """ seen: dict[bytes, int] = {} digests: dict[bytes, bytes] = {} for index, raw in enumerate(arr): token = _transition_key_token(raw, index) previous = seen.get(token) if previous is not None: - raise ValueError(f"{label} contains duplicate value at rows {previous} and {index}") + raise ValueError( + f"{label} contains duplicate value at rows {previous} and {index}" + ) from None seen[token] = index digest = hashlib.blake2s(token, digest_size=8, person=b"xykeyv1").digest() collision = digests.get(digest) if collision is not None and collision != token: - raise ValueError(f"{label} produced an identity digest collision") + raise ValueError(f"{label} produced an identity digest collision") from None digests[digest] = token raise AssertionError(f"{label} digest conflict did not reproduce") diff --git a/tests/test_animation.py b/tests/test_animation.py index 97fdd05c..e1099ee5 100644 --- a/tests/test_animation.py +++ b/tests/test_animation.py @@ -851,6 +851,46 @@ def test_a_duplicate_outranks_a_later_invalid_row(keys: list, rows: tuple[int, i _encode_transition_keys(keys, len(keys), "keys") +def test_the_superseded_token_error_is_not_chained_onto_the_duplicate() -> None: + """The prefix re-check runs inside an `except`, so without suppression the + error it deliberately discards would print *first*, above a 'During handling + of the above exception' banner — the traceback would lead with the one + message this is not supposed to report.""" + import traceback + + from xy.components import _encode_transition_keys + + try: + _encode_transition_keys(["a", "a", None], 3, "keys") + except ValueError as exc: + assert exc.__context__ is None or exc.__suppress_context__ + rendered = "".join(traceback.format_exception(exc)) + assert "During handling of the above exception" not in rendered + assert "is missing at row 2" not in rendered + assert "keys contains duplicate value at rows 0 and 1" in rendered + else: # pragma: no cover + raise AssertionError("expected a duplicate-key error") + + +def test_a_duplicate_outranks_a_non_valueerror_token_failure() -> None: + """Which row was first cannot depend on which exception the token raised. + + A key type whose `encode` raises something other than ValueError must not + smuggle a later row's failure past an earlier duplicate. + """ + from xy.components import _encode_transition_keys + + class BadStr(str): + def encode(self, *args: object, **kwargs: object) -> bytes: + raise RuntimeError("boom-encode") + + with pytest.raises(ValueError, match="keys contains duplicate value at rows 0 and 1"): + _encode_transition_keys(["a", "a", BadStr("z")], 3, "keys") + # With no earlier duplicate the original failure still propagates untouched. + with pytest.raises(RuntimeError, match="boom-encode"): + _encode_transition_keys(["a", "b", BadStr("z")], 3, "keys") + + def test_an_invalid_row_before_a_duplicate_still_wins() -> None: """The converse: nothing about the prefix re-check promotes a later duplicate over an earlier bad row.""" From 44005b25b2811c7d3d98b6cff04d5b60dc5b8cf5 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Sun, 26 Jul 2026 16:52:21 -0700 Subject: [PATCH 4/5] Test key uniqueness on growing prefixes, not once at the end MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- python/xy/components.py | 96 ++++++++++++++++++++++++++++------------- tests/test_animation.py | 45 +++++++++++++++++++ 2 files changed, 111 insertions(+), 30 deletions(-) diff --git a/python/xy/components.py b/python/xy/components.py index 732d26ac..60619386 100644 --- a/python/xy/components.py +++ b/python/xy/components.py @@ -3842,6 +3842,20 @@ def _resolve_color(data: Any, color: Any, *, context: Optional[str] = None) -> A raise +# First prefix at which `_encode_transition_keys` tests digest uniqueness, then +# doubling. Small enough that the common duplicate — an id column that is not +# actually unique — is caught in about a millisecond, large enough that the +# vectorized test is never run on a handful of rows where the per-row +# dictionaries it replaced would have been cheaper anyway. +_TRANSITION_KEY_CHECK_STRIDE = 4096 +# Growth factor between those prefixes. The intermediate tests are pure +# insurance, so they are kept far below the final full test's cost: at x8 they +# sum to about an eighth of it, where x2 would have tripled the total uniqueness +# work and cost the success path ~8%. The price is that a duplicate is seen +# within 8x rather than 2x the rows needed to reach it. +_TRANSITION_KEY_CHECK_GROWTH = 8 + + def _transition_key_token(value: Any, index: int) -> bytes: """Canonical, type-sensitive bytes for one supported stable key.""" if isinstance(value, np.generic): @@ -3998,37 +4012,59 @@ def _encode_transition_keys(value: Any, expected: int, label: str) -> np.ndarray if len(arr) != expected: raise ValueError(f"{label} must have length {expected}, got {len(arr)}") result = np.empty((expected, 2), dtype=np.uint32, order="F") - for index, raw in enumerate(arr): - try: - token = _transition_key_token(raw, index) - except Exception: - # Deferring the uniqueness test to after the walk also defers every - # duplicate behind every *later* bad row, which reorders the errors: - # `["a", "a", None]` reported the row-0/1 duplicate before this and - # must keep doing so. The rows already walked tokenized cleanly, so - # their digests are complete and a conflict among them belongs to an - # earlier row than this one. Error path only, and the prefix is - # re-walked below, so the message and the rows it names are the - # per-row bookkeeping's. - # - # Catching Exception, not ValueError: the rule is "the first bad row - # wins", and which exception the token raised does not change which - # row was first. A key type whose __format__/isoformat/encode raises - # something else would otherwise still mask the earlier duplicate. - if np.unique(result[:index].reshape(-1).view(np.uint64)).size != index: - _raise_transition_key_conflict(arr[:index], label) - raise - digest = hashlib.blake2s(token, digest_size=8, person=b"xykeyv1").digest() - result[index, 0] = int.from_bytes(digest[:4], "little") - result[index, 1] = int.from_bytes(digest[4:], "little") # Equal tokens hash equally, so distinct digests prove distinct keys *and* - # no collision: one vectorized uniqueness test stands in for the two - # per-row dictionaries this carried, which held a token and a digest bytes - # object for every row — hundreds of bytes of peak per 8-byte result row. - # A conflict is re-walked below so the message, and the rows it names, are - # exactly what the per-row bookkeeping produced. - if np.unique(result.reshape(-1).view(np.uint64)).size != expected: - _raise_transition_key_conflict(arr, label) + # no collision: a vectorized uniqueness test stands in for the two per-row + # dictionaries this carried, which held a token and a digest bytes object for + # every row — hundreds of bytes of peak per 8-byte result row. A conflict is + # re-walked by `_raise_transition_key_conflict` so the message, and the rows + # it names, are exactly what the per-row bookkeeping produced. + # + # The test runs on growing prefixes rather than once at the end. Testing only + # at the end would make a duplicate cost the entire walk plus a full re-walk: + # a non-unique id column, which is the ordinary form of this mistake, went + # from 16 ms and 32 MB peak on the dictionaries to 7.1 s and 65 MB — so the + # rewrite raised peak on a reachable path, the one thing it exists to lower. + # Prefixes bound both the wasted walk and the `np.unique` transient to the + # rows actually needed to see the duplicate. + # + # Blocked rather than checked per row on purpose: `arr[start:stop]` is a view + # of an object array, so the inner loop is exactly the loop it was, while one + # extra compare per row cost 8% of the success path in a 2M-row walk. + # + # `result` is Fortran-order because the native encoder's is and the payload + # ships `values[:, 0]` copy-free, so `reshape(-1)` below cannot be a view and + # each test copies the prefix it checks. That is 8 bytes per row against the + # dictionaries' hundreds, and the x8 growth keeps the copies to about an + # eighth of the final one on top of it; pairing the two u32 words any other + # way costs the same 8 bytes, and viewing the 2-D array directly raises. + start = 0 + block = _TRANSITION_KEY_CHECK_STRIDE + while start < expected: + stop = min(start + block, expected) + for index, raw in enumerate(arr[start:stop], start): + try: + token = _transition_key_token(raw, index) + except Exception: + # A later bad row must not mask an earlier duplicate: for + # `["a", "a", None]` the dictionaries reported the row-0/1 + # duplicate, and that is still the first thing wrong. Rows + # already walked tokenized cleanly, so their digests are + # complete and any conflict among them precedes this row. + # + # `except Exception`, not ValueError: which row was first does + # not depend on what the token raised, and a key type whose + # __format__/isoformat/encode raises something else would + # otherwise still slip past the earlier duplicate. + if np.unique(result[:index].reshape(-1).view(np.uint64)).size != index: + _raise_transition_key_conflict(arr[:index], label) + raise + digest = hashlib.blake2s(token, digest_size=8, person=b"xykeyv1").digest() + result[index, 0] = int.from_bytes(digest[:4], "little") + result[index, 1] = int.from_bytes(digest[4:], "little") + if np.unique(result[:stop].reshape(-1).view(np.uint64)).size != stop: + _raise_transition_key_conflict(arr[:stop], label) + start = stop + block *= _TRANSITION_KEY_CHECK_GROWTH return result diff --git a/tests/test_animation.py b/tests/test_animation.py index e1099ee5..4754268f 100644 --- a/tests/test_animation.py +++ b/tests/test_animation.py @@ -851,6 +851,51 @@ def test_a_duplicate_outranks_a_later_invalid_row(keys: list, rows: tuple[int, i _encode_transition_keys(keys, len(keys), "keys") +@pytest.mark.parametrize( + ("first", "second"), + [ + (0, 1), # both in the first block + (0, 4095), # first block, at its last row + (0, 4096), # spans the first block boundary + (4095, 4096), # straddles it + (4096, 4097), # both in the second block + (0, 32768), # spans the second boundary (stride x8) + (5000, 40000), # neither in the first block, different blocks + (0, 49999), # last row of all + ], +) +def test_a_duplicate_is_named_the_same_wherever_the_block_boundaries_fall( + first: int, second: int +) -> None: + """Uniqueness is tested on growing prefixes, so a duplicate whose two rows + land in different blocks must still be reported — and reported with the same + two row numbers a single trailing test would have produced.""" + from xy.components import _encode_transition_keys + + n = 50_000 + keys = [f"k-{i:09d}" for i in range(n)] + keys[second] = keys[first] + expected = f"keys contains duplicate value at rows {first} and {second}" + with pytest.raises(ValueError, match=re.escape(expected)): + _encode_transition_keys(keys, n, "keys") + + +def test_prefix_checks_do_not_change_the_encoding() -> None: + """The blocked walk must produce the same digests as one straight pass, at + sizes below, at, and above the first block boundary.""" + from xy.components import _encode_transition_keys + + for n in (1, 4095, 4096, 4097, 32769, 40000): + keys = [f"k-{i:09d}" for i in range(n)] + got = _encode_transition_keys(keys, n, "keys") + assert got.shape == (n, 2) and got.dtype == np.uint32 + # Every row is an independent hash of its own key, so a one-row encode + # of key i must equal row i of the bulk encode. + for probe in {0, n // 2, n - 1}: + single = _encode_transition_keys([keys[probe]], 1, "keys") + np.testing.assert_array_equal(single[0], got[probe]) + + def test_the_superseded_token_error_is_not_chained_onto_the_duplicate() -> None: """The prefix re-check runs inside an `except`, so without suppression the error it deliberately discards would print *first*, above a 'During handling From 152e1d18cfb2843323caffb9fbf1a79e952c4acd Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Sun, 26 Jul 2026 17:21:52 -0700 Subject: [PATCH 5/5] Pack prefix digests from the columns, not through a reshape MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- python/xy/components.py | 35 +++++++++++++++++++++++++++-------- tests/test_animation.py | 27 +++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 8 deletions(-) diff --git a/python/xy/components.py b/python/xy/components.py index 60619386..594f8f32 100644 --- a/python/xy/components.py +++ b/python/xy/components.py @@ -3995,6 +3995,27 @@ def _fixed_transition_key_values(value: Any) -> np.ndarray | None: return None +def _transition_key_digests(result: np.ndarray, stop: int) -> np.ndarray: + """The first `stop` rows' identity digests, one uint64 per row. + + `result` is Fortran-order — the native encoder's layout, so the payload ships + `values[:, 0]` copy-free — which makes `reshape(-1).view(np.uint64)` wrong + twice over. 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"; for longer prefixes it silently copies. Both + columns are contiguous in this order, so combining them is correct whatever + the layout and costs the same 8 bytes a row the reshape copy did. + + The two words are packed hi-word-first here rather than in the wire's + little-endian order. Only injectivity matters: equal packed values iff equal + (lo, hi) pairs, which is what the uniqueness test is asking. + """ + packed = result[:stop, 0].astype(np.uint64) + packed <<= np.uint64(32) + packed |= result[:stop, 1] + return packed + + def _encode_transition_keys(value: Any, expected: int, label: str) -> np.ndarray: fixed = _fixed_transition_key_values(value) if fixed is not None: @@ -4031,12 +4052,10 @@ def _encode_transition_keys(value: Any, expected: int, label: str) -> np.ndarray # of an object array, so the inner loop is exactly the loop it was, while one # extra compare per row cost 8% of the success path in a 2M-row walk. # - # `result` is Fortran-order because the native encoder's is and the payload - # ships `values[:, 0]` copy-free, so `reshape(-1)` below cannot be a view and - # each test copies the prefix it checks. That is 8 bytes per row against the - # dictionaries' hundreds, and the x8 growth keeps the copies to about an - # eighth of the final one on top of it; pairing the two u32 words any other - # way costs the same 8 bytes, and viewing the 2-D array directly raises. + # Each test packs the prefix it checks into one uint64 a row + # (`_transition_key_digests`) — 8 bytes per row against the dictionaries' + # hundreds, and the x8 growth keeps those temporaries to about an eighth of + # the final one on top of it. start = 0 block = _TRANSITION_KEY_CHECK_STRIDE while start < expected: @@ -4055,13 +4074,13 @@ def _encode_transition_keys(value: Any, expected: int, label: str) -> np.ndarray # not depend on what the token raised, and a key type whose # __format__/isoformat/encode raises something else would # otherwise still slip past the earlier duplicate. - if np.unique(result[:index].reshape(-1).view(np.uint64)).size != index: + if np.unique(_transition_key_digests(result, index)).size != index: _raise_transition_key_conflict(arr[:index], label) raise digest = hashlib.blake2s(token, digest_size=8, person=b"xykeyv1").digest() result[index, 0] = int.from_bytes(digest[:4], "little") result[index, 1] = int.from_bytes(digest[4:], "little") - if np.unique(result[:stop].reshape(-1).view(np.uint64)).size != stop: + if np.unique(_transition_key_digests(result, stop)).size != stop: _raise_transition_key_conflict(arr[:stop], label) start = stop block *= _TRANSITION_KEY_CHECK_GROWTH diff --git a/tests/test_animation.py b/tests/test_animation.py index 4754268f..9fdeafed 100644 --- a/tests/test_animation.py +++ b/tests/test_animation.py @@ -827,6 +827,33 @@ def test_transition_key_type_errors_still_report_their_row() -> None: _encode_transition_keys(["a", [], "c"], 3, "keys") +@pytest.mark.parametrize("n_good", [1, 2, 3, 17]) +def test_a_short_fallback_prefix_is_packable(n_good: int) -> None: + """Date keys always take the Python fallback, so this reaches the prefix + digest test at tiny lengths — where the Fortran-order buffer bites. + + `result[:1]` is a one-row slice of an F-order (N, 2) array, and numpy can + satisfy `reshape(-1)` on it with a strided *view*; re-viewing that as uint64 + raises "the last axis must be contiguous", replacing the key error with a + numpy internal one. The digests are packed from the two columns instead. + """ + from xy.components import _encode_transition_keys + + keys: list = [dt.date(2024, 1, 1) + dt.timedelta(days=i) for i in range(n_good)] + encoded = _encode_transition_keys(keys, n_good, "keys") + assert encoded.shape == (n_good, 2) + assert encoded[:, 0].flags.c_contiguous and encoded[:, 1].flags.c_contiguous + + # ...and the error path, which packs the completed prefix before re-raising. + with pytest.raises(ValueError, match=f"animation key is missing at row {n_good}"): + _encode_transition_keys([*keys, None], n_good + 1, "keys") + # ...and with a duplicate inside that prefix, which must win instead. + if n_good >= 2: + dup = [keys[0], *keys[:-1], None] + with pytest.raises(ValueError, match=r"keys contains duplicate value at rows 0 and 1"): + _encode_transition_keys(dup, len(dup), "keys") + + @pytest.mark.parametrize( ("keys", "rows"), [