JingsongLi commented on code in PR #82:
URL:
https://github.com/apache/paimon-vector-index/pull/82#discussion_r3844180787
##########
core/src/index.rs:
##########
@@ -1158,6 +1162,36 @@ pub struct VectorIndexTrainer {
}
impl VectorIndexTrainer {
+ pub fn from_options(options: &HashMap<String, String>) -> io::Result<Self>
{
+ let mut config_options = ConfigOptions::new(options)?;
+ let approximate_assignment = config_options
+ .optional("approximate-assignment")
+ .map(|value| parse_bool_option("approximate-assignment", &value))
+ .transpose()?;
+ config_options.values.remove("approximate-assignment");
+ let config = VectorIndexConfig::from_options(&config_options.values)?;
+ if approximate_assignment.is_some() && config.index_type() !=
IndexType::IvfPq {
+ return Err(invalid_input(
+ "approximate-assignment is only valid for IVF-PQ",
+ ));
+ }
+ match approximate_assignment {
+ Some(enabled) => Self::new_with_approximate_assignment(config,
enabled),
+ None => Self::new(config),
+ }
+ }
+
+ pub fn new_with_approximate_assignment(
+ config: VectorIndexConfig,
+ enabled: bool,
+ ) -> io::Result<Self> {
+ let mut trainer = Self::new(config)?;
Review Comment:
[P2] Reject unsupported configs in the typed constructor
`from_options` rejects `approximate-assignment` for every non-IVF-PQ index,
but this public typed constructor accepts Flat, SQ, RQ, or DiskANN configs and
silently ignores `enabled`. A caller applying the setting to a runtime-selected
config therefore receives `Ok` even though the requested behavior was not
honored. Please validate `config.index_type() == IndexType::IvfPq` here as well
and add a non-IVF-PQ regression test.
##########
core/src/vamana.rs:
##########
@@ -711,47 +719,76 @@ impl VamanaGraph {
query: &[f32],
search_list_size: usize,
) -> Vec<ScoredNode> {
+ let mut scratch = self.search_scratch(search_list_size);
Review Comment:
[P2] Clamp the public search width before allocating scratch
`search_scratch` preallocates several buffers directly from
`search_list_size`, even when the graph contains far fewer nodes. The previous
implementation remained bounded by `adjacency.len()`. On a small graph,
`greedy_search(..., usize::MAX)` now panics immediately with `capacity
overflow`, and less extreme oversized widths can cause avoidable large
allocations. Please clamp the effective width to the graph size (and handle an
empty graph) before constructing the scratch state.
##########
core/src/pq.rs:
##########
@@ -370,6 +376,84 @@ impl ProductQuantizer {
);
}
+ /// Blocked batch encode for the IVF-PQ add path.
+ ///
+ /// The nbits=8 path uses a transposed-codebook kernel: per sub-quantizer
+ /// the centroids are transposed once to `[dsub][ksub]` so the inner
+ /// distance loop is stride-1 over `ksub` and runs on SIMD (NEON/AVX2,
+ /// scalar fallback). This removes the per-vector-per-sub GEMM calls and
+ /// their distance-table memory traffic. Distances use the norms identity
+ /// `argmin_j (|c_j|^2 - 2 q·c_j)`; the row's own norm is constant per
+ /// argmin and dropped. Results match [`Self::encode_batch`] except for
+ /// ulp-level argmin ties caused by the different summation order, so use
+ /// this only where codes are freshly produced (index build), not where
+ /// byte-stable output is pinned.
+ pub(crate) fn encode_batch_blocked(&self, data: &[f32], n: usize, codes:
&mut [u8]) {
+ // Transposing the codebook costs O(d * ksub); skip it for tiny
+ // batches and for the 4-bit packed path, which keeps the original
+ // per-vector implementation.
+ if self.nbits == 8
+ && n >= ENCODE_TRANSPOSE_MIN_ROWS
+ && (0..self.m).all(|sub| self.chunk_dim(sub) >= 4)
+ {
+ self.encode_batch_8bit_transposed(data, n, codes);
+ return;
+ }
+ self.encode_batch(data, n, codes);
Review Comment:
[P2] Keep PQ scoring consistent across IVF-PQ batch boundaries
This fallback sends batches smaller than 32 rows back through
`encode_batch`, whose norm/dot L2 calculation can lose the distance difference
through cancellation, while larger batches use direct squared differences.
Since `IVFPQIndex::try_add` chunks at 32,768 rows, adding 32,799 identical rows
can encode the first 32,768 as code 1 and the final 31 as code 0 solely because
they land in the tail batch. A deterministic `d=4, m=1` reproduction with
`q=[100000000; 4]`, centroid 0 `[100000008; 4]`, and centroid 1 `[100000000;
4]` selects centroid 0 in the tail even though the f64 squared distances are
256 and 0. Please keep the transpose threshold as a layout optimization, but
use the numerically stable direct-difference scoring rule for the small-batch
path as well.
##########
core/src/ivfpq.rs:
##########
@@ -114,6 +162,24 @@ impl IVFPQIndex {
codes: vec![Vec::new(); nlist],
precomputed_table: Vec::new(),
fastscan_codes: Vec::new(),
+ assign_graph: None,
+ approximate_assignment: use_approximate_assignment(d, nlist),
Review Comment:
[P2] Document the new default and its compatibility escape hatch
This silently changes existing large IVF-PQ builds to approximate coarse
assignment once `dimension * nlist >= 1_000_000` and the row threshold is
reached, which can change persisted list placement and low-`nprobe` recall. The
public IVF-PQ option table and usage examples do not mention this default or
the new `approximate-assignment=false` override. Please document the threshold,
true/false/default semantics, recall trade-off, and the fact that this is
build-only state that is not serialized.
##########
core/src/ivfpq.rs:
##########
@@ -218,19 +350,59 @@ impl IVFPQIndex {
&data[offset * self.d..(offset + batch_n) * self.d],
&ids[offset..offset + batch_n],
batch_n,
- );
+ )?;
offset += batch_n;
}
+ Ok(())
}
- fn add_batch(&mut self, data: &[f32], ids: &[i64], n: usize) {
+ fn add_batch(&mut self, data: &[f32], ids: &[i64], n: usize) ->
io::Result<()> {
+ if self.approximate_assignment && self.assign_graph.is_none() {
+ self.auto_assignment_rows_seen =
self.auto_assignment_rows_seen.saturating_add(n);
Review Comment:
[P2] Do not amortize graph construction against rows that cannot benefit
This counter includes rows that were already assigned exactly, but those
rows cannot amortize a graph built later. For example, 16,383 rows added in
earlier calls followed by one final row causes the complete centroid graph to
be built for that one-row tail; the same total input in one batch uses the
graph for every row. This makes build cost and persisted assignment layout
depend heavily on caller batching. Please base activation on expected
remaining/total work, require enough current work to repay startup, or
share/buffer the graph across writers.
##########
core/src/ivfpq.rs:
##########
@@ -208,8 +280,68 @@ impl IVFPQIndex {
self.pq.train(&pq_train_data, n);
}
+ /// Build the build-only centroid graph for approximate assignment.
+ /// Skipped for small nlist, where an exact scan is cheap.
+ fn rebuild_assign_graph(&mut self) -> io::Result<()> {
+ self.assign_graph = None;
+ if !self.approximate_assignment {
+ return Ok(());
+ }
+ let params = crate::diskann::DiskAnnBuildParams {
+ max_degree: 12,
+ build_search_list_size: APPROX_ASSIGN_SEARCH_LIST,
+ alpha: 1.2,
+ seed: 42,
+ memory_budget_bytes: APPROX_ASSIGN_MEMORY_BUDGET_BYTES,
+ storage_layout: crate::diskann::DiskAnnStorageLayout::Compact,
+ raw_vector_encoding: crate::diskann::DiskAnnRawVectorEncoding::F32,
+ build_distance:
crate::diskann::DiskAnnBuildDistance::FullPrecision,
+ };
+ if let Err(error) = validate_assign_graph_memory_budget(
+ self.nlist,
+ params.max_degree,
+ params.build_search_list_size,
+ rayon::current_num_threads(),
+ params.memory_budget_bytes,
+ ) {
+ if self.approximate_assignment_explicit {
+ return Err(error);
+ }
+ emit_log(
+ LogLevel::Warn,
+ &format!("automatic IVF-PQ approximate assignment disabled:
{error}"),
+ );
+ self.approximate_assignment = false;
+ return Ok(());
+ }
+ match crate::vamana::VamanaGraph::build(
+ &self.quantizer_centroids,
+ self.nlist,
+ self.d,
+ params,
+ ) {
+ Ok(graph) => self.assign_graph = Some(graph),
+ Err(error) if self.approximate_assignment_explicit => return
Err(error),
+ Err(error) => {
+ emit_log(
+ LogLevel::Warn,
+ &format!(
+ "automatic IVF-PQ approximate assignment disabled
after graph build failed: {error}"
+ ),
+ );
+ self.approximate_assignment = false;
+ }
+ }
+ Ok(())
+ }
+
/// Add vectors in batches (Faiss-style: batch assign → batch residual →
batch encode).
pub fn add(&mut self, data: &[f32], ids: &[i64], n: usize) {
+ self.try_add(data, ids, n)
+ .expect("explicit IVF-PQ approximate assignment graph build
failed");
Review Comment:
[P2] Expose a fallible low-level add path
Explicit approximate assignment introduces recoverable graph-build errors,
including the fixed memory-budget rejection, but the public low-level `add` API
converts them into a panic while `try_add` remains crate-private.
`VectorIndexWriter`, FFI, and JNI correctly propagate the same error, so direct
`IVFPQIndex` callers are the only callers that cannot recover and may abort
under `panic=abort`. Please expose a public `Result`-returning add path, or
change `add` to return `io::Result<()>` if compatibility permits.
--
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]
