perf(vector): reuse prepared f32 PQ8 codebooks during partition search

[xiaguanglei]

Problem

An IVF_PQ query builds one query-to-codeword distance table for every probed
partition. The query residual changes per partition, but the PQ codebook does
not.

For an f32 PQ8 index, the existing table builders pass the original AoS
codebook to the generic batch distance API. That API prepares/transposes the
same 256 centroids again for every sub-vector and every probed partition. With
nprobes=32, a DIM=128 query therefore repeatedly prepares the same 128 KiB
codebook layout 32 times before bulk PQ-code scoring can begin. Once the Lance
index cache is warm, this repeated CPU setup is a large part of end-to-end
latency.

The cost grows with the number of probes and IVF partitions. This explains the
measured shape of the baseline: one-probe queries have no reusable work, while
high-probe searches spend a substantial fraction of their time preparing an
unchanged codebook.

Dot-product PQ had the same horizontal-reduction shape but no reusable prepared
target implementation. The current-format V3 partition loader also rebuilt
partition storage from metadata alone, so an index-level prepared L2 codebook
could not reach the partition search path.

Solution

• add a portable DotPrepared SoA target layout matching the existing
  L2Prepared design; LLVM can vectorize across PQ centroids on AArch64 and
  x86_64
• keep one current-format index-level ProductQuantizer and share its prepared
  targets with every loaded PQ partition through Arc
• initialize the active metric's prepared targets lazily with OnceLock, so
  partitions share one allocation and concurrent first use is safe
• keep the existing low-setup builder when the query is guaranteed to probe
  only one partition
• enable the prepared search path only for f32 PQ8; PQ4 has only 16 centroids
  and end-to-end testing did not show enough benefit to justify changing its
  setup path

The per-partition distance table is still computed from that partition's query
residual. The optimization removes only redundant codebook preparation; it
does not reuse a distance table across different residuals or change distance
semantics.

Scope

The fast path covers current-format f32 PQ8 storage:

• IVF_PQ and IVF_HNSW_PQ
• L2 and Dot
• Cosine through Lance's existing normalize-then-L2 conversion

PQ4, f16, f64, legacy partition loading, serialized metadata, and public APIs
remain on their existing paths.

End-to-end benchmark

Method

Measured on SIFT-1M, the industry-standard ANN benchmark
(corpus-texmex): 1M x 128-d f32 base, the
official 10k-query set, and the official exact top-100 ground truth. The exact
input files:

File	Rows / shape	SHA-256
sift_base.fvecs	1,000,000 x 128	21f66e2975057b5728ba56de1c825bac4f4d89d596609ae985741c6242631816
sift_query.fvecs	10,000 x 128	f7fc9be140accdfd64116c2fa2365ecdb69b8f084970c6b0532db5ff79ac8fdc
sift_groundtruth.ivecs	10,000 x 100	2b71de0a8d5a83e6a84eec3e23fb8b611d8801dd9b3a6cd62f070ab65ea65f4f

Environment and controls:

• x86_64 Skylake server (shared cloud host), repository release-with-debug
  profile
• one physical Lance dataset and index built by the baseline binary, then
  opened read-only by both saved binaries
• interleaved BCCBBCCB A/B protocol, warm cache, concurrency=1,
  median-of-medians aggregation
• result hashes and Recall@10 checked for identity on every cell
• benchmark harness kept on a separate local-only branch; it is not part of
  this implementation PR

IVF_PQ (SIFT-1M, nlist=256, PQ8, L2)

Index sift1m-ivf256-pq16-l2.lance (IVF_PQ, 8-bit PQ), interleaved BCCBBCCB
A/B, repetitions=6, warm cache, concurrency=1:

nprobes	Baseline	Candidate	Mean latency change	QPS change
1	4.326 ms	4.227 ms	-2.28%	+1.90%
4	5.218 ms	4.695 ms	-10.02%	+11.80%
16	7.592 ms	6.635 ms	-12.60%	+13.09%
32	10.600 ms	8.863 ms	-16.39%	+18.38%
64	13.484 ms	11.831 ms	-12.26%	+17.97%
128	15.852 ms	14.954 ms	-5.66%	+11.93%

Result hashes and Recall@10 matched in every cell. The win peaks at
nprobes=32 (-16.4%) and falls off at higher probe counts where partition I/O
becomes the larger share of the query — the shape the setup-cost model predicts:
one-probe queries have no reusable work (-2.3%, within noise), and each
additional probe amortizes the once-only codebook preparation.

