diff --git a/crates/integrations/datafusion/src/filter_pushdown.rs b/crates/integrations/datafusion/src/filter_pushdown.rs index cdbe6460d..6155062ec 100644 --- a/crates/integrations/datafusion/src/filter_pushdown.rs +++ b/crates/integrations/datafusion/src/filter_pushdown.rs @@ -89,6 +89,20 @@ pub(crate) fn analyze_filters( } } +pub(crate) fn is_safe_vector_prefilter(predicate: &Predicate) -> bool { + match predicate { + Predicate::Leaf { literals, .. } => !literals.iter().any(|literal| { + matches!(literal, Datum::Float(value) if value.is_nan()) + || matches!(literal, Datum::Double(value) if value.is_nan()) + }), + Predicate::And(children) | Predicate::Or(children) => { + children.iter().all(is_safe_vector_prefilter) + } + Predicate::Not(inner) => is_safe_vector_prefilter(inner), + Predicate::AlwaysTrue | Predicate::AlwaysFalse => true, + } +} + #[cfg(test)] pub(crate) fn build_pushed_predicate(filters: &[Expr], fields: &[DataField]) -> Option { analyze_filters(filters, fields, true).pushed_predicate @@ -912,6 +926,27 @@ mod tests { ); } + #[test] + fn test_vector_prefilter_rejects_nan_literals() { + let predicate = Predicate::Leaf { + column: "score".to_string(), + index: 0, + data_type: DataType::Float(FloatType::new()), + op: PredicateOperator::Eq, + literals: vec![Datum::Float(f32::NAN)], + }; + assert!(!is_safe_vector_prefilter(&predicate)); + + let finite = Predicate::Leaf { + column: "score".to_string(), + index: 0, + data_type: DataType::Float(FloatType::new()), + op: PredicateOperator::Eq, + literals: vec![Datum::Float(f32::INFINITY)], + }; + assert!(is_safe_vector_prefilter(&finite)); + } + #[test] fn test_negated_inexact_float_array_membership_falls_open() { use datafusion::functions_nested::expr_fn::array_has; diff --git a/crates/integrations/datafusion/src/lateral_vector_search.rs b/crates/integrations/datafusion/src/lateral_vector_search.rs index 30342b14e..674ef2afa 100644 --- a/crates/integrations/datafusion/src/lateral_vector_search.rs +++ b/crates/integrations/datafusion/src/lateral_vector_search.rs @@ -17,10 +17,14 @@ use std::any::Any; use std::cmp::Ordering; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::fmt; use std::hash::{Hash, Hasher}; -use std::sync::Arc; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering}; +use std::sync::{Arc, Mutex, Weak}; +use std::task::{Context, Poll}; +use std::time::Duration; use async_trait::async_trait; use datafusion::arrow::array::{ @@ -37,23 +41,28 @@ use datafusion::common::{ use datafusion::datasource::TableProvider; use datafusion::execution::context::{QueryPlanner, SessionState}; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; -use datafusion::logical_expr::{Expr, Extension, LogicalPlan, TableScan, UserDefinedLogicalNode}; +use datafusion::logical_expr::utils::{conjunction, split_conjunction}; +use datafusion::logical_expr::{ + Expr, Extension, Filter, LogicalPlan, Projection, TableScan, UserDefinedLogicalNode, +}; use datafusion::optimizer::{ApplyOrder, Optimizer, OptimizerConfig, OptimizerRule}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, - PlanProperties, + PlanProperties, RecordBatchStream, }; use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner}; use datafusion::prelude::SessionConfig; -use futures::{StreamExt, TryStreamExt}; -use paimon::spec::ROW_ID_FIELD_NAME; -use paimon::table::{RowRange, Table}; +use futures::{Stream, StreamExt, TryStreamExt}; +use paimon::spec::{Predicate, ROW_ID_FIELD_NAME}; +use paimon::table::{PreparedVectorSearchFilter, RowRange, Table}; use paimon::vector_search::SearchResult; +use tokio::sync::OnceCell; use crate::error::to_datafusion_error; +use crate::filter_pushdown::{analyze_filters, is_safe_vector_prefilter}; use crate::vector_search::LateralVectorSearchTableProvider; #[derive(Debug)] @@ -114,6 +123,84 @@ impl OptimizerRule for RewriteLateralVectorSearch { plan: LogicalPlan, _config: &dyn OptimizerConfig, ) -> DFResult> { + if let LogicalPlan::Filter(filter) = plan { + let (extension, projection) = match filter.input.as_ref() { + LogicalPlan::Extension(extension) => (extension, None), + LogicalPlan::Projection(projection) + if projection + .expr + .iter() + .all(|expr| matches!(expr, Expr::Column(_))) => + { + let LogicalPlan::Extension(extension) = projection.input.as_ref() else { + return Ok(Transformed::no(LogicalPlan::Filter(filter))); + }; + (extension, Some(projection)) + } + _ => return Ok(Transformed::no(LogicalPlan::Filter(filter))), + }; + let Some(node) = extension + .node + .as_any() + .downcast_ref::() + else { + return Ok(Transformed::no(LogicalPlan::Filter(filter))); + }; + let mut target_predicates = Vec::new(); + let mut residual_predicates = Vec::new(); + for conjunct in split_conjunction(&filter.predicate) { + if conjunct + .column_refs() + .iter() + .any(|column| node.input.schema().index_of_column(column).is_ok()) + { + residual_predicates.push(conjunct.clone()); + continue; + } + let analysis = analyze_filters( + std::slice::from_ref(conjunct), + node.target_table.schema().fields(), + true, + ); + match analysis.pushed_predicate { + Some(predicate) + if !analysis.requires_residual && is_safe_vector_prefilter(&predicate) => + { + target_predicates.push(predicate); + } + _ => { + residual_predicates.push(conjunct.clone()); + } + } + } + if target_predicates.is_empty() { + return Ok(Transformed::no(LogicalPlan::Filter(filter))); + } + let predicate = Predicate::and(target_predicates); + let predicate = match &node.filter { + Some(existing) => Predicate::and(vec![existing.clone(), predicate]), + None => predicate, + }; + let extension = LogicalPlan::Extension(Extension { + node: Arc::new(node.with_filter(predicate)), + }); + let rewritten = match projection { + Some(projection) => LogicalPlan::Projection(Projection::try_new_with_schema( + projection.expr.clone(), + Arc::new(extension), + Arc::clone(&projection.schema), + )?), + None => extension, + }; + let rewritten = match conjunction(residual_predicates) { + Some(predicate) => { + LogicalPlan::Filter(Filter::try_new(predicate, Arc::new(rewritten))?) + } + None => rewritten, + }; + return Ok(Transformed::yes(rewritten)); + } + let LogicalPlan::Join(join) = plan else { return Ok(Transformed::no(plan)); }; @@ -181,6 +268,7 @@ pub(crate) struct LateralVectorSearchNode { query_vector_expr: Expr, limit: usize, schema: DFSchemaRef, + filter: Option, } impl LateralVectorSearchNode { @@ -201,9 +289,16 @@ impl LateralVectorSearchNode { query_vector_expr, limit, schema, + filter: None, } } + fn with_filter(&self, filter: Predicate) -> Self { + let mut node = self.clone(); + node.filter = Some(filter); + node + } + fn target_table(&self) -> &Table { &self.target_table } @@ -253,8 +348,8 @@ impl UserDefinedLogicalNode for LateralVectorSearchNode { fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { write!( f, - "LateralVectorSearch: column={}, limit={}", - self.target_column, self.limit + "LateralVectorSearch: column={}, limit={}, filter={:?}", + self.target_column, self.limit, self.filter ) } @@ -274,6 +369,7 @@ impl UserDefinedLogicalNode for LateralVectorSearchNode { query_vector_expr: exprs.into_iter().next().unwrap(), limit: self.limit, schema: Arc::clone(&self.schema), + filter: self.filter.clone(), })) } @@ -284,6 +380,7 @@ impl UserDefinedLogicalNode for LateralVectorSearchNode { self.target_column.hash(&mut state); self.query_vector_expr.hash(&mut state); self.limit.hash(&mut state); + format!("{:?}", self.filter).hash(&mut state); } fn dyn_eq(&self, other: &dyn UserDefinedLogicalNode) -> bool { @@ -293,6 +390,7 @@ impl UserDefinedLogicalNode for LateralVectorSearchNode { && self.target_column == other.target_column && self.query_vector_expr == other.query_vector_expr && self.limit == other.limit + && self.filter == other.filter }) } @@ -331,7 +429,7 @@ impl ExtensionPlanner for LateralVectorSearchExtensionPlanner { logical_inputs[0].schema(), session_state, )?; - Ok(Some(Arc::new(LateralVectorSearchExec::new( + let mut exec = LateralVectorSearchExec::new( Arc::clone(&physical_inputs[0]), node.target_table().clone(), Arc::clone(node.target_schema()), @@ -339,7 +437,11 @@ impl ExtensionPlanner for LateralVectorSearchExtensionPlanner { query_vector_expr, node.limit(), Arc::new(node.schema().as_arrow().clone()), - )))) + ); + if let Some(filter) = &node.filter { + exec = exec.with_filter(filter.clone()); + } + Ok(Some(Arc::new(exec))) } } @@ -352,9 +454,251 @@ struct LateralVectorSearchExec { query_vector_expr: Arc, limit: usize, output_schema: ArrowSchemaRef, + filter: Option, + prepared_filter_cache: Arc, plan_properties: Arc, } +#[derive(Debug)] +struct ExecutionPreparedFilterEntry { + context: Weak, + prepared_filter: Arc>, + partition_count: usize, + completed_partitions: HashSet, + active_partition_leases: HashMap, +} + +#[derive(Debug, Default)] +struct ExecutionPreparedFilterCache { + // DataFusion passes the same TaskContext Arc to every partition of one + // execution. Keep the prepared filter alive for that TaskContext so + // sequential partitions resolve the same target snapshot. DataFusion 54 + // exposes no TaskContext drop hook, so one cache-level reaper observes the + // weak contexts and removes subset executions after their context dies. + entries: Mutex>, + reaper_running: AtomicBool, +} + +#[derive(Clone)] +struct ExecutionPreparedFilterLease { + prepared_filter: Arc>, + _completion: Arc, +} + +impl ExecutionPreparedFilterLease { + fn new( + cache: &Arc, + context: Arc, + prepared_filter: Arc>, + partition: usize, + ) -> Self { + Self { + prepared_filter: Arc::clone(&prepared_filter), + _completion: Arc::new(ExecutionPartitionCompletion { + cache: Arc::downgrade(cache), + context, + prepared_filter, + partition, + }), + } + } +} + +struct ExecutionPartitionCompletion { + cache: Weak, + context: Arc, + prepared_filter: Arc>, + partition: usize, +} + +impl Drop for ExecutionPartitionCompletion { + fn drop(&mut self) { + if let Some(cache) = self.cache.upgrade() { + cache.finish_partition(&self.context, &self.prepared_filter, self.partition); + } + } +} + +impl ExecutionPreparedFilterCache { + fn for_execution( + self: &Arc, + context: &Arc, + partition: usize, + partition_count: usize, + ) -> ExecutionPreparedFilterLease { + debug_assert!(partition < partition_count); + let mut entries = self + .entries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + entries.retain(|entry| entry.context.strong_count() > 0); + for entry in entries.iter_mut() { + let Some(entry_context) = entry.context.upgrade() else { + continue; + }; + if Arc::ptr_eq(&entry_context, context) && entry.partition_count == partition_count { + entry.completed_partitions.remove(&partition); + *entry.active_partition_leases.entry(partition).or_default() += 1; + let lease = ExecutionPreparedFilterLease::new( + self, + Arc::clone(context), + Arc::clone(&entry.prepared_filter), + partition, + ); + drop(entries); + self.ensure_context_reaper(); + return lease; + } + } + + let prepared_filter = Arc::new(OnceCell::new()); + let mut active_partition_leases = HashMap::new(); + active_partition_leases.insert(partition, 1); + entries.push(ExecutionPreparedFilterEntry { + context: Arc::downgrade(context), + prepared_filter: Arc::clone(&prepared_filter), + partition_count, + completed_partitions: HashSet::new(), + active_partition_leases, + }); + let lease = ExecutionPreparedFilterLease::new( + self, + Arc::clone(context), + prepared_filter, + partition, + ); + drop(entries); + self.ensure_context_reaper(); + lease + } + + fn ensure_context_reaper(self: &Arc) { + if self.reaper_running.swap(true, AtomicOrdering::AcqRel) { + return; + } + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + self.reaper_running.store(false, AtomicOrdering::Release); + return; + }; + let cache = Arc::downgrade(self); + runtime.spawn(async move { + loop { + tokio::time::sleep(Duration::from_millis(50)).await; + let Some(cache) = cache.upgrade() else { + return; + }; + if cache.prune_dead_contexts() { + cache.reaper_running.store(false, AtomicOrdering::Release); + let has_entries = !cache + .entries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_empty(); + if has_entries { + cache.ensure_context_reaper(); + } + return; + } + } + }); + } + + fn prune_dead_contexts(&self) -> bool { + let mut entries = self + .entries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + entries.retain(|entry| entry.context.strong_count() > 0); + entries.is_empty() + } + + fn finish_partition( + &self, + context: &Arc, + prepared_filter: &Arc>, + partition: usize, + ) { + let mut entries = self + .entries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let Some(entry_index) = entries + .iter() + .position(|entry| Arc::ptr_eq(&entry.prepared_filter, prepared_filter)) + else { + return; + }; + let entry = &mut entries[entry_index]; + let Some(active_leases) = entry.active_partition_leases.get_mut(&partition) else { + return; + }; + *active_leases -= 1; + if *active_leases == 0 { + entry.active_partition_leases.remove(&partition); + entry.completed_partitions.insert(partition); + } + if entry.active_partition_leases.is_empty() + && (Arc::strong_count(context) == 1 + || entry.completed_partitions.len() == entry.partition_count) + { + entries.remove(entry_index); + } + } +} + +struct ExecutionScopedStream { + schema: ArrowSchemaRef, + inner: Option, + lease: Option, +} + +impl ExecutionScopedStream { + fn new( + schema: ArrowSchemaRef, + inner: SendableRecordBatchStream, + lease: ExecutionPreparedFilterLease, + ) -> Self { + Self { + schema, + inner: Some(inner), + lease: Some(lease), + } + } + + fn finish(&mut self) { + drop(self.inner.take()); + drop(self.lease.take()); + } +} + +impl Stream for ExecutionScopedStream { + type Item = DFResult; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let Some(inner) = &mut this.inner else { + return Poll::Ready(None); + }; + let result = inner.as_mut().poll_next(cx); + if matches!(&result, Poll::Ready(None) | Poll::Ready(Some(Err(_)))) { + this.finish(); + } + result + } +} + +impl RecordBatchStream for ExecutionScopedStream { + fn schema(&self) -> ArrowSchemaRef { + Arc::clone(&self.schema) + } +} + +impl Drop for ExecutionScopedStream { + fn drop(&mut self) { + self.finish(); + } +} + impl LateralVectorSearchExec { fn new( input: Arc, @@ -380,11 +724,22 @@ impl LateralVectorSearchExec { query_vector_expr, limit, output_schema, + filter: None, + prepared_filter_cache: Arc::new(ExecutionPreparedFilterCache::default()), plan_properties, } } - async fn process_batch(&self, batch: RecordBatch) -> DFResult { + fn with_filter(mut self, filter: Predicate) -> Self { + self.filter = Some(filter); + self + } + + async fn process_batch( + &self, + batch: RecordBatch, + prepared_filter: Option<&OnceCell>, + ) -> DFResult { if batch.num_rows() == 0 { return empty_batch(self.output_schema.clone()); } @@ -398,17 +753,39 @@ impl LateralVectorSearchExec { return empty_batch(self.output_schema.clone()); } - let mut builder = self.target_table.new_batch_vector_search_builder(); - let results = builder + let prepared_filter = match &self.filter { + Some(filter) => { + let prepared_filter = prepared_filter.ok_or_else(|| { + DataFusionError::Internal( + "filtered lateral vector search is missing execution state".to_string(), + ) + })?; + let prepared = prepared_filter + .get_or_try_init(|| { + self.target_table + .prepare_vector_search_filter(filter.clone()) + }) + .await + .map_err(to_datafusion_error)?; + Some(prepared) + } + None => None, + }; + let target_table = prepared_filter + .map(PreparedVectorSearchFilter::table) + .unwrap_or(&self.target_table); + let mut builder = target_table.new_batch_vector_search_builder(); + builder .with_vector_column(&self.target_column) .with_query_vectors(query_vectors) - .with_limit(self.limit) - .execute() - .await - .map_err(to_datafusion_error)?; + .with_limit(self.limit); + if let Some(prepared_filter) = prepared_filter { + builder.with_prepared_filter(prepared_filter.clone()); + } + let results = builder.execute().await.map_err(to_datafusion_error)?; let (target_batch, target_row_id_to_index) = - read_target_rows(&self.target_table, &self.target_schema, &results).await?; + read_target_rows(target_table, &self.target_schema, &results).await?; let mut left_indices = Vec::new(); let mut right_indices = Vec::new(); @@ -452,8 +829,8 @@ impl DisplayAs for LateralVectorSearchExec { fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { write!( f, - "LateralVectorSearchExec: column={}, limit={}", - self.target_column, self.limit + "LateralVectorSearchExec: column={}, limit={}, filter={:?}", + self.target_column, self.limit, self.filter ) } } @@ -478,7 +855,7 @@ impl ExecutionPlan for LateralVectorSearchExec { if children.len() != 1 { return internal_err!("LateralVectorSearchExec expects one child"); } - Ok(Arc::new(Self::new( + let mut exec = Self::new( children.remove(0), self.target_table.clone(), Arc::clone(&self.target_schema), @@ -486,7 +863,12 @@ impl ExecutionPlan for LateralVectorSearchExec { Arc::clone(&self.query_vector_expr), self.limit, Arc::clone(&self.output_schema), - ))) + ); + if let Some(filter) = &self.filter { + exec = exec.with_filter(filter.clone()); + } + exec.prepared_filter_cache = Arc::clone(&self.prepared_filter_cache); + Ok(Arc::new(exec)) } fn execute( @@ -494,19 +876,39 @@ impl ExecutionPlan for LateralVectorSearchExec { partition: usize, context: Arc, ) -> DFResult { - let input = self.input.execute(partition, context)?; + let input = self.input.execute(partition, Arc::clone(&context))?; + let prepared_filter = self.filter.as_ref().map(|_| { + self.prepared_filter_cache.for_execution( + &context, + partition, + self.input.output_partitioning().partition_count(), + ) + }); + let prepared_filter_cell = prepared_filter + .as_ref() + .map(|lease| Arc::clone(&lease.prepared_filter)); let exec = self.clone(); let stream = input.then(move |batch| { let exec = exec.clone(); + let prepared_filter_cell = prepared_filter_cell.clone(); async move { let batch = batch?; - exec.process_batch(batch).await + exec.process_batch(batch, prepared_filter_cell.as_deref()) + .await } }); - Ok(Box::pin(RecordBatchStreamAdapter::new( + let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( self.output_schema.clone(), Box::pin(stream), - ))) + )); + Ok(match prepared_filter { + Some(lease) => Box::pin(ExecutionScopedStream::new( + self.output_schema.clone(), + stream, + lease, + )), + None => stream, + }) } fn partition_statistics(&self, _partition: Option) -> DFResult> { @@ -693,3 +1095,174 @@ fn empty_batch(schema: ArrowSchemaRef) -> DFResult { pub(crate) fn session_config() -> SessionConfig { SessionConfig::new().with_information_schema(true) } + +#[cfg(test)] +mod tests { + use super::*; + + struct ContextHoldingStream { + schema: ArrowSchemaRef, + _context: Arc, + } + + impl Stream for ContextHoldingStream { + type Item = DFResult; + + fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Pending + } + } + + impl RecordBatchStream for ContextHoldingStream { + fn schema(&self) -> ArrowSchemaRef { + Arc::clone(&self.schema) + } + } + + #[test] + fn prepared_filter_cache_releases_completed_execution() { + let cache = Arc::new(ExecutionPreparedFilterCache::default()); + let context = Arc::new(TaskContext::default()); + let first = cache.for_execution(&context, 0, 2); + let prepared_filter = Arc::downgrade(&first.prepared_filter); + + drop(first); + let second = cache.for_execution(&context, 1, 2); + assert!(Arc::ptr_eq( + &prepared_filter + .upgrade() + .expect("an unfinished execution should retain its prepared filter"), + &second.prepared_filter + )); + + drop(second); + assert!( + prepared_filter.upgrade().is_none(), + "finishing every partition should release the prepared filter even while the task context remains alive" + ); + assert!(cache + .entries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_empty()); + } + + #[test] + fn prepared_filter_cache_releases_subset_execution_after_context_drop() { + let cache = Arc::new(ExecutionPreparedFilterCache::default()); + let context = Arc::new(TaskContext::default()); + let lease = cache.for_execution(&context, 0, 4); + let prepared_filter = Arc::downgrade(&lease.prepared_filter); + + drop(context); + drop(lease); + + assert!( + prepared_filter.upgrade().is_none(), + "unstarted declared partitions must not retain a completed execution" + ); + assert!(cache + .entries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_empty()); + } + + #[tokio::test] + async fn prepared_filter_cache_releases_when_context_drops_after_stream() { + let cache = Arc::new(ExecutionPreparedFilterCache::default()); + let context = Arc::new(TaskContext::default()); + let lease = cache.for_execution(&context, 0, 4); + let prepared_filter = Arc::downgrade(&lease.prepared_filter); + + drop(lease); + assert!( + prepared_filter.upgrade().is_some(), + "the retained context must keep sequential partition reuse possible" + ); + drop(context); + + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while prepared_filter.upgrade().is_some() { + tokio::task::yield_now().await; + } + }) + .await + .expect("dropping the execution context must trigger cache cleanup"); + assert!(cache + .entries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_empty()); + } + + #[test] + fn execution_stream_cancellation_drops_inner_before_cache_lease() { + let cache = Arc::new(ExecutionPreparedFilterCache::default()); + let context = Arc::new(TaskContext::default()); + let lease = cache.for_execution(&context, 0, 4); + let prepared_filter = Arc::downgrade(&lease.prepared_filter); + let schema = Arc::new(datafusion::arrow::datatypes::Schema::empty()); + let inner: SendableRecordBatchStream = Box::pin(ContextHoldingStream { + schema: Arc::clone(&schema), + _context: Arc::clone(&context), + }); + let stream = ExecutionScopedStream::new(schema, inner, lease); + + drop(context); + drop(stream); + + assert!( + prepared_filter.upgrade().is_none(), + "cancellation must drop the child stream's context before releasing the cache lease" + ); + } + + #[test] + fn prepared_filter_cache_evicts_only_the_completed_execution() { + let cache = Arc::new(ExecutionPreparedFilterCache::default()); + let first_context = Arc::new(TaskContext::default()); + let second_context = Arc::new(TaskContext::default()); + let first = cache.for_execution(&first_context, 0, 1); + let second = cache.for_execution(&second_context, 0, 1); + let first_filter = Arc::downgrade(&first.prepared_filter); + let second_filter = Arc::downgrade(&second.prepared_filter); + + drop(first); + + assert!(first_filter.upgrade().is_none()); + assert!( + second_filter.upgrade().is_some(), + "completing one execution must not evict another execution's filter" + ); + assert_eq!( + cache + .entries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .len(), + 1 + ); + + drop(second); + assert!(second_filter.upgrade().is_none()); + } + + #[test] + fn prepared_filter_cache_waits_for_all_lease_clones() { + let cache = Arc::new(ExecutionPreparedFilterCache::default()); + let context = Arc::new(TaskContext::default()); + let lease = cache.for_execution(&context, 0, 1); + let lease_clone = lease.clone(); + let prepared_filter = Arc::downgrade(&lease.prepared_filter); + + drop(lease); + assert!( + prepared_filter.upgrade().is_some(), + "in-flight batch futures may retain a cloned lease" + ); + + drop(lease_clone); + assert!(prepared_filter.upgrade().is_none()); + } +} diff --git a/crates/integrations/datafusion/src/vector_search.rs b/crates/integrations/datafusion/src/vector_search.rs index 5a38942f7..f680800d6 100644 --- a/crates/integrations/datafusion/src/vector_search.rs +++ b/crates/integrations/datafusion/src/vector_search.rs @@ -44,11 +44,12 @@ use datafusion::prelude::SessionContext; use futures::{stream, TryStreamExt}; use paimon::catalog::Catalog; use paimon::spec::{ - BigIntType, CoreOptions, DataField, DataType, ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME, + BigIntType, CoreOptions, DataField, DataType, Predicate, ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME, }; use paimon::table::Table; use crate::error::to_datafusion_error; +use crate::filter_pushdown::{analyze_filters, is_safe_vector_prefilter}; use crate::runtime::{await_with_runtime, block_on_with_runtime}; use crate::table::{datafusion_read_fields, PaimonTableProvider}; use crate::table_function_args::{ @@ -262,10 +263,19 @@ impl TableProvider for VectorSearchTableProvider { &self, _state: &dyn Session, projection: Option<&Vec>, - _filters: &[Expr], + filters: &[Expr], limit: Option, ) -> DFResult> { let projected_schema = project_schema(&self.schema(), projection)?; + let filter_analysis = analyze_filters(filters, self.inner.table().schema().fields(), true); + if filter_analysis.requires_residual { + return Err(DataFusionError::Plan( + "vector_search cannot apply a partially translated scalar pre-filter".to_string(), + )); + } + let pushed_predicate = filter_analysis + .pushed_predicate + .filter(is_safe_vector_prefilter); // An outer `LIMIT 0` needs no rows. if limit == Some(0) { @@ -278,7 +288,7 @@ impl TableProvider for VectorSearchTableProvider { // small outer LIMIT doesn't read/materialize everything). All of this — search, // read and rank-order gather — runs at execution time in the exec's stream, so // planning / EXPLAIN stays cheap and the work is driven by the TaskContext. - Ok(Arc::new(VectorSearchExec::new( + let mut exec = VectorSearchExec::new( self.inner.table().clone(), self.column_name.clone(), self.query_vector.clone(), @@ -286,17 +296,36 @@ impl TableProvider for VectorSearchTableProvider { limit, projection.cloned(), projected_schema, - ))) + ); + if let Some(filter) = pushed_predicate { + exec = exec.with_filter(filter); + } + Ok(Arc::new(exec)) } fn supports_filters_pushdown( &self, filters: &[&Expr], ) -> DFResult> { - Ok(vec![ - TableProviderFilterPushDown::Unsupported; - filters.len() - ]) + let fields = self.inner.table().schema().fields(); + Ok(filters + .iter() + .map(|filter| { + let analysis = analyze_filters(std::slice::from_ref(*filter), fields, true); + if analysis + .pushed_predicate + .as_ref() + .is_some_and(is_safe_vector_prefilter) + && !analysis.requires_residual + { + // Keep DataFusion's residual filter as a correctness backstop + // while using the same predicate before vector Top-K. + TableProviderFilterPushDown::Inexact + } else { + TableProviderFilterPushDown::Unsupported + } + }) + .collect()) } } @@ -315,6 +344,7 @@ struct VectorSearchExec { output_limit: Option, projection: Option>, output_schema: ArrowSchemaRef, + filter: Option, plan_properties: Arc, } @@ -342,21 +372,45 @@ impl VectorSearchExec { output_limit, projection, output_schema, + filter: None, plan_properties, } } + fn with_filter(mut self, filter: Predicate) -> Self { + self.filter = Some(filter); + self + } + async fn compute_batch(&self) -> DFResult { + let prepared_filter = match &self.filter { + Some(filter) => Some( + self.table + .prepare_vector_search_filter(filter.clone()) + .await + .map_err(to_datafusion_error)?, + ), + None => None, + }; + let search_table = prepared_filter + .as_ref() + .map(|prepared| prepared.table()) + .unwrap_or(&self.table); + // Best-first row-ids from the index, searched at the full top-k so the ANN // recall is unchanged (data-evolution / global-index path; PK-vector tables are // unsupported here, as before). let mut search_result = await_with_runtime(async { - let mut builder = self.table.new_vector_search_builder(); + let mut builder = self.table.new_batch_vector_search_builder(); builder .with_vector_column(&self.column_name) - .with_query_vector(self.query_vector.clone()) + .with_query_vectors(vec![self.query_vector.clone()]) .with_limit(self.search_limit); - builder.execute_scored().await.map_err(to_datafusion_error) + if let Some(prepared) = &prepared_filter { + builder.with_prepared_filter(prepared.clone()); + } + let mut results = builder.execute().await.map_err(to_datafusion_error)?; + Ok::<_, DataFusionError>(results.remove(0)) }) .await?; @@ -376,10 +430,10 @@ impl VectorSearchExec { // Read the projected columns (+ internal `_ROW_ID`); the row-range scan yields // file order, realigned to relevance rank below. - let read_fields = projected_read_fields(&self.table, self.projection.as_ref())?; + let read_fields = projected_read_fields(search_table, self.projection.as_ref())?; let row_ranges = search_result.to_row_ranges().map_err(to_datafusion_error)?; let batches = await_with_runtime(async { - let mut read_builder = self.table.new_read_builder(); + let mut read_builder = search_table.new_read_builder(); read_builder .with_read_type(read_fields) .with_row_ranges(row_ranges); @@ -406,8 +460,8 @@ impl DisplayAs for VectorSearchExec { fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { write!( f, - "VectorSearchExec: column={}, search_limit={}, output_limit={:?}", - self.column_name, self.search_limit, self.output_limit + "VectorSearchExec: column={}, search_limit={}, output_limit={:?}, filter={:?}", + self.column_name, self.search_limit, self.output_limit, self.filter ) } } diff --git a/crates/integrations/datafusion/tests/read_tables.rs b/crates/integrations/datafusion/tests/read_tables.rs index e4fc47211..8756a6ce2 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -1727,6 +1727,7 @@ mod vector_search_tests { }; use datafusion::arrow::record_batch::RecordBatch; use datafusion::datasource::MemTable; + use datafusion::physical_plan::ExecutionPlanProperties; use paimon::catalog::Identifier; use paimon::spec::{ArrayType, DataType, FloatType, IntType, Schema, VarCharType}; use paimon::table::BranchManager; @@ -2279,6 +2280,17 @@ mod vector_search_tests { #[tokio::test] async fn test_vector_search_lateral_join_uses_query_vectors() { let (ctx, _catalog, _tmp) = create_java_vindex_vector_search_context().await; + ctx.sql( + "CALL sys.create_global_index( \ + table => 'default.test_java_vindex_vector', \ + index_column => 'id', \ + index_type => 'btree')", + ) + .await + .expect("BTree index build SQL should parse") + .collect() + .await + .expect("BTree index build SQL should execute"); let query_batch = build_vector_batch( vec![10, 20], vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]], @@ -2303,6 +2315,192 @@ mod vector_search_tests { let rows = extract_query_result_ids(&batches); assert_eq!(rows, vec![(10, 0), (10, 1), (20, 1), (20, 2)]); + + let filtered = ctx + .sql( + "SELECT q.id AS query_id, r.id AS result_id \ + FROM paimon.default.queries q \ + CROSS JOIN LATERAL vector_search('paimon.default.test_java_vindex_vector', 'embedding', q.embedding, 1) AS r \ + WHERE r.id = 2 \ + ORDER BY query_id, result_id", + ) + .await + .expect("filtered lateral vector_search SQL should parse") + .collect() + .await + .expect("filtered lateral vector_search query should execute"); + assert_eq!( + extract_query_result_ids(&filtered), + vec![(10, 2), (20, 2)], + "the target-side filter must be applied before each lateral Top-K" + ); + + let mixed_filter = ctx + .sql( + "SELECT q.id AS query_id, r.id AS result_id \ + FROM paimon.default.queries q \ + CROSS JOIN LATERAL vector_search('paimon.default.test_java_vindex_vector', 'embedding', q.embedding, 1) AS r \ + WHERE q.id = 10 AND r.id = 2 \ + ORDER BY query_id, result_id", + ) + .await + .expect("mixed-filter lateral vector_search SQL should parse") + .collect() + .await + .expect("mixed-filter lateral vector_search query should execute"); + assert_eq!( + extract_query_result_ids(&mixed_filter), + vec![(10, 2)], + "target-only conjuncts must be pre-filtered while left conjuncts remain residual" + ); + + let cross_side_filter = ctx + .sql( + "SELECT q.id AS query_id, r.id AS result_id \ + FROM paimon.default.queries q \ + CROSS JOIN LATERAL vector_search('paimon.default.test_java_vindex_vector', 'embedding', q.embedding, 1) AS r \ + WHERE r.id = 2 AND q.id = r.id * 5 \ + ORDER BY query_id, result_id", + ) + .await + .expect("cross-side filtered lateral vector_search SQL should parse") + .collect() + .await + .expect("cross-side filtered lateral vector_search query should execute"); + assert_eq!( + extract_query_result_ids(&cross_side_filter), + vec![(10, 2)], + "cross-side conjuncts must remain residual" + ); + } + + #[tokio::test] + async fn test_lateral_vector_search_sequential_partitions_share_scalar_prefilter_snapshot() { + let (ctx, catalog, _tmp) = create_java_vindex_vector_search_context().await; + ctx.sql( + "CALL sys.create_global_index( \ + table => 'default.test_java_vindex_vector', \ + index_column => 'id', \ + index_type => 'btree')", + ) + .await + .expect("initial BTree index build SQL should parse") + .collect() + .await + .expect("initial BTree index build should execute"); + + let first_query_batch = build_vector_batch(vec![10], vec![vec![1.0, 0.0, 0.0, 0.0]]); + let second_query_batch = build_vector_batch(vec![20], vec![vec![1.0, 0.0, 0.0, 0.0]]); + let query_table = MemTable::try_new( + first_query_batch.schema(), + vec![vec![first_query_batch], vec![second_query_batch]], + ) + .expect("Failed to create partitioned query vector table"); + ctx.register_temp_table("paimon.default.refresh_queries", Arc::new(query_table)) + .expect("Failed to register query vector table"); + + let dataframe = ctx + .sql( + "SELECT r.id \ + FROM paimon.default.refresh_queries q \ + CROSS JOIN LATERAL vector_search('paimon.default.test_java_vindex_vector', 'embedding', q.embedding, 1) AS r \ + WHERE r.id = 6", + ) + .await + .expect("filtered lateral vector_search SQL should parse"); + let physical_plan = dataframe + .create_physical_plan() + .await + .expect("physical plan should be created"); + assert_eq!( + physical_plan.output_partitioning().partition_count(), + 2, + "the regression requires sequential execution of two input partitions" + ); + + let task_context = ctx.ctx().task_ctx(); + let first = datafusion::physical_plan::common::collect( + physical_plan + .execute(0, Arc::clone(&task_context)) + .expect("first partition should start"), + ) + .await + .expect("first partition should finish"); + assert!(extract_ids_in_order(&first).is_empty()); + + let identifier = Identifier::new("default", "test_java_vindex_vector"); + let table = catalog + .get_table(&identifier) + .await + .expect("load vector target table"); + let write_builder = table + .new_write_builder() + .with_commit_user("test-user") + .expect("configure target append"); + let mut writer = write_builder.new_write().expect("create target writer"); + writer + .write_arrow_batch(&build_vector_batch(vec![6], vec![vec![1.0, 0.0, 0.0, 0.0]])) + .await + .expect("append target row"); + let messages = writer + .prepare_commit() + .await + .expect("prepare target append"); + write_builder + .new_commit() + .commit(messages) + .await + .expect("commit target append"); + + ctx.sql( + "CALL sys.create_global_index( \ + table => 'default.test_java_vindex_vector', \ + index_column => 'embedding', \ + index_type => 'ivf-flat', \ + options => 'ivf-flat.dimension=4,ivf-flat.nlist=1,ivf-flat.distance.metric=l2')", + ) + .await + .expect("incremental vector index build SQL should parse") + .collect() + .await + .expect("incremental vector index build should execute"); + ctx.sql( + "CALL sys.create_global_index( \ + table => 'default.test_java_vindex_vector', \ + index_column => 'id', \ + index_type => 'btree')", + ) + .await + .expect("incremental BTree index build SQL should parse") + .collect() + .await + .expect("incremental BTree index build should execute"); + + let second = datafusion::physical_plan::common::collect( + physical_plan + .execute(1, Arc::clone(&task_context)) + .expect("second partition should start"), + ) + .await + .expect("second partition should finish"); + assert!( + extract_ids_in_order(&second).is_empty(), + "sequential partitions in one execution must use the same scalar-filter snapshot" + ); + + drop(task_context); + let refreshed = datafusion::physical_plan::common::collect( + physical_plan + .execute(1, ctx.ctx().task_ctx()) + .expect("refreshed execution should start"), + ) + .await + .expect("refreshed execution should finish"); + assert_eq!( + extract_ids_in_order(&refreshed), + vec![6], + "a reused physical plan must refresh the scalar filter for a new execution" + ); } // Manual run with a local Lumina native library: @@ -2450,6 +2648,18 @@ mod vector_search_tests { .await .expect("vindex index build SQL should execute"); + ctx.sql( + "CALL sys.create_global_index( \ + table => 'default.vindex_build_query_e2e', \ + index_column => 'id', \ + index_type => 'btree')", + ) + .await + .expect("BTree index build SQL should parse") + .collect() + .await + .expect("BTree index build SQL should execute"); + let index_batches = ctx .sql("SELECT index_type, row_count, row_range_start, row_range_end, index_field_name FROM paimon.default.`vindex_build_query_e2e$table_indexes` WHERE index_type = 'ivf-flat'") .await @@ -2475,6 +2685,19 @@ mod vector_search_tests { .expect("vector_search query should execute"); let ids = extract_ids(&search_batches); assert_eq!(ids, vec![0, 1]); + + let filtered_batches = ctx + .sql("SELECT id FROM vector_search('paimon.default.vindex_build_query_e2e', 'embedding', '[1.0, 0.0]', 2) WHERE id >= 4") + .await + .expect("filtered vector_search SQL should parse") + .collect() + .await + .expect("filtered vector_search query should execute"); + assert_eq!( + extract_ids_in_order(&filtered_batches), + vec![5, 4], + "the scalar predicate must be applied before vector Top-K without losing rank order" + ); } } diff --git a/crates/paimon/src/io/file_io.rs b/crates/paimon/src/io/file_io.rs index 1257f6328..b15c1e2d1 100644 --- a/crates/paimon/src/io/file_io.rs +++ b/crates/paimon/src/io/file_io.rs @@ -63,11 +63,16 @@ pub(crate) trait FileIOProvider: std::fmt::Debug + Send + Sync { #[derive(Clone)] pub struct FileIO { + /// Private identity shared by clones, but not by independently built storage backends. + storage_lineage: Arc, storage: Arc, cache: Option>, provider: Option>, } +#[derive(Debug)] +struct StorageLineage; + impl std::fmt::Debug for FileIO { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("FileIO") @@ -79,6 +84,10 @@ impl std::fmt::Debug for FileIO { } impl FileIO { + pub(crate) fn shares_storage_lineage(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.storage_lineage, &other.storage_lineage) + } + /// Attach an externally managed block cache. /// /// `block_size` controls the aligned ranges presented to the cache. @@ -533,6 +542,7 @@ impl FileIOBuilder { let cache = self.cache.clone(); let storage = Storage::build(self)?; Ok(FileIO { + storage_lineage: Arc::new(StorageLineage), storage: Arc::new(storage), cache, provider: None, @@ -1254,6 +1264,16 @@ mod file_action_test { assert!(!file_io_2.exists(path).await.unwrap()); } + #[test] + fn test_storage_lineage_is_shared_only_by_file_io_clones() { + let file_io = setup_memory_file_io(); + let clone = file_io.clone(); + let independent = setup_memory_file_io(); + + assert!(file_io.shares_storage_lineage(&clone)); + assert!(!file_io.shares_storage_lineage(&independent)); + } + #[tokio::test] async fn test_get_status_fs() { let file_io = setup_fs_file_io(); diff --git a/crates/paimon/src/lumina/reader.rs b/crates/paimon/src/lumina/reader.rs index 25514a67f..00025176c 100644 --- a/crates/paimon/src/lumina/reader.rs +++ b/crates/paimon/src/lumina/reader.rs @@ -27,6 +27,65 @@ const MIN_SEARCH_LIST_SIZE: usize = 16; // C ABI returns int64_t -1 for invalid results, which casts to u64::MAX in Rust. const SENTINEL: u64 = u64::MAX; +trait LuminaSearch { + fn search( + &self, + query: &[f32], + n: i32, + k: i32, + distances: &mut [f32], + labels: &mut [u64], + options: &HashMap, + ) -> crate::Result<()>; + + #[allow(clippy::too_many_arguments)] + fn search_with_filter( + &self, + query: &[f32], + n: i32, + k: i32, + distances: &mut [f32], + labels: &mut [u64], + filter_ids: &[u64], + options: &HashMap, + ) -> crate::Result<()>; + + fn get_count(&self) -> crate::Result; +} + +impl LuminaSearch for LuminaSearcher { + fn search( + &self, + query: &[f32], + n: i32, + k: i32, + distances: &mut [f32], + labels: &mut [u64], + options: &HashMap, + ) -> crate::Result<()> { + LuminaSearcher::search(self, query, n, k, distances, labels, options) + } + + fn search_with_filter( + &self, + query: &[f32], + n: i32, + k: i32, + distances: &mut [f32], + labels: &mut [u64], + filter_ids: &[u64], + options: &HashMap, + ) -> crate::Result<()> { + LuminaSearcher::search_with_filter( + self, query, n, k, distances, labels, filter_ids, options, + ) + } + + fn get_count(&self) -> crate::Result { + LuminaSearcher::get_count(self) + } +} + fn ensure_search_list_size(search_options: &mut HashMap, top_k: usize) { if !search_options.contains_key("diskann.search.list_size") { let list_size = std::cmp::max((top_k as f64 * 1.5) as usize, MIN_SEARCH_LIST_SIZE); @@ -285,8 +344,8 @@ impl LuminaVectorGlobalIndexReader { } } -fn search_lumina( - searcher: &LuminaSearcher, +fn search_lumina( + searcher: &S, index_meta: &LuminaIndexMeta, search_options_base: &HashMap, vector_search: &VectorSearch, @@ -311,9 +370,9 @@ fn search_lumina( return Ok(None); } - let include_row_ids = &vector_search.include_row_ids; + let include_row_ids = vector_search.effective_include_row_ids(); - let (distances, labels) = if let Some(ref include_ids) = include_row_ids { + let (distances, labels) = if let Some(include_ids) = include_row_ids { let filter_id_list: Vec = include_ids.iter().collect(); if filter_id_list.is_empty() { return Ok(None); @@ -358,8 +417,8 @@ fn search_lumina( Ok(Some(id_to_scores)) } -fn search_lumina_batch( - searcher: &LuminaSearcher, +fn search_lumina_batch( + searcher: &S, index_meta: &LuminaIndexMeta, search_options_base: &HashMap, vector_searches: &[VectorSearch], @@ -367,23 +426,18 @@ fn search_lumina_batch( if vector_searches.is_empty() { return Ok(Vec::new()); } - if vector_searches - .iter() - .any(|vector_search| vector_search.include_row_ids.is_some()) - { - return vector_searches - .iter() - .map(|vector_search| { - search_lumina(searcher, index_meta, search_options_base, vector_search) - }) - .collect(); - } let limit = vector_searches[0].limit; - if vector_searches + let same_limit = vector_searches + .iter() + .all(|vector_search| vector_search.limit == limit); + let shared_filter = same_limit + .then(|| shared_batch_include_row_ids(vector_searches)) + .flatten(); + let has_filter = vector_searches .iter() - .any(|vector_search| vector_search.limit != limit) - { + .any(|vector_search| vector_search.effective_include_row_ids().is_some()); + if has_filter && shared_filter.is_none() || !same_limit { return vector_searches .iter() .map(|vector_search| { @@ -406,9 +460,18 @@ fn search_lumina_batch( } } + let filter_id_list = + shared_filter.map(|include_row_ids| include_row_ids.iter().collect::>()); + if filter_id_list.as_ref().is_some_and(Vec::is_empty) { + return Ok(vec![None; vector_searches.len()]); + } + let index_metric = index_meta.metric()?; let count = searcher.get_count()? as usize; - let effective_k = std::cmp::min(limit, count); + let effective_k = filter_id_list.as_ref().map_or_else( + || std::cmp::min(limit, count), + |ids| std::cmp::min(std::cmp::min(limit, count), ids.len()), + ); if effective_k == 0 { return Ok(vec![None; vector_searches.len()]); } @@ -422,14 +485,27 @@ fn search_lumina_batch( let mut labels = new_label_buffer(vector_searches.len() * effective_k); let mut search_opts: HashMap = search_options_base.clone(); ensure_search_list_size(&mut search_opts, effective_k); - searcher.search( - &query, - vector_searches.len() as i32, - effective_k as i32, - &mut distances, - &mut labels, - &search_opts, - )?; + if let Some(filter_ids) = filter_id_list { + search_opts.insert("search.thread_safe_filter".to_string(), "true".to_string()); + searcher.search_with_filter( + &query, + vector_searches.len() as i32, + effective_k as i32, + &mut distances, + &mut labels, + &filter_ids, + &search_opts, + )?; + } else { + searcher.search( + &query, + vector_searches.len() as i32, + effective_k as i32, + &mut distances, + &mut labels, + &search_opts, + )?; + } let mut results = Vec::with_capacity(vector_searches.len()); for query_index in 0..vector_searches.len() { @@ -450,6 +526,22 @@ fn search_lumina_batch( Ok(results) } +fn shared_batch_include_row_ids( + vector_searches: &[VectorSearch], +) -> Option<&std::sync::Arc> { + let first = vector_searches.first()?.shared_include_row_ids.as_ref()?; + vector_searches + .iter() + .skip(1) + .all(|vector_search| { + vector_search + .shared_include_row_ids + .as_ref() + .is_some_and(|include_row_ids| std::sync::Arc::ptr_eq(first, include_row_ids)) + }) + .then_some(first) +} + fn write_temp_index_file(stream: &mut S) -> crate::Result { stream .seek(SeekFrom::Start(0)) @@ -498,8 +590,247 @@ impl Drop for LuminaVectorGlobalIndexReader { #[cfg(test)] mod tests { use super::*; + use crate::lumina::{KEY_DIMENSION, KEY_DISTANCE_METRIC}; use crate::vector_search::GlobalIndexIOMeta; use std::io::Cursor; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex}; + + #[derive(Debug, PartialEq)] + struct FilteredSearchCall { + query: Vec, + n: i32, + k: i32, + filter_ids: Vec, + } + + struct RecordingSearcher { + count: u64, + count_calls: AtomicUsize, + unfiltered_calls: Mutex, i32, i32)>>, + filtered_calls: Mutex>, + } + + impl RecordingSearcher { + fn new(count: u64) -> Self { + Self { + count, + count_calls: AtomicUsize::new(0), + unfiltered_calls: Mutex::new(Vec::new()), + filtered_calls: Mutex::new(Vec::new()), + } + } + } + + impl LuminaSearch for RecordingSearcher { + fn search( + &self, + query: &[f32], + n: i32, + k: i32, + distances: &mut [f32], + labels: &mut [u64], + _options: &HashMap, + ) -> crate::Result<()> { + self.unfiltered_calls + .lock() + .expect("unfiltered call lock") + .push((query.to_vec(), n, k)); + for (index, (distance, label)) in + distances.iter_mut().zip(labels.iter_mut()).enumerate() + { + *distance = index as f32; + *label = index as u64; + } + Ok(()) + } + + fn search_with_filter( + &self, + query: &[f32], + n: i32, + k: i32, + distances: &mut [f32], + labels: &mut [u64], + filter_ids: &[u64], + _options: &HashMap, + ) -> crate::Result<()> { + self.filtered_calls + .lock() + .expect("filtered call lock") + .push(FilteredSearchCall { + query: query.to_vec(), + n, + k, + filter_ids: filter_ids.to_vec(), + }); + for (index, (distance, label)) in + distances.iter_mut().zip(labels.iter_mut()).enumerate() + { + *distance = index as f32; + *label = filter_ids[index % filter_ids.len()]; + } + Ok(()) + } + + fn get_count(&self) -> crate::Result { + self.count_calls.fetch_add(1, Ordering::Relaxed); + Ok(self.count) + } + } + + fn test_index_meta(dim: usize) -> LuminaIndexMeta { + LuminaIndexMeta::new(HashMap::from([ + (KEY_DIMENSION.to_string(), dim.to_string()), + (KEY_DISTANCE_METRIC.to_string(), "l2".to_string()), + ])) + } + + #[test] + fn test_shared_filter_uses_one_lumina_batch_search() { + let searcher = RecordingSearcher::new(10); + let shared_filter = Arc::new(roaring::RoaringTreemap::from_iter([2, 4, 6])); + let mut first = VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap(); + first.set_shared_include_row_ids(Arc::clone(&shared_filter)); + let mut second = VectorSearch::new(vec![0.0, 1.0], 2, "embedding".to_string()).unwrap(); + second.set_shared_include_row_ids(Arc::clone(&shared_filter)); + + let results = search_lumina_batch( + &searcher, + &test_index_meta(2), + &HashMap::new(), + &[first, second], + ) + .expect("shared filtered batch search should succeed"); + + assert_eq!(results.len(), 2); + assert!(searcher + .unfiltered_calls + .lock() + .expect("unfiltered call lock") + .is_empty()); + assert_eq!( + *searcher.filtered_calls.lock().expect("filtered call lock"), + vec![FilteredSearchCall { + query: vec![1.0, 0.0, 0.0, 1.0], + n: 2, + k: 2, + filter_ids: vec![2, 4, 6], + }] + ); + } + + #[test] + fn test_equal_but_distinct_filters_keep_scalar_fallback() { + let searcher = RecordingSearcher::new(10); + let mut first = VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap(); + first.set_shared_include_row_ids(Arc::new(roaring::RoaringTreemap::from_iter([2, 4, 6]))); + let mut second = VectorSearch::new(vec![0.0, 1.0], 2, "embedding".to_string()).unwrap(); + second.set_shared_include_row_ids(Arc::new(roaring::RoaringTreemap::from_iter([2, 4, 6]))); + + search_lumina_batch( + &searcher, + &test_index_meta(2), + &HashMap::new(), + &[first, second], + ) + .expect("distinct filtered searches should succeed"); + + let calls = searcher.filtered_calls.lock().expect("filtered call lock"); + assert_eq!(calls.len(), 2); + assert!(calls.iter().all(|call| call.n == 1)); + } + + #[test] + fn test_mixed_filters_and_limits_keep_scalar_fallback() { + let shared_filter = Arc::new(roaring::RoaringTreemap::from_iter([2, 4, 6])); + let mut filtered = VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap(); + filtered.set_shared_include_row_ids(Arc::clone(&shared_filter)); + let unfiltered = VectorSearch::new(vec![0.0, 1.0], 2, "embedding".to_string()).unwrap(); + let mixed_searcher = RecordingSearcher::new(10); + + search_lumina_batch( + &mixed_searcher, + &test_index_meta(2), + &HashMap::new(), + &[filtered, unfiltered], + ) + .expect("mixed filtered searches should succeed"); + + assert_eq!( + mixed_searcher + .filtered_calls + .lock() + .expect("filtered call lock") + .len(), + 1 + ); + assert_eq!( + mixed_searcher + .unfiltered_calls + .lock() + .expect("unfiltered call lock") + .len(), + 1 + ); + + let mut first = VectorSearch::new(vec![1.0, 0.0], 1, "embedding".to_string()).unwrap(); + first.set_shared_include_row_ids(Arc::clone(&shared_filter)); + let mut second = VectorSearch::new(vec![0.0, 1.0], 2, "embedding".to_string()).unwrap(); + second.set_shared_include_row_ids(shared_filter); + let differing_limit_searcher = RecordingSearcher::new(10); + + search_lumina_batch( + &differing_limit_searcher, + &test_index_meta(2), + &HashMap::new(), + &[first, second], + ) + .expect("differing-limit filtered searches should succeed"); + + let calls = differing_limit_searcher + .filtered_calls + .lock() + .expect("filtered call lock"); + assert_eq!(calls.len(), 2); + assert_eq!(calls.iter().map(|call| call.k).collect::>(), [1, 2]); + assert!(calls.iter().all(|call| call.n == 1)); + } + + #[test] + fn test_empty_shared_filter_skips_native_search() { + let searcher = RecordingSearcher::new(10); + let shared_filter = Arc::new(roaring::RoaringTreemap::new()); + let mut first = VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap(); + first.set_shared_include_row_ids(Arc::clone(&shared_filter)); + let mut second = VectorSearch::new(vec![0.0, 1.0], 2, "embedding".to_string()).unwrap(); + second.set_shared_include_row_ids(Arc::clone(&shared_filter)); + + let results = search_lumina_batch( + &searcher, + &test_index_meta(2), + &HashMap::new(), + &[first, second], + ) + .expect("empty shared filter should succeed"); + + assert_eq!(results, vec![None, None]); + assert_eq!( + searcher.count_calls.load(Ordering::Relaxed), + 0, + "an empty shared filter should avoid all native searcher calls" + ); + assert!(searcher + .unfiltered_calls + .lock() + .expect("unfiltered call lock") + .is_empty()); + assert!(searcher + .filtered_calls + .lock() + .expect("filtered call lock") + .is_empty()); + } #[test] fn test_convert_distance_to_score() { diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 3b0f447d4..ff99fefd2 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -154,13 +154,18 @@ pub use table_scan::TableScan; pub use table_update::TableUpdate; pub use table_write::TableWrite; pub use tag_manager::TagManager; -pub use vector_search_builder::{BatchVectorSearchBuilder, VectorSearchBuilder}; +pub use vector_search_builder::{ + BatchVectorSearchBuilder, PreparedVectorSearchFilter, VectorSearchBuilder, +}; pub use vindex_index_build_builder::VindexIndexBuildBuilder; pub use write_builder::WriteBuilder; use crate::catalog::{validate_branch_name, Identifier, DEFAULT_MAIN_BRANCH}; use crate::io::FileIO; -use crate::spec::{CoreOptions, DataField, Snapshot, TableSchema}; +use crate::spec::{ + CoreOptions, DataField, Snapshot, TableSchema, SCAN_SNAPSHOT_ID_OPTION, SCAN_TAG_NAME_OPTION, + SCAN_TIMESTAMP_MILLIS_OPTION, SCAN_VERSION_OPTION, SCAN_WATERMARK_OPTION, +}; use std::collections::HashMap; /// Table represents a table in the catalog. @@ -425,6 +430,50 @@ impl Table { } } + /// Create a read-only copy pinned to an already resolved snapshot. + /// + /// Replaces any selector that originally resolved the snapshot with an + /// explicit `scan.snapshot-id`, so every subsequent scan stage observes the + /// same snapshot. The snapshot's schema is loaded when it differs from the + /// current table schema. + pub(crate) async fn copy_with_resolved_snapshot(&self, snapshot: &Snapshot) -> Result { + let mut options = self.schema.options().clone(); + for selector in [ + SCAN_TIMESTAMP_MILLIS_OPTION, + SCAN_WATERMARK_OPTION, + SCAN_VERSION_OPTION, + SCAN_SNAPSHOT_ID_OPTION, + SCAN_TAG_NAME_OPTION, + ] { + options.remove(selector); + } + options.insert( + SCAN_SNAPSHOT_ID_OPTION.to_string(), + snapshot.id().to_string(), + ); + + let schema = if snapshot.schema_id() == self.schema.id() { + self.schema.copy_with_replaced_options(options) + } else { + self.schema_manager + .schema(snapshot.schema_id()) + .await? + .copy_with_replaced_options(options) + }; + Ok(Self { + file_io: self.file_io.clone(), + identifier: self.identifier.clone(), + location: self.location.clone(), + schema, + schema_manager: self.schema_manager.clone(), + branch: self.branch.clone(), + branch_reference: self.branch_reference, + rest_env: self.rest_env.clone(), + time_traveled: true, + travel_snapshot: Some(snapshot.clone()), + }) + } + /// Create a copy of this table with extra options merged in, switching to /// the schema of the time-travelled snapshot when the merged options /// select one. diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index aa1fdc6dc..1fc662b90 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -45,6 +45,7 @@ use crate::table::pk_vector_position_read::{ }; use crate::table::pk_vector_scan::{PkVectorScan, PkVectorScanPlan}; use crate::table::read_builder::resolve_projected_fields; +use crate::table::row_id_predicate::intersect_sorted_ranges; use crate::table::source::DataSplit; use crate::table::{ find_field_id_by_name, merge_row_ranges, ArrowRecordBatchStream, RowRange, Table, @@ -196,6 +197,34 @@ pub struct BatchVectorSearchBuilder<'a> { options: HashMap, projection: Option>, filter: Option, + include_row_ids: Option>, + prepared_filter: Option, +} + +/// A scalar vector pre-filter resolved once against one pinned snapshot. +/// +/// Reusing this value avoids repeating the same scalar-index/table read for +/// every input batch of a lateral vector query. +#[derive(Debug, Clone)] +pub struct PreparedVectorSearchFilter { + table: Table, + include_row_ids: Arc, +} + +impl PreparedVectorSearchFilter { + pub fn table(&self) -> &Table { + &self.table + } + + pub fn include_row_ids(&self) -> &Arc { + &self.include_row_ids + } +} + +fn same_vector_search_table(left: &Table, right: &Table) -> bool { + left.file_io().shares_storage_lineage(right.file_io()) + && left.location().trim_end_matches('/') == right.location().trim_end_matches('/') + && left.branch() == right.branch() } /// The primary-key vector route's search output plus the source context a later @@ -249,16 +278,14 @@ impl<'a> VectorSearchBuilder<'a> { self } - /// Attach a residual scalar predicate applied *after* vector recall on the - /// primary-key vector path: each recalled candidate file is re-read and only - /// rows satisfying `filter` survive, folded into the search so best-first - /// order and Top-K still hold. Mirrors Java `PrimaryKeyVectorRead`'s - /// residual-filter support. Only the primary-key vector path consumes it, and - /// only when the table exposes physical rows directly (deletion vectors - /// enabled without merge-on-read); otherwise the query fails loud. A query - /// that does not resolve to the primary-key vector path (no PK-vector index, - /// or a non-PK-vector column) also fails loud rather than silently ignoring - /// the filter. + /// Attach a scalar predicate applied before vector Top-K. + /// + /// On the primary-key vector path this remains a residual allow-list over + /// physical positions, mirroring Java `PrimaryKeyVectorRead`. On the + /// data-evolution/global-index path the predicate is evaluated through a + /// snapshot-pinned table read (which can use scalar global indexes such as + /// BTree), producing global row IDs that are localized for each vector-index + /// shard and passed to the vector backend as an include filter. /// /// The whole predicate is both pushed into the scan — where it prunes whole /// data files by their column stats — and applied per row as a residual over @@ -325,27 +352,16 @@ impl<'a> VectorSearchBuilder<'a> { } } - // The data-evolution (global-index) fall-through path cannot honor a - // residual filter — it never reads physical rows. Rather than silently - // drop the predicate and return unfiltered results, fail loud when a - // filter is set on a query that does not resolve to the primary-key - // vector path. - if self.filter.is_some() { - return Err(crate::Error::DataInvalid { - message: "vector search filter is only supported on the primary-key vector path" - .to_string(), - source: None, - }); - } - let mut batch_builder = BatchVectorSearchBuilder::new(self.table); - let mut results = batch_builder + batch_builder .with_vector_column(vector_column) .with_query_vectors(vec![query_vector.clone()]) .with_limit(limit) - .with_options(self.options.clone()) - .execute() - .await?; + .with_options(self.options.clone()); + if let Some(filter) = &self.filter { + batch_builder.with_filter(filter.clone()); + } + let mut results = batch_builder.execute().await?; debug_assert_eq!(results.len(), 1); Ok(results.remove(0)) @@ -397,7 +413,8 @@ impl<'a> VectorSearchBuilder<'a> { // Data-evolution (global-index) vector search: materialize rows from the // scored global row-ids and attach the unified score column. A non-vector // column or a set filter fails loud inside execute_scored below. - self.execute_de_vector_read().await + self.execute_de_vector_read(vector_column, query_vector, limit) + .await } /// Materialize the best-first data-evolution vector search hits into Arrow @@ -405,21 +422,20 @@ impl<'a> VectorSearchBuilder<'a> { /// subsequent row-range read materializes those rows, and each row's score is /// joined back by `_ROW_ID`. Output columns are the projected user table /// columns (all user columns by default) plus `__paimon_search_score`; `_ROW_ID` - /// is always hidden. A filter is unsupported here and fails loud inside - /// `execute_scored`. - async fn execute_de_vector_read(&self) -> crate::Result { + /// is always hidden. A scalar filter is applied before vector Top-K by the + /// snapshot-pinned scored search below. + async fn execute_de_vector_read( + &self, + vector_column: &str, + query_vector: &[f32], + limit: usize, + ) -> crate::Result { // Validate the target column exists and is a vector-bearing type before any // work. The data-evolution search returns an empty result for an unknown // field (its scored-path behavior), which would make a typo'd or scalar // column look like a normal empty read here — violating `execute_read`'s // fail-loud contract (a C/Doris caller would see EOF, not an input error). // Reject it up front instead. - let vector_column = - self.vector_column - .as_deref() - .ok_or_else(|| crate::Error::ConfigInvalid { - message: "Vector column must be set via with_vector_column()".to_string(), - })?; let field = self .table .schema() @@ -449,12 +465,25 @@ impl<'a> VectorSearchBuilder<'a> { }); } - let sr = self.execute_scored().await?; - // Resolve the projected user columns up front so an invalid projection // fails loud even when the result is empty. let mut read_type = self.resolve_materialize_read_type()?; + let Some(snapshot) = crate::table::time_travel::resolve_snapshot(self.table).await? else { + return Ok(Box::pin(stream::empty())); + }; + let pinned_table = self.table.copy_with_resolved_snapshot(&snapshot).await?; + let mut search_builder = pinned_table.new_vector_search_builder(); + search_builder + .with_vector_column(vector_column) + .with_query_vector(query_vector.to_vec()) + .with_limit(limit) + .with_options(self.options.clone()); + if let Some(filter) = &self.filter { + search_builder.with_filter(filter.clone()); + } + let sr = search_builder.execute_scored().await?; + if sr.is_empty() { return Ok(Box::pin(stream::empty())); } @@ -472,7 +501,7 @@ impl<'a> VectorSearchBuilder<'a> { read_type.push(row_id_data_field()); } - let mut read_builder = self.table.new_read_builder(); + let mut read_builder = pinned_table.new_read_builder(); read_builder .with_read_type(read_type) .with_row_ranges(ranges); @@ -1171,6 +1200,8 @@ impl<'a> BatchVectorSearchBuilder<'a> { options: HashMap::new(), projection: None, filter: None, + include_row_ids: None, + prepared_filter: None, } } @@ -1194,14 +1225,36 @@ impl<'a> BatchVectorSearchBuilder<'a> { self } - /// Attach a residual scalar predicate applied *after* vector recall on the - /// primary-key vector path, shared across every query in the batch. Mirrors - /// the single [`VectorSearchBuilder::with_filter`]: only the primary-key - /// vector path (via [`execute_read`](Self::execute_read)) consumes it, and only - /// when the table exposes physical rows directly (deletion vectors without - /// merge-on-read); otherwise the query fails loud. + /// Attach one scalar predicate shared by every query in the batch and applied + /// before vector Top-K. See [`VectorSearchBuilder::with_filter`] for the + /// primary-key and data-evolution execution semantics. pub fn with_filter(&mut self, filter: Predicate) -> &mut Self { self.filter = Some(filter); + self.include_row_ids = None; + self.prepared_filter = None; + self + } + + /// Attach a prepared scalar pre-filter together with the exact table + /// snapshot against which its row-ID allow-list was evaluated. + pub fn with_prepared_filter( + &mut self, + prepared_filter: PreparedVectorSearchFilter, + ) -> &mut Self { + self.prepared_filter = Some(prepared_filter); + self.filter = None; + self.include_row_ids = None; + self + } + + /// Attach a caller-managed row-ID allow-list. + /// + /// This low-level API does not bind the allow-list to a table snapshot. + /// Prefer [`Self::with_prepared_filter`] for scalar pre-filters. + pub fn with_include_row_ids(&mut self, include_row_ids: RoaringTreemap) -> &mut Self { + self.include_row_ids = Some(Arc::new(include_row_ids)); + self.filter = None; + self.prepared_filter = None; self } @@ -1217,11 +1270,31 @@ impl<'a> BatchVectorSearchBuilder<'a> { pub async fn execute(&self) -> crate::Result> { let timing_enabled = vector_search_timing_enabled(); let total_start = timing_enabled.then(Instant::now); - // Fail closed: like `execute_read` and the single-query builder, this - // returns data-derived row ids/scores outside `TableScan`/`TableRead`, - // so it must refuse a `query-auth.enabled` table before any fast path - // (an empty snapshot would otherwise return empty results and bypass it). - let core = CoreOptions::new(self.table.schema().options()); + // The builder target is authoritative for current auth/type policy. + // A prepared filter only pins a snapshot and may carry older options. + CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; + if let Some(prepared) = &self.prepared_filter { + if !same_vector_search_table(self.table, prepared.table()) { + return Err(crate::Error::DataInvalid { + message: format!( + "Prepared vector search filter belongs to a different table: builder target is '{}@{}', prepared filter target is '{}@{}'", + self.table.location(), + self.table.branch(), + prepared.table().location(), + prepared.table().branch(), + ), + source: None, + }); + } + } + // Check the pinned execution view as defense in depth before any fast + // path returns data-derived row ids/scores outside TableScan/TableRead. + let execution_table = self + .prepared_filter + .as_ref() + .map(PreparedVectorSearchFilter::table) + .unwrap_or(self.table); + let core = CoreOptions::new(execution_table.schema().options()); core.ensure_read_authorized()?; let vector_column = self.vector_column @@ -1270,20 +1343,7 @@ impl<'a> BatchVectorSearchBuilder<'a> { } } - // The data-evolution (global-index) fall-through path cannot honor a - // residual filter — it never reads physical rows. Rather than silently - // drop the predicate and return unfiltered results, fail loud when a - // filter is set on a batch that does not resolve to the primary-key - // vector path, mirroring the single-query builder. - if self.filter.is_some() { - return Err(crate::Error::DataInvalid { - message: "vector search filter is only supported on the primary-key vector path" - .to_string(), - source: None, - }); - } - - let vector_searches = query_vectors + let mut vector_searches = query_vectors .iter() .map(|vector| { VectorSearch::new(vector.clone(), limit, vector_column.to_string()) @@ -1291,11 +1351,19 @@ impl<'a> BatchVectorSearchBuilder<'a> { }) .collect::>>()?; - let snapshot_manager = self.table.snapshot_manager(); + if self + .prepared_filter + .as_ref() + .is_some_and(|prepared| prepared.include_row_ids().is_empty()) + { + return Ok(vec![SearchResult::empty(); vector_searches.len()]); + } + + let snapshot_manager = execution_table.snapshot_manager(); let setup = total_start.map_or(Duration::ZERO, |start| start.elapsed()); let snapshot_start = timing_enabled.then(Instant::now); - let snapshot = match crate::table::time_travel::resolve_snapshot(self.table).await? { + let snapshot = match crate::table::time_travel::resolve_snapshot(execution_table).await? { Some(s) => s, None => { let snapshot = snapshot_start.map_or(Duration::ZERO, |start| start.elapsed()); @@ -1317,12 +1385,42 @@ impl<'a> BatchVectorSearchBuilder<'a> { } }; let snapshot_elapsed = snapshot_start.map_or(Duration::ZERO, |start| start.elapsed()); + let pinned_table = match &self.prepared_filter { + Some(prepared) => prepared.table().clone(), + None => { + execution_table + .copy_with_resolved_snapshot(&snapshot) + .await? + } + }; + + if let Some(prepared) = &self.prepared_filter { + for search in &mut vector_searches { + search.set_shared_include_row_ids(Arc::clone(prepared.include_row_ids())); + } + } else if let Some(include_row_ids) = &self.include_row_ids { + if include_row_ids.is_empty() { + return Ok(vec![SearchResult::empty(); vector_searches.len()]); + } + for search in &mut vector_searches { + search.set_shared_include_row_ids(Arc::clone(include_row_ids)); + } + } else if let Some(filter) = &self.filter { + let include_row_ids = matching_row_ids_for_filter(&pinned_table, filter).await?; + if include_row_ids.is_empty() { + return Ok(vec![SearchResult::empty(); vector_searches.len()]); + } + let include_row_ids = Arc::new(include_row_ids); + for search in &mut vector_searches { + search.set_shared_include_row_ids(Arc::clone(&include_row_ids)); + } + } let manifest_start = timing_enabled.then(Instant::now); let index_entries = match snapshot.index_manifest() { Some(index_manifest_name) => { let manifest_path = snapshot_manager.manifest_path(index_manifest_name); - IndexManifest::read(self.table.file_io(), &manifest_path).await? + IndexManifest::read(execution_table.file_io(), &manifest_path).await? } None => Vec::new(), }; @@ -1331,11 +1429,11 @@ impl<'a> BatchVectorSearchBuilder<'a> { let evaluate_start = timing_enabled.then(Instant::now); let results = evaluate_batch_vector_search( VectorSearchEvaluation { - table: Some(self.table), - file_io: self.table.file_io(), - table_path: self.table.location(), - table_options: self.table.schema().options(), - schema_fields: self.table.schema().fields(), + table: Some(&pinned_table), + file_io: pinned_table.file_io(), + table_path: pinned_table.location(), + table_options: pinned_table.schema().options(), + schema_fields: pinned_table.schema().fields(), next_row_id: snapshot.next_row_id(), }, &index_entries, @@ -1522,6 +1620,84 @@ struct VectorSearchEvaluation<'a> { next_row_id: Option, } +async fn matching_row_ids_for_filter( + table: &Table, + filter: &Predicate, +) -> crate::Result { + let mut read_builder = table.new_read_builder(); + read_builder + .with_projection(&[ROW_ID_FIELD_NAME])? + .with_filter(filter.clone()); + let plan = read_builder.new_scan().plan().await?; + let read = read_builder.new_read()?; + let mut stream = read.to_arrow(plan.splits())?; + let mut row_ids = RoaringTreemap::new(); + while let Some(batch) = stream.try_next().await? { + let index = + batch + .schema() + .index_of(ROW_ID_FIELD_NAME) + .map_err(|_| crate::Error::DataInvalid { + message: format!( + "scalar vector pre-filter read is missing {ROW_ID_FIELD_NAME}" + ), + source: None, + })?; + let values = batch + .column(index) + .as_any() + .downcast_ref::() + .ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "scalar vector pre-filter {ROW_ID_FIELD_NAME} column is not Int64" + ), + source: None, + })?; + for row in 0..values.len() { + if values.is_null(row) { + return Err(crate::Error::DataInvalid { + message: format!( + "scalar vector pre-filter produced a null {ROW_ID_FIELD_NAME}" + ), + source: None, + }); + } + let row_id = values.value(row); + let row_id = u64::try_from(row_id).map_err(|_| crate::Error::DataInvalid { + message: format!( + "scalar vector pre-filter produced a negative {ROW_ID_FIELD_NAME}: {row_id}" + ), + source: None, + })?; + row_ids.insert(row_id); + } + } + Ok(row_ids) +} + +impl Table { + /// Resolve a scalar predicate once and pin all later vector-search/read + /// stages to the same snapshot. + pub async fn prepare_vector_search_filter( + &self, + filter: Predicate, + ) -> crate::Result { + CoreOptions::new(self.schema().options()).ensure_read_authorized()?; + let Some(snapshot) = crate::table::time_travel::resolve_snapshot(self).await? else { + return Ok(PreparedVectorSearchFilter { + table: self.clone(), + include_row_ids: Arc::new(RoaringTreemap::new()), + }); + }; + let table = self.copy_with_resolved_snapshot(&snapshot).await?; + let include_row_ids = matching_row_ids_for_filter(&table, &filter).await?; + Ok(PreparedVectorSearchFilter { + table, + include_row_ids: Arc::new(include_row_ids), + }) + } +} + #[derive(Default)] struct IndexSearchTiming { permit_wait: Duration, @@ -1634,6 +1810,40 @@ async fn evaluate_batch_vector_search( let index_search_limit = indexed_search_limit(max_limit, refine_factor)?; let vector_entry_count = vector_entries.len(); + let vector_search_plans = if let Some(include_row_ids) = + shared_batch_include_row_ids(vector_searches) + { + let ranges = vector_entries + .iter() + .map(|entry| { + let meta = entry.index_file.global_index_meta.as_ref().ok_or_else(|| { + crate::Error::DataInvalid { + message: format!( + "Vector index '{}' is missing global index metadata", + entry.index_file.file_name + ), + source: None, + } + })?; + Ok((meta.row_range_start, meta.row_range_end)) + }) + .collect::>>()?; + vector_entries + .iter() + .copied() + .zip(localize_shared_include_row_ids( + include_row_ids.as_ref(), + &ranges, + )?) + .filter_map(|(entry, local_filter)| local_filter.map(|filter| (entry, Some(filter)))) + .collect::>() + } else { + vector_entries + .iter() + .copied() + .map(|entry| (entry, None)) + .collect::>() + }; let mut permit_wait = Duration::ZERO; let mut file_reader_open = Duration::ZERO; let mut index_search = Duration::ZERO; @@ -1668,9 +1878,9 @@ async fn evaluate_batch_vector_search( Some(RangeReadLimiter::new(range_read_concurrency)), ) }; - let futures: Vec<_> = vector_entries + let futures: Vec<_> = vector_search_plans .into_iter() - .map(|entry| { + .map(|(entry, shared_local_filter)| { let range_read_limiter = range_read_limiter.clone(); let global_meta = entry.index_file.global_index_meta.as_ref().unwrap(); let backend = VectorIndexBackend::from_index_type(&entry.index_file.index_type) @@ -1696,6 +1906,37 @@ async fn evaluate_batch_vector_search( options.extend(search_options.clone()); let input = evaluation.file_io.new_input(&path); async move { + if let Some(local_filter) = shared_local_filter { + let local_filter = Arc::new(local_filter); + for vector_search in &mut vector_searches { + vector_search + .set_shared_include_row_ids(Arc::clone(&local_filter)); + } + } else { + for vector_search in &mut vector_searches { + if let Some(include_row_ids) = + vector_search.effective_include_row_ids() + { + vector_search.set_shared_include_row_ids(Arc::new( + localize_include_row_ids( + include_row_ids, + row_range_start, + row_range_end, + )?, + )); + } + } + } + if vector_searches.iter().all(|search| { + search + .effective_include_row_ids() + .is_some_and(|row_ids| row_ids.is_empty()) + }) { + return Ok(( + vec![SearchResult::empty(); vector_searches.len()], + IndexSearchTiming::default(), + )); + } let permit_start = timing_enabled.then(Instant::now); let permit = acquire_process_global_search_permit(concurrency).await?; let permit_wait = @@ -2569,7 +2810,7 @@ async fn maybe_rerank_indexed_batch_results( } let mut candidate_search = vector_search.clone(); - candidate_search.include_row_ids = Some(include_row_ids); + candidate_search.set_shared_include_row_ids(Arc::new(include_row_ids)); candidate_searches.push(candidate_search); candidate_results.push(candidates); } @@ -2653,6 +2894,113 @@ fn row_id_to_i64_for_range(row_id: u64) -> crate::Result { }) } +fn shared_batch_include_row_ids(vector_searches: &[VectorSearch]) -> Option<&Arc> { + let first = vector_searches.first()?.shared_include_row_ids.as_ref()?; + vector_searches + .iter() + .skip(1) + .all(|search| { + search + .shared_include_row_ids + .as_ref() + .is_some_and(|include_row_ids| Arc::ptr_eq(first, include_row_ids)) + }) + .then_some(first) +} + +fn prune_raw_ranges_by_include_row_ids( + raw_ranges: &[RowRange], + vector_searches: &[VectorSearch], +) -> crate::Result> { + if vector_searches + .iter() + .any(|search| search.effective_include_row_ids().is_none()) + { + return Ok(raw_ranges.to_vec()); + } + + let include_ranges = + if let Some(include_row_ids) = shared_batch_include_row_ids(vector_searches) { + sorted_row_ids_to_row_ranges(include_row_ids.iter())? + } else { + let mut union = RoaringTreemap::new(); + for include_row_ids in vector_searches + .iter() + .filter_map(VectorSearch::effective_include_row_ids) + { + for row_id in include_row_ids.iter() { + union.insert(row_id); + } + } + sorted_row_ids_to_row_ranges(union.iter())? + }; + Ok(intersect_sorted_ranges(raw_ranges, &include_ranges)) +} + +fn localize_include_row_ids( + include_row_ids: &RoaringTreemap, + row_range_start: i64, + row_range_end: i64, +) -> crate::Result { + let start = u64::try_from(row_range_start).map_err(|_| crate::Error::DataInvalid { + message: format!("Negative vector index row range start: {row_range_start}"), + source: None, + })?; + let end = u64::try_from(row_range_end).map_err(|_| crate::Error::DataInvalid { + message: format!("Negative vector index row range end: {row_range_end}"), + source: None, + })?; + let mut localized = RoaringTreemap::new(); + for row_id in include_row_ids.iter() { + if row_id >= start && row_id <= end { + localized.insert(row_id - start); + } + } + Ok(localized) +} + +fn localize_shared_include_row_ids( + include_row_ids: &RoaringTreemap, + ranges: &[(i64, i64)], +) -> crate::Result>> { + let mut validated_ranges = Vec::with_capacity(ranges.len()); + for (index, &(start, end)) in ranges.iter().enumerate() { + if start < 0 || end < start { + return Err(crate::Error::DataInvalid { + message: format!("Invalid vector index row range [{start}, {end}]"), + source: None, + }); + } + validated_ranges.push((start as u64, end as u64, index)); + } + validated_ranges.sort_unstable_by_key(|(start, _, _)| *start); + + let mut localized = (0..ranges.len()) + .map(|_| RoaringTreemap::new()) + .collect::>(); + let mut active = Vec::::new(); + let mut next_range = 0usize; + for row_id in include_row_ids.iter() { + while next_range < validated_ranges.len() && validated_ranges[next_range].0 <= row_id { + active.push(next_range); + next_range += 1; + } + active.retain(|range_index| validated_ranges[*range_index].1 >= row_id); + for range_index in &active { + let (start, _, original_index) = validated_ranges[*range_index]; + localized[original_index].insert(row_id - start); + } + if next_range == validated_ranges.len() && active.is_empty() { + break; + } + } + + Ok(localized + .into_iter() + .map(|filter| (!filter.is_empty()).then_some(filter)) + .collect()) +} + async fn detail_data_ranges_for_table(table: &Table) -> crate::Result> { let plan = table .new_read_builder() @@ -2937,6 +3285,10 @@ async fn read_raw_batch_vector_search( if raw_ranges.is_empty() { return Ok((vec![SearchResult::empty(); vector_searches.len()], None)); } + let raw_ranges = prune_raw_ranges_by_include_row_ids(raw_ranges, vector_searches)?; + if raw_ranges.is_empty() { + return Ok((vec![SearchResult::empty(); vector_searches.len()], None)); + } let field_name = &vector_searches[0].field_name; if vector_searches @@ -2954,7 +3306,7 @@ async fn read_raw_batch_vector_search( let mut read_builder = table.new_read_builder(); read_builder .with_projection(&[field_name.as_str(), ROW_ID_FIELD_NAME])? - .with_row_ranges(raw_ranges.to_vec()); + .with_row_ranges(raw_ranges); let plan = read_builder.new_scan().plan().await?; let plan_elapsed = plan_start.map_or(Duration::ZERO, |start| start.elapsed()); let split_count = plan.splits().len(); @@ -3023,15 +3375,22 @@ async fn read_raw_batch_vector_search( struct RawScoringPlan { all_query_indices: Vec, + shared_filter_groups: Vec, candidate_query_indices: HashMap>, query_l2_squared_norms: Vec, dense_query_dimension: Option, dense_query_matrix: Option>, } +struct SharedRawFilterGroup { + include_row_ids: Arc, + query_indices: Vec, +} + impl RawScoringPlan { fn new(vector_searches: &[VectorSearch], metric: RawVectorMetric) -> Self { let mut all_query_indices = Vec::new(); + let mut shared_filter_groups = Vec::new(); let mut candidate_query_indices: HashMap> = HashMap::new(); let query_l2_squared_norms = vector_searches .iter() @@ -3045,16 +3404,23 @@ impl RawScoringPlan { }) .collect(); - for (query_index, vector_search) in vector_searches.iter().enumerate() { - if let Some(include_row_ids) = &vector_search.include_row_ids { - for row_id in include_row_ids.iter() { - candidate_query_indices - .entry(row_id) - .or_default() - .push(query_index); + if let Some(include_row_ids) = shared_batch_include_row_ids(vector_searches) { + shared_filter_groups.push(SharedRawFilterGroup { + include_row_ids: Arc::clone(include_row_ids), + query_indices: (0..vector_searches.len()).collect(), + }); + } else { + for (query_index, vector_search) in vector_searches.iter().enumerate() { + if let Some(include_row_ids) = vector_search.effective_include_row_ids() { + for row_id in include_row_ids.iter() { + candidate_query_indices + .entry(row_id) + .or_default() + .push(query_index); + } + } else { + all_query_indices.push(query_index); } - } else { - all_query_indices.push(query_index); } } @@ -3077,6 +3443,7 @@ impl RawScoringPlan { Self { all_query_indices, + shared_filter_groups, candidate_query_indices, query_l2_squared_norms, dense_query_dimension, @@ -3335,6 +3702,20 @@ fn collect_raw_batch_vector_batch( )?; } } + for group in &scoring_plan.shared_filter_groups { + if group.include_row_ids.contains(row_id) { + for &query_index in &group.query_indices { + offer_raw_vector_score( + raw_row, + query_index, + metric, + vector_searches, + scoring_plan, + top_k_out, + )?; + } + } + } } if !dense_row_ids.is_empty() { @@ -3707,6 +4088,14 @@ mod tests { } fn vector_test_table() -> Table { + vector_test_table_at("memory:/vector_test") + } + + fn vector_test_table_at(location: &str) -> Table { + vector_test_table_with_file_io(FileIOBuilder::new("memory").build().unwrap(), location) + } + + fn vector_test_table_with_file_io(file_io: FileIO, location: &str) -> Table { let schema = Schema::builder() .column("id", DataType::Int(IntType::new())) .column( @@ -3716,9 +4105,9 @@ mod tests { .build() .unwrap(); Table::new( - FileIOBuilder::new("memory").build().unwrap(), + file_io, Identifier::new("default", "vector_test"), - "memory:/vector_test".to_string(), + location.to_string(), TableSchema::new(0, &schema), None, ) @@ -3747,6 +4136,101 @@ mod tests { assert_eq!(find_field_id_by_name(&fields, "nonexistent"), None); } + #[test] + fn shared_include_filter_is_localized_once_per_index_shard() { + let include_row_ids = RoaringTreemap::from_iter([101, 205, 999]); + let localized = localize_shared_include_row_ids( + &include_row_ids, + &[(100, 109), (200, 209), (300, 309)], + ) + .unwrap(); + + assert_eq!( + localized[0].as_ref().unwrap().iter().collect::>(), + vec![1] + ); + assert_eq!( + localized[1].as_ref().unwrap().iter().collect::>(), + vec![5] + ); + assert!(localized[2].is_none(), "an empty shard must be skipped"); + } + + #[test] + fn shared_batch_include_filter_requires_the_same_arc() { + let shared = Arc::new(RoaringTreemap::from_iter([1, 2, 3])); + let mut shared_searches = vec![ + VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap(), + VectorSearch::new(vec![0.0, 1.0], 2, "embedding".to_string()).unwrap(), + ]; + for search in &mut shared_searches { + search.set_shared_include_row_ids(Arc::clone(&shared)); + } + let detected = shared_batch_include_row_ids(&shared_searches).unwrap(); + assert!(Arc::ptr_eq(detected, &shared)); + + let mut equal_but_distinct = shared_searches.clone(); + equal_but_distinct[1] + .set_shared_include_row_ids(Arc::new(RoaringTreemap::from_iter([1, 2, 3]))); + assert!(shared_batch_include_row_ids(&equal_but_distinct).is_none()); + + let mut owned = shared_searches; + owned[1] = owned[1] + .clone() + .with_include_row_ids(RoaringTreemap::from_iter([1, 2, 3])); + assert!(shared_batch_include_row_ids(&owned).is_none()); + } + + #[test] + fn shared_raw_filter_does_not_expand_row_query_associations() { + let shared = Arc::new(RoaringTreemap::from_iter(0..1_000)); + let mut searches = (0..128) + .map(|_| VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap()) + .collect::>(); + for search in &mut searches { + search.set_shared_include_row_ids(Arc::clone(&shared)); + } + + let plan = RawScoringPlan::new(&searches, RawVectorMetric::L2); + let expanded_associations = plan + .candidate_query_indices + .values() + .map(Vec::len) + .sum::(); + + assert_eq!( + expanded_associations, 0, + "one shared bitmap must stay O(B + Q), not expand to O(B * Q)" + ); + assert_eq!(plan.shared_filter_groups.len(), 1); + assert!(Arc::ptr_eq( + &plan.shared_filter_groups[0].include_row_ids, + &shared + )); + assert_eq!(plan.shared_filter_groups[0].query_indices.len(), 128); + } + + #[test] + fn shared_raw_filter_prunes_unindexed_ranges_before_reading() { + let shared = Arc::new(RoaringTreemap::from_iter([7, 1_000, 1_001, 900_000])); + let mut searches = (0..128) + .map(|_| VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap()) + .collect::>(); + for search in &mut searches { + search.set_shared_include_row_ids(Arc::clone(&shared)); + } + + let raw_ranges = vec![RowRange::new(0, 999_999)]; + assert_eq!( + prune_raw_ranges_by_include_row_ids(&raw_ranges, &searches).unwrap(), + vec![ + RowRange::new(7, 7), + RowRange::new(1_000, 1_001), + RowRange::new(900_000, 900_000), + ] + ); + } + #[test] fn test_raw_vector_score_matches_java_metric_semantics() { let l2 = compute_raw_vector_score(&[1.0, 2.0], &[1.0, 4.0], RawVectorMetric::L2); @@ -4608,6 +5092,34 @@ mod tests { ); } + #[tokio::test] + async fn prepared_filter_cannot_bypass_builder_target_query_auth() { + let source = vector_test_table(); + let prepared = source + .prepare_vector_search_filter(id_gt_filter(&source, 0)) + .await + .unwrap(); + let target = source.copy_with_options(HashMap::from([( + "query-auth.enabled".to_string(), + "true".to_string(), + )])); + + let err = target + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0]]) + .with_limit(1) + .with_prepared_filter(prepared) + .execute() + .await + .expect_err("a stale prepared filter must not bypass current target authorization"); + + assert!( + matches!(err, crate::Error::Unsupported { ref message } if message.contains("query-auth.enabled")), + "builder target authorization must remain authoritative, got: {err:?}" + ); + } + fn pk_data_file(name: &str, row_count: i64, first_row_id: Option) -> DataFileMeta { DataFileMeta { file_name: name.to_string(), @@ -5529,6 +6041,14 @@ mod tests { .await .unwrap(); assert!(built > 0, "DE fixture must build a global vector index"); + let built = table + .new_sorted_global_index_build_builder() + .with_index_column("id") + .with_index_type("btree") + .execute() + .await + .unwrap(); + assert!(built > 0, "DE fixture must build a scalar BTree index"); table } @@ -6007,14 +6527,13 @@ mod tests { } #[tokio::test] - async fn execute_scored_filter_on_non_pk_vector_path_fails_loud() { - // No PK-vector index configured, so `execute_scored` would fall through to - // the data-evolution path, which never consumes the filter. Silently - // returning unfiltered rows is a wrong-read; the query must fail loud - // instead. + async fn execute_scored_filter_on_empty_de_path_returns_empty() { + // No PK-vector index and no snapshot: the request follows the + // data-evolution path. Scalar pre-filter support must not turn an empty + // table into an error. let table = pk_vector_table(&[]); let filter = id_gt_filter(&table, 2); - let err = table + let result = table .new_vector_search_builder() .with_vector_column("embedding") .with_query_vector(vec![1.0]) @@ -6022,13 +6541,8 @@ mod tests { .with_filter(filter) .execute_scored() .await - .map(|_| ()) - .expect_err("filter on the non-PK-vector path must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } - if message.contains("only supported on the primary-key vector path")), - "unexpected error: {err:?}" - ); + .expect("an empty data-evolution search with a filter should succeed"); + assert!(result.is_empty()); } #[tokio::test] @@ -6894,27 +7408,207 @@ mod tests { } #[tokio::test] - async fn de_execute_read_with_filter_fails_loud() { - // A filter on the data-evolution path is unsupported (the DE path never - // reads physical rows), so execute_read must fail loud rather than drop the - // predicate. The guard lives in execute_scored. + async fn resolved_vector_snapshot_can_be_reused_by_all_read_stages() { + let table = de_vector_table().await; + let snapshot = crate::table::time_travel::resolve_snapshot(&table) + .await + .unwrap() + .unwrap(); + let pinned = table.copy_with_resolved_snapshot(&snapshot).await.unwrap(); + + assert_eq!( + pinned.travel_snapshot().map(|snapshot| snapshot.id()), + Some(snapshot.id()) + ); + let options = CoreOptions::new(pinned.schema().options()); + let selector = options.try_time_travel_selector().unwrap().unwrap(); + assert!(matches!( + selector, + crate::spec::TimeTravelSelector::SnapshotId { + value, + option_name: crate::spec::SCAN_SNAPSHOT_ID_OPTION, + } if value == snapshot.id().to_string() + )); + } + + #[tokio::test] + async fn de_execute_read_applies_scalar_filter_before_top_k() { + // Row id=1 is the closest vector to [1, 0], but the scalar filter excludes + // it. Filter-before-Top-K must return the best rows among ids > 1 instead + // of recalling id=1 first and filtering it after the search. let table = de_vector_table().await; let filter = id_gt_filter(&table, 1); - let err = table + let mut stream = table .new_vector_search_builder() .with_vector_column("embedding") .with_query_vector(vec![1.0, 0.0]) - .with_limit(3) + .with_limit(2) .with_filter(filter) .execute_read() .await - .map(|_| ()) - .expect_err("DE read with a filter must fail loud"); + .expect("DE vector search should support a scalar pre-filter"); + + let mut ids = Vec::new(); + while let Some(batch) = stream.try_next().await.unwrap() { + let id = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + ids.extend((0..id.len()).map(|row| id.value(row))); + } + + assert_eq!(ids, vec![3, 2]); + } + + #[tokio::test] + async fn de_scalar_filter_with_no_matching_rows_returns_empty() { + let table = de_vector_table().await; + let filter = id_gt_filter(&table, 99); + + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(2) + .with_filter(filter.clone()) + .execute_scored() + .await + .unwrap(); + assert!(result.is_empty()); + + let results = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0], vec![0.0, 1.0]]) + .with_limit(2) + .with_filter(filter) + .execute() + .await + .unwrap(); + assert_eq!(results.len(), 2); + assert!(results.iter().all(SearchResult::is_empty)); + } + + #[tokio::test] + async fn prepared_de_scalar_filter_can_be_reused_by_batch_search() { + let table = de_vector_table().await; + let prepared = table + .prepare_vector_search_filter(id_gt_filter(&table, 1)) + .await + .unwrap(); + let results = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0], vec![0.0, 1.0]]) + .with_limit(2) + .with_prepared_filter(prepared) + .execute() + .await + .unwrap(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].row_ids, vec![2, 1]); + assert_eq!(results[1].row_ids, vec![1, 2]); + } + + #[tokio::test] + async fn prepared_filter_from_different_table_is_rejected() { + let prepared = PreparedVectorSearchFilter { + table: vector_test_table_at("memory:/prepared_filter_source"), + include_row_ids: Arc::new(RoaringTreemap::from_iter([1])), + }; + let target = vector_test_table_at("memory:/prepared_filter_target"); + + let error = target + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0]]) + .with_limit(1) + .with_prepared_filter(prepared) + .execute() + .await + .expect_err("a prepared filter must not retarget the builder to another table"); + assert!( - matches!(err, crate::Error::DataInvalid { .. }), - "unexpected error: {err:?}" + error.to_string().contains("different table"), + "unexpected error: {error}" + ); + } + + #[tokio::test] + async fn prepared_filter_from_same_location_but_different_file_io_is_rejected() { + let location = "memory:/prepared_filter_shared_location"; + let source = + vector_test_table_with_file_io(FileIOBuilder::new("memory").build().unwrap(), location); + let prepared = source + .prepare_vector_search_filter(id_gt_filter(&source, 0)) + .await + .unwrap(); + let target = + vector_test_table_with_file_io(FileIOBuilder::new("memory").build().unwrap(), location); + + let error = target + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0]]) + .with_limit(1) + .with_prepared_filter(prepared) + .execute() + .await + .expect_err("a prepared filter must stay bound to its FileIO lineage"); + + assert!( + error.to_string().contains("different table"), + "unexpected error: {error}" ); } + + #[tokio::test] + async fn de_scalar_filter_applies_to_unindexed_raw_fallback() { + let table = de_vector_table().await; + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut vector_builder = + ListBuilder::new(Float32Builder::new()).with_field(element_field.clone()); + vector_builder.values().append_value(1.0); + vector_builder.values().append_value(0.0); + vector_builder.append(true); + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("embedding", ArrowDataType::List(element_field), true), + ])), + vec![ + Arc::new(Int32Array::from(vec![4])) as ArrayRef, + Arc::new(vector_builder.finish()) as ArrayRef, + ], + ) + .unwrap(); + let mut writer = TableWrite::new(&table, "test-user".to_string()).unwrap(); + writer.write_arrow_batch(&batch).await.unwrap(); + let messages = writer.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let table = table.copy_with_options(HashMap::from([ + ("vector-index.search-mode".to_string(), "full".to_string()), + ("scalar-index.search-mode".to_string(), "full".to_string()), + ])); + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(1) + .with_filter(id_gt_filter(&table, 3)) + .execute_scored() + .await + .unwrap(); + + assert_eq!(result.row_ids, vec![3]); + } } /// Tests for [`residual_positions_by_file`]: the residual predicate is applied at diff --git a/crates/paimon/src/vector_search.rs b/crates/paimon/src/vector_search.rs index 1625243e6..e1a3de10d 100644 --- a/crates/paimon/src/vector_search.rs +++ b/crates/paimon/src/vector_search.rs @@ -17,6 +17,7 @@ use std::cmp::Ordering; use std::collections::{BinaryHeap, HashMap}; +use std::sync::Arc; #[derive(Clone)] pub struct VectorSearch { @@ -25,6 +26,7 @@ pub struct VectorSearch { pub field_name: String, pub options: HashMap, pub include_row_ids: Option, + pub(crate) shared_include_row_ids: Option>, } impl VectorSearch { @@ -53,6 +55,7 @@ impl VectorSearch { field_name, options: HashMap::new(), include_row_ids: None, + shared_include_row_ids: None, }) } @@ -63,8 +66,23 @@ impl VectorSearch { pub fn with_include_row_ids(mut self, include_row_ids: roaring::RoaringTreemap) -> Self { self.include_row_ids = Some(include_row_ids); + self.shared_include_row_ids = None; self } + + pub(crate) fn set_shared_include_row_ids( + &mut self, + include_row_ids: Arc, + ) { + self.include_row_ids = None; + self.shared_include_row_ids = Some(include_row_ids); + } + + pub(crate) fn effective_include_row_ids(&self) -> Option<&roaring::RoaringTreemap> { + self.shared_include_row_ids + .as_deref() + .or(self.include_row_ids.as_ref()) + } } impl std::fmt::Display for VectorSearch { diff --git a/crates/paimon/src/vindex/reader.rs b/crates/paimon/src/vindex/reader.rs index 8f8eab35f..418b43013 100644 --- a/crates/paimon/src/vindex/reader.rs +++ b/crates/paimon/src/vindex/reader.rs @@ -26,7 +26,7 @@ use paimon_vindex_core::io::{ReadRequest, SeekRead, SeekReadCapabilities}; use std::collections::BinaryHeap; use std::collections::HashMap; use std::io; -use std::sync::{Condvar, Mutex}; +use std::sync::{Arc, Condvar, Mutex}; use std::time::{Duration, Instant}; const DEFAULT_NPROBE: usize = 16; @@ -421,16 +421,39 @@ fn search_vindex( Ok(Some(id_to_scores)) } -#[derive(Clone, PartialEq, Eq)] +#[derive(Clone)] struct PreparedSearch { params: VectorSearchParams, - filter_bytes: Option>, + filter_bytes: Option>, } +impl PreparedSearch { + fn same_batch_group(&self, other: &Self) -> bool { + self.params == other.params + && match (&self.filter_bytes, &other.filter_bytes) { + (Some(left), Some(right)) => { + Arc::ptr_eq(left, right) || left.as_ref() == right.as_ref() + } + (None, None) => true, + _ => false, + } + } +} + +#[cfg(test)] fn prepare_search( metadata: &VectorIndexMetadata, options: &HashMap, vector_search: &VectorSearch, +) -> crate::Result> { + prepare_search_with_shared_filter(metadata, options, vector_search, None) +} + +fn prepare_search_with_shared_filter( + metadata: &VectorIndexMetadata, + options: &HashMap, + vector_search: &VectorSearch, + shared_filter_bytes: Option<&Arc<[u8]>>, ) -> crate::Result> { if vector_search.vector.len() != metadata.dimension { return Err(crate::Error::DataInvalid { @@ -472,19 +495,15 @@ fn prepare_search( ), }; - let filter_bytes = if let Some(include_ids) = &vector_search.include_row_ids { + let filter_bytes = if let Some(include_ids) = vector_search.effective_include_row_ids() { if include_ids.is_empty() { return Ok(None); } top_k = top_k.min(include_ids.len() as usize); - let mut bytes = Vec::new(); - include_ids - .serialize_into(&mut bytes) - .map_err(|e| crate::Error::DataInvalid { - message: format!("Failed to serialize vector search row-id filter: {}", e), - source: Some(Box::new(e)), - })?; - Some(bytes) + Some(match shared_filter_bytes { + Some(filter_bytes) => Arc::clone(filter_bytes), + None => serialize_row_id_filter(include_ids)?, + }) } else { None }; @@ -496,6 +515,55 @@ fn prepare_search( })) } +fn serialize_row_id_filter(include_ids: &roaring::RoaringTreemap) -> crate::Result> { + let mut bytes = Vec::new(); + include_ids + .serialize_into(&mut bytes) + .map_err(|e| crate::Error::DataInvalid { + message: format!("Failed to serialize vector search row-id filter: {}", e), + source: Some(Box::new(e)), + })?; + Ok(Arc::from(bytes)) +} + +fn shared_batch_include_row_ids( + vector_searches: &[VectorSearch], +) -> Option<&Arc> { + let first = vector_searches.first()?.shared_include_row_ids.as_ref()?; + vector_searches + .iter() + .skip(1) + .all(|search| { + search + .shared_include_row_ids + .as_ref() + .is_some_and(|include_row_ids| Arc::ptr_eq(first, include_row_ids)) + }) + .then_some(first) +} + +fn prepare_batch_searches( + metadata: &VectorIndexMetadata, + options: &HashMap, + vector_searches: &[VectorSearch], +) -> crate::Result>> { + let shared_filter_bytes = shared_batch_include_row_ids(vector_searches) + .filter(|include_row_ids| !include_row_ids.is_empty()) + .map(|include_row_ids| serialize_row_id_filter(include_row_ids)) + .transpose()?; + vector_searches + .iter() + .map(|search| { + prepare_search_with_shared_filter( + metadata, + options, + search, + shared_filter_bytes.as_ref(), + ) + }) + .collect() +} + fn execute_scalar_search( reader: &mut VIndexReader, vector_search: &VectorSearch, @@ -537,11 +605,17 @@ fn search_batch_vindex( ..VindexBatchStats::default() }); - for (index, search) in vector_searches.iter().enumerate() { - let Some(prepared) = prepare_search(metadata, options, search)? else { + for (index, prepared) in prepare_batch_searches(metadata, options, vector_searches)? + .into_iter() + .enumerate() + { + let Some(prepared) = prepared else { continue; }; - if let Some((_, indices)) = groups.iter_mut().find(|(key, _)| key == &prepared) { + if let Some((_, indices)) = groups + .iter_mut() + .find(|(key, _)| key.same_batch_group(&prepared)) + { indices.push(index); } else { groups.push((prepared, vec![index])); @@ -1032,7 +1106,7 @@ mod tests { }; let prepared = PreparedSearch { params: VectorSearchParams::new(10, 16), - filter_bytes: Some(vec![0; 128]), + filter_bytes: Some(Arc::from(vec![0; 128])), }; let chunk_size = native_batch_chunk_size(&metadata, &prepared, 1); let full_chunk = native_batch_chunk_working_set_bytes(&metadata, &prepared, chunk_size); @@ -1043,6 +1117,38 @@ mod tests { assert!(final_chunk < full_chunk); } + #[test] + fn batch_preparation_serializes_shared_filter_once() { + let metadata = VectorIndexMetadata { + index_type: paimon_vindex_core::index::IndexType::IvfFlat, + dimension: TEST_DIMENSION, + nlist: 16, + metric: MetricType::L2, + total_vectors: 1_000_000, + pq_m: None, + pq_bits: None, + rq_bits: None, + diskann: None, + }; + let shared_filter = Arc::new(roaring::RoaringTreemap::from_iter(0..100_000)); + let mut searches = vec![query(); 128]; + for search in &mut searches { + search.set_shared_include_row_ids(Arc::clone(&shared_filter)); + } + + let prepared = prepare_batch_searches(&metadata, &HashMap::new(), &searches).unwrap(); + let first = prepared[0] + .as_ref() + .and_then(|search| search.filter_bytes.as_ref()) + .expect("shared filter should be serialized"); + assert!(prepared.iter().all(|search| { + search + .as_ref() + .and_then(|search| search.filter_bytes.as_ref()) + .is_some_and(|filter| Arc::ptr_eq(first, filter)) + })); + } + #[test] fn native_batch_memory_pool_admits_only_available_bytes() { let pool = NativeBatchMemoryPool::new(64); @@ -1149,7 +1255,7 @@ mod tests { ) .unwrap() .unwrap(); - assert!(automatic != explicit); + assert_ne!(automatic.params, explicit.params); let explicit_params = explicit.params; assert_eq!( explicit_params.search_width, diff --git a/crates/paimon/tests/pk_vector_batch_test.rs b/crates/paimon/tests/pk_vector_batch_test.rs index 87e4ca59d..65e16c38b 100644 --- a/crates/paimon/tests/pk_vector_batch_test.rs +++ b/crates/paimon/tests/pk_vector_batch_test.rs @@ -730,14 +730,13 @@ async fn empty_snapshot_still_rejects_zero_limit() { ); } -/// A filter set on a batch `execute()` (the scored / data-evolution path) must -/// fail loud rather than silently drop the predicate: that path never reads -/// physical rows, so it cannot honor a residual filter. Mirrors the single-query -/// `execute_scored` guard. +/// A filter set on a batch `execute()` (the scored / data-evolution path) is +/// accepted even when the snapshot is empty. The scalar pre-filter is evaluated +/// before vector Top-K, and batch result arity is preserved when no rows match. // Gated off Windows for the same `file://` tempdir reason as `pk_vector_baseline_test`. #[cfg(not(windows))] #[tokio::test] -async fn batch_execute_with_filter_on_non_pk_vector_table_fails_loud() { +async fn batch_execute_with_filter_on_empty_non_pk_vector_table_returns_empty() { let tmp = tempfile::tempdir().expect("create temp dir"); let location = format!("file://{}", tmp.path().display()); let file_io = FileIOBuilder::new("file").build().unwrap(); @@ -770,18 +769,16 @@ async fn batch_execute_with_filter_on_non_pk_vector_table_fails_loud() { .greater_or_equal("id", Datum::Int(1)) .expect("build filter on id"); + let queries = vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]]; let mut batch = table.new_batch_vector_search_builder(); - let err = batch + let results = batch .with_vector_column(VECTOR_COLUMN) - .with_query_vectors(vec![vec![1.0, 0.0, 0.0, 0.0]]) + .with_query_vectors(queries.clone()) .with_limit(3) .with_filter(filter) .execute() .await - .expect_err("a filter on the data-evolution batch path must fail loud"); - assert!( - err.to_string() - .contains("only supported on the primary-key vector path"), - "expected a filter-unsupported error, got: {err}" - ); + .expect("the data-evolution batch path must accept scalar pre-filters"); + assert_eq!(results.len(), queries.len()); + assert!(results.iter().all(|result| result.is_empty())); } diff --git a/docs/src/sql.md b/docs/src/sql.md index bac91b7d8..5919c4b20 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -1339,6 +1339,32 @@ The function performs ANN search across all matching vector index files for the target column, merges results, and returns the top-k rows ordered by relevance score. If no matching index is found, an empty result is returned. +### Scalar Pre-Filters + +Add a `WHERE` clause to restrict the rows considered by vector Top-K. The +predicate is evaluated before the vector index selects its nearest neighbors: + +```sql +SELECT id, event_time +FROM vector_search( + 'paimon.my_db.items', + 'embedding', + '[1.0, 0.0, 0.0, 0.0]', + 10 +) +WHERE event_time >= TIMESTAMP '2026-08-01 00:00:00'; +``` + +On data-evolution tables, Paimon resolves the predicate to matching global row +IDs using a snapshot-pinned table read. Scalar global indexes such as BTree can +narrow this read. The matching global row IDs are intersected with each vector +index shard and passed to the vector backend as its row filter, so an excluded +nearest neighbor does not consume one of the requested Top-K positions. + +Only predicates that can be translated completely to Paimon predicates are +pushed into vector search. DataFusion keeps its residual filter for ordinary +`vector_search` queries as an additional correctness check. + ### Refine / Rerank Vector index search can optionally refine ANN results by reading the raw vectors @@ -1421,6 +1447,28 @@ ORDER BY query_id, result_id; The query-vector column must have Arrow type `List` or `FixedSizeList`. Null query-vector rows produce no joined results, and null elements inside a vector are rejected. The lateral form returns the left row joined with the top-k matching rows from the target Paimon table for that row's query vector. +Fully translatable target-table predicates are also applied before each lateral +Top-K: + +```sql +SELECT q.id AS query_id, r.id AS result_id +FROM paimon.my_db.queries q +CROSS JOIN LATERAL vector_search( + 'paimon.my_db.items', + 'embedding', + q.embedding, + 10 +) AS r +WHERE r.event_time >= TIMESTAMP '2026-08-01 00:00:00' +ORDER BY query_id, result_id; +``` + +For conjunctions, target-only predicates such as `r.event_time >= ...` are +pushed into vector search. Predicates that reference the left relation or both +sides remain normal join-result filters. Unsupported or inexact target +predicates also remain post-Top-K residual filters, so they may return fewer +than the requested number of rows. + ### Supported Metrics The distance metric is configured at index creation time via table options: