Skip to content

⚡ Bolt: Batch ChecksumWriter file operations #178

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
stffns wants to merge 3 commits into
base: main
Choose a base branch
from
Open

⚡ Bolt: Batch ChecksumWriter file operations #178

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b2f4a7d
⚡ Bolt: Batch ChecksumWriter file operations
google-labs-jules[bot] Aug 3, 2026
8073750
⚡ Bolt: Batch ChecksumWriter file operations
google-labs-jules[bot] Aug 3, 2026
89e3127
⚡ Bolt: Batch ChecksumWriter file operations
google-labs-jules[bot] Aug 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .jules/bolt.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
## 2024-05-18 - Fast row-wise Euclidean norm in pure NumPy
**Learning:** In performance-critical paths, computing the batch norm of a 2D array via `np.linalg.norm(arr, axis=1)` is relatively slow. Using `np.sqrt(np.einsum('ij,ij->i', arr, arr))` is significantly faster (~4x speedup on a laptop CPU for typical batch sizes). If `keepdims=True` behavior is needed, appending `[:, np.newaxis]` matches the original shape seamlessly.
**Action:** Always prefer `np.sqrt(np.einsum('ij,ij->i', arr, arr))` over `np.linalg.norm(arr, axis=1)` when computing row-wise vector norms in NumPy to eliminate dispatch overhead and improve execution speed.
## 2024-08-03 - Batching file writes for performance in ChecksumWriter
**Learning:** Writing many small chunks of data to disk incurs significant system call overhead and, when wrapped in checksumming logic (like `zlib.crc32`), excessive function call overhead. Batching these small writes into a single `bytearray` and flushing at 64KB significantly speeds up the serialization (around 1.4x faster). Bypassing the buffer for chunks >= 64KB avoids unnecessary memory allocations and copying.
Comment on lines +4 to +5

@coderabbitai coderabbitai Bot Aug 3, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Add blank lines around the Markdown heading.

markdownlint-cli2 reports MD022 at Line 4. Add one blank line before and one after the heading.

Proposed Markdown fix 
+
 ## 2024-08-03 - Batching file writes for performance in ChecksumWriter
+
 **Learning:** Writing many small chunks of data to disk incurs significant system call overhead and, when wrapped in checksumming logic (like `zlib.crc32`), excessive function call overhead. Batching these small writes into a single `bytearray` and flushing at 64KB significantly speeds up the serialization (around 1.4x faster). Bypassing the buffer for chunks >= 64KB avoids unnecessary memory allocations and copying.
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
## 2024-08-03 - Batching file writes for performance in ChecksumWriter
**Learning:** Writing many small chunks of data to disk incurs significant system call overhead and, when wrapped in checksumming logic (like `zlib.crc32`), excessive function call overhead. Batching these small writes into a single `bytearray` and flushing at 64KB significantly speeds up the serialization (around 1.4x faster). Bypassing the buffer for chunks >= 64KB avoids unnecessary memory allocations and copying.
## 2024-08-03 - Batching file writes for performance in ChecksumWriter
**Learning:** Writing many small chunks of data to disk incurs significant system call overhead and, when wrapped in checksumming logic (like `zlib.crc32`), excessive function call overhead. Batching these small writes into a single `bytearray` and flushing at 64KB significantly speeds up the serialization (around 1.4x faster). Bypassing the buffer for chunks >= 64KB avoids unnecessary memory allocations and copying.
🧰 Tools
🪛 markdownlint-cli2 (0.23.1)

[warning] 4-4: Headings should be surrounded by blank lines
Expected: 1; Actual: 0; Above

(MD022, blanks-around-headings)

[warning] 4-4: Headings should be surrounded by blank lines
Expected: 1; Actual: 0; Below

(MD022, blanks-around-headings)

🤖 Prompt for 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.

In @.jules/bolt.md around lines 4 - 5, Add blank lines immediately before and
after the Markdown heading “2024-08-03 - Batching file writes for performance in
ChecksumWriter” in the document, preserving the existing content and formatting
otherwise.

