diff --git a/README.md b/README.md index f4dccb2..02617a7 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,56 @@ Queries accept at most 128 subvectors. Both `num_vectors * k` and `refine_factor * k` must be at most 100,000 to bound plan expansion and candidate allocation. The existing single-vector API and its defaults are unchanged. +## Batch nearest-neighbor search + +Use `lance_scanner_nearest_batch` or `Scanner::nearest_batch` to search a +`FixedSizeList` column with independent query vectors: + +```cpp +const float queries[] = {1.0f, 0.0f, 0.0f, 1.0f}; +auto scanner = dataset.scan(); +scanner.nearest_batch("embedding", queries, 2, 2, 10) + .use_index(false); +``` + +Here there are two queries of dimension two, each returning up to ten rows. +The result includes the zero-based `query_index` (`int32`) and `_distance`, +alongside the requested dataset columns. A batch of one query still includes +`query_index`. Duplicate query vectors retain distinct indices; queries with no +matching rows produce no result rows. Result row order is not part of the API contract. + +This differs from `nearest_multivector`: that API scores one logical query +against a `List>` column and returns one ranked result set. +Batch input requires a fixed-size vector column and rejects a dataset with a +reserved `query_index` column. Float16, Float32, Float64, Int8, and UInt8 are accepted; +the query element type must match the column. UInt8 uses the Hamming metric. +Floating-point query values must be finite. Values are copied before the setter +returns, so the caller can immediately release its input buffer. + +Configure the query before scanning. Successful nearest-query setters replace +the previous query; an invalid batch request leaves it unchanged. FTS and +nearest queries remain mutually exclusive. Distance ranges currently require a +single-vector query; replacing it with a batch clears its bounds. All queries +share the scanner's metric, filter, index selection, and tuning parameters. +Existing prefilter and postfilter semantics apply; postfiltering can return +fewer than `k` rows. + +The scanner's global `limit` and `offset` settings are rejected in batch mode, +in either configuration order: they cannot represent independent result windows. +For per-query pagination, request enough candidates with `k` and apply the window +separately to each `query_index` in the caller. + +Bounds are 128 queries, 64 MiB of copied query values, and 100,000 candidates for +`num_queries * k * refine_factor` (use 1 when refinement is unset). Explicit +refinement must be positive. These bounds do not cap total index/payload memory; +callers should also configure scanner resource controls. Arrow streams use the +existing ownership and cancellation contract, including early stream release. + +Execution delegates to the pinned Lance scanner. Flat search can share data +scans. Eligible IVF queries can share partition scans; HNSW, refinement, adaptive +partition probing, and other ineligible plans use Lance's per-query fallback. +The batch API does not guarantee shared execution or a particular speedup. + ## Building There are four supported entry points; pick whichever matches your toolchain. diff --git a/include/lance/lance.h b/include/lance/lance.h index d29fc0b..552ee38 100644 --- a/include/lance/lance.h +++ b/include/lance/lance.h @@ -2131,7 +2131,7 @@ int32_t lance_scanner_nearest( * Call after lance_scanner_nearest and before starting the scan. Both bounds * are copied before returning; NULL means unbounded on that side. Passing NULL * for both clears the range. A successful replacement nearest query also clears - * the range. Multi-vector queries and scans without nearest are rejected. + * the range. Batch/multi-vector queries and scans without nearest are rejected. * * Bounds must be finite; negative distances are allowed. When both are present, * lower_bound must be strictly smaller than upper_bound. Invalid calls leave the @@ -2150,6 +2150,33 @@ int32_t lance_scanner_set_distance_range( const float* upper_bound ); +/** + * Set independent k-NN queries on a FixedSizeList column. + * query_data contains dimension * num_queries aligned, row-major elements, copied + * before returning. Supported element types: Float16/32/64, Int8, and UInt8 (Hamming). + * Floating-point values must be finite and the type/dimension must match the column. + * + * Every query returns up to k rows with a zero-based int32 query_index column and + * _distance, including batches of one query. Duplicate queries retain separate + * indices. A dataset column named query_index is rejected. No output row order + * is guaranteed. Index, metric, filter and search settings are shared by all queries. + * + * dimension and k must be positive; num_queries must be in 1..=128. Copied values + * are bounded to 64 MiB; num_queries * k * refine_factor must be <= 100000, using + * 1 when refinement is unset. Explicit refinement must be positive. These are + * request bounds, not a total index/payload-memory limit. + * + * Must be configured before scanning. Replaces a previous nearest query atomically; + * an invalid request preserves it. Mutually exclusive with FTS and global scanner + * limit/offset, in either configuration order. Apply per-query windows in the caller. + * Streams retain the existing ownership, error and early-release/cancellation contract. + * Returns 0 on success, -1 on error; use lance_last_error_code/message for details. + */ +int32_t lance_scanner_nearest_batch( + LanceScanner* scanner, const char* column, const void* query_data, + size_t dimension, size_t num_queries, LanceDataType element_type, uint32_t k +); + /** * Set one multi-vector query on a List> column. * Inner vectors must be non-nullable and contain no null elements; the outer list may be nullable. diff --git a/include/lance/lance.hpp b/include/lance/lance.hpp index 882059a..edcd39f 100644 --- a/include/lance/lance.hpp +++ b/include/lance/lance.hpp @@ -1636,6 +1636,22 @@ class Scanner { return *this; } + /// Independent queries; each returns up to k rows tagged with query_index. + Scanner& nearest_batch(const std::string& column, const void* query_data, + size_t dimension, size_t num_queries, + LanceDataType element_type, uint32_t k) { + if (lance_scanner_nearest_batch(handle_.get(), column.c_str(), query_data, + dimension, num_queries, element_type, k) != 0) + check_error(); + return *this; + } + + /// Float32 batch convenience overload; input values are copied before returning. + Scanner& nearest_batch(const std::string& column, const float* query_data, + size_t dimension, size_t num_queries, uint32_t k) { + return nearest_batch(column, query_data, dimension, num_queries, LANCE_DTYPE_FLOAT32, k); + } + /// One multi-vector query, copied from dimension * num_vectors row-major elements. Scanner& nearest_multivector(const std::string& column, const void* query_data, size_t dimension, size_t num_vectors, diff --git a/src/scanner.rs b/src/scanner.rs index 61717bf..c73cc0d 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -144,7 +144,53 @@ pub struct LanceScanner { schema: Option, } +#[derive(Clone, Copy, PartialEq, Eq)] +enum NearestMode { + Single, + MultiVector, + Batch { num_queries: usize }, +} + +const MAX_BATCH_QUERIES: usize = 128; +const MAX_BATCH_QUERY_BYTES: usize = 64 * 1024 * 1024; +const MAX_BATCH_CANDIDATES: usize = 100_000; + +fn batch_query_value_count(dimension: usize, num_queries: usize, width: usize) -> Result { + // Check both products before constructing an FFI slice, including on 32-bit targets. + let overflow = || { + let message = format!( + "batch query byte size overflows usize: dimension={dimension}, num_queries={num_queries}, element_width={width}" + ); + lance_core::Error::invalid_input_source(message.into()) + }; + let count = dimension.checked_mul(num_queries).ok_or_else(overflow)?; + let bytes = count.checked_mul(width).ok_or_else(overflow)?; + if bytes > MAX_BATCH_QUERY_BYTES { + let message = format!( + "batch query byte size {bytes} exceeds limit {MAX_BATCH_QUERY_BYTES}: dimension={dimension}, num_queries={num_queries}, element_width={width}" + ); + return Err(lance_core::Error::invalid_input_source(message.into())); + } + Ok(count) +} + +fn validate_batch_candidates(num_queries: usize, k: u32, refine_factor: Option) -> Result<()> { + let refine = refine_factor.unwrap_or(1) as usize; + if refine == 0 + || num_queries + .checked_mul(k as usize) + .and_then(|n| n.checked_mul(refine)) + .is_none_or(|n| n > MAX_BATCH_CANDIDATES) + { + return Err(lance_core::Error::invalid_input_source(format!( + "batch num_queries ({num_queries}) * k ({k}) * refine_factor ({refine}) must be in 1..={MAX_BATCH_CANDIDATES}" + ).into())); + } + Ok(()) +} + struct NearestQuery { + mode: NearestMode, column: String, query: arrow_array::ArrayRef, k: u32, @@ -371,12 +417,20 @@ impl LanceScanner { if let Some(handling) = &self.blob_handling { scanner.blob_handling(handling.clone()); } - let multi_vector = self.nearest.as_ref().is_some_and(|query| { - matches!( - query.query.data_type(), - arrow_schema::DataType::FixedSizeList(_, _) - ) - }); + // Both batch and multi-vector inputs are matrices, but only multi-vector + // queries may use the custom row-level scoring and global result window. + let multi_vector = self + .nearest + .as_ref() + .is_some_and(|query| query.mode == NearestMode::MultiVector); + if let Some(NearestQuery { + mode: NearestMode::Batch { num_queries }, + k, + .. + }) = &self.nearest + { + validate_batch_candidates(*num_queries, *k, self.refine_factor)?; + } if self.limit.is_some() || self.offset.is_some() { scanner.limit(self.limit, self.offset)?; if multi_vector { @@ -1041,6 +1095,14 @@ unsafe fn scanner_set_limit_inner(scanner: *mut LanceScanner, limit: i64) -> Res )); } let s = unsafe { &mut *scanner }; + if s.nearest + .as_ref() + .is_some_and(|q| matches!(q.mode, NearestMode::Batch { .. })) + { + return Err(lance_core::Error::invalid_input_source( + "global limit is unsupported for batch nearest; apply windows per query_index in the caller".into(), + )); + } s.limit = Some(limit); Ok(0) } @@ -1061,6 +1123,14 @@ unsafe fn scanner_set_offset_inner(scanner: *mut LanceScanner, offset: i64) -> R )); } let s = unsafe { &mut *scanner }; + if s.nearest + .as_ref() + .is_some_and(|q| matches!(q.mode, NearestMode::Batch { .. })) + { + return Err(lance_core::Error::invalid_input_source( + "global offset is unsupported for batch nearest; apply windows per query_index in the caller".into(), + )); + } s.offset = Some(offset); Ok(0) } @@ -2502,7 +2572,32 @@ macro_rules! scanner_set_u32 { }; } -scanner_set_u32!(lance_scanner_set_refine_factor, refine_factor); +/// Configure refinement without exceeding a batch query's candidate budget. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_scanner_set_refine_factor( + scanner: *mut LanceScanner, + value: u32, +) -> i32 { + scanner_poison_check!(scanner, -1); + scanner_ffi_try!( + scanner, + (|| -> Result { + let s = unsafe { scanner.as_mut() } + .ok_or_else(|| lance_core::Error::invalid_input_source("scanner is NULL".into()))?; + if let Some(NearestQuery { + mode: NearestMode::Batch { num_queries }, + k, + .. + }) = &s.nearest + { + s.ensure_scan_not_started("refine_factor")?; + validate_batch_candidates(*num_queries, *k, Some(value))?; + } + s.refine_factor = Some(value); + Ok(0) + })() + ) +} scanner_set_u32!(lance_scanner_set_ef, ef); /// Set inclusive lower and exclusive upper distance bounds on a single-vector query. @@ -2536,14 +2631,11 @@ unsafe fn scanner_set_distance_range_inner( .nearest .as_mut() .ok_or_else(|| invalid("distance_range requires nearest() to be configured".into()))?; - // Multi-vector scores aggregate subvectors; filtering each subvector's - // candidates would not enforce the requested range on the final score. - if matches!( - nearest.query.data_type(), - arrow_schema::DataType::FixedSizeList(_, _) - ) { + // Range bounds are currently a single-query contract. Both batch and + // multi-vector queries use matrices, so check the explicit mode. + if nearest.mode != NearestMode::Single { return Err(invalid( - "distance_range does not support multi-vector queries".into(), + "distance_range requires a single-vector nearest query".into(), )); } let lower_bound = unsafe { lower_bound.as_ref() }.copied(); @@ -2913,6 +3005,7 @@ unsafe fn scanner_nearest_inner( let query = unsafe { decode_query_values(query_data, query_len, element_type)? }; s.nearest = Some(NearestQuery { + mode: NearestMode::Single, column: column_str.to_string(), query, k, @@ -2922,17 +3015,164 @@ unsafe fn scanner_nearest_inner( Ok(0) } +/// Set independent nearest-neighbor queries copied from a row-major matrix. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_scanner_nearest_batch( + scanner: *mut LanceScanner, + column: *const c_char, + query_data: *const c_void, + dimension: usize, + num_queries: usize, + element_type: i32, + k: u32, +) -> i32 { + scanner_poison_check!(scanner, -1); + scanner_ffi_try!(scanner, unsafe { + nearest_batch_inner( + scanner, + column, + query_data, + dimension, + num_queries, + element_type, + k, + ) + }) +} + +unsafe fn nearest_batch_inner( + scanner: *mut LanceScanner, + column: *const c_char, + query_data: *const c_void, + dimension: usize, + num_queries: usize, + element_type: i32, + k: u32, +) -> Result { + use arrow_array::{Array, Float16Array, Float32Array, Float64Array}; + use arrow_schema::{DataType, Field}; + let invalid = |message: String| 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".into(), + )); + } + let s = unsafe { &mut *scanner }; + s.ensure_scan_not_started("nearest_batch")?; + if dimension == 0 || dimension > i32::MAX as usize || k == 0 { + return Err(invalid(format!( + "dimension ({dimension}) and k ({k}) must be positive; dimension must fit int32" + ))); + } + if num_queries == 0 || num_queries > MAX_BATCH_QUERIES { + return Err(invalid(format!( + "num_queries ({num_queries}) must be in 1..={MAX_BATCH_QUERIES}" + ))); + } + if s.fts_query.is_some() || s.fts_context.is_some() { + return Err(invalid( + "batch nearest and full-text search are mutually exclusive".into(), + )); + } + if s.limit.is_some() || s.offset.is_some() { + return Err(invalid("batch nearest does not support global limit/offset; apply windows per query_index in the caller".into())); + } + validate_batch_candidates(num_queries, k, s.refine_factor)?; + let (data_type, width) = match element_type { + 0 => (DataType::Float32, 4), + 1 => (DataType::Float16, 2), + 2 => (DataType::Float64, 8), + 3 => (DataType::UInt8, 1), + 4 => (DataType::Int8, 1), + _ => { + return Err(invalid(format!( + "batch element_type ({element_type}) must be float16, float32, float64, uint8, or int8" + ))); + } + }; + let count = batch_query_value_count(dimension, num_queries, width)?; + let column = unsafe { helpers::parse_c_string(column)? }.unwrap(); + if s.dataset.schema().field("query_index").is_some() { + return Err(invalid( + "batch nearest reserves the dataset column name query_index".into(), + )); + } + let field = s + .dataset + .schema() + .field(column) + .ok_or_else(|| invalid(format!("batch vector column {column:?} does not exist")))?; + match field.data_type() { + DataType::FixedSizeList(element, dim) + if dim == dimension as i32 && *element.data_type() == data_type => {} + actual => { + return Err(invalid(format!( + "batch column {column:?} must be FixedSizeList<{data_type}, {dimension}>, got {actual}" + ))); + } + } + let values = unsafe { decode_query_values(query_data, count, element_type)? }; + let finite = match values.data_type() { + DataType::Float16 => values + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .all(|v| v.is_some_and(|v| v.is_finite())), + DataType::Float32 => values + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .all(|v| v.is_some_and(f32::is_finite)), + DataType::Float64 => values + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .all(|v| v.is_some_and(f64::is_finite)), + DataType::UInt8 | DataType::Int8 => true, + actual => { + // A future decoder type must return an input error instead of poisoning the scanner. + return Err(invalid(format!( + "unsupported batch query data type {actual} for element_type={element_type}" + ))); + } + }; + if !finite { + return Err(invalid( + "batch query values must be finite and non-null".into(), + )); + } + // Keep the matrix shape even for one query: Lance uses it to emit query_index. + let query = arrow_array::FixedSizeListArray::try_new( + Arc::new(Field::new("item", data_type, false)), + dimension as i32, + values, + None, + )?; + s.nearest = Some(NearestQuery { + mode: NearestMode::Batch { num_queries }, + column: column.to_string(), + query: Arc::new(query), + k, + lower_bound: None, + upper_bound: None, + }); + Ok(0) +} + unsafe fn decode_query_values( query_data: *const c_void, query_len: usize, element_type: i32, ) -> Result { - let dtype = match element_type { - 0 => LanceDataType::Float32, - 1 => LanceDataType::Float16, - 2 => LanceDataType::Float64, - 3 => LanceDataType::UInt8, - 4 => LanceDataType::Int8, + let (dtype, alignment) = match element_type { + 0 => (LanceDataType::Float32, std::mem::align_of::()), + 1 => (LanceDataType::Float16, std::mem::align_of::()), + 2 => (LanceDataType::Float64, std::mem::align_of::()), + 3 => (LanceDataType::UInt8, std::mem::align_of::()), + 4 => (LanceDataType::Int8, std::mem::align_of::()), _ => { return Err(lance_core::Error::invalid_input_source( format!("invalid element_type: {}", element_type).into(), @@ -2940,6 +3180,17 @@ unsafe fn decode_query_values( } }; + // from_raw_parts requires alignment even when the caller supplies raw bytes. + // Validate here so single, multi-vector, and batch queries share the same guard. + if !(query_data as usize).is_multiple_of(alignment) { + return Err(lance_core::Error::invalid_input_source( + format!( + "query_data must be aligned to {alignment} bytes for element_type={element_type}" + ) + .into(), + )); + } + let query: arrow_array::ArrayRef = match dtype { LanceDataType::Float32 => { let slice = unsafe { std::slice::from_raw_parts(query_data as *const f32, query_len) }; @@ -3069,6 +3320,7 @@ unsafe fn nearest_multivector_inner( None, )?; s.nearest = Some(NearestQuery { + mode: NearestMode::MultiVector, column: column.to_string(), query: Arc::new(query), k, @@ -3246,6 +3498,136 @@ mod tests { use lance::io::exec::PreFilterSource; use lance_index::{IndexType, scalar::InvertedIndexParams}; + #[test] + fn batch_query_sizes_distinguish_overflow_from_byte_limit() { + for (dimension, num_queries, width) in [(usize::MAX, 2, 1), (usize::MAX / 2 + 1, 1, 2)] { + let error = batch_query_value_count(dimension, num_queries, width) + .unwrap_err() + .to_string(); + for detail in [ + "overflows usize".to_string(), + format!("dimension={dimension}"), + format!("num_queries={num_queries}"), + format!("element_width={width}"), + ] { + assert!(error.contains(&detail), "{error} lacks {detail}"); + } + } + assert_eq!( + batch_query_value_count(MAX_BATCH_QUERY_BYTES / 8, 1, 8).unwrap(), + MAX_BATCH_QUERY_BYTES / 8 + ); + let error = batch_query_value_count(MAX_BATCH_QUERY_BYTES / 8 + 1, 1, 8) + .unwrap_err() + .to_string(); + assert!(error.contains("exceeds limit"), "{error}"); + assert!(!error.contains("overflows"), "{error}"); + } + + #[test] + fn query_decoder_rejects_unaligned_typed_buffers() { + let storage = [0u64; 4]; + let unaligned = unsafe { storage.as_ptr().cast::().add(1) }.cast(); + for element_type in [0, 1, 2] { + let error = unsafe { decode_query_values(unaligned, 2, element_type) } + .unwrap_err() + .to_string(); + assert!(error.contains("query_data must be aligned"), "{error}"); + } + } + + #[test] + fn batch_matrix_does_not_enter_multivector_scoring() { + use arrow_array::{Array, FixedSizeListArray, Float32Array}; + use futures::TryStreamExt; + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().join("batch.lance"); + let element = Arc::new(Field::new("item", DataType::Float32, false)); + let vectors = FixedSizeListArray::try_new( + element.clone(), + 2, + Arc::new(Float32Array::from(vec![0., 0., 1., 0., 4., 0., 5., 0.])), + None, + ) + .unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("embedding", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![0, 1, 2, 3])), + Arc::new(vectors), + ], + ) + .unwrap(); + block_on(Dataset::write( + arrow::record_batch::RecordBatchIterator::new(vec![Ok(batch)], schema), + uri.to_str().unwrap(), + None, + )) + .unwrap(); + let (dataset, scanner) = open_dataset_and_scanner(uri.to_str().unwrap()); + unsafe { + (*scanner).nearest = Some(NearestQuery { + mode: NearestMode::Batch { num_queries: 2 }, + column: "embedding".to_string(), + query: Arc::new( + FixedSizeListArray::try_new( + element, + 2, + Arc::new(Float32Array::from(vec![0., 0., 4., 0.])), + None, + ) + .unwrap(), + ), + k: 2, + lower_bound: None, + upper_bound: None, + }); + (*scanner).use_index = Some(false); + let result = block_on(async { + (*scanner) + .build_scanner()? + .try_into_stream() + .await? + .try_collect::>() + .await + }); + lance_scanner_close(scanner); + lance_dataset_close(dataset); + let batches = + result.expect("independent batch queries must not use multivector scoring"); + let mut rows = Vec::new(); + for batch in batches { + let queries = batch + .column_by_name("query_index") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let distances = batch + .column_by_name("_distance") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + rows.push((queries.value(row), ids.value(row), distances.value(row))); + } + } + rows.sort_by_key(|row| (row.0, row.1)); + assert_eq!(rows, vec![(0, 0, 0.), (0, 1, 1.), (1, 2, 0.), (1, 3, 1.)]); + } + } + /// Write a 3-row dataset to a tempdir, returning (tempdir, uri). fn create_test_dataset() -> (tempfile::TempDir, String) { let tmp = tempfile::tempdir().unwrap(); diff --git a/tests/batch_nearest_test.rs b/tests/batch_nearest_test.rs new file mode 100644 index 0000000..1db0122 --- /dev/null +++ b/tests/batch_nearest_test.rs @@ -0,0 +1,712 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::ffi::{CStr, CString}; +use std::ptr; +use std::sync::Arc; + +use arrow::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream}; +use arrow::record_batch::RecordBatchIterator; +use arrow_array::{ + Array, FixedSizeListArray, Float32Array, Int32Array, RecordBatch, RecordBatchReader, +}; +use arrow_schema::{DataType, Field, Schema}; +use lance::Dataset; +use lance::dataset::{WriteMode, WriteParams}; +use lance_c::*; + +fn c(s: &str) -> CString { + CString::new(s).unwrap() +} +fn ok(code: i32) { + if code != 0 { + unsafe { + let msg = lance_last_error_message(); + let message = CStr::from_ptr(msg).to_string_lossy().into_owned(); + lance_free_string(msg); + panic!("{message}"); + } + } +} + +fn fixture() -> (tempfile::TempDir, CString) { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("batch.lance"); + lance_c::runtime::block_on(async { + for fragment in 0..2 { + let ids: Vec = (fragment * 64..(fragment + 1) * 64).collect(); + let vectors = FixedSizeListArray::try_new( + Arc::new(Field::new("item", DataType::Float32, false)), + 2, + Arc::new(Float32Array::from( + ids.iter().flat_map(|i| [*i as f32, 1.]).collect::>(), + )), + None, + ) + .unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("embedding", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(ids)), Arc::new(vectors)], + ) + .unwrap(); + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + path.to_str().unwrap(), + Some(WriteParams { + mode: if fragment == 0 { + WriteMode::Create + } else { + WriteMode::Append + }, + ..Default::default() + }), + ) + .await + .unwrap(); + } + }); + (dir, c(path.to_str().unwrap())) +} + +struct Scan(*mut LanceScanner); +impl Drop for Scan { + fn drop(&mut self) { + unsafe { lance_scanner_close(self.0) } + } +} +struct Data(*mut LanceDataset); +impl Drop for Data { + fn drop(&mut self) { + unsafe { lance_dataset_close(self.0) } + } +} +impl Data { + fn open(uri: &CString) -> Self { + let ds = unsafe { lance_dataset_open(uri.as_ptr(), ptr::null(), 0) }; + assert!(!ds.is_null()); + Self(ds) + } + fn scan(&self) -> Scan { + let scan = unsafe { lance_scanner_new(self.0, ptr::null(), ptr::null()) }; + assert!(!scan.is_null()); + Scan(scan) + } +} + +fn configure(scan: &Scan, q: &[f32], k: u32) { + ok(unsafe { + lance_scanner_nearest_batch( + scan.0, + c("embedding").as_ptr(), + q.as_ptr().cast(), + 2, + q.len() / 2, + 0, + k, + ) + }); +} +fn read(scan: &Scan) -> Vec<(i32, i32, f32)> { + let mut stream = FFI_ArrowArrayStream::empty(); + ok(unsafe { lance_scanner_to_arrow_stream(scan.0, &mut stream) }); + let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream) }.unwrap(); + assert_eq!( + reader + .schema() + .field_with_name("query_index") + .unwrap() + .data_type(), + &DataType::Int32 + ); + let mut result = Vec::new(); + for batch in reader { + let batch = batch.unwrap(); + let qs = batch + .column_by_name("query_index") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let ds = batch + .column_by_name("_distance") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + result.push((qs.value(i), ids.value(i), ds.value(i))); + } + } + result.sort_by(|a, b| a.0.cmp(&b.0).then(a.2.total_cmp(&b.2)).then(a.1.cmp(&b.1))); + result +} +fn expected(q: &[f32], k: usize, minimum_id: i32) -> Vec<(i32, i32, f32)> { + let mut result = Vec::new(); + for (index, q) in q.as_chunks::<2>().0.iter().enumerate() { + let mut rows = (minimum_id..128) + .map(|id| { + ( + index as i32, + id, + (id as f32 - q[0]).powi(2) + (1. - q[1]).powi(2), + ) + }) + .collect::>(); + rows.sort_by(|a, b| a.2.total_cmp(&b.2).then(a.1.cmp(&b.1))); + result.extend(rows.into_iter().take(k)); + } + result +} + +#[test] +fn batch_nearest_flat_preserves_independent_queries_and_copies_input() { + let (_dir, uri) = fixture(); + let data = Data::open(&uri); + for count in [1, 3] { + let scan = data.scan(); + let mut q = vec![2.25, 1., 100.25, 1., 2.25, 1.]; + q.truncate(count * 2); + let want = expected(&q, 3, 0); + configure(&scan, &q, 3); + q.fill(f32::NAN); + ok(unsafe { lance_scanner_set_use_index(scan.0, false) }); + ok(unsafe { lance_scanner_set_batch_size(scan.0, 2) }); + ok(unsafe { lance_scanner_set_strict_batch_size(scan.0, true) }); + assert_eq!(read(&scan), want); + } +} + +#[test] +fn batch_nearest_prefilter_empty_results_and_projection() { + let (_dir, uri) = fixture(); + let data = Data::open(&uri); + let q = [2.25, 1., 100.25, 1.]; + for minimum in [80, 128] { + let filter = c(&format!("id >= {minimum}")); + let column = c("id"); + let projection = [column.as_ptr(), ptr::null()]; + let scan = Scan(unsafe { lance_scanner_new(data.0, projection.as_ptr(), filter.as_ptr()) }); + configure(&scan, &q, 3); + ok(unsafe { lance_scanner_set_use_index(scan.0, false) }); + ok(unsafe { lance_scanner_set_prefilter(scan.0, true) }); + assert_eq!(read(&scan), expected(&q, 3, minimum)); + } +} + +#[test] +fn batch_nearest_indexed_shared_and_refined_fallback() { + let (_dir, uri) = fixture(); + let data = Data::open(&uri); + let params = LanceVectorIndexParams { + index_type: LanceVectorIndexType::IvfFlat, + metric: LanceMetricType::L2, + num_partitions: 2, + num_sub_vectors: 0, + num_bits: 0, + max_iterations: 0, + hnsw_m: 0, + hnsw_ef_construction: 0, + sample_rate: 0, + }; + ok(unsafe { + lance_dataset_create_vector_index( + data.0, + c("embedding").as_ptr(), + ptr::null(), + ¶ms, + true, + ) + }); + let q = [2.25, 1., 100.25, 1., 2.25, 1.]; + for refine in [None, Some(2)] { + let scan = data.scan(); + configure(&scan, &q, 3); + ok(unsafe { lance_scanner_set_nprobes(scan.0, 2) }); + if let Some(r) = refine { + ok(unsafe { lance_scanner_set_refine_factor(scan.0, r) }); + } + assert_eq!(read(&scan), expected(&q, 3, 0)); + } +} + +#[test] +fn batch_nearest_respects_selected_index_segments() { + use lance::index::{DatasetIndexExt, vector::VectorIndexParams}; + use lance_index::IndexType; + use lance_linalg::distance::MetricType; + let (_dir, uri) = fixture(); + let segments = lance_c::runtime::block_on(async { + let mut ds = Dataset::open(uri.to_str().unwrap()).await.unwrap(); + let params = VectorIndexParams::ivf_flat(1, MetricType::L2); + let mut segments = Vec::new(); + for fragment in 0..2 { + segments.push( + ds.create_index_builder(&["embedding"], IndexType::Vector, ¶ms) + .name("embedding_idx".into()) + .fragments(vec![fragment]) + .execute_uncommitted() + .await + .unwrap(), + ); + } + let uuids = segments + .iter() + .map(|s| *s.uuid.as_bytes()) + .collect::>(); + ds.commit_existing_index_segments("embedding_idx", "embedding", segments) + .await + .unwrap(); + uuids + }); + let data = Data::open(&uri); + let queries = [2.25, 1., 100.25, 1.]; + for (selected, domain) in [(vec![0], 0..64), (vec![1], 64..128), (vec![0, 1], 0..128)] { + let scan = data.scan(); + configure(&scan, &queries, 3); + let uuids = selected + .iter() + .flat_map(|i| segments[*i]) + .collect::>(); + ok(unsafe { lance_scanner_set_index_segments(scan.0, uuids.as_ptr(), selected.len()) }); + ok(unsafe { lance_scanner_set_nprobes(scan.0, 1) }); + let mut want = Vec::new(); + for (index, q) in queries.as_chunks::<2>().0.iter().enumerate() { + let mut rows = domain + .clone() + .map(|id| { + ( + index as i32, + id, + (id as f32 - q[0]).powi(2) + (1. - q[1]).powi(2), + ) + }) + .collect::>(); + rows.sort_by(|a, b| a.2.total_cmp(&b.2).then(a.1.cmp(&b.1))); + want.extend(rows.into_iter().take(3)); + } + assert_eq!(read(&scan), want, "selected segments: {selected:?}"); + } +} + +#[test] +fn batch_nearest_size_error_identifies_the_request() { + let (_dir, uri) = fixture(); + let data = Data::open(&uri); + let scan = data.scan(); + let q = [2.25f32, 1., 100.25, 1.]; + configure(&scan, &q, 2); + // Rejection must happen before reading the oversized input buffer. + assert_eq!( + unsafe { + lance_scanner_nearest_batch( + scan.0, + c("embedding").as_ptr(), + q.as_ptr().cast(), + 10_000_000, + 2, + 0, + 2, + ) + }, + -1 + ); + let message = unsafe { + let ptr = lance_last_error_message(); + let message = CStr::from_ptr(ptr).to_string_lossy().into_owned(); + lance_free_string(ptr); + message + }; + for detail in [ + "dimension=10000000", + "num_queries=2", + "element_width=4", + "80000000", + "67108864", + ] { + assert!(message.contains(detail), "{message} lacks {detail}"); + } + ok(unsafe { lance_scanner_set_use_index(scan.0, false) }); + assert_eq!(read(&scan), expected(&q, 2, 0)); +} + +#[test] +fn batch_nearest_int8_matches_independent_l2_queries() { + let values = vec![-4i8, -2, -1, 1, 3, 2, 7, 4]; + let (_dir, uri) = typed_fixture(Arc::new(arrow_array::Int8Array::from(values)), false); + let data = Data::open(&uri); + let scan = data.scan(); + let queries = [-3i8, -2, 6, 3]; + ok(unsafe { + lance_scanner_nearest_batch( + scan.0, + c("embedding").as_ptr(), + queries.as_ptr().cast(), + 2, + 2, + 4, + 2, + ) + }); + ok(unsafe { lance_scanner_set_use_index(scan.0, false) }); + assert_eq!( + read(&scan), + vec![(0, 0, 1.), (0, 1, 13.), (1, 3, 2.), (1, 2, 10.)] + ); +} + +#[test] +fn batch_nearest_rejects_invalid_input_atomically() { + let (_dir, uri) = fixture(); + let data = Data::open(&uri); + let scan = data.scan(); + let q = [2.25f32, 1., 100.25, 1.]; + configure(&scan, &q, 2); + for (dim, count, dtype, k) in [ + (0, 2, 0, 2), + (2, 0, 0, 2), + (3, 1, 0, 2), + (2, 2, 1, 2), + (2, 2, 4, 2), + (2, 2, 99, 2), + (2, 2, 0, 0), + (usize::MAX, 2, 0, 2), + (2, usize::MAX, 0, 2), + (2, 129, 0, 2), + (i32::MAX as usize, 1, 0, 2), + (2, 2, 0, 100001), + ] { + assert_ne!( + unsafe { + lance_scanner_nearest_batch( + scan.0, + c("embedding").as_ptr(), + q.as_ptr().cast(), + dim, + count, + dtype, + k, + ) + }, + 0 + ); + } + for bad in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + assert_ne!( + unsafe { + lance_scanner_nearest_batch( + scan.0, + c("embedding").as_ptr(), + [bad, 1.].as_ptr().cast(), + 2, + 1, + 0, + 2, + ) + }, + 0 + ); + } + assert_ne!( + unsafe { + lance_scanner_nearest_batch( + scan.0, + c("missing").as_ptr(), + q.as_ptr().cast(), + 2, + 2, + 0, + 2, + ) + }, + 0 + ); + assert_ne!( + unsafe { + lance_scanner_nearest_batch(scan.0, c("id").as_ptr(), q.as_ptr().cast(), 2, 2, 0, 2) + }, + 0 + ); + ok(unsafe { lance_scanner_set_use_index(scan.0, false) }); + assert_eq!(read(&scan), expected(&q, 2, 0)); + assert_ne!( + unsafe { + lance_scanner_nearest_batch( + scan.0, + c("embedding").as_ptr(), + q.as_ptr().cast(), + 2, + 2, + 0, + 2, + ) + }, + 0 + ); +} + +#[test] +fn batch_nearest_rejects_global_windows_in_both_configuration_orders() { + let (_dir, uri) = fixture(); + let data = Data::open(&uri); + let q = [2.25f32, 1., 100.25, 1.]; + for offset in [false, true] { + let scan = data.scan(); + configure(&scan, &q, 2); + assert_ne!( + unsafe { + if offset { + lance_scanner_set_offset(scan.0, 1) + } else { + lance_scanner_set_limit(scan.0, 1) + } + }, + 0 + ); + ok(unsafe { lance_scanner_set_use_index(scan.0, false) }); + assert_eq!(read(&scan), expected(&q, 2, 0)); + let scan = data.scan(); + ok(unsafe { + if offset { + lance_scanner_set_offset(scan.0, 1) + } else { + lance_scanner_set_limit(scan.0, 1) + } + }); + assert_ne!( + unsafe { + lance_scanner_nearest_batch( + scan.0, + c("embedding").as_ptr(), + q.as_ptr().cast(), + 2, + 2, + 0, + 2, + ) + }, + 0 + ); + } +} + +#[test] +fn batch_nearest_refinement_budget_and_query_replacement() { + let (_dir, uri) = fixture(); + let data = Data::open(&uri); + let q = [2.25f32, 1., 100.25, 1.]; + let scan = data.scan(); + configure(&scan, &q, 2); + assert_ne!(unsafe { lance_scanner_set_refine_factor(scan.0, 0) }, 0); + assert_ne!(unsafe { lance_scanner_set_refine_factor(scan.0, 25001) }, 0); + ok(unsafe { + lance_scanner_nearest(scan.0, c("embedding").as_ptr(), q.as_ptr().cast(), 2, 0, 2) + }); + ok(unsafe { lance_scanner_set_limit(scan.0, 1) }); + let scan = data.scan(); + ok(unsafe { lance_scanner_set_refine_factor(scan.0, 25001) }); + assert_ne!( + unsafe { + lance_scanner_nearest_batch( + scan.0, + c("embedding").as_ptr(), + q.as_ptr().cast(), + 2, + 2, + 0, + 2, + ) + }, + 0 + ); +} + +fn typed_fixture(values: Arc, reserved_name: bool) -> (tempfile::TempDir, CString) { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("typed.lance"); + let vectors = FixedSizeListArray::try_new( + Arc::new(Field::new("item", values.data_type().clone(), false)), + 2, + values, + None, + ) + .unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new( + if reserved_name { "query_index" } else { "id" }, + DataType::Int32, + false, + ), + Field::new("embedding", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![0, 1, 2, 3])), + Arc::new(vectors), + ], + ) + .unwrap(); + lance_c::runtime::block_on(Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + path.to_str().unwrap(), + None, + )) + .unwrap(); + (dir, c(path.to_str().unwrap())) +} + +#[test] +fn batch_nearest_supports_float16_float64_and_binary_queries() { + use arrow_array::{Float16Array, Float64Array, UInt8Array}; + let f16s: Vec = [0., 0., 1., 1., 2., 2., 3., 3.] + .into_iter() + .map(half::f16::from_f32) + .collect(); + let f64s = [0f64, 0., 1., 1., 2., 2., 3., 3.]; + let bytes = [0u8, 0, 1, 1, 15, 15, 255, 255]; + for (values, query, dtype, metric) in [ + ( + Arc::new(Float16Array::from(f16s.clone())) as Arc, + f16s.as_ptr().cast(), + 1, + 0, + ), + ( + Arc::new(Float64Array::from(f64s.to_vec())) as Arc, + f64s.as_ptr().cast(), + 2, + 0, + ), + ( + Arc::new(UInt8Array::from(bytes.to_vec())) as Arc, + bytes.as_ptr().cast(), + 3, + 3, + ), + ] { + let (_dir, uri) = typed_fixture(values, false); + let data = Data::open(&uri); + let scan = data.scan(); + ok(unsafe { + lance_scanner_nearest_batch(scan.0, c("embedding").as_ptr(), query, 2, 4, dtype, 1) + }); + ok(unsafe { lance_scanner_set_metric(scan.0, metric) }); + ok(unsafe { lance_scanner_set_use_index(scan.0, false) }); + assert_eq!( + read(&scan), + vec![(0, 0, 0.), (1, 1, 0.), (2, 2, 0.), (3, 3, 0.)] + ); + } +} + +#[test] +fn batch_nearest_rejects_reserved_query_index_column() { + let (_dir, uri) = typed_fixture(Arc::new(Float32Array::from(vec![0.; 8])), true); + let data = Data::open(&uri); + let scan = data.scan(); + assert_ne!( + unsafe { + lance_scanner_nearest_batch( + scan.0, + c("embedding").as_ptr(), + [0f32, 0.].as_ptr().cast(), + 2, + 1, + 0, + 1, + ) + }, + 0 + ); +} + +#[test] +fn batch_nearest_exported_stream_owns_query_and_survives_handle_close() { + let (_dir, uri) = fixture(); + let mut stream = FFI_ArrowArrayStream::empty(); + { + let data = Data::open(&uri); + let scan = data.scan(); + configure(&scan, &[2.25, 1., 100.25, 1.], 16); + ok(unsafe { lance_scanner_set_use_index(scan.0, false) }); + ok(unsafe { lance_scanner_set_batch_size(scan.0, 1) }); + ok(unsafe { lance_scanner_to_arrow_stream(scan.0, &mut stream) }); + } + let mut reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream) }.unwrap(); + assert!(reader.next().unwrap().unwrap().num_rows() > 0); + // Releasing an unfinished exported stream exercises the existing cancellation path. + drop(reader); +} + +#[test] +fn batch_nearest_and_full_text_search_are_mutually_exclusive() { + let (_dir, uri) = fixture(); + let data = Data::open(&uri); + let q = [2.25f32, 1.]; + let scan = data.scan(); + ok(unsafe { lance_scanner_full_text_search(scan.0, c("term").as_ptr(), ptr::null(), 0) }); + assert_ne!( + unsafe { + lance_scanner_nearest_batch( + scan.0, + c("embedding").as_ptr(), + q.as_ptr().cast(), + 2, + 1, + 0, + 1, + ) + }, + 0 + ); + let scan = data.scan(); + configure(&scan, &q, 1); + assert_ne!( + unsafe { lance_scanner_full_text_search(scan.0, c("term").as_ptr(), ptr::null(), 0) }, + 0 + ); +} + +#[test] +fn batch_nearest_respects_fragment_scope() { + let (_dir, uri) = fixture(); + let data = Data::open(&uri); + let scan = data.scan(); + let q = [2.25, 1., 100.25, 1.]; + configure(&scan, &q, 3); + ok(unsafe { lance_scanner_set_fragment_ids(scan.0, [1u64].as_ptr(), 1) }); + ok(unsafe { lance_scanner_set_prefilter(scan.0, true) }); + ok(unsafe { lance_scanner_set_use_index(scan.0, false) }); + assert_eq!(read(&scan), expected(&q, 3, 64)); +} + +#[test] +fn batch_replacement_clears_single_query_range_and_rejects_new_bounds() { + let (_dir, uri) = fixture(); + let data = Data::open(&uri); + let scan = data.scan(); + let q = [2.25f32, 1., 100.25, 1.]; + ok(unsafe { + lance_scanner_nearest(scan.0, c("embedding").as_ptr(), q.as_ptr().cast(), 2, 0, 3) + }); + // This range excludes every L2 result; it must not leak into the replacement batch. + ok(unsafe { lance_scanner_set_distance_range(scan.0, ptr::null(), &0.0) }); + configure(&scan, &q, 3); + assert_eq!( + unsafe { lance_scanner_set_distance_range(scan.0, ptr::null(), &0.0) }, + -1 + ); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + ok(unsafe { lance_scanner_set_use_index(scan.0, false) }); + assert_eq!(read(&scan), expected(&q, 3, 0)); +} diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 2a00f5e..a51dea6 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -17276,3 +17276,28 @@ fn test_blob_close_keeps_the_pending_error_readable() { close_blob_handles(&rest); unsafe { lance_dataset_close(ds) }; } + +#[test] +fn test_scanner_nearest_batch_null_arguments() { + let (_dir, uri) = create_vector_dataset(8, 2); + unsafe { + let ds = lance_dataset_open(c_str(&uri).as_ptr(), ptr::null(), 0); + let scan = lance_scanner_new(ds, ptr::null(), ptr::null()); + let query = [0f32, 1.]; + for (scanner, column, data) in [ + (ptr::null_mut(), c_str("embedding"), query.as_ptr().cast()), + (scan, c_str("embedding"), ptr::null()), + ] { + assert_ne!( + lance_scanner_nearest_batch(scanner, column.as_ptr(), data, 2, 1, 0, 1), + 0 + ); + } + assert_ne!( + lance_scanner_nearest_batch(scan, ptr::null(), query.as_ptr().cast(), 2, 1, 0, 1), + 0 + ); + lance_scanner_close(scan); + lance_dataset_close(ds); + } +} diff --git a/tests/cpp/test_c_api.c b/tests/cpp/test_c_api.c index ce59a40..efad5ed 100644 --- a/tests/cpp/test_c_api.c +++ b/tests/cpp/test_c_api.c @@ -1391,6 +1391,72 @@ static void test_distance_range(const char *uri) { printf("OK\n"); } +static void check_batch_nearest_results(struct ArrowArrayStream *stream) { + struct ArrowSchema schema = {0}; + ASSERT(stream->get_schema(stream, &schema) == 0, "batch nearest result"); + int query_column = -1, id_column = -1, distance_column = -1; + for (int64_t i = 0; i < schema.n_children; ++i) { + if (strcmp(schema.children[i]->name, "query_index") == 0) { + query_column = (int)i; + } + if (strcmp(schema.children[i]->name, "id") == 0) id_column = (int)i; + if (strcmp(schema.children[i]->name, "_distance") == 0) distance_column = (int)i; + } + ASSERT(query_column >= 0, "batch nearest result"); + ASSERT(strcmp(schema.children[query_column]->format, "i") == 0, "batch nearest result"); + ASSERT(id_column >= 0 && distance_column >= 0, "missing result columns"); + schema.release(&schema); + int counts[2] = {0, 0}; + int seen[2] = {0, 0}; + while (1) { + struct ArrowArray batch = {0}; + ASSERT(stream->get_next(stream, &batch) == 0, "batch nearest result"); + if (!batch.release) break; + struct ArrowArray *queries = batch.children[query_column]; + const int32_t *values = (const int32_t *)queries->buffers[1]; + struct ArrowArray *ids = batch.children[id_column]; + struct ArrowArray *distances = batch.children[distance_column]; + const int32_t *id_values = (const int32_t *)ids->buffers[1]; + const float *distance_values = (const float *)distances->buffers[1]; + for (int64_t i = 0; i < batch.length; ++i) { + int32_t index = values[queries->offset + i]; + ASSERT(index >= 0 && index < 2, "batch nearest result"); + int32_t id = id_values[ids->offset + i]; + int32_t first_id = index == 0 ? 1 : 19; + ASSERT(id >= first_id && id <= first_id + 1, "wrong per-query neighbor"); + seen[index] |= 1 << (id - first_id); + float expected = (id == 1 || id == 20) ? 0.0f : 0.08f; + float error = distance_values[distances->offset + i] - expected; + ASSERT(error > -0.0001f && error < 0.0001f, "wrong per-query distance"); + ++counts[index]; + } + batch.release(&batch); + } + ASSERT(counts[0] == 2 && counts[1] == 2, "batch nearest result"); + ASSERT(seen[0] == 3 && seen[1] == 3, "missing per-query neighbor"); + stream->release(stream); +} + +static void test_batch_nearest(const char *uri) { + float queries[16]; + for (int i = 0; i < 8; ++i) { + queries[i] = 0.1f + i; + queries[8 + i] = 2.0f + i; + } + LanceDataset *dataset = lance_dataset_open(uri, NULL, 0); + ASSERT(dataset != NULL, "open batch dataset"); + LanceScanner *scanner = lance_scanner_new(dataset, NULL, NULL); + ASSERT(scanner != NULL, "open batch scanner"); + ASSERT(lance_scanner_nearest_batch(scanner, "embedding", queries, 8, 2, LANCE_DTYPE_FLOAT32, 2) == 0, "batch nearest"); + ASSERT(lance_scanner_set_use_index(scanner, false) == 0, "flat batch search"); + ASSERT(lance_scanner_set_batch_size(scanner, 1) == 0, "small output batches"); + struct ArrowArrayStream stream = {0}; + ASSERT(lance_scanner_to_arrow_stream(scanner, &stream) == 0, "export batch stream"); + lance_scanner_close(scanner); + lance_dataset_close(dataset); + check_batch_nearest_results(&stream); +} + int main(int argc, char **argv) { if (argc < 4) { fprintf(stderr, "Usage: %s \n", argv[0]); @@ -1402,6 +1468,7 @@ int main(int argc, char **argv) { const char *blob_uri = argv[3]; printf("Running C API tests with dataset: %s\n", uri); + test_batch_nearest(uri); test_open_and_metadata(uri); test_shared_session(uri); test_scan(uri); diff --git a/tests/cpp/test_cpp_api.cpp b/tests/cpp/test_cpp_api.cpp index e82582d..e2d541e 100644 --- a/tests/cpp/test_cpp_api.cpp +++ b/tests/cpp/test_cpp_api.cpp @@ -1271,6 +1271,75 @@ static void test_distance_range(const std::string& uri) { PASS(); } +static void check_batch_nearest_results(struct ArrowArrayStream *stream) { + ArrowSchema schema{}; + assert(stream->get_schema(stream, &schema) == 0); + int query_column = -1, id_column = -1, distance_column = -1; + for (int64_t i = 0; i < schema.n_children; ++i) { + if (strcmp(schema.children[i]->name, "query_index") == 0) { + query_column = (int)i; + } + if (strcmp(schema.children[i]->name, "id") == 0) id_column = (int)i; + if (strcmp(schema.children[i]->name, "_distance") == 0) distance_column = (int)i; + } + assert(query_column >= 0); + assert(strcmp(schema.children[query_column]->format, "i") == 0); + assert(id_column >= 0 && distance_column >= 0); + schema.release(&schema); + int counts[2] = {0, 0}; + int seen[2] = {0, 0}; + while (1) { + ArrowArray batch{}; + assert(stream->get_next(stream, &batch) == 0); + if (!batch.release) break; + struct ArrowArray *queries = batch.children[query_column]; + const int32_t *values = (const int32_t *)queries->buffers[1]; + struct ArrowArray *ids = batch.children[id_column]; + struct ArrowArray *distances = batch.children[distance_column]; + const int32_t *id_values = (const int32_t *)ids->buffers[1]; + const float *distance_values = (const float *)distances->buffers[1]; + for (int64_t i = 0; i < batch.length; ++i) { + int32_t index = values[queries->offset + i]; + assert(index >= 0 && index < 2); + int32_t id = id_values[ids->offset + i]; + int32_t first_id = index == 0 ? 1 : 19; + assert(id >= first_id && id <= first_id + 1); + seen[index] |= 1 << (id - first_id); + float expected = (id == 1 || id == 20) ? 0.0f : 0.08f; + float error = distance_values[distances->offset + i] - expected; + assert(error > -0.0001f && error < 0.0001f); + ++counts[index]; + } + batch.release(&batch); + } + assert(counts[0] == 2 && counts[1] == 2); + assert(seen[0] == 3 && seen[1] == 3); + stream->release(stream); +} + +static void test_batch_nearest(const std::string& uri) { + TEST(test_batch_nearest); + float queries[16]; + for (int i = 0; i < 8; ++i) { + queries[i] = 0.1f + i; + queries[8 + i] = 2.0f + i; + } + auto dataset = lance::Dataset::open(uri); + for (bool typed : {false, true}) { + auto scanner = dataset.scan(); + if (typed) { + scanner.nearest_batch("embedding", queries, 8, 2, LANCE_DTYPE_FLOAT32, 2); + } else { + scanner.nearest_batch("embedding", queries, 8, 2, 2); + } + scanner.use_index(false).batch_size(1); + ArrowArrayStream stream = {}; + scanner.to_arrow_stream(&stream); + check_batch_nearest_results(&stream); + } + PASS(); +} + int main(int argc, char** argv) { if (argc < 4) { fprintf(stderr, "Usage: %s \n", argv[0]); @@ -1297,6 +1366,7 @@ int main(int argc, char** argv) { test_restore_to_current(uri); test_error_exception(uri); test_index_lifecycle(uri); + test_batch_nearest(uri); test_nearest_smoke(uri); test_multivector_rejects_flat_column(uri); test_index_segments_smoke(uri); diff --git a/tests/multivector_test.rs b/tests/multivector_test.rs index 4223502..8905b86 100644 --- a/tests/multivector_test.rs +++ b/tests/multivector_test.rs @@ -992,3 +992,36 @@ fn multivector_distance_range_is_rejected_instead_of_filtering_subvectors() { lance_dataset_close(ds); } } + +#[test] +fn batch_nearest_rejects_multivector_column_without_replacing_query() { + let (_dir, uri) = fixture(); + unsafe { + let dataset = lance_dataset_open(uri.as_ptr(), ptr::null(), 0); + assert!(!dataset.is_null()); + let scan = lance_scanner_new(dataset, ptr::null(), ptr::null()); + assert!(!scan.is_null()); + let column = CString::new("vectors").unwrap(); + let query = [1f32, 0.]; + assert_eq!( + lance_scanner_nearest_multivector( + scan, + column.as_ptr(), + query.as_ptr().cast(), + 2, + 1, + 0, + 1 + ), + 0 + ); + assert_ne!( + lance_scanner_nearest_batch(scan, column.as_ptr(), query.as_ptr().cast(), 2, 1, 0, 1), + 0 + ); + assert_eq!(lance_scanner_set_use_index(scan, false), 0); + assert_eq!(collect(scan), vec![(1, 0.)]); + lance_scanner_close(scan); + lance_dataset_close(dataset); + } +}