Gabriel39 commented on code in PR #68028:
URL: https://github.com/apache/doris/pull/68028#discussion_r4022025541
##########
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:
Fixed in c3c89e3440. Undefined cosine pairs are ignored; if any query
subvector has no defined stored match, the row is filtered before TopK instead
of failing the scan. Valid subvectors in a mixed zero/nonzero row remain
eligible. Added native exact/refined tests for zero rows alongside valid rows
and zero queries, plus BE and SQL regression coverage over an indexed fragment
and an appended all-zero fragment. The original failure was reproduced before
the fix; all 12 native tests and 52 BE reader tests now pass. Synced to
lance-format/lance-c#83 (aabf748).
##########
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:
Fixed in c3c89e3440. The scorer now uses Lance's parse_field_path and
traverses StructArray fields, preserving the existing direct-column fast path.
The pinned KNN helper itself is private. Added exact and indexed-refinement
tests for payload.vectors, payload.`vectors.with.dot`, and
payload.inner.`vectors.with.dot`. The new test reproduced the missing
List-column error before the fix and now passes on both Lance pins. Synced the
same implementation and tests to lance-format/lance-c#83 (aabf748).
--
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]
