JunRuiLee commented on code in PR #108:
URL:
https://github.com/apache/paimon-vector-index/pull/108#discussion_r4056911480
##########
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:
Thanks, fixed in e404a85. Non-L2 probing now uses thresholded parallel
iteration, retaining checked distances and query order. Added 1/4-worker and
overflow regression tests.
##########
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:
Thanks, fixed in e404a85. Removed subtable copy/restore; the scanner applies
the offset to the first selected entry, preserving rounding and signed zero.
Kept the rounding test and added regression coverage.
--
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]
