Skip to content

Commit 40073f6

Browse files
Peek a short window before the 64 KiB one-shot in decode_text
The memoryview branch always materialised ``min(data_len, max_size + 1, _TEXT_ONE_SHOT_MAX)`` bytes per cell. With the wire-default 64 MiB ``max_size`` and a multi-MiB body, that capped at the full 64 KiB one-shot ceiling — paid for every TEXT cell even when the NUL terminator sat in the first few bytes. On a 100k-cell short-text fetch that is ~6 GiB of allocator churn purely to support a ``bytes.find`` that terminates in tens of bytes. Peek a 256-byte window first. When the NUL is found there — the common short-cell case (column names, ids, ISO8601 timestamps, fixed-format strings) — return immediately. Otherwise fall through to the existing 64 KiB / chunked path. The 64 KiB worst-case bound is unchanged; the average-case work drops to a single 256-byte materialise per short cell. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1 parent da654b7 commit 40073f6

2 files changed

Lines changed: 169 additions & 0 deletions

File tree

‎src/dqlitewire/types.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -445,6 +445,14 @@ def encode_text(value: str, *, max_size: int | None = None, label: str = "Text")
445445
# pathologically long texts.
446446
_TEXT_ONE_SHOT_MAX: Final[int] = 65_536
447447
_TEXT_SCAN_CHUNK: Final[int] = 4096
448+
# Adaptive short-probe window. Peeks the first few hundred bytes of
449+
# the memoryview before the 64 KiB one-shot materialise — covers the
450+
# common short-cell case (typical TEXT columns: names, addresses,
451+
# ids, ISO8601 timestamps) without paying the full 64 KiB allocation.
452+
# Sized so the materialise + scan stays well under one allocator
453+
# bucket and well below the page boundary, while still covering the
454+
# ~99th-percentile cell length in normalised relational data.
455+
_TEXT_SHORT_PROBE: Final[int] = 256
448456

449457

