From 350f03b654a08e9fca27b47f9dbf9644a9b89353 Mon Sep 17 00:00:00 2001 From: shyjsarah <44659226+shyjsarah@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:16:11 -0700 Subject: [PATCH 1/4] feat(table): support scalar pre-filter for data-evolution vector search --- crates/paimon/src/table/mod.rs | 53 +- .../paimon/src/table/vector_search_builder.rs | 587 +++++++++++++++--- crates/paimon/src/vector_search.rs | 7 +- 3 files changed, 551 insertions(+), 96 deletions(-) diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 35f1e45a0..1859b7bb0 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -153,13 +153,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. @@ -424,6 +429,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..abbcb80ed 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -196,6 +196,27 @@ pub struct BatchVectorSearchBuilder<'a> { options: HashMap, projection: Option>, filter: Option, + include_row_ids: 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 + } } /// The primary-key vector route's search output plus the source context a later @@ -249,16 +270,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 +344,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 +405,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 @@ -407,19 +416,18 @@ impl<'a> VectorSearchBuilder<'a> { /// 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 { + 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 +457,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 +493,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 +1192,7 @@ impl<'a> BatchVectorSearchBuilder<'a> { options: HashMap::new(), projection: None, filter: None, + include_row_ids: None, } } @@ -1194,14 +1216,32 @@ 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 + } + + /// Reuse row IDs from a previously prepared scalar pre-filter. + /// + /// The builder's table must be [`PreparedVectorSearchFilter::table`] (or an + /// equivalent copy pinned to the same snapshot). + 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 + } + + /// Reuse a shared row-ID allow-list without copying its bitmap. + pub fn with_shared_include_row_ids( + &mut self, + include_row_ids: Arc, + ) -> &mut Self { + self.include_row_ids = Some(include_row_ids); + self.filter = None; self } @@ -1270,20 +1310,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()) @@ -1317,6 +1344,25 @@ impl<'a> BatchVectorSearchBuilder<'a> { } }; let snapshot_elapsed = snapshot_start.map_or(Duration::ZERO, |start| start.elapsed()); + let pinned_table = self.table.copy_with_resolved_snapshot(&snapshot).await?; + + 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.include_row_ids = Some(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.include_row_ids = Some(Arc::clone(&include_row_ids)); + } + } let manifest_start = timing_enabled.then(Instant::now); let index_entries = match snapshot.index_manifest() { @@ -1331,11 +1377,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 +1568,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 +1758,44 @@ 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 shared_include_row_ids = + vector_searches[0] + .include_row_ids + .as_ref() + .filter(|include_row_ids| { + vector_searches + .iter() + .all(|search| search.include_row_ids.as_ref() == Some(*include_row_ids)) + }); + let vector_search_plans = if let Some(include_row_ids) = shared_include_row_ids { + 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, &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 +1830,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 +1858,34 @@ 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.include_row_ids = Some(Arc::clone(&local_filter)); + } + } else { + for vector_search in &mut vector_searches { + if let Some(include_row_ids) = vector_search.include_row_ids.as_ref() { + vector_search.include_row_ids = + Some(Arc::new(localize_include_row_ids( + include_row_ids, + row_range_start, + row_range_end, + )?)); + } + } + } + if vector_searches.iter().all(|search| { + search + .include_row_ids + .as_ref() + .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 +2759,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.include_row_ids = Some(Arc::new(include_row_ids)); candidate_searches.push(candidate_search); candidate_results.push(candidates); } @@ -2653,6 +2843,70 @@ fn row_id_to_i64_for_range(row_id: u64) -> crate::Result { }) } +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() @@ -3747,6 +4001,26 @@ 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 test_raw_vector_score_matches_java_metric_semantics() { let l2 = compute_raw_vector_score(&[1.0, 2.0], &[1.0, 4.0], RawVectorMetric::L2); @@ -5529,6 +5803,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 +6289,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 +6303,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,26 +7170,155 @@ 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"); - assert!( - matches!(err, crate::Error::DataInvalid { .. }), - "unexpected error: {err:?}" - ); + .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 = prepared + .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_shared_include_row_ids(Arc::clone(prepared.include_row_ids())) + .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 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]); } } diff --git a/crates/paimon/src/vector_search.rs b/crates/paimon/src/vector_search.rs index 1625243e6..0eda3afad 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 { @@ -24,7 +25,7 @@ pub struct VectorSearch { pub limit: usize, pub field_name: String, pub options: HashMap, - pub include_row_ids: Option, + pub include_row_ids: Option>, } impl VectorSearch { @@ -62,7 +63,7 @@ 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.include_row_ids = Some(Arc::new(include_row_ids)); self } } @@ -335,7 +336,7 @@ mod tests { assert_eq!(cloned.limit, vector_search.limit); assert_eq!(cloned.field_name, vector_search.field_name); assert_eq!(cloned.options, vector_search.options); - assert_eq!(cloned.include_row_ids.as_ref(), Some(&include_row_ids)); + assert_eq!(cloned.include_row_ids.as_deref(), Some(&include_row_ids)); } #[test] From b545eef8b4632ba417d3d782dd6e62c5d556d6a3 Mon Sep 17 00:00:00 2001 From: shyjsarah <44659226+shyjsarah@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:16:12 -0700 Subject: [PATCH 2/4] feat(datafusion): push scalar filters into vector search --- .../datafusion/src/filter_pushdown.rs | 35 ++++ .../datafusion/src/lateral_vector_search.rs | 166 ++++++++++++++++-- .../datafusion/src/vector_search.rs | 84 +++++++-- .../datafusion/tests/read_tables.rs | 93 ++++++++++ docs/src/sql.md | 48 +++++ 5 files changed, 393 insertions(+), 33 deletions(-) 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..96abad545 100644 --- a/crates/integrations/datafusion/src/lateral_vector_search.rs +++ b/crates/integrations/datafusion/src/lateral_vector_search.rs @@ -37,7 +37,10 @@ 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}; @@ -49,11 +52,13 @@ use datafusion::physical_plan::{ 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 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 +119,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 +264,7 @@ pub(crate) struct LateralVectorSearchNode { query_vector_expr: Expr, limit: usize, schema: DFSchemaRef, + filter: Option, } impl LateralVectorSearchNode { @@ -201,9 +285,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 +344,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 +365,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 +376,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 +386,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 +425,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 +433,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,6 +450,8 @@ struct LateralVectorSearchExec { query_vector_expr: Arc, limit: usize, output_schema: ArrowSchemaRef, + filter: Option, + prepared_filter: Arc>, plan_properties: Arc, } @@ -380,10 +480,17 @@ impl LateralVectorSearchExec { query_vector_expr, limit, output_schema, + filter: None, + prepared_filter: Arc::new(OnceCell::new()), plan_properties, } } + fn with_filter(mut self, filter: Predicate) -> Self { + self.filter = Some(filter); + self + } + async fn process_batch(&self, batch: RecordBatch) -> DFResult { if batch.num_rows() == 0 { return empty_batch(self.output_schema.clone()); @@ -398,17 +505,35 @@ impl LateralVectorSearchExec { return empty_batch(self.output_schema.clone()); } - let mut builder = self.target_table.new_batch_vector_search_builder(); - let results = builder + let (target_table, include_row_ids) = match &self.filter { + Some(filter) => { + let prepared = self + .prepared_filter + .get_or_try_init(|| { + self.target_table + .prepare_vector_search_filter(filter.clone()) + }) + .await + .map_err(to_datafusion_error)?; + ( + prepared.table(), + Some(Arc::clone(prepared.include_row_ids())), + ) + } + None => (&self.target_table, None), + }; + 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(include_row_ids) = include_row_ids { + builder.with_shared_include_row_ids(include_row_ids); + } + 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 +577,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 +603,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 +611,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 = Arc::clone(&self.prepared_filter); + Ok(Arc::new(exec)) } fn execute( diff --git a/crates/integrations/datafusion/src/vector_search.rs b/crates/integrations/datafusion/src/vector_search.rs index 5a38942f7..4b5d78cab 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 = search_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_shared_include_row_ids(Arc::clone(prepared.include_row_ids())); + } + 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..baed71741 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -2279,6 +2279,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 +2314,63 @@ 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" + ); } // Manual run with a local Lumina native library: @@ -2450,6 +2518,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 +2555,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/docs/src/sql.md b/docs/src/sql.md index 34a033e96..f524cfc4e 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -1318,6 +1318,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 @@ -1400,6 +1426,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: From 6c951a8cf00a8770ba1e9cd03671b7c85747f9c2 Mon Sep 17 00:00:00 2001 From: shyjsarah <44659226+shyjsarah@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:19:40 -0700 Subject: [PATCH 3/4] perf(vector-search): share scalar filter bitmaps --- crates/paimon/src/lumina/reader.rs | 6 ++-- .../paimon/src/table/vector_search_builder.rs | 30 ++++++++++--------- crates/paimon/src/vector_search.rs | 23 ++++++++++++-- crates/paimon/src/vindex/reader.rs | 2 +- 4 files changed, 40 insertions(+), 21 deletions(-) diff --git a/crates/paimon/src/lumina/reader.rs b/crates/paimon/src/lumina/reader.rs index 25514a67f..a77a5006f 100644 --- a/crates/paimon/src/lumina/reader.rs +++ b/crates/paimon/src/lumina/reader.rs @@ -311,9 +311,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); @@ -369,7 +369,7 @@ fn search_lumina_batch( } if vector_searches .iter() - .any(|vector_search| vector_search.include_row_ids.is_some()) + .any(|vector_search| vector_search.effective_include_row_ids().is_some()) { return vector_searches .iter() diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index abbcb80ed..d68ea1c00 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -1351,7 +1351,7 @@ impl<'a> BatchVectorSearchBuilder<'a> { return Ok(vec![SearchResult::empty(); vector_searches.len()]); } for search in &mut vector_searches { - search.include_row_ids = Some(Arc::clone(include_row_ids)); + 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?; @@ -1360,7 +1360,7 @@ impl<'a> BatchVectorSearchBuilder<'a> { } let include_row_ids = Arc::new(include_row_ids); for search in &mut vector_searches { - search.include_row_ids = Some(Arc::clone(&include_row_ids)); + search.set_shared_include_row_ids(Arc::clone(&include_row_ids)); } } @@ -1760,12 +1760,11 @@ async fn evaluate_batch_vector_search( let vector_entry_count = vector_entries.len(); let shared_include_row_ids = vector_searches[0] - .include_row_ids - .as_ref() + .effective_include_row_ids() .filter(|include_row_ids| { vector_searches .iter() - .all(|search| search.include_row_ids.as_ref() == Some(*include_row_ids)) + .all(|search| search.effective_include_row_ids() == Some(*include_row_ids)) }); let vector_search_plans = if let Some(include_row_ids) = shared_include_row_ids { let ranges = vector_entries @@ -1861,24 +1860,27 @@ async fn evaluate_batch_vector_search( if let Some(local_filter) = shared_local_filter { let local_filter = Arc::new(local_filter); for vector_search in &mut vector_searches { - vector_search.include_row_ids = Some(Arc::clone(&local_filter)); + 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.include_row_ids.as_ref() { - vector_search.include_row_ids = - Some(Arc::new(localize_include_row_ids( + 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 - .include_row_ids - .as_ref() + .effective_include_row_ids() .is_some_and(|row_ids| row_ids.is_empty()) }) { return Ok(( @@ -2759,7 +2761,7 @@ async fn maybe_rerank_indexed_batch_results( } let mut candidate_search = vector_search.clone(); - candidate_search.include_row_ids = Some(Arc::new(include_row_ids)); + candidate_search.set_shared_include_row_ids(Arc::new(include_row_ids)); candidate_searches.push(candidate_search); candidate_results.push(candidates); } @@ -3300,7 +3302,7 @@ impl RawScoringPlan { .collect(); for (query_index, vector_search) in vector_searches.iter().enumerate() { - if let Some(include_row_ids) = &vector_search.include_row_ids { + 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) diff --git a/crates/paimon/src/vector_search.rs b/crates/paimon/src/vector_search.rs index 0eda3afad..e1a3de10d 100644 --- a/crates/paimon/src/vector_search.rs +++ b/crates/paimon/src/vector_search.rs @@ -25,7 +25,8 @@ pub struct VectorSearch { pub limit: usize, pub field_name: String, pub options: HashMap, - pub include_row_ids: Option>, + pub include_row_ids: Option, + pub(crate) shared_include_row_ids: Option>, } impl VectorSearch { @@ -54,6 +55,7 @@ impl VectorSearch { field_name, options: HashMap::new(), include_row_ids: None, + shared_include_row_ids: None, }) } @@ -63,9 +65,24 @@ impl VectorSearch { } pub fn with_include_row_ids(mut self, include_row_ids: roaring::RoaringTreemap) -> Self { - self.include_row_ids = Some(Arc::new(include_row_ids)); + 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 { @@ -336,7 +353,7 @@ mod tests { assert_eq!(cloned.limit, vector_search.limit); assert_eq!(cloned.field_name, vector_search.field_name); assert_eq!(cloned.options, vector_search.options); - assert_eq!(cloned.include_row_ids.as_deref(), Some(&include_row_ids)); + assert_eq!(cloned.include_row_ids.as_ref(), Some(&include_row_ids)); } #[test] diff --git a/crates/paimon/src/vindex/reader.rs b/crates/paimon/src/vindex/reader.rs index 8f8eab35f..2fa1dc60e 100644 --- a/crates/paimon/src/vindex/reader.rs +++ b/crates/paimon/src/vindex/reader.rs @@ -472,7 +472,7 @@ 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); } From 1a5bc2d0f6bcc0a1126a2b36cc6dafb6d16dbb03 Mon Sep 17 00:00:00 2001 From: shyjsarah <44659226+shyjsarah@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:58:48 -0700 Subject: [PATCH 4/4] test(vector-search): update data-evolution filter expectation --- .../paimon/src/table/vector_search_builder.rs | 4 ++-- crates/paimon/tests/pk_vector_batch_test.rs | 23 ++++++++----------- 2 files changed, 12 insertions(+), 15 deletions(-) diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index d68ea1c00..19fcede86 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -414,8 +414,8 @@ 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`. + /// 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, 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())); }