github-actions[bot] commented on code in PR #68028:
URL: https://github.com/apache/doris/pull/68028#discussion_r4021940567
##########
thirdparty/patches/lance-c-0.1.9-multivector.patch:
##########
@@ -0,0 +1,1328 @@
+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,45 @@
+ 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 {
++ use datafusion::physical_expr::{PhysicalSortExpr, expressions};
++ use datafusion::physical_plan::{
++ coalesce_partitions::CoalescePartitionsExec,
limit::GlobalLimitExec,
++ sorts::sort::SortExec,
++ };
++ let plan =
crate::multivector::rewrite(self.scanner.create_plan().await?)?;
++ let sort = PhysicalSortExpr {
++ expr: expressions::col("_distance", plan.schema().as_ref())?,
++ options: arrow::compute::SortOptions {
++ descending: false,
++ nulls_first: false,
++ },
++ };
++ // Fragment-scoped Lance plans can reorder candidate batches
during payload take.
++ // Apply the result window only after restoring distance order
across all partitions.
++ // The nearest plan already bounds the candidate rows by k.
++ let sorted = Arc::new(SortExec::new(
++ [sort].into(),
++ Arc::new(CoalescePartitionsExec::new(plan)),
++ ));
++ let plan = Arc::new(GlobalLimitExec::new(sorted, 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 +2052,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 +2105,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,363 @@
++// 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,
++ })
++}
++
++#[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<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))
++ .min_by(f32::total_cmp)
++ .ok_or_else(|| invalid("cannot score an empty multi-vector
row"))?;
++ score += best as f64;
++ }
++ let score = score as f32;
++ if !score.is_finite() {
Review Comment:
[P1] Do not let one finite zero-norm cosine row abort the scan. Both query
and stored-value validation accept `[[0,0]]`, but Lance's cosine kernel derives
NaN by dividing by the zero norm. The ordinary KNN scorer masks an undefined
row; this branch instead returns `multi-vector distance is not finite`, so one
zero stored subvector makes the whole exact scan (and indexed refinement) fail.
Please either preserve row-level filtering or reject zero-norm cosine
subvectors consistently up front, and cover a valid row alongside a zero row.
##########
thirdparty/patches/lance-c-0.1.9-multivector.patch:
##########
@@ -0,0 +1,1328 @@
+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,45 @@
+ 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 {
++ use datafusion::physical_expr::{PhysicalSortExpr, expressions};
++ use datafusion::physical_plan::{
++ coalesce_partitions::CoalescePartitionsExec,
limit::GlobalLimitExec,
++ sorts::sort::SortExec,
++ };
++ let plan =
crate::multivector::rewrite(self.scanner.create_plan().await?)?;
++ let sort = PhysicalSortExpr {
++ expr: expressions::col("_distance", plan.schema().as_ref())?,
++ options: arrow::compute::SortOptions {
++ descending: false,
++ nulls_first: false,
++ },
++ };
++ // Fragment-scoped Lance plans can reorder candidate batches
during payload take.
++ // Apply the result window only after restoring distance order
across all partitions.
++ // The nearest plan already bounds the candidate rows by k.
++ let sorted = Arc::new(SortExec::new(
++ [sort].into(),
++ Arc::new(CoalescePartitionsExec::new(plan)),
++ ));
++ let plan = Arc::new(GlobalLimitExec::new(sorted, 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 +2052,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 +2105,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,363 @@
++// 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,
++ })
++}
++
++#[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<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))
++ .min_by(f32::total_cmp)
++ .ok_or_else(|| invalid("cannot score an empty multi-vector
row"))?;
++ score += best as f64;
++ }
++ let score = score as f32;
++ if !score.is_finite() {
++ return Err(invalid("multi-vector distance is not finite"));
++ }
++ Ok(score)
++}
++
++fn exact_batch(
++ batch: RecordBatch,
++ query: ArrayRef,
++ column: &str,
++ metric: DistanceType,
++ schema: SchemaRef,
++) -> Result<RecordBatch> {
++ if batch.num_rows() == 0 {
++ return Ok(RecordBatch::new_empty(schema));
++ }
++ let vectors = batch
++ .column_by_name(column)
Review Comment:
[P2] Preserve nested field-path resolution here. The new API validates the
column with Lance's field-path-aware `Schema::field`, and the pinned planner
accepts paths such as `payload.vectors`; the original KNN scorer then descends
the struct path. `RecordBatch::column_by_name()` only sees top-level names, so
every accepted nested-column exact search fails here, and indexed search fails
during refinement. Reuse the pinned path resolver (including quoted dotted
names), or reject nested paths synchronously, and add exact/refined 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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
