github-actions[bot] commented on code in PR #68028:
URL: https://github.com/apache/doris/pull/68028#discussion_r4022237009
##########
thirdparty/patches/lance-c-0.1.9-multivector.patch:
##########
@@ -0,0 +1,1713 @@
+diff --git a/src/lib.rs b/src/lib.rs
+--- a/src/lib.rs
++++ b/src/lib.rs
+@@ -39,6 +39,7 @@
+ mod index_model;
+ mod index_segment;
+ mod merge_insert;
++mod multivector;
+ mod restore;
+ pub mod runtime;
+ mod scanner;
+diff --git a/src/scanner.rs b/src/scanner.rs
+--- a/src/scanner.rs
++++ b/src/scanner.rs
+@@ -226,8 +226,18 @@
+ if let Some(cols) = &self.columns {
+ scanner.project(cols)?;
+ }
++ let multi_vector = self.nearest.as_ref().is_some_and(|query| {
++ matches!(
++ query.query.data_type(),
++ arrow_schema::DataType::FixedSizeList(_, _)
++ )
++ });
+ if self.limit.is_some() || self.offset.is_some() {
+ scanner.limit(self.limit, self.offset)?;
++ if multi_vector {
++ // Retain Lance's window validation, but defer truncation
until the final sort.
++ scanner.limit(None, None)?;
++ }
+ }
+ if let Some(bs) = self.batch_size {
+ scanner.batch_size(bs);
+@@ -261,7 +271,27 @@
+ if let Some(np) = self.nprobes {
+ scanner.nprobes(np as usize);
+ }
+- if let Some(rf) = self.refine_factor {
++ if multi_vector {
++ if matches!(
++ self.metric_override,
++ Some(crate::index::LanceMetricType::Hamming)
++ ) {
++ return Err(lance_core::Error::invalid_input_source(
++ "multi-vector queries support only l2, cosine, and
dot metrics".into(),
++ ));
++ }
++ let refine = self.refine_factor.unwrap_or(1);
++ if refine == 0
++ || n.k as usize
++ > crate::multivector::MAX_QUERY_VECTOR_CANDIDATES /
refine as usize
++ {
++ return Err(lance_core::Error::invalid_input_source(
++ "multi-vector refined candidate count must be in
1..=100000".into(),
++ ));
++ }
++ // Validate actual stored values and refine candidate scores
before TopK.
++ scanner.refine(refine);
++ } else if let Some(rf) = self.refine_factor {
+ scanner.refine(rf);
+ }
+ if let Some(ef) = self.ef {
+@@ -269,6 +299,9 @@
+ }
+ if let Some(m) = self.metric_override {
+ scanner.distance_metric(m.to_distance());
++ } else if multi_vector {
++ // Resolve the same default on indexed and uncovered
fragments.
++
scanner.distance_metric(lance_linalg::distance::DistanceType::L2);
+ }
+ if let Some(ui) = self.use_index {
+ scanner.use_index(ui);
+@@ -300,6 +333,12 @@
+ Ok(PreparedScanner {
+ scanner,
+ distributed_fts,
++ multi_vector_window: multi_vector.then_some((
++ self.offset.unwrap_or(0) as usize,
++ self.limit.map(|n| n as usize),
++ )),
++ batch_size: self.batch_size,
++ scan_statistics_callback: self.scan_statistics_callback.clone(),
+ })
+ }
+ }
+@@ -314,10 +353,26 @@
+ struct PreparedScanner {
+ scanner: lance::dataset::scanner::Scanner,
+ distributed_fts: Option<PreparedFtsExecution>,
++ multi_vector_window: Option<(usize, Option<usize>)>,
++ batch_size: Option<usize>,
++ scan_statistics_callback: Option<ExecutionStatsCallback>,
+ }
+
+ impl PreparedScanner {
+ async fn try_into_stream(self) -> Result<DatasetRecordBatchStream> {
++ if let Some((offset, limit)) = self.multi_vector_window {
++ let plan =
crate::multivector::rewrite(self.scanner.create_plan().await?)?;
++ let plan = crate::multivector::apply_result_window(plan, offset,
limit)?;
++ let stream = lance_datafusion::exec::execute_plan(
++ plan,
++ lance_datafusion::exec::LanceExecutionOptions {
++ batch_size: self.batch_size,
++ execution_stats_callback: self.scan_statistics_callback,
++ ..Default::default()
++ },
++ )?;
++ return Ok(DatasetRecordBatchStream::new(stream));
++ }
+ let Some(distributed_fts) = self.distributed_fts else {
+ return self.scanner.try_into_stream().await;
+ };
+@@ -1978,6 +2033,21 @@
+ }
+ let column_str = unsafe { helpers::parse_c_string(column)? }.unwrap();
+
++ let query = unsafe { decode_query_values(query_data, query_len,
element_type)? };
++
++ s.nearest = Some(NearestQuery {
++ column: column_str.to_string(),
++ query,
++ k,
++ });
++ Ok(0)
++}
++
++unsafe fn decode_query_values(
++ query_data: *const c_void,
++ query_len: usize,
++ element_type: i32,
++) -> Result<arrow_array::ArrayRef> {
+ let dtype = match element_type {
+ 0 => LanceDataType::Float32,
+ 1 => LanceDataType::Float16,
+@@ -2016,9 +2086,112 @@
+ }
+ };
+
++ Ok(query)
++}
++
++/// Set one multi-vector query, supplied as a row-major matrix of
floating-point values.
++/// The caller must supply dimension * num_vectors aligned elements matching
the column type.
++#[unsafe(no_mangle)]
++pub unsafe extern "C" fn lance_scanner_nearest_multivector(
++ scanner: *mut LanceScanner,
++ column: *const c_char,
++ query_data: *const c_void,
++ dimension: usize,
++ num_vectors: usize,
++ element_type: i32,
++ k: u32,
++) -> i32 {
++ scanner_poison_check!(scanner, -1);
++ scanner_ffi_try!(scanner, unsafe {
++ nearest_multivector_inner(
++ scanner,
++ column,
++ query_data,
++ dimension,
++ num_vectors,
++ element_type,
++ k,
++ )
++ },)
++}
++
++unsafe fn nearest_multivector_inner(
++ scanner: *mut LanceScanner,
++ column: *const c_char,
++ query_data: *const c_void,
++ dimension: usize,
++ num_vectors: usize,
++ element_type: i32,
++ k: u32,
++) -> Result<i32> {
++ use arrow_schema::{DataType, Field};
++ let invalid = |message: &str|
lance_core::Error::invalid_input_source(message.into());
++ if scanner.is_null() || column.is_null() || query_data.is_null() {
++ return Err(invalid("scanner, column, and query_data must not be
NULL"));
++ }
++ if dimension == 0 || dimension > i32::MAX as usize || num_vectors == 0 ||
k == 0 {
++ return Err(invalid(
++ "dimension, num_vectors, and k must be positive; dimension must
fit int32",
++ ));
++ }
++ if num_vectors > crate::multivector::MAX_QUERY_VECTORS
++ || num_vectors > crate::multivector::MAX_QUERY_VECTOR_CANDIDATES / k
as usize
++ {
++ return Err(invalid(
++ "multi-vector query exceeds 128 subvectors or 100000
subvector-candidates",
++ ));
++ }
++ let (data_type, width) = match element_type {
++ 0 => (DataType::Float32, 4),
++ 1 => (DataType::Float16, 2),
++ 2 => (DataType::Float64, 8),
++ _ => {
++ return Err(invalid(
++ "multi-vector queries require float16, float32, or float64",
++ ));
++ }
++ };
++ let count = dimension
++ .checked_mul(num_vectors)
++ .filter(|count| *count <= isize::MAX as usize / width)
++ .ok_or_else(|| invalid("query matrix byte size overflows"))?;
++ let s = unsafe { &mut *scanner };
++ if s.fts_query.is_some() || s.fts_context.is_some() {
++ return Err(invalid(
++ "nearest and full-text search are mutually exclusive",
++ ));
++ }
++ let column = unsafe { helpers::parse_c_string(column)? }.unwrap();
++ let field = s
++ .dataset
++ .schema()
++ .field(column)
++ .ok_or_else(|| invalid("multi-vector column does not exist"))?;
++ match field.data_type() {
++ DataType::List(child) if !child.is_nullable() => match
child.data_type() {
++ DataType::FixedSizeList(element, dim)
++ if *dim == dimension as i32 && *element.data_type() ==
data_type => {}
++ _ => return Err(invalid("multi-vector dimension/type mismatch")),
++ },
++ _ => {
++ return Err(invalid(
++ "multi-vector column must be List of non-nullable
FixedSizeList",
++ ));
++ }
++ }
++ // A primitive array is interpreted as one vector by Lance. Preserve
matrix shape even
++ // for a single subvector. Lance does not preserve element nullability in
its schema.
++ let values = unsafe { decode_query_values(query_data, count,
element_type)? };
++ crate::multivector::validate_query(values.as_ref())?;
++ let query = arrow_array::FixedSizeListArray::try_new(
++ Arc::new(Field::new("item", data_type, false)),
++ dimension as i32,
++ values,
++ None,
++ )?;
+ s.nearest = Some(NearestQuery {
+- column: column_str.to_string(),
+- query,
++ column: column.to_string(),
++ query: Arc::new(query),
+ k,
+ });
+ Ok(0)
+diff --git a/src/multivector.rs b/src/multivector.rs
+--- /dev/null
++++ b/src/multivector.rs
+@@ -0,0 +1,514 @@
++// SPDX-License-Identifier: Apache-2.0
++// SPDX-FileCopyrightText: Copyright The Lance Authors
++
++//! Correct multi-vector scoring before the pinned Lance plan's candidate
limits.
++
++use std::collections::HashMap;
++use std::sync::Arc;
++
++use arrow_array::types::{Float16Type, Float32Type, Float64Type};
++use arrow_array::{
++ Array, ArrayRef, ArrowPrimitiveType, BooleanArray, FixedSizeListArray,
Float32Array, ListArray,
++ RecordBatch, UInt64Array,
++};
++use arrow_schema::{DataType, SchemaRef};
++use datafusion::error::{DataFusionError, Result};
++use datafusion::execution::context::TaskContext;
++use datafusion::physical_plan::{
++ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
SendableRecordBatchStream,
++ stream::RecordBatchStreamAdapter,
++};
++use futures::{StreamExt, TryStreamExt, stream};
++use lance::io::exec::KNNVectorDistanceExec;
++use lance_linalg::distance::{Cosine, DistanceType, Dot, L2};
++
++// Lance creates one ANN branch per query vector,
++// each overfetching 10 * k candidates before scoring; wire bytes alone
cannot bound this work.
++pub(crate) const MAX_QUERY_VECTORS: usize = 128;
++pub(crate) const MAX_QUERY_VECTOR_CANDIDATES: usize = 100_000;
++
++fn invalid(message: impl Into<String>) -> DataFusionError {
++ DataFusionError::Execution(message.into())
++}
++
++/// Rewrite inside TopK/refinement, before any score can discard a candidate.
++pub(crate) fn rewrite(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn
ExecutionPlan>> {
++ let children = plan
++ .children()
++ .into_iter()
++ .map(|child| rewrite(child.clone()))
++ .collect::<Result<Vec<_>>>()?;
++ let plan = if children.is_empty() {
++ plan
++ } else {
++ plan.with_new_children(children)?
++ };
++ let mode = if let Some(exact) =
plan.downcast_ref::<KNNVectorDistanceExec>() {
++ if exact.is_batch {
++ return Err(invalid(
++ "expected one logical multi-vector query, not batch queries",
++ ));
++ }
++ Some(Scoring::Exact {
++ query: exact.query.clone(),
++ column: exact.column.clone(),
++ metric: exact.distance_type,
++ })
++ // This pinned Lance node is not publicly re-exported, so match its
stable plan name.
++ } else if plan.name() == "MultivectorScoringExec" {
++ Some(Scoring::Indexed)
++ } else {
++ None
++ };
++ Ok(match mode {
++ Some(mode) => Arc::new(MultiVectorScoreExec {
++ original: plan,
++ mode,
++ }),
++ None => plan,
++ })
++}
++
++/// Apply the final distance-ordered window without invalidating output
batching.
++pub(crate) fn apply_result_window(
++ plan: Arc<dyn ExecutionPlan>,
++ offset: usize,
++ limit: Option<usize>,
++) -> Result<Arc<dyn ExecutionPlan>> {
++ use datafusion::physical_expr::{PhysicalSortExpr, expressions};
++ use datafusion::physical_plan::{
++ coalesce_partitions::CoalescePartitionsExec, limit::GlobalLimitExec,
sorts::sort::SortExec,
++ };
++ if plan
++ .downcast_ref::<lance_datafusion::exec::StrictBatchSizeExec>()
++ .is_some()
++ {
++ // Offset can split a previously strict batch. Keep Lance's final
rechunker
++ // outside the window, preserving its resolved batch size, including
defaults.
++ let input = apply_result_window(plan.children()[0].clone(), offset,
limit)?;
++ return plan.with_new_children(vec![input]);
++ }
++ let sort = PhysicalSortExpr {
++ expr: expressions::col("_distance", plan.schema().as_ref())?,
++ options: arrow::compute::SortOptions {
++ descending: false,
++ nulls_first: false,
++ },
++ };
++ // Fragment-scoped payload takes can reorder batches. Restore distance
order
++ // before the window; the nearest plan already bounds candidate rows by k.
++ let sorted = Arc::new(SortExec::new(
++ [sort].into(),
++ Arc::new(CoalescePartitionsExec::new(plan)),
++ ));
++ Ok(Arc::new(GlobalLimitExec::new(sorted, offset, limit)))
++}
++
++#[derive(Clone, Debug)]
++enum Scoring {
++ Exact {
++ query: ArrayRef,
++ column: String,
++ metric: DistanceType,
++ },
++ Indexed,
++}
++
++#[derive(Debug)]
++struct MultiVectorScoreExec {
++ original: Arc<dyn ExecutionPlan>,
++ mode: Scoring,
++}
++
++impl DisplayAs for MultiVectorScoreExec {
++ fn fmt_as(&self, _: DisplayFormatType, f: &mut std::fmt::Formatter) ->
std::fmt::Result {
++ write!(f, "MultiVectorScore: {}", self.original.name())
++ }
++}
++
++impl ExecutionPlan for MultiVectorScoreExec {
++ fn name(&self) -> &str {
++ "MultiVectorScoreExec"
++ }
++ fn properties(&self) -> &Arc<PlanProperties> {
++ self.original.properties()
++ }
++ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
++ self.original.children()
++ }
++ fn required_input_distribution(&self) ->
Vec<datafusion::physical_expr::Distribution> {
++ self.original.required_input_distribution()
++ }
++ fn with_new_children(
++ self: Arc<Self>,
++ children: Vec<Arc<dyn ExecutionPlan>>,
++ ) -> Result<Arc<dyn ExecutionPlan>> {
++ Ok(Arc::new(Self {
++ original: self.original.clone().with_new_children(children)?,
++ mode: self.mode.clone(),
++ }))
++ }
++ fn execute(
++ &self,
++ partition: usize,
++ context: Arc<TaskContext>,
++ ) -> Result<SendableRecordBatchStream> {
++ let schema = self.schema();
++ match &self.mode {
++ Scoring::Exact {
++ query,
++ column,
++ metric,
++ } => {
++ let input = self.children()[0].execute(partition, context)?;
++ let query = query.clone();
++ let column = column.clone();
++ let metric = *metric;
++ let output_schema = schema.clone();
++ let output = input
++ .map(move |batch| {
++ let query = query.clone();
++ let column = column.clone();
++ let schema = output_schema.clone();
++ async move {
++ let batch = batch?;
++ tokio::task::spawn_blocking(move || {
++ exact_batch(batch, query, &column, metric,
schema)
++ })
++ .await
++ .map_err(|e|
DataFusionError::External(Box::new(e)))?
++ }
++ })
++
.buffered(lance_core::utils::tokio::get_num_compute_intensive_cpus());
++ Ok(Box::pin(RecordBatchStreamAdapter::new(schema, output)))
++ }
++ Scoring::Indexed => {
++ let inputs = self
++ .children()
++ .into_iter()
++ .map(|child| child.execute(partition, context.clone()))
++ .collect::<Result<Vec<_>>>()?;
++ let output_schema = schema.clone();
++ let output =
++ stream::once(async move { indexed_batch(inputs,
output_schema).await });
++ Ok(Box::pin(RecordBatchStreamAdapter::new(schema, output)))
++ }
++ }
++ }
++}
++
++fn row_distance<T: ArrowPrimitiveType>(
++ query: &dyn Array,
++ vectors: &FixedSizeListArray,
++ metric: DistanceType,
++) -> Result<Option<f32>>
++where
++ T::Native: L2 + Cosine + Dot + Into<f64>,
++{
++ let q = query
++ .as_any()
++ .downcast_ref::<arrow_array::PrimitiveArray<T>>()
++ .ok_or_else(|| invalid("multi-vector query element type mismatch"))?;
++ let values = vectors
++ .values()
++ .as_any()
++ .downcast_ref::<arrow_array::PrimitiveArray<T>>()
++ .ok_or_else(|| invalid("multi-vector stored element type mismatch"))?;
++ if vectors.null_count() != 0
++ || values.null_count() != 0
++ || values
++ .values()
++ .iter()
++ .any(|v| !Into::<f64>::into(*v).is_finite())
++ {
++ return Err(invalid(
++ "multi-vector stored subvectors must contain only finite,
non-null elements",
++ ));
++ }
++ let dimension = vectors.value_length() as usize;
++ let distance = metric.func();
++ // Subtracting each small distance from 1 rounds it away before TopK. Sum
minima
++ // directly, using f64 only for the accumulator; the base kernels and
output remain f32.
++ let mut score = 0.0f64;
++ for query_vector in q.values().chunks_exact(dimension) {
++ let best = values
++ .values()
++ .chunks_exact(dimension)
++ .map(|vector| distance(query_vector, vector))
++ // Finite zero-norm vectors have undefined cosine distance.
Ignore those
++ // pairs; a query with no defined match masks this row, not the
whole scan.
++ .filter(|distance| !distance.is_nan())
Review Comment:
[P2] Do not abort valid scans on finite-distance overflow
This filter removes NaN but retains infinities. Lance's pinned kernels
return f32 distances, so documented-valid finite operands can still overflow:
with L2, query `[0]` and stored `[f32::MAX]` produce `+Inf`; even individually
finite pair distances can overflow when their sum is cast back to f32. The
later check then errors the whole batch, even if another row has a normal score
and should rank. Dot has the analogous case. This is distinct from non-finite
inputs and zero-norm cosine. Please treat a nonrepresentable row score as no
match, or reject/bound such operands consistently, and add exact/refined
overflow cases.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
