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..594f8f32 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): @@ -3981,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: @@ -3998,22 +4033,85 @@ 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") + # Equal tokens hash equally, so distinct digests prove distinct keys *and* + # 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. + # + # 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: + 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(_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(_transition_key_digests(result, stop)).size != stop: + _raise_transition_key_conflict(arr[:stop], label) + start = stop + block *= _TRANSITION_KEY_CHECK_GROWTH + return result + + +def _raise_transition_key_conflict(arr: np.ndarray, label: str) -> None: + """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 - 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..9fdeafed 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 @@ -770,3 +771,212 @@ 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") + + +@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"), + [ + 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") + + +@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 + 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 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 + + 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