IVF_HNSW_PQ (SIFT-1M, nlist=256, PQ8, L2)

Secondary check on sift1m-ivf256-hnswpq-l2.lance (IVF_HNSW_PQ, same PQ8 fast
path):

nprobes	Mean latency change
8	-2.96%
32	-6.73%
128	-9.03%

HNSW_PQ shows a smaller, positive win because HNSW graph traversal adds fixed
work outside PQ table construction.

Prewarm path (prewarm_index API)

The query warm-up path above exercises load_partition during query
execution. A separate API, prewarm_index, eagerly materializes partition
storage into the index cache. Before this PR, that path called
try_from_batch_with_remapper, which rebuilt every prewarmed partition from
metadata alone, so each partition owned an empty prepared-target cache and
re-transposed the codebook on first use — i.e. the prewarm path silently
missed the prepared-codebook optimization entirely.

This PR fixes materialize_partition_for_prewarm to pass the retained
index-level quantizer through (try_from_batch_with_quantizer), so prewarmed
partitions share the same lazily initialized prepared targets as query-loaded
partitions. A regression test
(test_prewarm_materialization_shares_prepared_targets) locks this in.

IVF_PQ prewarm path, SIFT-1M, nlist=256, PQ8, L2 — --cache prewarm
(partitions loaded via prewarm_index), interleaved BCCBBCCB,
repetitions=6, warmup_rounds=1:

nprobes	Baseline	Candidate	Mean latency change	QPS change
1	4.080 ms	3.993 ms	-2.14%	+3.69%
4	5.089 ms	4.679 ms	-8.05%	+8.55%
16	7.796 ms	6.800 ms	-12.78%	+14.42%
32	11.207 ms	9.051 ms	-19.24%	+24.14%
64	17.222 ms	13.725 ms	-20.31%	+24.87%
128	25.570 ms	20.601 ms	-19.43%	+19.29%

Result hashes and Recall@10 matched in every cell. The prewarm-path win
(peaking ~-20%) is larger than the query-path win (-16.4%) because the
un-fixed prewarm path was fully on the legacy rebuild path, whereas the
query path was already partially optimized.

Memory and compatibility

Prepared targets are lazy and only the active metric can allocate one set. A
DIM=128/PQ8 index adds one shared 128 KiB f32 codebook layout after the first
multi-probe use, regardless of how many partitions are loaded. DeepSizeOf
uses Arc-aware accounting so shared targets are not charged once per
partition.

There are no file-format, protobuf, serialized metadata, dependency, or public
API changes. Legacy readers and writers are untouched.

Correctness coverage

• compare prepared and AoS L2/Dot distance tables element-by-element
• compare DotPrepared with the generic Dot implementation across empty, odd,
  short, PQ-like, and large shapes
• verify loaded current-format partitions share one lazy target allocation for
  L2 and Dot
• verify only Product Quantizer is retained by current-format partition
  storage; Flat, SQ, and RQ keep their existing ownership and memory behavior
• verify a one-probe calculator reproduces the legacy table without
  initializing the lazy targets
• verify PQ4 never installs prepared targets
• preserve legacy non-divisible-dimension search behavior
• compare scalar per-ID and bulk PQ distance APIs for L2 and Dot

Validation

cargo fmt --all -- --check
cargo check -p lance-index --tests
cargo check -p lance --tests
cargo test -p lance-index test_distance_with_legacy_truncated_dimension
cargo test -p lance-index test_four_bit_pq_keeps_low_setup_distance_path
cargo test -p lance-index test_partition_storages_share_prepared_targets
cargo test -p lance-index test_only_product_quantizer_is_retained_for_partition_storage
cargo test -p lance test_build_ivf_sq
cargo test -p lance test_build_ivf_rq
cargo clippy -p lance-index --tests -- -D warnings
cargo clippy -p lance --tests -- -D warnings

Earlier full-suite validation of the same PQ8 implementation also passed
cargo test -p lance-index, cargo test -p lance,
cargo clippy --all --tests --benches -- -D warnings, and
cargo check --workspace --tests --benches.

[lance-gatekeeper]

✅ Gate recommendation: approve.

The rebase preserves the previously reviewed patch exactly. Shared prepared targets still reach both on-demand and prewarm partition loads on the updated base, with no new finding.