shyjsarah commented on code in PR #82:
URL:
https://github.com/apache/paimon-vector-index/pull/82#discussion_r3869883629
##########
core/src/ivfpq.rs:
##########
@@ -228,9 +337,16 @@ impl IVFPQIndex {
// L2/IP without OPQ borrows the caller's batch instead of copying it.
let processed = self.preprocess_queries(data, n);
- let assignments =
- kmeans::find_nearest_batch(&processed, n,
&self.quantizer_centroids, self.nlist, d);
-
+ // Both branches return the exact nearest centroid; the projection only
+ // prunes the scan (see `projected_assign`).
+ let assignments = match &self.coarse_projection {
+ Some(projection) => {
+ projection.assign(&processed, n, &self.quantizer_centroids,
self.nlist)
+ }
+ None => {
+ kmeans::find_nearest_batch(&processed, n,
&self.quantizer_centroids, self.nlist, d)
Review Comment:
**[Major — correctness]** These two branches do not actually have the same
“exact nearest centroid” contract. `projection.assign` exact-checks candidates
using direct squared differences, while `find_nearest_batch` uses the f32
norm-minus-dot GEMM path; batch query probing uses that GEMM convention as
well. On valid finite vectors translated around `1e8`, I reproduced Enabled
assigning the true list while Disabled/query probing selected another list. An
index built with Enabled then missed its own two inserted vectors for a batch
query with `nprobe=1` (returned `[-1, -1]`). Please use one stable distance/tie
contract for both build and query, or conservatively exact-refine the GEMM
result. A regression should cover Enabled/Disabled/Auto, scalar and batch
paths, and serialized readers.
##########
core/src/projected_assign.rs:
##########
@@ -0,0 +1,1195 @@
+// 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;
+/// 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;
+ }
+ // 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);
+
+ 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);
Review Comment:
**[Major — resource usage]** Auto reaches `top_subspace` before checking
whether projected assignment is profitable, and the fit currently has no
explicit memory/work budget. For `d=nlist=8192`, Auto chooses `block=2728`; the
key PCA buffers can exceed 1 GiB and the eight iterations are roughly `6e12`
FLOPs. A small training input can still reach this because coarse k-means
repeats rows when `n <= nlist`. Please estimate checked peak bytes/work before
allocation, cap `block*nlist` and `block^2`, and skip/fail safely before
entering PCA when the configured budget would be exceeded.
##########
core/src/projected_assign.rs:
##########
@@ -0,0 +1,1195 @@
+// 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;
+/// 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;
+ }
+ // 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);
+
+ 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 calibration_rows = calibration_rows.min(calibration.len() /
d.max(1));
+ 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];
+ let (block_rows, _) =
+ crate::kmeans::assignment_block_plan(n, dp, nlist,
rayon::current_num_threads());
+ let block_rows = block_rows.min(MAX_BLOCK_ROWS);
+ let evaluations = out
+ .par_chunks_mut(block_rows)
Review Comment:
**[Major — resource usage]** `assignment_block_plan` returns a `parallel`
flag specifically to enforce a serial fallback when the aggregate scratch
budget would be exceeded, but this path discards it and always uses
`par_chunks_mut`. Each worker independently allocates the `rows*nlist` score
matrix, `nlist` bounds, and potentially `nlist` candidates. For
`nlist=1,048,576` and 32 workers, score+bounds are about 384 MiB aggregate and
worst-case candidates add about 512 MiB, bypassing the intended budget. Please
honor `parallel=false` and include all per-worker scratch—not only the score
matrix—in the planner.
--
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]