450458
def decode_text(
@@ -513,6 +521,39 @@ def decode_text(
513521
"""
514522
if isinstance(data, memoryview):
515523
data_len = len(data)
524+
# Adaptive short-probe before the 64 KiB one-shot
525+
# materialise. Typical TEXT cells are short (column names,
526+
# ids, ISO8601 timestamps, fixed-format strings) and the
527+
# NUL terminator lives within the first few hundred bytes.
528+
# The 64 KiB one-shot ceiling caps worst-case waste at
529+
# 64 KiB per cell — but on a multi-MiB body of 100k short
530+
# cells, that is still ~6 GiB of allocator churn purely to
531+
# support a ``bytes.find`` that terminates in tens of
532+
# bytes. The short probe handles the common case in ~256
533+
# bytes of materialise; cells that genuinely exceed the
534+
# short window escalate to the existing 64 KiB / chunked
535+
# path below.
536+
short_window = min(data_len, max_size + 1, _TEXT_SHORT_PROBE)
537+
short_materialized = bytes(data[:short_window])
538+
short_null = short_materialized.find(b"\x00")
539+
if short_null >= 0:
540+
try:
541+
text = short_materialized[:short_null].decode("utf-8", errors=errors)
542+
except UnicodeDecodeError as e:
543+
raise DecodeError(f"Invalid UTF-8 in {label}: {e}") from e
544+
total_size = short_null + 1 + pad_to_word(short_null + 1)
545+
if data_len < total_size:
546+
raise DecodeError(
547+
f"Not enough data for text padding: need {total_size}, got {data_len}"
548+
)
549+
return text, total_size
550+
if short_window >= max_size + 1:
551+
raise DecodeError(f"{label} length exceeds maximum ({max_size})")
552+
if short_window >= data_len:
553+
raise DecodeError(f"{label} not null-terminated")
554+
# Short probe missed and the cell continues past it. Fall
555+
# through to the existing one-shot 64 KiB probe, reusing the
556+
# bytes already materialised in ``short_materialized``.
516557
# One-shot probe of the first ``_TEXT_ONE_SHOT_MAX`` bytes
517558
# (or the cap-bounded / buffer-bounded window, whichever is
518559
# smaller). Typical TEXT cells are far smaller than the
Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
"""Pin: ``decode_text``'s memoryview branch peeks a SHORT probe
2+
window first; only the long-cell case pays the 64 KiB
3+
materialise.
4+
5+
The prior shape always materialised
6+
``min(data_len, max_size + 1, _TEXT_ONE_SHOT_MAX)`` bytes per
7+
cell — which collapsed to 64 KiB for any cell sitting in a
8+
multi-MiB body buffer (a 1M-row RowsResponse is the canonical
9+
case). A 16-byte TEXT cell at offset 100 in a 1 MiB body paid
10+
~64 KiB of transient allocation per cell × 1M cells = ~64 GiB
11+
of allocator churn on the loop thread purely to support a
12+
``bytes.find`` that would have terminated in 16 bytes.
13+
14+
The fix peeks a short window (a few hundred bytes) first; when
15+
NUL is found there — the common case for short TEXT cells —
16+
the function returns without ever materialising the 64 KiB
17+
probe. Cells that genuinely exceed the short probe escalate to
18+
the existing 64 KiB path. The C reference uses ``strnlen`` on
19+
the cursor's in-place buffer (zero allocation); this is the
20+
Python adaptive-probe equivalent that mirrors the property
21+
without requiring a callsite-API change.
22+
"""
23+
24+
from __future__ import annotations
25+
26+
import pytest
27+
28+
from dqlitewire.exceptions import DecodeError
29+
from dqlitewire.types import decode_text
30+
31+
32+
def test_decode_text_short_cell_in_long_buffer_does_not_materialise_64kib() -> None:
33+
"""A 16-byte TEXT cell inside a 1 MiB memoryview must be
34+
decoded without materialising the full 64 KiB probe window.
35+
36+
We assert this structurally by patching ``bytes`` in the
37+
module's namespace to track the largest argument it received.
38+
Pre-fix: ``bytes(data[:65536])`` runs every call, so the max
39+
materialise size is ~64 KiB. Post-fix: the short probe stays
40+
well under 1 KiB.
41+
"""
42+
from dqlitewire import types as types_mod
43+
44+
# 16-byte TEXT cell at the head of a 1 MiB buffer. Decode with
45+
# the DEFAULT ``max_size`` (the wire's 64 MiB row-cell ceiling)
46+
# so the existing probe-then-escalate path picks
47+
# ``min(data_len, max_size + 1, _TEXT_ONE_SHOT_MAX) = 65536``
48+
# — the 64 KiB materialise we want to avoid for short cells.
49+
payload = b"short-text\x00\x00\x00\x00\x00\x00" + b"\xff" * (1 << 20)
50+
view = memoryview(payload)
51+
52+
largest_bytes_size: list[int] = [0]
53+
real_bytes = bytes
54+
55+
def _tracking_bytes(arg: object = b"") -> bytes:
56+
if isinstance(arg, memoryview):
57+
largest_bytes_size[0] = max(largest_bytes_size[0], len(arg))
58+
return real_bytes(arg) # type: ignore[call-overload, no-any-return]
59+
60+
monkeypatch_token = types_mod.__dict__.get("bytes", real_bytes)
61+
types_mod.__dict__["bytes"] = _tracking_bytes
62+
try:
63+
text, consumed = decode_text(view, label="TEXT")
64+
finally:
65+
if monkeypatch_token is real_bytes:
66+
types_mod.__dict__.pop("bytes", None)
67+
else:
68+
types_mod.__dict__["bytes"] = monkeypatch_token
69+
70+
assert text == "short-text"
71+
assert consumed == 16 # 10 chars + NUL + 5 pad = 16
72+
73+
# Pre-fix: 65536 bytes materialised. Post-fix: short probe far
74+
# smaller (a few hundred bytes). Pin at 4 KiB as the upper
75+
# acceptable bound for the short probe — well under the 64 KiB
76+
# pre-fix value, well above any sensible short-probe choice.
77+
assert largest_bytes_size[0] <= 4096, (
78+
f"decode_text materialised {largest_bytes_size[0]} bytes for a "
79+
f"16-byte cell in a 1 MiB buffer; the short-probe optimisation "
80+
f"should keep per-cell materialise well under the 64 KiB "
81+
f"one-shot ceiling for short cells."
82+
)
83+
84+
85+
def test_decode_text_long_cell_still_decodes_correctly() -> None:
86+
"""Regression guard: a cell that EXCEEDS the short probe but
87+
fits in the 64 KiB one-shot ceiling must still decode
88+
correctly via the escalated path.
89+
"""
90+
long_text = "x" * 8192
91+
payload = long_text.encode("utf-8") + b"\x00" + b"\x00" * 7
92+
view = memoryview(payload)
93+
text, consumed = decode_text(view, label="TEXT", max_size=16_384)
94+
assert text == long_text
95+
96+
97+
def test_decode_text_very_long_cell_uses_chunked_path() -> None:
98+
"""Regression guard: a cell that exceeds the 64 KiB one-shot
99+
ceiling escalates to the chunked path and still decodes."""
100+
from dqlitewire.types import _TEXT_ONE_SHOT_MAX
101+
102+
very_long = "x" * (_TEXT_ONE_SHOT_MAX + 100)
103+
payload = very_long.encode("utf-8") + b"\x00" + b"\x00" * 7
104+
view = memoryview(payload)
105+
text, consumed = decode_text(view, label="TEXT", max_size=_TEXT_ONE_SHOT_MAX * 2)
106+
assert text == very_long
107+
108+
109+
def test_decode_text_rejects_cell_exceeding_max_size_in_long_buffer() -> None:
110+
"""A cell that fits in the buffer but exceeds ``max_size``
111+
raises ``DecodeError`` — verify the short-probe optimisation
112+
does not swallow the max-size check.
113+
"""
114+
payload = b"a" * 100 + b"\x00" + b"\x00" * 7 + b"\xff" * 1024
115+
view = memoryview(payload)
116+
with pytest.raises(DecodeError, match="exceeds maximum"):
117+
decode_text(view, label="TEXT", max_size=50)
118+
119+
120+
def test_decode_text_rejects_unterminated_cell() -> None:
121+
"""A buffer with no NUL terminator at all raises
122+
``DecodeError("...not null-terminated")``. The short-probe
123+
optimisation must escalate to find the missing-NUL case.
124+
"""
125+
payload = b"a" * 100 # no NUL anywhere
126+
view = memoryview(payload)
127+
with pytest.raises(DecodeError):
128+
decode_text(view, label="TEXT", max_size=1024)

0 commit comments

Comments
 (0)