jerry-024 commented on code in PR #108:
URL:
https://github.com/apache/paimon-vector-index/pull/108#discussion_r4056656926
##########
core/src/ivfpq.rs:
##########
@@ -1962,39 +1984,61 @@ fn scan_pq_range_list<C: Collector + Send>(
return Ok(());
}
let scan_query = |scratch: &mut PqRangeScratch, query: &mut
PqRangeQuery<C>| {
+ let cached_ip_table =
+ query.query_index < cached_queries && query.metric ==
MetricType::InnerProduct;
+ let table_centroid = centroid.filter(|_| !cached_ip_table);
let table = if query.query_index < cached_queries {
if query.table_list.is_none()
- || (centroid.is_some() && query.table_list != Some(list_id))
+ || (table_centroid.is_some() && query.table_list !=
Some(list_id))
{
build_pq_range_table(
pq,
&query.query,
- centroid,
+ query.metric,
+ table_centroid,
&mut scratch.residual_query,
&mut query.table,
);
query.table_list = Some(list_id);
}
- &query.table
+ &mut query.table
} else {
build_pq_range_table(
pq,
&query.query,
+ query.metric,
centroid,
&mut scratch.residual_query,
&mut scratch.table,
);
- &scratch.table
+ &mut scratch.table
};
- scan_pq_range_codes(
+ let ip_offset = centroid
+ .filter(|_| cached_ip_table)
+ .map(|centroid| -fvec_inner_product(&query.query, centroid));
+ if let Some(offset) = ip_offset {
Review Comment:
**[Ponytail][simplification]** The cached residual-IP path copies, mutates,
and restores the entire first `ksub` subtable just to add one per-list scalar,
which also requires special cleanup after collector errors. Since
`scan_pq_range_codes` has only this caller, pass `ip_offset` into it and start
the distance accumulation with that value instead. This removes
`first_subtable`, the backup/restore work, and the restore-on-error test; keep
the existing rounding test to guard the numerical behavior.
##########
core/src/range.rs:
##########
@@ -135,17 +136,76 @@ impl DistanceBand {
}
}
-/// Only L2 is certified so far. Cosine and inner-product bands return
-/// `Unsupported`, which means "we cannot serve this request, please fall
back",
-/// not "the call has a bug".
-pub(crate) fn ensure_certified_metric(metric: MetricType) -> io::Result<()> {
+pub(crate) fn prepare_range_queries(
+ queries: &[f32],
+ dimension: usize,
+ metric: MetricType,
+) -> io::Result<Cow<'_, [f32]>> {
+ if metric != MetricType::Cosine {
+ return Ok(Cow::Borrowed(queries));
+ }
+ let mut normalized = queries.to_vec();
+ for query in normalized.chunks_exact_mut(dimension) {
+ let norm = fvec_normalize(query);
+ if !norm.is_finite() || query.iter().any(|value| !value.is_finite()) {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidData,
+ "non-finite normalized query",
+ ));
+ }
+ }
+ Ok(Cow::Owned(normalized))
+}
+
+pub(crate) fn range_probe_lists(
+ queries: &[f32],
+ centroids: &[f32],
+ dimension: usize,
+ nlist: usize,
+ nprobe: usize,
+ metric: MetricType,
+) -> io::Result<Vec<Vec<usize>>> {
if metric == MetricType::L2 {
- return Ok(());
+ return Ok(kmeans::find_topk_batch(
+ queries,
+ queries.len() / dimension,
+ centroids,
+ nlist,
+ dimension,
+ nprobe,
+ )
+ .0);
+ }
+ if centroids.iter().any(|value| !value.is_finite()) {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidData,
+ "non-finite IVF centroid",
+ ));
+ }
+ queries
Review Comment:
**[Major][performance]** This non-L2 branch runs `find_topk_checked`
serially for every query, so Flat and SQ range searches do `O(nq * nlist *
dimension)` coarse-assignment work on the caller thread before list I/O;
`nprobe` and filtering cannot reduce it. For example, 1,024 queries, 4,096
lists, and 128 dimensions require about 537 million scalar component
evaluations. Please reuse the thresholded `par_chunks_exact` pattern already
used by PQ/RQ here, while retaining `find_topk_checked` so non-finite
query-centroid distances remain rejected and collecting results in query order.
--
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]
