Skip to content

Add freeze() to drop per-shard codes/residuals for read-only indexes#45

Merged
raphaelsty merged 5 commits into
mainfrom
freeze-index
May 28, 2026
Merged

Add freeze() to drop per-shard codes/residuals for read-only indexes#45
raphaelsty merged 5 commits into
mainfrom
freeze-index

Conversation

@raphaelsty

@raphaelsty raphaelsty commented May 28, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Adds FastPlaid.freeze() which deletes the per-shard {i}.codes.npy / {i}.residuals.npy files and sets "frozen": true in metadata.json. The merged file remains and becomes the sole storage, cutting on-disk usage roughly in half.
  • Adds FastPlaid.unfreeze() as the inverse: slices merged_codes.npy / merged_residuals.npy back into per-shard files using doclens.{i}.json, then drops the frozen key. The roundtrip is byte-identical — every file matches the original.
  • Teaches the loader to skip the chunk scan when an index is frozen and mmap the merged file directly (slightly faster reload, less peak RAM).
  • update() and delete() on a frozen index raise RuntimeError — both paths operate per shard and can't run without re-sharding (call unfreeze() first).

Motivation

After indexing, FastPlaid keeps two copies of the same compressed data on disk:

  • per-shard {i}.codes.npy / {i}.residuals.npy (used by update.rs / delete.rs)
  • a merged merged_codes.npy / merged_residuals.npy (used at search time via StridedTensor)

For workloads that won't mutate the index again, the shards are dead weight. freeze() makes that explicit and reclaims the space without any Rust changes. unfreeze() is there as an escape hatch if you change your mind.

Byte-identical roundtrip

unfreeze() reconstructs shards via a small helper that matches tch::Tensor::write_npy byte-for-byte:

  • Rust-style shape formatting (R,C,) (trailing comma, no spaces) instead of np.save's (R, C).
  • Explicit < byte order even for 1-byte types (np.save writes |).
  • Header padded so the prologue (10 + header) is a multiple of 16, matching tch's alignment.

metadata.json is written with indent=2 to match serde_json::to_writer_pretty. The frozen key is fully removed on unfreeze() so the file matches its pre-freeze bytes exactly.

What's untouched

  • Indexing path (create.rs) and the chunk-then-merge flow are unchanged. Peak RAM during create is preserved.
  • Search hot path is identical — same merged mmap, same StridedTensor.
  • Non-frozen indexes behave exactly as before.

Test plan

  • test_freeze_removes_shards_and_search_still_works — search results identical before and after freeze(), shard files gone, merged files retained.
  • test_freeze_persists_across_reload — a frozen index reopens correctly in a new FastPlaid instance and returns the same results.
  • test_freeze_update_raisesupdate() on a frozen index raises RuntimeError.
  • test_freeze_delete_raisesdelete() on a frozen index raises RuntimeError.
  • test_freeze_idempotent — calling freeze() twice is a no-op.
  • test_freeze_unfreeze_roundtrip_byte_identical — create a multi-shard index, copy the directory, freeze+unfreeze the original, assert every file is byte-identical to the snapshot via filecmp.cmp(shallow=False).
  • test_unfreeze_restores_update_capability — after unfreeze(), update() works again.
  • test_unfreeze_idempotent — calling unfreeze() on a non-frozen index is a no-op.
  • Existing TestBasicCreateAndSearch and TestUpdate suites still pass.

After indexing, FastPlaid keeps both the per-shard files
({i}.codes.npy, {i}.residuals.npy) used by update/delete and a merged
file used at search time, roughly doubling on-disk storage. For a
read-only index the shards are redundant.

freeze() ensures the merged file is up to date, deletes every per-shard
codes/residuals file, and writes "frozen": true into metadata.json. The
loader skips the chunk scan when frozen and mmaps the merged file
directly. Subsequent update/delete calls raise; search is unaffected.
unfreeze() slices merged_codes.npy / merged_residuals.npy back into
per-shard {i}.codes.npy / {i}.residuals.npy using each shard's
doclens.{i}.json, then drops the frozen flag from metadata.json so
update/delete work again.

Shard .npy files are written via a small helper that matches
tch::Tensor::write_npy byte-for-byte (Rust-style shape formatting,
explicit byte order, 16-byte prologue alignment). metadata.json is
re-pretty-printed with two-space indent to match serde_json's
to_writer_pretty output.

Adds test_freeze_unfreeze_roundtrip_byte_identical which creates a
multi-shard index, copies the directory, freezes+unfreezes the
original, and asserts every file (except the regenerated merged_*
manifests) is byte-identical to the snapshot.
Python text-mode open() translates \n to \r\n on Windows, breaking
byte equality against the Rust serde_json::to_writer_pretty output
(which always uses \n). Pass newline="\n" explicitly so the
freeze/unfreeze roundtrip is byte-identical on all platforms.
@raphaelsty raphaelsty merged commit 6fe755c into main May 28, 2026
13 checks passed
@raphaelsty raphaelsty deleted the freeze-index branch May 28, 2026 16:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant