leaves12138 commented on code in PR #734:
URL: https://github.com/apache/paimon-rust/pull/734#discussion_r3820405657
##########
crates/paimon/src/table/vector_search_builder.rs:
##########
@@ -3284,14 +3337,160 @@ fn collect_raw_batch_vector_batch(
}
}
+ if !dense_row_ids.is_empty() {
+ let query_matrix = scoring_plan
+ .dense_query_matrix
+ .as_deref()
+ .expect("dense query dimensions were validated above");
+ let dimension = dense_dimension.expect("dense rows require dense
queries");
+ let queries_per_chunk = (RAW_SCORE_MATRIX_TARGET_ELEMENTS /
dense_row_ids.len())
+ .max(1)
+ .min(scoring_plan.all_query_indices.len());
+ for (query_chunk_index, query_indices) in scoring_plan
+ .all_query_indices
+ .chunks(queries_per_chunk)
+ .enumerate()
+ {
+ let query_start = query_chunk_index * queries_per_chunk *
dimension;
+ let query_end = query_start + query_indices.len() * dimension;
+ let scores = compute_raw_vector_score_matrix(
+ &dense_vectors,
+ dense_row_ids.len(),
+ &query_matrix[query_start..query_end],
+ query_indices.len(),
+ dimension,
+ &scoring_plan.query_l2_squared_norms,
+ query_indices,
+ metric,
+ )?;
+ for (matrix_query_index, &query_index) in
query_indices.iter().enumerate() {
+ let query_scores = &scores[matrix_query_index *
dense_row_ids.len()
+ ..(matrix_query_index + 1) * dense_row_ids.len()];
+ top_k_out[query_index].offer_many(
+ dense_row_ids
+ .iter()
+ .zip(query_scores)
+ .map(|(&row_id, &score)| RawScoredRow { row_id, score
}),
+ );
+ }
+ }
+ }
+
+ Ok(())
+}
+
+fn ensure_raw_vector_dimension(stored_len: usize, query_len: usize) ->
crate::Result<()> {
+ if stored_len != query_len {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "Query vector dimension mismatch: raw row has {}, but query
has {}",
+ stored_len, query_len
+ ),
+ source: None,
+ });
+ }
Ok(())
}
+#[allow(clippy::too_many_arguments)]
+fn compute_raw_vector_score_matrix(
+ stored_vectors: &[f32],
+ row_count: usize,
+ query_vectors: &[f32],
+ query_count: usize,
+ dimension: usize,
+ query_l2_squared_norms: &[f32],
+ query_indices: &[usize],
+ metric: RawVectorMetric,
+) -> crate::Result<Vec<f32>> {
+ let score_count =
+ row_count
+ .checked_mul(query_count)
+ .ok_or_else(|| crate::Error::DataInvalid {
+ message: "Vector raw search score matrix is too
large".to_string(),
+ source: None,
+ })?;
+ debug_assert_eq!(stored_vectors.len(), row_count * dimension);
+ debug_assert_eq!(query_vectors.len(), query_count * dimension);
+ debug_assert_eq!(query_indices.len(), query_count);
+
+ let mut scores = vec![0.0; score_count];
+ // Query × stored-vector^T produces a query-major score matrix. Each
query's
+ // scores are contiguous, which feeds partial Top-K without strided reads.
+ sgemm_a_bt(
+ query_count,
+ row_count,
+ dimension,
+ 1.0,
+ query_vectors,
+ stored_vectors,
+ 0.0,
+ &mut scores,
+ );
+ if metric == RawVectorMetric::InnerProduct {
+ return Ok(scores);
+ }
+
+ let stored_l2_squared_norms = stored_vectors
+ .chunks_exact(dimension)
+ .map(|vector| vector.iter().map(|value| value * value).sum::<f32>())
+ .collect::<Vec<_>>();
+ for (matrix_query_index, &query_index) in query_indices.iter().enumerate()
{
+ for (row_index, &stored_l2_squared_norm) in
stored_l2_squared_norms.iter().enumerate() {
+ let score = &mut scores[matrix_query_index * row_count +
row_index];
+ let query_l2_squared_norm = query_l2_squared_norms[query_index];
+ *score = match metric {
+ RawVectorMetric::L2 => {
+ if !stored_l2_squared_norm.is_finite() ||
!query_l2_squared_norm.is_finite() {
+ let stored =
+ &stored_vectors[row_index * dimension..(row_index
+ 1) * dimension];
+ let query = &query_vectors
+ [matrix_query_index *
dimension..(matrix_query_index + 1) * dimension];
+ let squared_distance = query
+ .iter()
+ .zip(stored)
+ .map(|(query_value, stored_value)| {
+ let difference = query_value - stored_value;
+ difference * difference
+ })
+ .sum::<f32>();
+ 1.0 / (1.0 + squared_distance)
+ } else {
+ let squared_distance =
Review Comment:
The `||a||² + ||b||² - 2·a·b` reconstruction is not numerically equivalent
to the previous scalar L2 loop for large finite components. I reproduced this
with four queries and 128-D vectors: for an exact `[1e10; 128]` match, this
path returns about `4.44e-16` instead of `1.0`; a vector differing by 1024 in
one component receives the same matrix score, while the scalar path returns
about `9.54e-7`. This can change Top-K results, not just add a small error.
Please use a numerically stable path or fallback when cancellation can
dominate, and add large-finite-value coverage.
##########
crates/paimon/src/table/vector_search_builder.rs:
##########
@@ -3372,7 +3562,7 @@ fn compute_raw_vector_score_from_values(
dot += q * stored;
norm_b += stored * stored;
}
- let denominator = query_l2_norm * norm_b.sqrt();
+ let denominator = (query_l2_squared_norm * norm_b).sqrt();
Review Comment:
Taking `sqrt(query_norm_squared * stored_norm_squared)` introduces an
overflow that the previous scalar implementation avoided. For an identical
`[1e15, 0]` query and stored vector, both squared norms are finite (`1e30`),
but their product becomes infinity, so this code returns cosine `0` instead of
`1`. This affects the scalar path too, not only matrix scoring. Please compute
the two square roots before multiplying (as before), or otherwise avoid the
intermediate product overflow, and add a regression test.
--
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]
