JingsongLi commented on code in PR #90:
URL:
https://github.com/apache/paimon-vector-index/pull/90#discussion_r3930422959
##########
core/src/pq.rs:
##########
@@ -455,6 +478,70 @@ impl ProductQuantizer {
);
}
+ /// Transposed codebook: per sub a `[dsub][ksub]` block at a uniform
+ /// stride of `max_dsub * ksub`, so sub lookup stays O(1) for balanced
+ /// non-uniform chunks. Returns the table and the per-sub stride.
+ fn build_transposed_codebook(&self) -> (Vec<f32>, usize) {
+ let m = self.m;
+ let ksub = self.ksub;
+ let max_dsub = (0..m).map(|sub|
self.chunk_dim(sub)).max().unwrap_or(0);
+ let sub_stride = max_dsub
+ .checked_mul(ksub)
+ .expect("transposed codebook stride overflows usize");
+ let mut transposed = vec![0.0f32; m * sub_stride];
+ for sub in 0..m {
+ let dsub = self.chunk_dim(sub);
+ let c_base = self.centroid_chunk_base(sub);
+ let dst = &mut transposed[sub * sub_stride..sub * sub_stride +
dsub * ksub];
+ for j in 0..ksub {
+ for k in 0..dsub {
+ dst[k * ksub + j] = self.centroids[c_base + j * dsub + k];
+ }
+ }
+ }
+ (transposed, sub_stride)
+ }
+
+ /// Transposed-codebook encode for nbits=8. Rows are split into blocks
+ /// only for parallelism; each row is encoded independently.
+ fn encode_batch_8bit_transposed(&self, data: &[f32], n: usize, codes: &mut
[u8]) {
+ let d = self.d;
+ let m = self.m;
+ let ksub = self.ksub;
+ let cs = self.code_size();
+ debug_assert_eq!(cs, m);
+
+ let (transposed, sub_stride) = self.build_transposed_codebook();
Review Comment:
Could we avoid rebuilding the entire transposed codebook for unamortized
tiny batches? build_transposed_codebook allocates and copies d × 256 floats on
every call (768 KiB at d=768), even when n=1. The parent used the canonical
path below max(32, 4 × threads); on native arm64 with d=768, m=96, dsub=8, an
isolated release probe measured auto at 209 µs versus 95 µs for the
parent-equivalent canonical path (2.20× slower). At n=7 auto was already
faster, so this is limited to row-at-a-time ingestion. A centroid-major
direct-L2 single-row kernel with the same dimension-order mul_add, NaN, and tie
semantics would preserve split invariance without paying the transpose;
alternatively, cache the transpose with reliable invalidation for the publicly
mutable centroids.
##########
docs/api.html:
##########
@@ -50,6 +50,7 @@ <h2>Shared lifecycle</h2>
<div class="flow" aria-label="Unified API lifecycle"><div
class="flow-step"><small>01</small><strong>Create a Trainer<br>Parse and
validate options</strong></div><div
class="flow-step"><small>02</small><strong>Submit one or more<br>training
batches</strong></div><div class="flow-step"><small>03</small><strong>Finish
training and<br>create a one-shot Writer</strong></div><div
class="flow-step"><small>04</small><strong>Add row IDs / vectors<br>and write
the file</strong></div><div class="flow-step"><small>05</small><strong>Detect
file magic<br>and execute searches</strong></div></div>
<ul><li>Vectors are contiguous <code>f32</code> values; length must
equal <code>vector_count × dimension</code>.</li><li>Training data may arrive
in batches. Every IVF trainer keeps a deterministic reservoir of at most
<code>max(65,536, 64 × resolved nlist)</code> vectors. DiskANN starts from a
50,000-row cap and lowers it when necessary so the retained sample, optional
cosine-normalized copy, codebook, and parallel PQ-training scratch fit
<code>diskann.memory-budget-bytes</code>. Sampling is independent of batch
boundaries.</li><li>The Python and Java one-shot <code>train</code> helpers
infer <code>dimension</code> from the matrix and use its row count for
automatic <code>nlist</code>. When the matrix is only a sample, pass the final
corpus size as <code>expected-vector-count</code>. Streaming Trainer APIs
require a concrete dimension before their first batch.</li><li>A Writer may
receive production vectors in multiple batches. Row-ID count must equal vector
count.</li>
<li>Readers expose metadata, single-query search, batch search, and
Roaring64-filtered variants.</li><li>Files carry their type and resolved model
sections. Callers do not pass index options again when opening a
Reader.</li></ul>
<div class="callout warning"><strong>IVF coarse assignment is
approximate by default for large centroid matrices</strong>When <code>dimension
× nlist ≥ 1,000,000</code>, <code>ivf.coarse-assignment=auto</code> uses a
Vamana graph while training and adding vectors. Search still selects lists by
exact centroid distance, so graph assignment can lower recall at small
<code>nprobe</code> and does not guarantee that a vector is found by a
self-query with <code>nprobe=1</code>. Set
<code>ivf.coarse-assignment=exact</code> to disable the graph, preserve exact
nearest-centroid assignment, and avoid graph startup cost for small non-empty
batches. Empty batches never build the graph.</div>
+ <div class="callout warning"><strong>IVF-PQ encoding defaults to
auto</strong><code>ivf.pq-encoding=auto</code> uses the transposed direct-L2
encoder on x86 with AVX2+FMA and on AArch64, and the blocked SGEMM
expanded-form encoder otherwise. Codes can differ across these backends; NaN
and high-dynamic-range inputs also follow the selected backend's arithmetic.
Set <code>ivf.pq-encoding=canonical</code> to reproduce
<code>ProductQuantizer::encode_batch</code> on the same CPU and runtime
backend; it is substantially slower. Neither mode promises byte-identical codes
across CPU feature sets. The choice affects builds only, is not stored, and
does not change the index format or search path.</div>
Review Comment:
The documented dispatch is broader than the implementation. Auto uses
transposed/SGEMM only for 8-bit, ksub=256 shapes whose every chunk has dsub>=4;
dsub=1/2/3 falls back to canonical even on AVX2+FMA/AArch64, and a non-finite
codebook on the no-fast-kernel branch also uses canonical. For example,
dimension=128 with pq.m=64 is valid but auto already equals canonical, so the
arithmetic and performance guidance here is incorrect. Could we qualify this
callout—and the matching IVF-PQ and release docs—with the actual gates and
canonical fallback?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
