Gabriel39 commented on code in PR #68028:
URL: https://github.com/apache/doris/pull/68028#discussion_r4021778984
##########
thirdparty/patches/lance-c-0.1.9-multivector.patch:
##########
@@ -0,0 +1,508 @@
+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);
+@@ -300,6 +310,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 +330,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 = 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 +2029,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 +2082,104 @@
+ }
+ };
+
++ 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",
++ ));
++ }
++ 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)
Review Comment:
Fixed in cd8255cb4c. Exact scoring/refinement validates actual primitive
validity and finiteness before calling the distance kernels. The reconstructed
nullable schema flag remains accepted, but actual null/NaN/Infinity values fail
execution. Added fixtures written through Arrow/Lance and native, BE and SQL
regressions, including distance-only projection.
--
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]