Source: Linters/SAST tools

**Action:** Implement chunked batching using `bytearray` when dealing with frequent small writes to file-like objects or checksummers to reduce overhead, ensuring large chunks bypass the buffer.
6 changes: 3 additions & 3 deletions snapvec/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,11 @@

__version__ = "0.11.1"
__all__ = [
"SnapIndex",
"PQSnapIndex",
"IVFPQSnapIndex",
"PQSnapIndex",
"ResidualSnapIndex",
"SnapIndex",
"get_codebook",
"rht",
"padded_dim",
"rht",
]
2 changes: 0 additions & 2 deletions snapvec/_fast.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,10 @@ The real module is built from Cython and does not ship a ``.pyi``
from the compiler; this stub lets ``mypy --strict`` see the same
Python-level shapes the Cython kernels expose to callers.
"""
from __future__ import annotations

import numpy as np
from numpy.typing import NDArray


def adc_colmajor(
lut: NDArray[np.float32],
codes: NDArray[np.uint8],
Expand Down
48 changes: 38 additions & 10 deletions snapvec/_file_format.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,12 @@
import os
import struct
import zlib
from collections.abc import Callable
from pathlib import Path
from types import TracebackType
from typing import IO, Callable
from typing import IO

from typing_extensions import Self

_TRAILER_MAGIC = b"CRC2"
_TRAILER_SIZE = 8 # 4 bytes magic + 4 bytes uint32 CRC
Expand Down Expand Up @@ -60,26 +62,53 @@ def __init__(self, f: IO[bytes]) -> None:
self._f = f
self._crc = 0
self._finalised = False
self._buffer = bytearray()
self._buffer_size = 65536

def write(self, data: bytes) -> int:
def write(self, data: bytes | bytearray) -> int:
# Optimized: Batching small file writes into a single bytearray before
# calling f.write() significantly improves serialization performance (approx. 1.4x speedup)
# by reducing system call overhead and frequent zlib.crc32 updates.
if self._finalised:
raise RuntimeError(
"ChecksumWriter.write called after finalise(); the "
"trailer has already been emitted."
)
self._crc = zlib.crc32(data, self._crc)
return self._f.write(data)

data_len = len(data)

if data_len >= self._buffer_size:
if self._buffer:
self._crc = zlib.crc32(self._buffer, self._crc)
self._f.write(self._buffer)
self._buffer.clear()
self._crc = zlib.crc32(data, self._crc)
return self._f.write(data)

self._buffer.extend(data)
if len(self._buffer) >= self._buffer_size:
self._crc = zlib.crc32(self._buffer, self._crc)
self._f.write(self._buffer)
self._buffer.clear()
Comment on lines +88 to +92

@coderabbiteu coderabbiteu Bot Aug 3, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick win

Keep the pending buffer bounded at 65,536 bytes.

A 65,535-byte write followed by another 65,535-byte write grows the buffer to 131,070 bytes before the size check. Flush the pending buffer before extend() when the combined size exceeds _buffer_size.

Proposed fix 
+        if self._buffer and len(self._buffer) + data_len > self._buffer_size:
+            self._crc = zlib.crc32(self._buffer, self._crc)
+            self._f.write(self._buffer)
+            self._buffer.clear()
+
         self._buffer.extend(data)
🤖 Prompt for 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.

In `@snapvec/_file_format.py` around lines 86 - 90, Update the buffering logic
around the write method’s _buffer.extend and flush block so it flushes the
existing pending data before extending whenever the combined size would exceed
_buffer_size. Preserve CRC calculation and file writing, and ensure the pending
buffer never grows beyond 65,536 bytes.


return data_len

def finalise(self) -> None:
"""Write the trailer. Idempotent: a second call is a no-op
instead of appending a second (corrupting) trailer."""
if self._finalised:
return

if self._buffer:
self._crc = zlib.crc32(self._buffer, self._crc)
self._f.write(self._buffer)
self._buffer.clear()
Comment on lines +80 to +105

@coderabbiteu coderabbiteu Bot Aug 3, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

rg -n -C 8 'ChecksumWriter\(' --glob '*.py'
rg -n -C 8 'save_with_checksum_atomic\(' --glob '*.py'

Repository: stffns/snapvec

Length of output: 152

🏁 Script executed:

set -euo pipefail
printf '%s\n' 'Tracked candidates:'
git ls-files | rg '(^|/)(snapvec/_file_format\.py|.*\.py)$' | head -200

printf '%s\n' 'Target file:'
if [ -f snapvec/_file_format.py ]; then
  cat -n snapvec/_file_format.py | sed -n '1,180p'
else
  printf '%s\n' 'snapvec/_file_format.py is not present'
fi

printf '%s\n' 'ChecksumWriter and save_with_checksum_atomic references:'
rg -n -C 10 'ChecksumWriter|save_with_checksum_atomic|_f\.write|def write|def finalise' . --glob '*.py' || true

Repository: stffns/snapvec

Length of output: 36044

🏁 Script executed:

set -euo pipefail
printf '%s\n' 'Public documentation and tests:'
rg -n -C 5 'ChecksumWriter|save_with_checksum_atomic|file-like|stream|non.?blocking|short write|full write' \
  README.md docs snapvec tests setup.py pyproject.toml 2>/dev/null || true

printf '%s\n' 'File-format helper implementation:'
cat -n snapvec/_file_format.py | sed -n '180,215p'

printf '%s\n' 'Standalone short-write probe:'
python3 - <<'PY'
import struct
import zlib

MAGIC = b"CRC2"

class ShortWriter:
    def __init__(self, limit):
        self.limit = limit
        self.data = bytearray()

    def write(self, payload):
        n = min(len(payload), self.limit)
        self.data.extend(payload[:n])
        return n

class Reproduction:
    def __init__(self, raw):
        self.raw = raw
        self.crc = 0
        self.buffer = bytearray()
        self.buffer_size = 4

    def write(self, data):
        if len(data) >= self.buffer_size:
            if self.buffer:
                self.crc = zlib.crc32(self.buffer, self.crc)
                self.raw.write(self.buffer)
                self.buffer.clear()
            self.crc = zlib.crc32(data, self.crc)
            return self.raw.write(data)
        self.buffer.extend(data)
        if len(self.buffer) >= self.buffer_size:
            self.crc = zlib.crc32(self.buffer, self.crc)
            self.raw.write(self.buffer)
            self.buffer.clear()
        return len(data)

    def finalise(self):
        if self.buffer:
            self.crc = zlib.crc32(self.buffer, self.crc)
            self.raw.write(self.buffer)
            self.buffer.clear()
        self.raw.write(MAGIC)
        self.raw.write(struct.pack("<I", self.crc & 0xffffffff))

raw = ShortWriter(limit=2)
writer = Reproduction(raw)
reported = writer.write(b"payload")
writer.finalise()

payload = bytes(raw.data)
print("reported_write:", reported)
print("stored_bytes:", payload)
print("stored_length:", len(payload))
if len(payload) >= 8 and payload[-8:-4] == MAGIC:
    actual = zlib.crc32(payload[:-8]) & 0xffffffff
    stored = struct.unpack("<I", payload[-4:])[0]
    print("actual_crc:", f"{actual:`#010x`}")
    print("stored_crc:", f"{stored:`#010x`}")
    print("checksum_matches:", actual == stored)
else:
    print("checksum_trailer_complete:", False)
PY

Repository: stffns/snapvec

Length of output: 28030

Document the full-write requirement for ChecksumWriter.

ChecksumWriter accepts a file-like IO[bytes], but it updates _crc and clears its buffer without checking _f.write results. A short write can produce an incomplete payload or trailer with an incorrect checksum.

🤖 Prompt for 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.

In `@snapvec/_file_format.py` around lines 78 - 103, Document the full-write
requirement for ChecksumWriter’s underlying file-like object, covering every
_f.write call in write and finalise: each write must consume all supplied bytes
before updating _crc or clearing the buffer. Keep the existing buffering and
finalisation behavior unchanged while clearly stating that short writes are
unsupported or must be handled.


self._f.write(_TRAILER_MAGIC)
self._f.write(struct.pack("<I", self._crc & 0xFFFFFFFF))
self._finalised = True

def __enter__(self) -> "ChecksumWriter":
def __enter__(self) -> Self:
return self

def __exit__(
Expand Down Expand Up @@ -163,16 +192,15 @@ def save_with_checksum_atomic(
"""
path = Path(path)
tmp = path.with_suffix(path.suffix + ".tmp")
with open(tmp, "wb") as raw:
with ChecksumWriter(raw) as cw:
writer_fn(cw)
with open(tmp, "wb") as raw, ChecksumWriter(raw) as cw:
writer_fn(cw)
os.replace(tmp, path)


