shyjsarah commented on code in PR #82:
URL:
https://github.com/apache/paimon-vector-index/pull/82#discussion_r3877894034
##########
core/src/kmeans.rs:
##########
@@ -747,10 +786,10 @@ pub fn find_topk_batch(
d: usize,
nprobe: usize,
) -> (Vec<Vec<usize>>, Vec<Vec<f32>>) {
- let centroid_norms = (0..k)
- .map(|c| fvec_norm_l2sqr(¢roids[c * d..(c + 1) * d]))
- .collect::<Vec<_>>();
- find_topk_batch_with_centroid_norms(queries, nq, centroids,
¢roid_norms, k, d, nprobe)
+ (0..nq)
Review Comment:
**[Major — resource usage]** The stable batch probing fix now runs one
`find_topk` allocation per active Rayon query. Each call allocates `Vec<(f32,
usize)>` with capacity `nlist` (16 bytes per centroid on 64-bit targets), with
no aggregate query scratch budget. At `nlist=1,048,576` and 32 active workers
this can keep roughly 512 MiB of tuple buffers live. A bounded allocator
reproduction at `k=262,144, nq=32, threads=32` measured about 109 MiB peak on
this head versus 39 MiB before the change. Please keep the direct-distance
contract but bound/reuse storage—for example an O(nprobe) heap per query, a
capped per-worker buffer, or query blocking based on `workers * nlist *
tuple_size`.
##########
core/src/kmeans.rs:
##########
@@ -705,16 +735,10 @@ pub(crate) fn find_nearest_batch(
k: usize,
d: usize,
) -> Vec<usize> {
- if n == 0 {
- return Vec::new();
- }
- if n == 1 {
- return vec![find_nearest(&data[..d], centroids, k, d)];
- }
-
- let mut assignments = vec![0usize; n];
- assign_clusters_fast(data, n, d, centroids, k, &mut assignments, 0.0);
- assignments
+ (0..n)
Review Comment:
**[Major — performance regression]** Replacing the blocked batch assignment
with independent row scans fixes cancellation, but the four-centroid kernel
does not recover the lost batch-level reuse. In an 8-worker release comparison
at `d=768, nlist=4096, n=2048`, the previous blocked path had a 38.1 ms median
versus 59.6 ms here (about 56% higher latency / 36% lower throughput).
One-worker batch probing also regressed from 1.88 to 3.58 ms and from 5.47 to
8.58 ms on two tested shapes. Since this is a build-performance PR and the
exact fallback is used whenever Auto is unavailable/rejected, please add a
numerically stable blocked direct-distance kernel (with a small-batch
crossover), or exact-refine an error-bounded GEMM candidate set.
##########
core/src/projected_assign.rs:
##########
@@ -0,0 +1,1259 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Exact coarse assignment through a low-dimensional projection.
+//!
+//! `IVFPQIndex::add` must find the nearest coarse centroid of every row. The
+//! exact scan is one `n × nlist × d` GEMM, which dominates `add` once PQ
+//! encoding is vectorized. This module keeps the result exact but does most of
+//! the work in a `d'`-dimensional PCA subspace of the centroids:
+//!
+//! 1. project the row and centroids with a contractive PCA projection `P`
+//! (`d' × d`);
+//! 2. for every centroid compute a conservative lower bound on the true
+//! squared distance: the projected distance plus the reverse triangle
+//! inequality on the components outside the subspace
+//! (`|x-c|² >= |P₀(x-c)|² + (|P₀⊥x| - |P₀⊥c|)²`), with every f32
+//! projection, GEMM and f64 rounding error accounted for explicitly;
+//! 3. evaluate exact distances in ascending bound order and stop as soon as
+//! the next bound cannot beat the best exact distance (branch-and-bound).
+//!
+//! Step 2 is a valid lower bound for any contractive `P`, so the assignment
+//! never depends on how well the PCA converged; `P` only decides how many
+//! exact evaluations step 3 needs. Rows without low-dimensional structure
+//! degrade to checking every centroid, i.e. the exact scan plus one small
+//! GEMM. At `train`, candidate widths are therefore timed on a sample of
+//! training vectors, including projection, bound construction, candidate
+//! collection, and exact checks. Automatic mode keeps the fastest projection
+//! only when it is clearly faster than an exact scan of the same sample.
+//! Without a sample, forced mode takes the widest candidate and automatic
+//! mode leaves the exact scan in place.
+//!
+//! The projection is a training artifact: it is derived from the centroids,
+//! shared by `IVFPQIndex::from_trained`, and never serialized.
+
+use crate::blas::{dgemm_a_bt, sgemm_a_bt};
+use crate::distance::{fvec_l2sqr, fvec_l2sqr_four};
+use nalgebra::{DMatrix, SymmetricEigen};
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+use rayon::prelude::*;
+use std::time::{Duration, Instant};
+
+/// Automatic mode only times widths with `d' * MAX_AUTO_DP_DIVISOR <= d`.
+const MAX_AUTO_DP_DIVISOR: usize = 3;
+/// Calibration samples below this count as absent.
+const MIN_CALIBRATION_ROWS: usize = 256;
+/// Rows of the training data used to calibrate `d'`.
+pub(crate) const CALIBRATION_ROWS: usize = 2048;
+/// With a calibration sample, automatic mode keeps the projection only when
+/// its measured time is below this fraction of the exact scan.
+const MAX_AUTO_COST_FRACTION: f64 = 0.85;
+/// Block subspace iterations; the bound stays valid regardless of convergence.
+const SUBSPACE_ITERATIONS: usize = 8;
+const SUBSPACE_SEED: u64 = 0x7a5e_c7ed;
+/// Keep automatic PCA fitting bounded before allocating its work matrices.
+const MAX_AUTO_PCA_BYTES: u128 = 256 * 1024 * 1024;
+const MAX_AUTO_PCA_FLOPS: u128 = 100_000_000_000;
+/// Smallest projection worth a GEMM; also the rounding unit for `d'`.
+const MIN_DP: usize = 8;
+/// Rows per parallel block; the projected score matrix is `rows × nlist`.
+const MAX_BLOCK_ROWS: usize = 1024;
+/// The projection is re-orthonormalized in f64 and scaled by
+/// `1 - ORTHO_MARGIN`, so its Gram matrix is `(1 - ORTHO_MARGIN)² · I` up to
+/// rounding. `orthonormal_f64` verifies (Gershgorin, including the f64
+/// rounding of the Gram entries) that every eigenvalue of `P Pᵀ` lies within
+/// `GRAM_MARGIN` of that diagonal, and declines the projection otherwise, so
+/// the constant singular-value bounds below hold for every stored
+/// projection. Gram-Schmidt typically lands within ~1e-15.
+const ORTHO_MARGIN: f64 = 1e-9;
+const GRAM_MARGIN: f64 = 1e-9;
+/// Bounds on the squared singular values of the stored projection (see
+/// `ORTHO_MARGIN`): `|P v|² / SIGMA_MAX_SQ <= |P₀ v|² <= |P v|² /
SIGMA_MIN_SQ`
+/// for the orthonormal basis `P₀` of its row space.
+const SIGMA_MAX_SQ: f64 = 1.0;
+const SIGMA_MIN_SQ: f64 = (1.0 - ORTHO_MARGIN) * (1.0 - ORTHO_MARGIN) * (1.0 -
GRAM_MARGIN);
+/// Candidates evaluated per branch-and-bound stage before re-pruning.
+const STAGE: usize = 32;
+
+/// How `IVFPQIndex::add` chooses between the exact centroid scan and the
+/// projected branch-and-bound. Both produce the exact nearest centroid.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
+pub enum ProjectedAssignment {
+ /// Use the projection when the centroids are compressible enough for it
+ /// to pay off; otherwise scan.
+ #[default]
+ Auto,
+ /// Always build and use the projection (still exact).
+ Enabled,
+ /// Always scan.
+ Disabled,
+}
+
+/// Contractive projection of the coarse centroids plus the per-centroid data
+/// the lower bound needs.
+#[derive(Clone, Debug)]
+pub struct CoarseProjection {
+ d: usize,
+ dp: usize,
+ explained_variance: f64,
+ /// Centroid mean (f64), subtracted before projecting rows and centroids.
+ /// Translation cancels in the projected term but not in the residual
+ /// term, whose norms must not be dominated by a common offset.
+ mean: Vec<f64>,
+ /// `dp × d`, scaled so its operator norm does not exceed one. Stored in
+ /// f64 (exactly the f32 values used for the singular-value bounds) so the
+ /// row projection carries only f64 rounding.
+ proj: Vec<f64>,
+ /// `nlist × dp`: projected centroids.
+ cents_p: Vec<f32>,
+ /// `|cents_p[c]|²`.
+ cents_p_norms: Vec<f64>,
+ /// `|cents_p[c]|`, cached to keep square roots out of assignment's hot
loop.
+ cents_p_norms_sqrt: Vec<f64>,
+ /// Upper bound on the f32 GEMM error in each projected centroid.
+ cents_p_errors: Vec<f64>,
+ /// Relative error factors of the `rows × nlist` f32 GEMM and of the f64
+ /// arithmetic in the bound, fixed by `dp`.
+ gemm_error_factor: f64,
+ f64_error_factor: f64,
+ /// Interval for `|P₀⊥ c|`, the norm of each centroid's component outside
+ /// the projected subspace, including every rounding error above.
+ cents_res_lo: Vec<f64>,
+ cents_res_hi: Vec<f64>,
+}
+
+impl CoarseProjection {
+ /// Fit a projection to `nlist` centroids of dimension `d`.
+ ///
+ /// `calibration` holds `calibration_rows` sample vectors in the centroid
+ /// space (typically training vectors). Candidate widths up to `d / 3`
+ /// (`d / 2` when forced) are timed on that sample and the fastest wins;
+ /// automatic mode keeps the projection only when it also beats the exact
+ /// scan on the sample by `MAX_AUTO_COST_FRACTION`. Without a sample,
+ /// forced mode takes the widest width and automatic mode builds nothing.
+ ///
+ /// Returns `None` when the projection is not worth it (`force == false`)
+ /// or cannot be built (too few centroids, zero variance).
+ pub(crate) fn train(
+ cents: &[f32],
+ nlist: usize,
+ d: usize,
+ force: bool,
+ calibration: &[f32],
+ calibration_rows: usize,
+ ) -> Option<Self> {
+ if nlist == 0 || d == 0 {
+ return None;
+ }
+ let calibration_rows = calibration_rows.min(calibration.len() / d);
+ if !force && calibration_rows < MIN_CALIBRATION_ROWS {
+ return None;
+ }
+ // Largest d' auto mode may pick, rounded down so the gate below holds;
+ // forced mode searches up to d / 2 (beyond that the projected GEMM
+ // costs about as much as the scan it is meant to replace).
+ let auto_limit = (d / MAX_AUTO_DP_DIVISOR) / MIN_DP * MIN_DP;
+ let block = if force {
+ (d / 2).max(MIN_DP).min(d)
+ } else {
+ auto_limit
+ };
+ if !force && (block < MIN_DP || nlist < 2 * block) {
+ return None;
+ }
+ let block = block.min(nlist);
+ if !force && !auto_pca_within_budget(nlist, d, block) {
+ return None;
+ }
+
+ let mean = column_mean(cents, nlist, d);
+ let centered: Vec<f32> = cents
+ .iter()
+ .enumerate()
+ .map(|(i, v)| v - mean[i % d] as f32)
+ .collect();
+ let total_variance = centered
+ .iter()
+ .map(|v| (*v as f64) * (*v as f64))
+ .sum::<f64>()
+ / nlist as f64;
+ if total_variance.is_nan() || total_variance <= 0.0 {
+ return None;
+ }
+
+ let (basis, eigenvalues) = top_subspace(¢ered, nlist, d, block);
+ let mut projection = if calibration_rows >= MIN_CALIBRATION_ROWS {
+ Self::select_width_by_time(
+ cents,
+ &mean,
+ &basis,
+ nlist,
+ d,
+ block,
+ force,
+ &calibration[..calibration_rows * d],
+ calibration_rows,
+ )?
+ } else if force {
+ // No sample to time against: take the widest projection.
+ Self::from_basis(cents, &mean, &basis, nlist, d, block)?
+ } else {
+ return None;
+ };
+ projection.explained_variance =
+ eigenvalues[..projection.dp].iter().sum::<f64>() / total_variance;
+ Some(projection)
+ }
+
+ /// Width minimizing elapsed assignment time on the calibration sample.
+ #[allow(clippy::too_many_arguments)]
+ fn select_width_by_time(
+ cents: &[f32],
+ mean: &[f64],
+ basis: &[f32],
+ nlist: usize,
+ d: usize,
+ block: usize,
+ force: bool,
+ calibration: &[f32],
+ calibration_rows: usize,
+ ) -> Option<Self> {
+ let exact_elapsed = if force {
+ None
+ } else {
+ let started = Instant::now();
+ let _ =
+ crate::kmeans::find_nearest_batch(calibration,
calibration_rows, cents, nlist, d);
+ Some(started.elapsed())
+ };
+ let mut best: Option<(Duration, Self)> = None;
+ let mut widths: Vec<usize> = [8usize, 4, 2]
+ .iter()
+ .map(|div| (block / div).div_ceil(MIN_DP) * MIN_DP)
+ .chain([(block * 3 / 4).div_ceil(MIN_DP) * MIN_DP, block])
+ .filter(|&w| w >= MIN_DP && w <= block)
+ .collect();
+ if force && widths.is_empty() {
+ widths.push(block);
+ }
+ widths.sort_unstable();
+ widths.dedup();
+ for dp in widths.into_iter().rev() {
+ let Some(candidate) = Self::from_basis(cents, mean, basis, nlist,
d, dp) else {
+ continue;
+ };
+ let started = Instant::now();
+ let _ = candidate.assign_with_stats(calibration, calibration_rows,
cents, nlist);
+ let elapsed = started.elapsed();
+ if best.as_ref().is_none_or(|(time, _)| elapsed <= *time) {
+ best = Some((elapsed, candidate));
+ }
+ }
+ let (elapsed, candidate) = best?;
+ if let Some(exact_elapsed) =
+ exact_elapsed.filter(|exact_elapsed| !is_fast_enough(elapsed,
*exact_elapsed))
+ {
+ crate::logging::emit_log(
+ crate::logging::LogLevel::Info,
+ &format!(
+ "IVF-PQ projected assignment not used: best width d'={}
measured {:.0}% of exact scan time",
+ candidate.dp,
+ elapsed.as_secs_f64() / exact_elapsed.as_secs_f64() * 100.0
+ ),
+ );
+ return None;
+ }
+ Some(candidate)
+ }
+
+ fn from_basis(
+ cents: &[f32],
+ mean: &[f64],
+ basis: &[f32],
+ nlist: usize,
+ d: usize,
+ dp: usize,
+ ) -> Option<Self> {
+ let proj = orthonormal_f64(&basis[..dp * d], dp, d)?;
+ let mean = mean.to_vec();
+ // Center and project in f64, then round to f32 for the row×centroid
+ // GEMM. `cent_norms` are `|c - mean|²` in f64.
+ let (cents_p, cent_norms) = project_rows(cents, nlist, d, &proj, dp,
&mean);
+ let cents_p_norms: Vec<f64> = (0..nlist)
+ .map(|c| norm_l2sqr_f64(¢s_p[c * dp..(c + 1) * dp]))
+ .collect();
+ let cents_p_norms_sqrt: Vec<f64> = cents_p_norms.iter().map(|norm|
norm.sqrt()).collect();
+ let cents_p_errors: Vec<f64> = (0..nlist)
+ .map(|c| {
+ projection_error(
+ centered_norm_upper(cent_norms[c], d),
+ cents_p_norms_sqrt[c],
+ d,
+ dp,
+ )
+ })
+ .collect();
+ let (cents_res_lo, cents_res_hi): (Vec<f64>, Vec<f64>) = (0..nlist)
+ .map(|c| residual_interval(cent_norms[c], d, cents_p_norms[c],
cents_p_errors[c]))
+ .unzip();
+ Some(Self {
+ d,
+ dp,
+ explained_variance: 0.0,
+ mean,
+ proj,
+ cents_p,
+ cents_p_norms,
+ cents_p_norms_sqrt,
+ cents_p_errors,
+ gemm_error_factor: 2.0 *
gamma_f32(dp.saturating_mul(2).saturating_add(1)),
+ f64_error_factor:
gamma_f64(dp.saturating_mul(2).saturating_add(8)),
+ cents_res_lo,
+ cents_res_hi,
+ })
+ }
+
+ pub fn dimension(&self) -> usize {
+ self.dp
+ }
+
+ pub fn explained_variance(&self) -> f64 {
+ self.explained_variance
+ }
+
+ /// Exact nearest centroid of every row (ties: smallest centroid index).
+ /// `data` must already be in the centroid space (normalized / rotated).
+ pub(crate) fn assign(&self, data: &[f32], n: usize, cents: &[f32], nlist:
usize) -> Vec<usize> {
+ self.assign_with_stats(data, n, cents, nlist).0
+ }
+
+ /// `assign` plus the total number of exact distance evaluations, for
+ /// tests and benchmarks.
+ fn assign_with_stats(
+ &self,
+ data: &[f32],
+ n: usize,
+ cents: &[f32],
+ nlist: usize,
+ ) -> (Vec<usize>, usize) {
+ let d = self.d;
+ let dp = self.dp;
+ debug_assert_eq!(self.cents_p.len(), nlist * dp);
+ let mut out = vec![0usize; n];
+ // Per worker: bounds (f64) and candidates (f64 + u32), plus projected
+ // rows and their f64 norms. The planner already counts the score
matrix.
+ let fixed_scratch = nlist.saturating_mul(
+ (std::mem::size_of::<f64>() + std::mem::size_of::<(f64, u32)>())
+ .div_ceil(std::mem::size_of::<f32>()),
+ );
+ let row_scratch = dp.saturating_add(2);
Review Comment:
**[Minor — incomplete scratch accounting]** Honoring `parallel=false` fixes
the original worker-amplification case, but `row_scratch = dp + 2` omits
transient buffers held by `project_rows`: `rows*d` f64 centered data, `rows*dp`
f64 output, and the f32 projection allocated while those buffers are still
live. The actual row cost is approximately `2*d + 3*dp + 2` f32 units. For
`d=8192, nlist=dp=512, threads=32`, the planner models about 16 MiB but can
select a parallel plan whose projection phase reaches roughly 271 MiB. Please
account for the projection and score phases explicitly (or move to reusable
caller-owned buffers), and add an end-to-end planner regression rather than
testing only the supplied estimates.
--
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]
