Skip to content

perf(py): zero-copy contiguous input + GIL release for pq_encode_batch - #108

Draft
konjoinfinity wants to merge 3 commits into
mainfrom
claude/upbeat-mccarthy-dnb6un
Draft

perf(py): zero-copy contiguous input + GIL release for pq_encode_batch#108
konjoinfinity wants to merge 3 commits into
mainfrom
claude/upbeat-mccarthy-dnb6un

Conversation

@konjoinfinity

Copy link
Copy Markdown
Collaborator

Summary

  • OPTIMIZATION_OPPORTUNITIES.md item 10 flagged pq_encode_batch as always copying its [N, D] vectors input with .iter().copied().collect(), even when the numpy array was already C-contiguous (the common case). Verified this was still true and fixed it: now borrows directly via .as_slice() when contiguous, mirroring the quantize_int8_batch idiom already established elsewhere in this file, and only falls back to an owned copy for the (rare) non-contiguous case.
  • Also releases the GIL around the pq_encode_into call — it's already rayon-parallel internally but previously ran with the GIL held, unlike quantize_int8_batch/train which already release it.
  • centroids ([M, K, sub_dim]) is left as a direct copy: it's orders of magnitude smaller than vectors for realistic N, and PQCodebook must own its centroid buffer regardless, so there's no copy to avoid there.
  • While digging through the roadmap I found items 7 and 8 in the same section (search_batch_arrays, vectorized query normalization) had also already shipped but were still listed as open — corrected those too so a future pass doesn't re-chase them.

Measured (Rust-level microbenchmark of the copy step in isolation — this container has no pytest/numpy to drive an end-to-end Python benchmark): copying a [200,000, 768] f32 array (614 MB) took a stable ~220–230ms across repeated trials, vs. ~935ms for the actual pq_encode_into work at the same scale (M=96, K=256) — the removed copy is ~20–25% of total call time at this scale, and a larger fraction for smaller batches where the fixed copy cost dominates more.

Type of change

  • Performance improvement
  • Documentation update

Checklist

  • cargo test passes locally (238 vectro_lib tests; pq_encode_into's numerics are untouched, only input ownership changed)
  • cargo clippy -- -D warnings reports no errors (verified against both the project's baseline flags and the full CI gate flags: -D warnings -D clippy::pedantic -D clippy::unwrap_used -D clippy::expect_used -D clippy::panic -D clippy::todo -D clippy::dbg_macro)
  • cargo fmt --check passes on the files touched by this PR
  • pytest python/tests/ — not run (no pytest/numpy installed in this container); no Python source changed, only the Rust FFI binding's internals
  • ruff check python/ — not applicable, no Python changed
  • No hardcoded absolute paths
  • No embedding/eval/model-weight files staged
  • Scoped to one logical concern (the pq_encode_batch fix + doc reconciliation for the same section)
  • Before/after timing included above
  • No new binary format / no serialization change

Related issues

Continuation of the plan following #99 / #100 / #104 / #105 / #106 / #107.


Generated by Claude Code

claude added 3 commits July 2, 2026 01:36
pq_encode_batch always copied its [N, D] vectors input with
.iter().copied().collect(), even when the numpy array was already
C-contiguous. Now borrows directly via .as_slice() in that case
(mirroring quantize_int8_batch's established idiom) and only copies for
the rare non-contiguous case. Also releases the GIL around the
pq_encode_into call (already rayon-parallel internally, but previously
held the GIL), matching the same fix already shipped elsewhere.
centroids stays a direct copy -- it's orders of magnitude smaller than
vectors and PQCodebook must own its buffer regardless.

Measured (Rust-level copy microbenchmark; no pytest/numpy available in
this container for an end-to-end Python benchmark): the avoided copy
was ~220-230ms of a ~935ms pq_encode_into call at [200000, 768],
M=96, K=256 (~20-25% of total call time).

Also reconciles OPTIMIZATION_OPPORTUNITIES.md's FFI/marshalling section:
items 7 and 8 (search_batch_arrays, vectorized query normalization)
turned out to already be shipped; corrected to prevent re-chasing them.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EhXzwfPHcfRj7NdSaSeasB
Split the varr.as_slice().unwrap_or_else(...) call onto two lines to
match rustfmt's line-width preference, flagged by the gates job's
repo:fmt-check (now correctly matching content post kiban v1.1.5,
unlike the earlier false-positive net-new bug this session worked
around).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EhXzwfPHcfRj7NdSaSeasB
The previous fmt fix reverted the as_slice().unwrap_or_else() split
based on a manual line-width miscount -- rustfmt genuinely wants that
chain split (confirmed via a direct rustfmt --check run against this
file), so restore it. Also multi-line the adjacent PQCodebook struct
literal, which rustfmt's struct_lit_width heuristic (not the general
line-width limit) wants split regardless of the ~90-char line fitting
under 100.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EhXzwfPHcfRj7NdSaSeasB
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.

2 participants