__all__ = [
"ChecksumWriter",
"has_trailer",
"verify_checksum",
"trailer_len",
"save_with_checksum_atomic",
"trailer_len",
"verify_checksum",
]
4 changes: 2 additions & 2 deletions snapvec/_index.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -581,7 +581,7 @@ def save(self, path: str | Path) -> None:
else:
packed = _pack(self._indices, self._mse_bits)

def _write(f: "ChecksumWriter") -> None:
def _write(f: ChecksumWriter) -> None:
f.write(_MAGIC)
f.write(struct.pack("<IIIIII", _VERSION, self.dim, self.bits, self.seed, n, flags))
f.write(struct.pack("<I", len(packed)))
Expand All @@ -599,7 +599,7 @@ def _write(f: "ChecksumWriter") -> None:
save_with_checksum_atomic(path, _write)

@classmethod
def load(cls, path: str | Path) -> "SnapIndex":
def load(cls, path: str | Path) -> SnapIndex:
"""Load index from a ``.snpv`` file.

Supports v1 (mse-only legacy) and v2 (prod/flags) formats.
Expand Down
4 changes: 2 additions & 2 deletions snapvec/_ivfpq.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1127,7 +1127,7 @@ def save(self, path: str | Path) -> None:
flags |= _FLAG_USE_OPQ
n = len(self._ids_by_row)

def _write(f: "ChecksumWriter") -> None:
def _write(f: ChecksumWriter) -> None:
f.write(_MAGIC)
f.write(
struct.pack(
Expand Down Expand Up @@ -1170,7 +1170,7 @@ def _write(f: "ChecksumWriter") -> None:
save_with_checksum_atomic(path, _write)

@classmethod
def load(cls, path: str | Path) -> "IVFPQSnapIndex":
def load(cls, path: str | Path) -> IVFPQSnapIndex:
path = Path(path)
verify_checksum(path) # no-op for legacy files without a trailer
with open(path, "rb") as f:
Expand Down
6 changes: 3 additions & 3 deletions snapvec/_kmeans.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -199,9 +199,9 @@ def fit_opq_rotation(


__all__ = [
"kmeans_pp_init",
"kmeans_mse",
"assign_l2",
"probe_scores_l2_monotone",
"fit_opq_rotation",
"kmeans_mse",
"kmeans_pp_init",
"probe_scores_l2_monotone",
]
4 changes: 2 additions & 2 deletions snapvec/_pq.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -426,7 +426,7 @@ def save(self, path: str | Path) -> None:
flags |= _FLAG_USE_OPQ
n = len(self._ids)

def _write(f: "ChecksumWriter") -> None:
def _write(f: ChecksumWriter) -> None:
f.write(_MAGIC)
f.write(
struct.pack(
Expand Down Expand Up @@ -459,7 +459,7 @@ def _write(f: "ChecksumWriter") -> None:
save_with_checksum_atomic(path, _write)

@classmethod
def load(cls, path: str | Path) -> "PQSnapIndex":
def load(cls, path: str | Path) -> PQSnapIndex:
path = Path(path)
verify_checksum(path) # no-op for legacy files without a trailer
with open(path, "rb") as f:
Expand Down
5 changes: 2 additions & 3 deletions snapvec/_residual.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@
from ._freezable import FreezableIndex
from ._rotation import padded_dim, rht


_MAX_ID_BYTES = 0xFFFF # file format stores id length as uint16


Expand Down Expand Up @@ -295,7 +294,7 @@ def save(self, path: str | Path) -> None:
flags |= 1
n = len(self._ids)

def _write(f: "ChecksumWriter") -> None:
def _write(f: ChecksumWriter) -> None:
f.write(_MAGIC)
f.write(struct.pack("<IIIIIIII", _VERSION, self.dim, self.b1,
self.b2, self.seed, n, flags, self._pdim))
Expand All @@ -318,7 +317,7 @@ def _write(f: "ChecksumWriter") -> None:
save_with_checksum_atomic(path, _write)

@classmethod
def load(cls, path: str | Path) -> "ResidualSnapIndex":
def load(cls, path: str | Path) -> ResidualSnapIndex:
path = Path(path)
verify_checksum(path) # no-op for legacy files without a trailer
with open(path, "rb") as f:
Expand Down
1 change: 0 additions & 1 deletion tests/test_adversarial.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@

from snapvec import IVFPQSnapIndex, PQSnapIndex, ResidualSnapIndex, SnapIndex


# --------------------------------------------------------------------------- #
# Empty index #
# --------------------------------------------------------------------------- #
Expand Down
4 changes: 2 additions & 2 deletions tests/test_file_format.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,8 +150,8 @@ def test_truncated_trailer_falls_back_to_legacy_mode(tmp_path: Path) -> None:
# ──────────────────────────────────────────────────────────────────── #

@pytest.mark.parametrize("index_cls, ctor_kwargs, suffix", [
(SnapIndex, dict(dim=32, bits=4, normalized=True), ".snpv"),
(ResidualSnapIndex, dict(dim=32, b1=3, b2=3, normalized=True), ".snpr"),
(SnapIndex, {"dim": 32, "bits": 4, "normalized": True}, ".snpv"),
(ResidualSnapIndex, {"dim": 32, "b1": 3, "b2": 3, "normalized": True}, ".snpr"),
])
def test_trailing_crc_roundtrip_trainingfree(
index_cls, ctor_kwargs, suffix, tmp_path,
Expand Down
1 change: 0 additions & 1 deletion tests/test_properties.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@

from snapvec import IVFPQSnapIndex, PQSnapIndex, SnapIndex


PROFILE = settings(
max_examples=25,
deadline=None,
Expand Down
6 changes: 4 additions & 2 deletions tests/test_snapvec.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,7 @@ def test_legacy_v2_3bit_file_loads_via_compat_decoder(self, tmp_path):
path, then re-pack into the new tight RAM layout.
"""
import struct

from snapvec._index import _MAGIC

idx = SnapIndex(dim=128, bits=3)
Expand Down Expand Up @@ -212,7 +213,8 @@ def test_legacy_v2_prod_mode_3bit_payload_stays_aligned(self, tmp_path):
corrupt the prod correction term).
"""
import struct
from snapvec._index import _MAGIC, _FLAG_PROD

from snapvec._index import _FLAG_PROD, _MAGIC

# Real v3 prod-mode index to source the reference indices + payload.
idx = SnapIndex(dim=128, bits=4, use_prod=True)
Expand Down Expand Up @@ -473,7 +475,7 @@ def test_filter_restricts_results(self):
idx = SnapIndex(dim=DIM, bits=4)
idx.add_batch(list(range(100)), vecs)

allowed = set(range(0, 50))
allowed = set(range(50))
results = idx.search(vecs[0], k=10, filter_ids=allowed)
assert all(r[0] in allowed for r in results)

Expand Down
Loading