diff --git a/bindings/c/src/tests.rs b/bindings/c/src/tests.rs index 31b62de27..45899ab4e 100644 --- a/bindings/c/src/tests.rs +++ b/bindings/c/src/tests.rs @@ -3175,6 +3175,89 @@ fn vector_search_pk_table_read_matches_rust() { } } +/// The bucket-split terminal marshals an array of buffers, which the single-split +/// terminal does not: a caller passing a null array, a zero count, or a null +/// entry has made an input error, and it must be reported as one rather than +/// reaching the decoder as corrupt data. +#[test] +fn vector_search_bucket_splits_reject_malformed_input() { + let path = "memory:/vsearch_pk_split_args"; + let (query, vectors) = pk_fixture_smoke(); + let table = build_pk_vector_table(path, &vectors); + let handle = unsafe { wrap_table(table) }; + + unsafe { + // No splits at all. + let builder = c_vector_builder(handle, VECTOR_COLUMN, &query, 3, ptr::null_mut()); + let result = paimon_vector_search_builder_execute_read_for_bucket_splits( + builder, + ptr::null(), + ptr::null(), + 0, + ); + paimon_vector_search_builder_free(builder); + assert!(!result.error.is_null(), "a null split array must error"); + paimon_error_free(result.error); + + // A count that does not match the (absent) arrays. + let builder = c_vector_builder(handle, VECTOR_COLUMN, &query, 3, ptr::null_mut()); + let result = paimon_vector_search_builder_execute_read_for_bucket_splits( + builder, + ptr::null(), + ptr::null(), + 1, + ); + paimon_vector_search_builder_free(builder); + assert!(!result.error.is_null(), "a null split array must error"); + paimon_error_free(result.error); + + // A null entry inside an otherwise valid array. + let bytes: Vec = vec![1, 2, 3, 4]; + let ptrs: [*const u8; 2] = [bytes.as_ptr(), ptr::null()]; + let lens: [usize; 2] = [bytes.len(), 0]; + let builder = c_vector_builder(handle, VECTOR_COLUMN, &query, 3, ptr::null_mut()); + let result = paimon_vector_search_builder_execute_read_for_bucket_splits( + builder, + ptrs.as_ptr(), + lens.as_ptr(), + 2, + ); + paimon_vector_search_builder_free(builder); + assert!(!result.error.is_null(), "a null split entry must error"); + paimon_error_free(result.error); + + unwrap_table(handle); + } +} + +/// Split bytes come from outside the process, so a buffer that is not a split +/// must surface as an error, not a panic across the ABI boundary. +#[test] +fn vector_search_bucket_splits_reject_corrupt_bytes() { + let path = "memory:/vsearch_pk_split_corrupt"; + let (query, vectors) = pk_fixture_smoke(); + let table = build_pk_vector_table(path, &vectors); + let handle = unsafe { wrap_table(table) }; + + unsafe { + let garbage: Vec = vec![0xAB; 64]; + let ptrs: [*const u8; 1] = [garbage.as_ptr()]; + let lens: [usize; 1] = [garbage.len()]; + let builder = c_vector_builder(handle, VECTOR_COLUMN, &query, 3, ptr::null_mut()); + let result = paimon_vector_search_builder_execute_read_for_bucket_splits( + builder, + ptrs.as_ptr(), + lens.as_ptr(), + 1, + ); + paimon_vector_search_builder_free(builder); + assert!(!result.error.is_null(), "garbage bytes must error"); + assert!(result.reader.is_null()); + paimon_error_free(result.error); + unwrap_table(handle); + } +} + #[test] fn vector_search_append_table_read_matches_rust() { let path = "memory:/vsearch_append_read"; diff --git a/bindings/c/src/vector_search.rs b/bindings/c/src/vector_search.rs index 80d79efed..90ba0afb2 100644 --- a/bindings/c/src/vector_search.rs +++ b/bindings/c/src/vector_search.rs @@ -344,6 +344,117 @@ pub unsafe extern "C" fn paimon_vector_search_builder_execute_read( } } +/// Run the vector search over bucket splits a Java planner produced, and stream +/// the materialized rows. +/// +/// This is the entry point for an engine that plans in Paimon Java and executes +/// here: the planner emits one `BucketVectorSearchSplit` per bucket and ships its +/// bytes to a worker, which calls this. `splits` points at `count` buffers and +/// `split_lens` at their lengths; both arrays must hold `count` entries. The +/// buffers are only read for the duration of the call. +/// +/// The splits are the plan -- their payload files, per-file row ranges and pinned +/// snapshot are used as given, and the table's index manifest is not read. Search, +/// optional refine, Top-K and materialization are the same as +/// `paimon_vector_search_builder_execute_read`, so the output is the projected +/// user columns plus `__paimon_search_score`, best-first. The Top-K is local to +/// the splits passed in; a caller distributing one call per bucket merges the +/// per-bucket results itself. +/// +/// Only a primary-key vector column can be read this way; a data-evolution table +/// returns an error rather than an answer from a different plan. Consume via +/// `paimon_record_batch_reader_next` and free with +/// `paimon_record_batch_reader_free`. +/// +/// # Safety +/// `b` must be a valid pointer from `paimon_table_new_vector_search_builder`, or +/// null (returns an error result). `splits` and `split_lens` must each point at +/// `count` valid entries, and each `splits[i]` at `split_lens[i]` readable bytes. +#[no_mangle] +pub unsafe extern "C" fn paimon_vector_search_builder_execute_read_for_bucket_splits( + b: *mut paimon_vector_search_builder, + splits: *const *const u8, + split_lens: *const usize, + count: usize, +) -> paimon_result_record_batch_reader { + if let Err(e) = check_non_null(b, "b") { + return paimon_result_record_batch_reader { + reader: std::ptr::null_mut(), + error: e, + }; + } + if splits.is_null() || split_lens.is_null() || count == 0 { + return paimon_result_record_batch_reader { + reader: std::ptr::null_mut(), + error: paimon_error::new( + PaimonErrorCode::InvalidInput, + "paimon_vector_search_builder_execute_read_for_bucket_splits: null or empty splits" + .to_string(), + ), + }; + } + + let ptrs = std::slice::from_raw_parts(splits, count); + let lens = std::slice::from_raw_parts(split_lens, count); + let mut buffers: Vec<&[u8]> = Vec::with_capacity(count); + for (i, (&ptr, &len)) in ptrs.iter().zip(lens).enumerate() { + // A null or empty buffer cannot be a split, and reaching the decoder with + // one would report it as corrupt data rather than as the caller's error. + if ptr.is_null() || len == 0 { + return paimon_result_record_batch_reader { + reader: std::ptr::null_mut(), + error: paimon_error::new( + PaimonErrorCode::InvalidInput, + format!( + "paimon_vector_search_builder_execute_read_for_bucket_splits: \ + split {i} is null or empty" + ), + ), + }; + } + buffers.push(std::slice::from_raw_parts(ptr, len)); + } + + let state = &*((*b).inner as *const VectorSearchState); + let mut builder = state.table.new_vector_search_builder(); + if let Some(col) = &state.vector_column { + builder.with_vector_column(col); + } + if let Some(v) = &state.query_vector { + builder.with_query_vector(v.clone()); + } + if let Some(limit) = state.limit { + builder.with_limit(limit); + } + if !state.options.is_empty() { + builder.with_options(state.options.clone()); + } + if let Some(f) = &state.filter { + builder.with_filter(f.clone()); + } + if let Some(cols) = &state.projection { + let col_refs: Vec<&str> = cols.iter().map(String::as_str).collect(); + builder.with_projection(&col_refs); + } + + match runtime().block_on(builder.execute_read_for_bucket_splits(&buffers)) { + Ok(stream) => { + let reader = Box::new(stream); + let wrapper = Box::new(paimon_record_batch_reader { + inner: Box::into_raw(reader) as *mut c_void, + }); + paimon_result_record_batch_reader { + reader: Box::into_raw(wrapper), + error: std::ptr::null_mut(), + } + } + Err(e) => paimon_result_record_batch_reader { + reader: std::ptr::null_mut(), + error: paimon_error::from_paimon(e), + }, + } +} + // --- C ABI signature guards ------------------------------------------------- // // These symbols are called across the FFI boundary with fixed argument counts: @@ -387,3 +498,10 @@ const _: unsafe extern "C" fn(*mut paimon_vector_search_builder) = const _: unsafe extern "C" fn( *mut paimon_vector_search_builder, ) -> paimon_result_record_batch_reader = paimon_vector_search_builder_execute_read; +const _: unsafe extern "C" fn( + *mut paimon_vector_search_builder, + *const *const u8, + *const usize, + usize, +) -> paimon_result_record_batch_reader = + paimon_vector_search_builder_execute_read_for_bucket_splits; diff --git a/crates/paimon/src/lumina/reader.rs b/crates/paimon/src/lumina/reader.rs index 00025176c..6301516e1 100644 --- a/crates/paimon/src/lumina/reader.rs +++ b/crates/paimon/src/lumina/reader.rs @@ -122,6 +122,27 @@ fn compare_scores(a: f32, b: f32) -> std::cmp::Ordering { } } +/// Turn an include-set into the dense id array Lumina's filtered search takes. +/// +/// Mirrors Java `LuminaVectorGlobalIndexReader.toScopedIds`: an include-set above +/// `Integer.MAX_VALUE` is refused BEFORE the array is allocated, because that array +/// is 8 bytes per id and the set can hold one id per live row of the segment. +fn to_scoped_ids(include_row_ids: &roaring::RoaringTreemap) -> crate::Result> { + let cardinality = include_row_ids.len(); + if cardinality > i32::MAX as u64 { + return Err(crate::Error::DataInvalid { + message: format!( + "include_row_ids cardinality ({cardinality}) exceeds {}", + i32::MAX + ), + source: None, + }); + } + let mut ids = Vec::with_capacity(cardinality as usize); + ids.extend(include_row_ids.iter()); + Ok(ids) +} + /// Allocate the label buffer for one native search, filled with [`SENTINEL`]. /// /// [`SENTINEL`] is the "no result" marker (the C ABI's `-1`), and @@ -373,7 +394,7 @@ fn search_lumina( let include_row_ids = vector_search.effective_include_row_ids(); let (distances, labels) = if let Some(include_ids) = include_row_ids { - let filter_id_list: Vec = include_ids.iter().collect(); + let filter_id_list: Vec = to_scoped_ids(include_ids)?; if filter_id_list.is_empty() { return Ok(None); } @@ -460,14 +481,26 @@ 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) { + // An include-set that permits nothing can match nothing, whatever the index + // holds. Answered from the set itself, so no native call is made at all. + if shared_filter.is_some_and(|filter| filter.is_empty()) { return Ok(vec![None; vector_searches.len()]); } - let index_metric = index_meta.metric()?; + // Java's order, which decides what an empty or unreadable index reports: + // `index.size()` first, then the dense filter, and the metric only on the way + // out. An empty index has nothing to return whatever the filter says, so + // converting first would report an oversized filter instead, and reading the + // metric first would report a malformed one. let count = searcher.get_count()? as usize; + if count == 0 { + return Ok(vec![None; vector_searches.len()]); + } + let filter_id_list = shared_filter + .map(|include_row_ids| to_scoped_ids(include_row_ids)) + .transpose()?; + let index_metric = index_meta.metric()?; + 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()), @@ -589,6 +622,25 @@ impl Drop for LuminaVectorGlobalIndexReader { #[cfg(test)] mod tests { + + #[test] + fn to_scoped_ids_refuses_a_set_larger_than_the_dense_array_can_index() { + // Java's own limit for this exact quantity at this exact seam + // (`LuminaVectorGlobalIndexReader.toScopedIds`). The array is 8 bytes per id, + // and the set can hold one id per live row of the segment, so the refusal has + // to come BEFORE the allocation. + let mut small = roaring::RoaringTreemap::new(); + small.insert_range(0..=9u64); + assert_eq!(to_scoped_ids(&small).unwrap().len(), 10); + + let mut over = roaring::RoaringTreemap::new(); + over.insert_range(0..=(i32::MAX as u64 + 1)); + let error = to_scoped_ids(&over) + .map(|_| ()) + .expect_err("a set this size cannot be handed to the backend"); + assert!(error.to_string().contains("exceeds"), "{error}"); + } + use super::*; use crate::lumina::{KEY_DIMENSION, KEY_DISTANCE_METRIC}; use crate::vector_search::GlobalIndexIOMeta; @@ -797,6 +849,46 @@ mod tests { assert!(calls.iter().all(|call| call.n == 1)); } + #[test] + fn an_empty_index_returns_empty_before_the_filter_is_sized() { + // Java's order: `index.size()` decides first, so an index holding nothing + // returns nothing whatever the filter says. Sizing the dense filter first + // would turn this into an oversized-filter error instead. + let searcher = RecordingSearcher::new(0); + let mut oversized = roaring::RoaringTreemap::new(); + oversized.insert_range(0..=(i32::MAX as u64 + 1)); + let shared_filter = Arc::new(oversized); + 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("an empty index reports no hits, not an oversized filter"); + + assert_eq!(results, vec![None, None]); + assert_eq!( + searcher.count_calls.load(Ordering::Relaxed), + 1, + "the index size is what the answer came from" + ); + 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_empty_shared_filter_skips_native_search() { let searcher = RecordingSearcher::new(10); diff --git a/crates/paimon/src/table/pk_vector_bucket_split.rs b/crates/paimon/src/table/pk_vector_bucket_split.rs index c82826757..0b10e6fa5 100644 --- a/crates/paimon/src/table/pk_vector_bucket_split.rs +++ b/crates/paimon/src/table/pk_vector_bucket_split.rs @@ -239,8 +239,20 @@ impl BucketVectorSearchSplit { // leave no way to say which of them a range belongs to; a bucket cannot // hold such a pair and the planner rejects it, but these bytes are // untrusted. + // + // The row count is checked here for EVERY file, not only for the ones the + // message goes on to list ranges for -- the range reader below only sees the + // files it lists. `bucket_search` re-checks the ACTIVE files it is about to + // read, so this is the fail-fast copy plus the only check a file that never + // becomes active gets. let mut row_counts: HashMap<&str, i64> = HashMap::new(); for file in data_split.data_files() { + if file.row_count < 0 { + return Err(data_invalid(format!( + "data file {} has a negative row count: {}", + file.file_name, file.row_count + ))); + } if row_counts .insert(file.file_name.as_str(), file.row_count) .is_some() @@ -301,8 +313,11 @@ fn nested_error(error: crate::Error) -> crate::Error { /// file. Java writes ranges its planner produced and re-checks nothing, so the /// checks here are what a reader of untrusted bytes needs rather than a mirror /// of the writer: `RowRange` cannot represent a descending pair at all, and a -/// range outside its file would read rows that are not there. The file's own row -/// count is checked first, so a forged one cannot lift that bound. +/// range outside its file would read rows that are not there. The bound is only as +/// good as the count it is taken against: every file's count was checked +/// non-negative when the bucket was indexed above, but a forged LARGER positive +/// count does lift this bound, so nothing downstream may size an allocation from +/// what it admits. fn read_row_ranges( cur: &mut &[u8], file_name: &str, @@ -313,15 +328,6 @@ fn read_row_ranges( "row ranges reference data file not present in the bucket split: {file_name}" )) })?; - // A negative row count is forged by construction. Rejecting it, rather than - // skipping the bound check for it, is what keeps the bound below meaningful: - // otherwise a forged count would lift it entirely. - if row_count < 0 { - return Err(data_invalid(format!( - "data file {file_name} has a negative row count: {row_count}" - ))); - } - let count = read_count(cur, "row range")?; let mut ranges: Vec = Vec::new(); for _ in 0..count { @@ -998,6 +1004,35 @@ mod tests { assert_error_contains(&bytes, "data file data-1.orc has a negative row count: -1"); } + /// The same forged count, on a message that lists NO ranges at all. The range + /// reader only ever sees the files the message lists, and `bucket_search` + /// re-checks only the files that become ACTIVE, so for anything else this is the + /// only check there is. + #[test] + fn rejects_a_negative_row_count_on_a_file_with_no_listed_ranges() { + let mut bytes = golden_without_row_ranges(); + bytes[DATA_FILE_ROW_COUNT_OFFSET..DATA_FILE_ROW_COUNT_OFFSET + 8] + .copy_from_slice(&crate::spec::DataFileMeta::ROW_COUNT_UNKNOWN.to_le_bytes()); + assert_error_contains(&bytes, "data file data-1.orc has a negative row count: -1"); + } + + /// The golden with its row-range section removed: the no-pre-filter shape Java + /// emits when its planner narrowed nothing. + fn golden_without_row_ranges() -> Vec { + let mut bytes = golden(); + let count_offset = bytes.len() - RANGE_SECTION_BYTES - 4; + bytes.truncate(count_offset + 4); + bytes[count_offset..count_offset + 4].copy_from_slice(&0i32.to_be_bytes()); + bytes + } + + #[test] + fn a_message_that_lists_no_ranges_decodes_to_no_entries() { + let split = BucketVectorSearchSplit::deserialize(&golden_without_row_ranges()).unwrap(); + assert!(split.row_ranges_by_file().is_empty()); + assert_eq!(split.data_split().data_files().len(), 1); + } + #[test] fn rejects_negative_row_range() { let mut bytes = golden(); diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs b/crates/paimon/src/table/pk_vector_orchestrator.rs index ac0d3ab35..ca38b08fa 100644 --- a/crates/paimon/src/table/pk_vector_orchestrator.rs +++ b/crates/paimon/src/table/pk_vector_orchestrator.rs @@ -26,8 +26,6 @@ use std::cmp::Ordering; use std::collections::HashMap; use std::sync::Arc; -use roaring::RoaringTreemap; - use crate::deletion_vector::DeletionVector; use crate::spec::BinaryRow; use crate::table::data_file_reader::DataFileReader; @@ -40,6 +38,7 @@ use crate::vindex::pkvector::bucket::{ }; use crate::vindex::pkvector::metric::{java_float_compare, VectorSearchMetric}; use crate::vindex::pkvector::result::PkVectorSearchResult; +use crate::vindex::pkvector::FileRowSelections; fn data_invalid(message: impl Into) -> crate::Error { crate::Error::DataInvalid { @@ -356,12 +355,13 @@ impl PkVectorOrchestrator { /// and split so a caller can build a reader keyed to the specific split/file. /// `skip_exact_fallback` forwards to `bucket_search`. /// - /// `residual_by_split`, when present, carries one per-file allow-list of - /// physical row positions per split (indexed parallel to `splits`): only - /// positions listed for a file may survive that bucket's search. A file - /// absent from its split's map (or mapped to an empty set) contributes no - /// candidates. `None` applies no residual filtering. The slice must have the - /// same length as `splits`. + /// `row_selections_by_split`, when present, carries one per-file row selection + /// per split (indexed parallel to `splits`), in the three states of + /// [`FileRowSelection`]: a file with NO entry is unrestricted, an empty entry + /// contributes no candidates, and a non-empty one limits which of its rows may. + /// A selection is either interval `Ranges` (from an engine's bucket split) or + /// `Positions` (from a residual data predicate). `None` restricts nothing at + /// all. The slice must have the same length as `splits`. /// /// This is the single-query wrapper over /// [`search_candidates_batch`](Self::search_candidates_batch): it searches the @@ -391,7 +391,7 @@ impl PkVectorOrchestrator { + Sync), search_options: &HashMap, skip_exact_fallback: bool, - residual_by_split: Option<&[HashMap]>, + row_selections_by_split: Option<&[FileRowSelections]>, concurrency: usize, ) -> crate::Result { let mut results = self @@ -405,7 +405,7 @@ impl PkVectorOrchestrator { exact_file_search, search_options, skip_exact_fallback, - residual_by_split, + row_selections_by_split, concurrency, ) .await?; @@ -422,8 +422,8 @@ impl PkVectorOrchestrator { /// lists get their own cross-bucket global Top-K. No query's candidates bleed /// into another's (independent per-query heaps). /// - /// The residual allow-list depends only on the filter and the plan, not the - /// query vector, so the SAME `residual_by_split` slice is shared across every + /// The row selections depend only on the filter and the plan, not the + /// query vector, so the SAME `row_selections_by_split` slice is shared across every /// query. Input-shape validation (positive limits, non-empty query, residual /// count) is applied per query / once as appropriate. /// @@ -457,7 +457,7 @@ impl PkVectorOrchestrator { + Sync), search_options: &HashMap, skip_exact_fallback: bool, - residual_by_split: Option<&[HashMap]>, + row_selections_by_split: Option<&[FileRowSelections]>, concurrency: usize, ) -> crate::Result> { // Eager input-shape validation (Java checkArgument parity). @@ -475,10 +475,10 @@ impl PkVectorOrchestrator { return Err(data_invalid("vector search query must not be empty")); } } - if let Some(per_split) = residual_by_split { + if let Some(per_split) = row_selections_by_split { if per_split.len() != splits.len() { return Err(data_invalid( - "residual range map count does not match split count", + "row selection map count does not match split count", )); } } @@ -528,7 +528,8 @@ impl PkVectorOrchestrator { ) }, ); - let residual_ranges = residual_by_split.map(|per_split| &per_split[split_index]); + let row_selections = + row_selections_by_split.map(|per_split| &per_split[split_index]); let per_query = bucket_search_batch( ann_searcher, &split.ann_segments, @@ -541,7 +542,7 @@ impl PkVectorOrchestrator { limit, search_options, skip_exact_fallback, - residual_ranges, + row_selections, concurrency, search_budget, ) @@ -1263,7 +1264,7 @@ mod e2e_tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _residual_ranges: Option<&HashMap>, + _row_selections: Option<&FileRowSelections>, ) -> crate::Result>> { Ok(queries.iter().map(|_| self.hits.clone()).collect()) } @@ -1795,10 +1796,13 @@ mod e2e_tests { ); // Allow only positions 0 and 2 for "r.mosaic"; pos1 (the best hit) is // excluded by the residual. - let mut allowed = RoaringTreemap::new(); + let mut allowed = roaring::RoaringTreemap::new(); allowed.insert(0); allowed.insert(2); - let residual_by_split = vec![HashMap::from([("r.mosaic".to_string(), allowed)])]; + let row_selections_by_split: Vec = vec![HashMap::from([( + "r.mosaic".to_string(), + crate::vindex::pkvector::FileRowSelection::Positions(allowed), + )])]; let opts = HashMap::new(); let result = PkVectorOrchestrator::new(make_reader(file_io, table_path)) .search_candidates( @@ -1811,7 +1815,7 @@ mod e2e_tests { &factory, &opts, false, - Some(&residual_by_split), + Some(&row_selections_by_split), 1, ) .await @@ -1852,8 +1856,7 @@ mod e2e_tests { }; let factory = unreachable_split_search(); // Two residual maps for a single split. - let residual_by_split: Vec> = - vec![HashMap::new(), HashMap::new()]; + let row_selections_by_split: Vec = vec![HashMap::new(), HashMap::new()]; let opts = HashMap::new(); let err = PkVectorOrchestrator::new(make_reader(file_io, table_path)) .search_candidates( @@ -1866,7 +1869,7 @@ mod e2e_tests { &factory, &opts, false, - Some(&residual_by_split), + Some(&row_selections_by_split), 1, ) .await diff --git a/crates/paimon/src/table/pk_vector_scan.rs b/crates/paimon/src/table/pk_vector_scan.rs index 21d70d1dc..39815ee5f 100644 --- a/crates/paimon/src/table/pk_vector_scan.rs +++ b/crates/paimon/src/table/pk_vector_scan.rs @@ -24,8 +24,6 @@ use std::collections::{BTreeMap, HashMap, HashSet}; use indexmap::IndexMap; -use roaring::RoaringTreemap; - use crate::spec::{ should_read_pk_index_source, BinaryRow, DataFileMeta, FileKind, IndexManifest, Predicate, PrimaryKeyIndexSourceFile, PrimaryKeyIndexSourceMeta, @@ -53,37 +51,6 @@ struct UnresolvedAnnSegment { /// (`BinaryRow` is not hashable) paired with the bucket number. type BucketKey = (Vec, i32); -/// Expand inclusive row ranges into the positions they allow. Only the search -/// kernel's membership tests need this; a read is limited by the ranges themselves. -pub(super) fn positions_in_ranges(ranges: &[RowRange]) -> crate::Result { - let mut positions = RoaringTreemap::new(); - for range in ranges { - let from = u64::try_from(range.from()) - .map_err(|_| data_invalid("row range bound must not be negative"))?; - let to = u64::try_from(range.to()) - .map_err(|_| data_invalid("row range bound must not be negative"))?; - positions.insert_range(from..=to); - } - Ok(positions) -} - -/// The whole of a file, for a file the message left unrestricted. -/// -/// A zero-row file gets no range rather than an empty one: `RowRange` is inclusive, -/// so it cannot express "nothing". An unknown count (`DataFileMeta::ROW_COUNT_UNKNOWN`, -/// or anything else negative) is rejected rather than read as "nothing", which would -/// silently drop the file from the search. Decoding only checks the count of a file -/// the message lists ranges for, so this is where an omitted one is checked. -fn whole_file_range(row_count: i64) -> crate::Result> { - match row_count { - 0 => Ok(Vec::new()), - count if count > 0 => Ok(vec![RowRange::new(0, count - 1)]), - count => Err(data_invalid(format!( - "data file row count must be known and non-negative, got {count}" - ))), - } -} - fn data_invalid(message: impl Into) -> crate::Error { crate::Error::DataInvalid { message: message.into(), @@ -248,16 +215,18 @@ pub(crate) struct PkVectorScanPlan { // at all (never written), which also yields empty `splits`. pub snapshot_id: i64, pub splits: Vec, - // Per-split allow-list of physical rows, indexed parallel to `splits`: only the - // rows listed for a data file may produce candidates from it. Ranges rather than - // materialized positions, because this is what a read is limited to — expanding - // a whole-file range of a large file into positions costs memory no reader needs. + // Per-split allow-list of physical rows, indexed parallel to `splits`, holding + // ONLY the files the engine's split actually restricted. Three states per file, + // as in Java `rowRangesByFile`: absent means every row is readable, an empty + // list means none is, and a non-empty list means exactly those. Ranges rather + // than materialized positions, because this is what a read is limited to, and + // because the row counts these arrive with are untrusted — sizing an allow-list + // from one is unbounded work. // Each list is normalized: sorted, non-overlapping, inclusive, file-local. // Populated when the plan was built from engine-supplied bucket splits, which - // carry row ranges the engine's own planner already resolved. `None` for a plan - // read from this table's index manifest, which places no positional restriction - // of its own -- distinct from `Some` of an empty allow-list, which permits - // nothing. + // carry row ranges the engine's own planner already resolved -- possibly an + // empty map, when that planner narrowed nothing. `None` for a plan read from + // this table's index manifest, which places no positional restriction at all. pub physical_row_ranges_by_split: Option>>>, } @@ -414,9 +383,6 @@ impl<'a> PkVectorScan<'a> { /// Mirrors what Java's `PrimaryKeyVectorRead` does with a /// `BucketVectorSearchSplit`: search the payloads the split names, over the rows /// the split allows. - // Entry point for engine-supplied splits; no in-tree caller reads a plan from - // them yet, and the tests drive `plan_from_bucket_splits` directly. - #[allow(dead_code)] pub(crate) fn plan_for_bucket_vector_splits( &self, splits: Vec, @@ -451,7 +417,6 @@ impl<'a> PkVectorScan<'a> { /// The `Table`-independent core of [`PkVectorScan::plan_for_bucket_vector_splits`], /// so planning from engine-supplied splits is testable the same way planning from a /// manifest is. -#[cfg_attr(not(test), allow(dead_code))] fn plan_from_bucket_splits( index_type: &str, vector_field_id: i32, @@ -577,15 +542,16 @@ fn plan_from_bucket_splits( index_file_in_data_file_dir, )?; - // Normalize the row ranges against the planned splits, which are grouped by - // bucket and so may be ordered differently from the input. + // Re-key the row ranges against the planned splits, which are grouped by bucket + // and so may be ordered differently from the input. // - // A file the message lists is restricted to the positions it lists. A file it - // omits is unrestricted: Java records ranges only for the files its own - // pre-filter narrowed, and leaves the rest out. The search kernel reads a - // missing entry as "no rows allowed", the opposite meaning, so the omission - // has to be turned into an explicit full-file range here rather than passed - // through. + // Only what the message actually listed is carried. A file it omits gets NO + // entry, which is what "every row" is spelled as throughout the search kernel + // and in Java (`rowRangesByFile.get(file) == null`). Synthesizing an explicit + // `[0, row_count - 1]` for it instead would look equivalent and is not: the + // row count arrives on the wire, and sizing an allow-list from an untrusted + // number is unbounded work. An explicitly empty list stays empty — that is the + // separate "no rows" state. let physical_row_ranges_by_split = splits .iter() .map(|split| { @@ -597,18 +563,15 @@ fn plan_from_bucket_splits( .data_split .data_files() .iter() - .map(|file| { - let allowed = match listed.and_then(|ranges| ranges.get(&file.file_name)) { - // Decoding checks each range's bounds but not their order or - // whether they overlap, and a read needs them normalized. - Some(ranges) => merge_row_ranges(ranges.clone()), - None => whole_file_range(file.row_count)?, - }; - Ok((file.file_name.clone(), allowed)) + .filter_map(|file| { + let ranges = listed.and_then(|ranges| ranges.get(&file.file_name))?; + // Decoding checks each range's bounds but not their order or + // whether they overlap, and a read needs them normalized. + Some((file.file_name.clone(), merge_row_ranges(ranges.clone()))) }) - .collect::>>>() + .collect::>>() }) - .collect::>>()?; + .collect::>(); Ok(PkVectorScanPlan { snapshot_id, @@ -1268,6 +1231,12 @@ mod tests { const BUCKET_SPLIT_GOLDEN: &[u8] = include_bytes!("goldens/bucket_vector_search_split_v1.bin"); + /// The Java-planned fixture #771's read test uses: one bucket, no pre-filter, + /// so `rangeFileCount == 0`. Provenance is in + /// `tests/pk_vector_bucket_split_read_test.rs`. + const BUCKET_SPLIT_NO_PREFILTER: &[u8] = + include_bytes!("../../testdata/pkvector_split/bucket_split_0.bin"); + fn int_partition(value: i32) -> BinaryRow { let mut builder = crate::spec::BinaryRowBuilder::new(1); builder.write_int(0, value); @@ -1333,13 +1302,14 @@ mod tests { } /// The positions a file's normalized ranges allow, for assertions that read - /// better as a row list than as ranges. - fn allowed(map: &HashMap>, file: &str) -> Vec { + /// better as a row list than as ranges. Test-only: the production path never + /// expands a range into positions. + fn allowed(map: &HashMap>, file: &str) -> Vec { map.get(file) .map(|ranges| { - positions_in_ranges(ranges) - .expect("planned ranges are in range") + ranges .iter() + .flat_map(|range| range.from()..=range.to()) .collect() }) .unwrap_or_default() @@ -1373,17 +1343,103 @@ mod tests { } #[test] - fn rejects_an_unknown_row_count_on_an_unlisted_file() { - // A file the message lists no ranges for is read as "the whole file", which - // needs a real row count. `ROW_COUNT_UNKNOWN` is -1, and reading that as "no - // rows" would drop the file from the search without a word; the decoder only - // checks the count of files it does carry ranges for. - let error = whole_file_range(DataFileMeta::ROW_COUNT_UNKNOWN) - .map(|_| ()) - .expect_err("an unknown row count cannot stand in for the whole file"); - assert!(error.to_string().contains("must be known"), "{error}"); - assert!(whole_file_range(0).unwrap().is_empty()); - assert_eq!(whole_file_range(3).unwrap(), vec![RowRange::new(0, 2)]); + fn an_unlisted_file_is_left_out_rather_than_expanded_from_its_row_count() { + // The row count comes off the wire. Turning an unlisted file into an explicit + // [0, row_count - 1] hands that number to the allow-list materialization, and + // i64::MAX does not come back. Java never writes the entry: absence already + // means "every row", so nothing has to be sized from the count. + let split = engine_split( + 11, + 0, + BinaryRow::new(0), + vec![dfm("d0", i64::MAX, 5, Some(1))], + vec![engine_payload(gim(2, 5, &[("d0", i64::MAX)]))], + &[], + ); + let plan = plan_from_bucket_splits("ivf-pq", 2, None, "/tbl", false, vec![split]).unwrap(); + let selections = plan + .physical_row_ranges_by_split + .expect("split-driven plan"); + assert!( + !selections[0].contains_key("d0"), + "an unlisted file must stay absent (unrestricted), not be expanded" + ); + } + + #[test] + fn the_java_fixture_lists_no_ranges_so_the_plan_restricts_nothing() { + // The committed Java-planned fixture is the no-pre-filter shape an engine + // ships by default. Nothing in it is restricted, so the plan must carry no + // per-file entry at all -- that absence is what lets the search skip the + // filtered ANN path. + let split = BucketVectorSearchSplit::deserialize(BUCKET_SPLIT_NO_PREFILTER).unwrap(); + assert!( + split.row_ranges_by_file().is_empty(), + "fixture provenance: Java recorded no pre-filter ranges" + ); + // Derived from the fixture itself so the test cannot drift from it. + let payload = &split.payload_files()[0]; + let index_type = payload.index_type().to_string(); + let field_id = payload.global_index_meta().index_field_id; + + let plan = plan_from_bucket_splits(&index_type, field_id, None, "/tbl", false, vec![split]) + .unwrap(); + let selections = plan + .physical_row_ranges_by_split + .expect("split-driven plan"); + assert_eq!(selections.len(), 1); + assert!( + selections[0].is_empty(), + "a split that lists no ranges restricts no file" + ); + } + + #[test] + fn the_java_fixture_reaches_the_ann_layer_with_no_mask_at_all() { + // The chain the two tests above only cover in halves: the committed + // Java-planned fixture, through planning, into the function that decides + // whether the ANN backend is handed an `include_row_ids` filter. `None` here + // is what makes Lumina call `search` rather than `search_with_filter`. + // Result equality against the manifest route cannot see this difference -- + // both return the same rows for this fixture. + let split = BucketVectorSearchSplit::deserialize(BUCKET_SPLIT_NO_PREFILTER).unwrap(); + let payload = &split.payload_files()[0]; + let index_type = payload.index_type().to_string(); + let field_id = payload.global_index_meta().index_field_id; + let source_meta = + PrimaryKeyIndexSourceMeta::from_global_index_meta(payload.global_index_meta()).unwrap(); + + let plan = plan_from_bucket_splits(&index_type, field_id, None, "/tbl", false, vec![split]) + .unwrap(); + let selections: crate::vindex::pkvector::FileRowSelections = plan + .physical_row_ranges_by_split + .expect("split-driven plan") + .remove(0) + .into_iter() + .map(|(file, ranges)| { + ( + file, + crate::vindex::pkvector::FileRowSelection::Ranges(ranges), + ) + }) + .collect(); + + let active: HashSet = source_meta + .source_files() + .iter() + .map(|file| file.file_name().to_string()) + .collect(); + assert!( + crate::vindex::pkvector::ann::build_live_row_ids( + source_meta.source_files(), + &active, + &HashMap::new(), + Some(&selections), + ) + .unwrap() + .is_none(), + "the no-pre-filter fixture must not put the backend on the filtered path" + ); } #[test] @@ -1474,8 +1530,9 @@ mod tests { .physical_row_ranges_by_split .expect("split-driven plan"); - assert!(allowed(&ranges[0], "d0").is_empty()); - assert_eq!(allowed(&ranges[0], "d1"), vec![0, 1, 2]); + // Listed empty: present and excluded. Unlisted: absent, meaning every row. + assert_eq!(ranges[0].get("d0").map(Vec::as_slice), Some(&[][..])); + assert!(!ranges[0].contains_key("d1")); } #[test] @@ -1542,8 +1599,8 @@ mod tests { let ranges = plan .physical_row_ranges_by_split .expect("split-driven plan"); - // Unaffected: the whole file stays readable. - assert_eq!(allowed(&ranges[0], "d0"), vec![0, 1, 2, 3]); + // Unaffected: the file stays unlisted, which is "every row". + assert!(!ranges[0].contains_key("d0")); } #[test] diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 447d69953..93767012f 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -33,6 +33,7 @@ use crate::table::global_index_scanner::{ unindexed_ranges_for_global_index_entries, RowRangeIndex, }; use crate::table::index_file_path::IndexFileLocation; +use crate::table::pk_vector_bucket_split::BucketVectorSearchSplit; use crate::table::pk_vector_data_file_reader::{ append_batch_vectors, DataFilePkVectorReaderFactory, }; @@ -44,7 +45,7 @@ use crate::table::pk_vector_orchestrator::{ use crate::table::pk_vector_position_read::{ PkVectorPositionRead, PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN, }; -use crate::table::pk_vector_scan::{positions_in_ranges, PkVectorScan, PkVectorScanPlan}; +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; @@ -60,6 +61,7 @@ use crate::vindex::pkvector::ann::{AnnSegmentSource, PkVectorAnnSearcher, Vindex use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment, ExactFileSearchFuture}; use crate::vindex::pkvector::exact::validate_query; use crate::vindex::pkvector::metric::VectorSearchMetric; +use crate::vindex::pkvector::{FileRowSelection, FileRowSelections}; use crate::vindex::range_reader::{RangeIoStats, RangeReadLimiter, VindexFileReader}; use crate::vindex::reader::VindexVectorGlobalIndexReader; use crate::vindex::{is_vindex_index_type, vector_search_timing_enabled, VindexVectorIndexOptions}; @@ -416,6 +418,131 @@ impl<'a> VectorSearchBuilder<'a> { .await } + /// Run this search over bucket splits an engine planned elsewhere, and + /// materialize the hits. + /// + /// The unit of work is Java's `BucketVectorSearchSplit` byte form: a planner + /// running in Paimon Java enumerates one split per bucket -- a bucket is never + /// divided, because the ANN current-segment decision needs the bucket's whole + /// active file set -- and ships each to a worker that calls this. The splits + /// are the plan: their payload files, their per-file row ranges and the + /// snapshot they pin are used as given, and this table's index manifest is not + /// read. + /// + /// Everything after planning is the ordinary primary-key vector read, so + /// search, optional refine, local Top-K and materialization stay identical to + /// [`execute_read`](Self::execute_read): output is the projected user columns + /// plus `__paimon_search_score`, best-first. The Top-K is local to the supplied + /// splits; a caller distributing one call per bucket merges the per-bucket + /// results itself. + /// + /// Only a primary-key vector column can be read this way. The data-evolution + /// route plans through the global index rather than through bucket splits, so + /// it is rejected rather than silently answered from a different plan. + pub async fn execute_read_for_bucket_splits( + &self, + split_bytes: &[&[u8]], + ) -> crate::Result { + // Fail closed: returns data outside `TableScan`/`TableRead`. + let core = CoreOptions::new(self.table.schema().options()); + core.ensure_read_authorized()?; + 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 query_vector = + self.query_vector + .as_ref() + .ok_or_else(|| crate::Error::ConfigInvalid { + message: "Query vector must be set via with_query_vector()".to_string(), + })?; + let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid { + message: "Limit must be set via with_limit()".to_string(), + })?; + + let pk_col = if core.primary_key_vector_index_enabled() { + let targets_pk_column = core + .primary_key_vector_index_columns() + .ok() + .is_some_and(|cols| cols.iter().any(|c| c == vector_column)); + if targets_pk_column { + core.primary_key_vector_index_column()? + } else { + return Err(bucket_split_route_error(vector_column)); + } + } else { + return Err(bucket_split_route_error(vector_column)); + }; + + // Decoding is the trust boundary: these bytes come from outside the + // process. Reject an empty request here rather than let it reach planning + // as "no splits", which cannot pin a snapshot. + if split_bytes.is_empty() { + return Err(crate::Error::DataInvalid { + message: "bucket-split read requires at least one split".to_string(), + source: None, + }); + } + let splits = split_bytes + .iter() + .map(|bytes| BucketVectorSearchSplit::deserialize(bytes)) + .collect::>>()?; + + // Resolve the query parameters (and reject a query the search cannot answer + // correctly) before planning, exactly as the manifest route does. + let params = resolve_pk_vector_search_params( + self.table, + &self.options, + self.filter.as_ref(), + &core, + &pk_col, + &[query_vector.as_slice()], + limit, + )?; + let plan = PkVectorScan::new( + self.table, + params.field_id, + params.index_type.clone(), + self.filter.clone(), + ) + .plan_for_bucket_vector_splits(splits)?; + + // Resolve the materialization read-type up front so an invalid projection + // fails loud even when the plan is empty and no rows will be read. + let read_type = self.resolve_materialize_read_type()?; + + let mut candidates = search_pk_candidates_batch_with_plan( + self.table, + &self.options, + self.filter.as_ref(), + &core, + &pk_col, + &[query_vector.as_slice()], + limit, + &plan, + ¶ms, + ) + .await?; + debug_assert_eq!(candidates.len(), 1); + let candidates = candidates.remove(0); + + // A separate, predicate-free materialization reader projecting the user + // columns (the search reader projects only the vector column). + let materialize_reader = DataFileReader::new( + self.table.file_io().clone(), + self.table.schema_manager().clone(), + self.table.schema().id(), + self.table.schema().fields().to_vec(), + read_type, + Vec::new(), + ); + + Self::materialize_candidates(candidates, &plan.splits, params.metric, &materialize_reader) + .await + } + /// Materialize the best-first data-evolution vector search hits into Arrow /// rows. The global-index search returns global `_ROW_ID`s and their scores; a /// subsequent row-range read materializes those rows, and each row's score is @@ -779,6 +906,18 @@ struct PkVectorSearchParams { indexed_limit: usize, } +/// A bucket split is a primary-key vector plan. The data-evolution route plans +/// through the global index instead, so answering it here would silently use a +/// different plan than the caller supplied. +fn bucket_split_route_error(vector_column: &str) -> crate::Error { + crate::Error::DataInvalid { + message: format!( + "bucket-split read requires a primary-key vector column, but '{vector_column}' is not one" + ), + source: None, + } +} + /// Resolve the query-level parameters and reject a query the search cannot answer /// correctly, before any planning or read happens. fn resolve_pk_vector_search_params( @@ -914,19 +1053,30 @@ fn resolve_pk_vector_search_params( /// rows an engine-supplied plan restricts each file to, and the positions a residual /// data predicate leaves behind. /// -/// Both sides list what is permitted, and both read a file's absence as "no rows -/// allowed", so combining them intersects files as well as positions. Either side -/// alone passes through unchanged; neither side means no positional restriction. +/// The two sides read a file's ABSENCE differently, and the merge has to respect +/// both readings: +/// +/// * The plan lists only what the engine's split narrowed, so an absent file is +/// unrestricted -- Java's `rowRangesByFile.get(file) == null`. +/// * The residual is exhaustive over the files a search can read from +/// (`residual_positions_by_file` registers every active file, empty when nothing +/// passed), so once a residual exists its silence about a file means "no rows". /// -/// This is where the plan's ranges become positions: the search kernel tests -/// membership, while a read is limited by the ranges themselves. When the residual -/// was evaluated over those same ranges the intersection cannot remove anything, and -/// is kept as the invariant that says so. +/// So: with no residual, a file the plan omits stays absent and unrestricted. With a +/// residual, a file it omits is excluded even if the plan restricted it, and a file +/// both describe keeps the intersection. Absent from BOTH is unrestricted, which is +/// what lets the ANN backend search unfiltered. +/// +/// The plan's ranges stay ranges. Expanding them into positions would be work sized +/// by row counts that arrived on the wire; where an intersection is genuinely needed +/// the residual positions — bounded by the rows its own read returned — are filtered +/// BY the ranges instead. When the residual was evaluated over those same ranges the +/// intersection cannot remove anything, and is kept as the invariant that says so. fn intersect_row_allow_lists( physical: Option<&[HashMap>]>, residual: Option>>, split_count: usize, -) -> crate::Result>>> { +) -> crate::Result>> { if let Some(maps) = physical { if maps.len() != split_count { return Err(crate::Error::DataInvalid { @@ -938,45 +1088,79 @@ fn intersect_row_allow_lists( }); } } + if let Some(maps) = residual.as_ref() { + if maps.len() != split_count { + return Err(crate::Error::DataInvalid { + message: format!( + "residual carries {} row allow-lists for {split_count} splits", + maps.len() + ), + source: None, + }); + } + } match (physical, residual) { - (None, residual) => Ok(residual), + (None, None) => Ok(None), + (None, Some(residual)) => Ok(Some( + residual + .into_iter() + .map(|per_file| { + per_file + .into_iter() + .map(|(file, positions)| (file, FileRowSelection::Positions(positions))) + .collect() + }) + .collect(), + )), (Some(physical), None) => Ok(Some( physical .iter() .map(|per_file| { per_file .iter() - .map(|(file, ranges)| Ok((file.clone(), positions_in_ranges(ranges)?))) - .collect::>>() + .map(|(file, ranges)| { + (file.clone(), FileRowSelection::Ranges(ranges.clone())) + }) + .collect() }) - .collect::>>()?, + .collect(), )), (Some(physical), Some(residual)) => { - if residual.len() != split_count { - return Err(crate::Error::DataInvalid { - message: format!( - "residual carries {} row allow-lists for {split_count} splits", - residual.len() - ), - source: None, - }); - } Ok(Some( physical .iter() .zip(residual) - .map(|(physical, residual)| { - physical - .iter() - .filter(|(file, _)| residual.contains_key(file.as_str())) - .map(|(file, ranges)| { - let allowed = positions_in_ranges(ranges)?; - let kept = &residual[file.as_str()]; - Ok((file.clone(), allowed & kept)) - }) - .collect::>>() + .map(|(physical, mut residual)| { + let mut merged: FileRowSelections = HashMap::new(); + for (file, ranges) in physical { + let range_selection = FileRowSelection::Ranges(ranges.clone()); + let selection = match residual.remove(file.as_str()) { + // Both restrict: keep the positions the ranges also + // allow. Filtering the positions (bounded by the read) + // by the ranges never expands the ranges. + Some(positions) => FileRowSelection::Positions( + positions + .iter() + .filter(|position| range_selection.contains(*position)) + .collect(), + ), + // The residual is exhaustive over the files the search + // can read from -- `residual_positions_by_file` + // registers every active file, empty when nothing + // passed. Its silence about a file therefore means "no + // rows", NOT "unrestricted", and must stay fail-closed + // here even though the plan has something to say. + None => FileRowSelection::Positions(RoaringTreemap::new()), + }; + merged.insert(file.clone(), selection); + } + // Whatever the residual restricted and the plan did not. + merged.extend(residual.into_iter().map(|(file, positions)| { + (file, FileRowSelection::Positions(positions)) + })); + merged }) - .collect::>>()?, + .collect(), )) } } @@ -1231,7 +1415,7 @@ async fn search_pk_raw_candidates_batch_with_plan( // built from engine-supplied bucket splits carries the physical positions each // file is limited to; a plan read from the index manifest carries none. Both // sides list what is permitted, so combining them is an intersection. - let residual_by_split = intersect_row_allow_lists( + let row_selections_by_split = intersect_row_allow_lists( plan.physical_row_ranges_by_split.as_deref(), residual_by_split, plan.splits.len(), @@ -1309,7 +1493,7 @@ async fn search_pk_raw_candidates_batch_with_plan( &factory, &search_options, skip_exact_fallback, - residual_by_split.as_deref(), + row_selections_by_split.as_deref(), concurrency, ) .await?; @@ -2495,18 +2679,23 @@ fn is_vector_global_index_file(index_file: &IndexFileMeta) -> bool { /// from the selection the read was limited to. This needs no `_ROW_ID` and no /// `first_row_id` — real primary-key tables never write one. /// -/// `allowed_rows` is the plan's per-file physical selection, when it has one. The -/// residual is evaluated over exactly those rows: an engine-supplied bucket split -/// can restrict a huge file to a handful of ranges, and reading the whole file only -/// to discard everything outside them afterwards would defeat the split. With no -/// selection every physical row is scanned, as before. +/// `allowed_rows` is the plan's per-file physical selection, keyed by data-file +/// name, with the plan's three states: a file it does not list is unrestricted and +/// the whole file is scanned; an empty range list excludes the file, which is +/// registered empty without a read; a non-empty list is scanned over exactly those +/// ranges, because an engine-supplied bucket split can restrict a huge file to a +/// handful of ranges and reading all of it to discard the rest would defeat the +/// split. /// -/// Every *active* data file in the split gets an entry, possibly empty. The -/// bucket search treats an absent entry and an empty entry identically (the file -/// contributes no candidates), so the empty entries only make the map cover every -/// active file. Non-active files (e.g. level-0 files the bucket search excludes) -/// are skipped entirely: they are never searched, so re-reading them would be -/// wasted IO. +/// Every *active* data file in the split gets an entry in the RESULT, possibly +/// empty, and that exhaustiveness is load-bearing. The search kernel reads a file's +/// absence from its selections as "unrestricted", so an active file missing here +/// would reach the search with no predicate applied at all -- the residual would be +/// silently dropped for it. (The merge below reads a residual's silence about a +/// file the PLAN listed as exclusion, so only a file both omit falls through, which +/// is exactly the case this exhaustiveness rules out.) Non-active files (e.g. +/// level-0 files the bucket search excludes) are skipped entirely: they are never +/// searched, so re-reading them would be wasted IO. /// /// `reader` must be predicate-free and project the residual columns; /// `residual.file_fields` are the fields the residual leaf indices point into @@ -2527,18 +2716,15 @@ async fn residual_positions_by_file( if !active_names.contains(file_meta.file_name.as_str()) { continue; } - let selection = match allowed_rows { - // A plan that lists nothing for a file permits nothing from it, whether - // the list is empty or the file is absent: both sides of the eventual - // intersection read absence that way. Registering it empty says so and - // costs no read. - Some(by_file) => match by_file.get(&file_meta.file_name) { - Some(ranges) if !ranges.is_empty() => Some(ranges.clone()), - _ => { - out.entry(file_meta.file_name.clone()).or_default(); - continue; - } - }, + // A file the plan lists an EMPTY range list for permits nothing; registering + // it empty says so and costs no read. A file the plan does not list at all + // is unrestricted, so the residual is evaluated over the whole file. + let selection = match allowed_rows.and_then(|by_file| by_file.get(&file_meta.file_name)) { + Some(ranges) if ranges.is_empty() => { + out.entry(file_meta.file_name.clone()).or_default(); + continue; + } + Some(ranges) => Some(ranges.clone()), None => None, }; let data_fields = reader.derive_data_fields(file_meta).await?; @@ -8102,30 +8288,39 @@ mod residual_positions_tests { #[tokio::test] async fn test_residual_does_not_read_a_file_the_plan_excludes() { - // A file the plan lists no rows for is registered empty and never opened. The - // empty entry is what tells the search the file contributes nothing; an - // absent one would mean the same, but then the map would not cover the split. + // An EMPTY range list is how a plan says "no rows of this file": it is + // registered empty and never opened. Absence means the opposite -- the plan + // narrowed nothing there -- so the residual reads the whole file. let (reader, split, active) = build_reader_and_split( "memory:/rpf_plan_excludes", &[("part-0.mosaic", vec![1, 2, 3], 0)], ) .await; - for allowed in [ - HashMap::from([("part-0.mosaic".to_string(), Vec::new())]), - HashMap::new(), - ] { - let map = residual_positions_by_file( - &reader, - &split, - &active, - &residual_id_gt(0), - Some(&allowed), - ) - .await - .unwrap(); - assert!(map.contains_key("part-0.mosaic")); - assert!(sorted(&map["part-0.mosaic"]).is_empty()); - } + + let excluded = HashMap::from([("part-0.mosaic".to_string(), Vec::new())]); + let map = residual_positions_by_file( + &reader, + &split, + &active, + &residual_id_gt(0), + Some(&excluded), + ) + .await + .unwrap(); + assert!(map.contains_key("part-0.mosaic")); + assert!(sorted(&map["part-0.mosaic"]).is_empty()); + + let unrestricted = HashMap::new(); + let map = residual_positions_by_file( + &reader, + &split, + &active, + &residual_id_gt(0), + Some(&unrestricted), + ) + .await + .unwrap(); + assert_eq!(sorted(&map["part-0.mosaic"]), vec![0, 1, 2]); } #[tokio::test] @@ -8290,10 +8485,17 @@ mod residual_positions_tests { .collect() } - fn listed(map: &HashMap, file: &str) -> Vec { - map.get(file) - .map(|positions| positions.iter().collect()) - .unwrap_or_default() + /// The positions a merged selection allows, expanded for readable assertions. + /// Test-only: the production path never expands a range. + fn listed(map: &FileRowSelections, file: &str) -> Vec { + match map.get(file) { + None => Vec::new(), + Some(FileRowSelection::Positions(positions)) => positions.iter().collect(), + Some(FileRowSelection::Ranges(ranges)) => ranges + .iter() + .flat_map(|range| (range.from() as u64)..=(range.to() as u64)) + .collect(), + } } #[test] @@ -8308,6 +8510,12 @@ mod residual_positions_tests { .unwrap() .expect("a plan restriction survives on its own"); assert_eq!(listed(&only_physical[0], "d0"), vec![1, 2]); + // Still intervals. Expanding them here is the unbounded step the plan side + // must never take, and the positions above cannot tell the two apart. + assert!( + matches!(only_physical[0]["d0"], FileRowSelection::Ranges(_)), + "the plan's ranges must reach the search as ranges" + ); let residual = vec![allow_list(&[("d0", &[3])])]; let only_residual = intersect_row_allow_lists(None, Some(residual), 1) @@ -8317,17 +8525,53 @@ mod residual_positions_tests { } #[test] - fn both_sides_intersect_and_a_file_either_omits_is_dropped() { - // `d0`: both list positions, so only the shared ones survive. `d1`: the - // residual kept nothing there, and its absence means "no rows", so the file - // must not come back unrestricted from the plan side. + fn both_sides_intersect_and_the_residual_stays_fail_closed() { + // `d0`: both restrict it, so only the shared positions survive. `d1`: the + // residual says nothing about it. The residual registers EVERY file the + // search can read from, so its silence is "no rows" -- the plan's ranges + // must not resurrect the file, and neither may its absence make it + // unrestricted. let physical = vec![range_allow_list(&[("d0", &[1, 2, 3]), ("d1", &[0, 1])])]; let residual = vec![allow_list(&[("d0", &[2, 3, 4])])]; let combined = intersect_row_allow_lists(Some(&physical), Some(residual), 1) .unwrap() .expect("both sides restrict"); assert_eq!(listed(&combined[0], "d0"), vec![2, 3]); + assert!( + combined[0]["d1"].is_excluded(), + "a file the residual omits must stay excluded" + ); + } + + #[test] + fn a_file_neither_side_restricts_stays_absent() { + // Absence is how "every row" is spelled. A merged map must not invent an + // entry for a file no one narrowed, or the ANN backend takes the filtered + // path for a query that filters nothing. + let physical = vec![range_allow_list(&[("d0", &[1])])]; + let combined = intersect_row_allow_lists(Some(&physical), None, 1) + .unwrap() + .expect("the plan restricts d0"); assert!(!combined[0].contains_key("d1")); + + let residual = vec![allow_list(&[("d0", &[1])])]; + let combined = intersect_row_allow_lists(Some(&physical), Some(residual), 1) + .unwrap() + .expect("both restrict d0"); + assert!(!combined[0].contains_key("d1")); + assert!(!combined[0].contains_key("d2")); + } + + #[test] + fn a_plan_that_restricts_nothing_produces_an_empty_selection_map() { + // The no-pre-filter split: the plan carries a map with no entries at all, + // and that must survive the merge as an empty map (which the ANN layer reads + // as "nothing to mask"), not become a per-file all-permitting mask. + let physical = vec![HashMap::new()]; + let combined = intersect_row_allow_lists(Some(&physical), None, 1) + .unwrap() + .expect("a split-driven plan is always Some"); + assert!(combined[0].is_empty()); } #[test] diff --git a/crates/paimon/src/vindex/pkvector/ann.rs b/crates/paimon/src/vindex/pkvector/ann.rs index e2be9f870..6940da4ee 100644 --- a/crates/paimon/src/vindex/pkvector/ann.rs +++ b/crates/paimon/src/vindex/pkvector/ann.rs @@ -25,6 +25,7 @@ use super::bucket::BucketAnnSegment; use super::data_invalid; use super::metric::{java_float_compare, VectorSearchMetric}; use super::result::PkVectorSearchResult; +use super::{FileRowSelection, FileRowSelections}; use crate::deletion_vector::DeletionVector; use crate::spec::{ PrimaryKeyIndexSourceFile as PkVectorSourceFile, @@ -42,21 +43,60 @@ use crate::vindex::range_reader::VindexFileReader; /// longer readable in this snapshot). Deletion vectors are applied only to active /// sources. /// -/// `residual_ranges` (when `Some`) restricts each source file to the physical row -/// positions allowed by a residual predicate on the data columns: `key = file -/// name`, `value = allowed physical positions`. A file with no entry (or an empty -/// entry) has no allowed rows and contributes nothing. When `residual_ranges` is -/// `Some`, a mask is always required (the residual can only narrow the live set), -/// so the result is always `Some`. Mirrors Java `rowRangesByFile`. +/// The most live row ids one ANN segment's mask may hold. /// -/// Returns `None` only when there is no residual, every source file is active, AND -/// no deletion vector is relevant — nothing to mask. Otherwise returns the masked -/// live ids. +/// Mirrors Java `LuminaVectorGlobalIndexReader.toScopedIds`, which refuses an +/// include-set above `Integer.MAX_VALUE` before allocating the dense array the +/// backend is handed. The number is Paimon's own for this exact quantity at this +/// exact seam, not one picked here. +/// +/// The limit is on what goes INTO the mask, charged before each insertion, so +/// nothing is ever allocated for a span that could not be handed over anyway. That +/// is what makes it a bound rather than a complaint: on the bucket-split route the +/// source row counts arrive on the wire (`PrimaryKeyIndexSourceMeta` is decoded from +/// the split's own `GlobalIndexMeta`, and only its sign is checked), and once +/// anything makes a mask necessary an unrestricted file is inserted wholesale at +/// whatever size that count claims -- `i64::MAX` never returns. +/// +/// Because it counts insertions, a segment that CLAIMS an impossible number of rows +/// but puts few or none in the mask -- all sources inactive, an explicitly empty +/// selection, a narrow range -- is not turned away. It is marginally stricter than +/// Java in one direction: Java checks the set after subtracting deletion vectors, +/// while the allocation being bounded here happens as rows go in. +const MAX_LIVE_ROW_IDS: u64 = i32::MAX as u64; + +/// Charge `rows` against what is left of [`MAX_LIVE_ROW_IDS`], before they are +/// inserted. +fn charge_live_rows(remaining: &mut u64, rows: u64) -> crate::Result<()> { + *remaining = remaining.checked_sub(rows).ok_or_else(|| { + data_invalid(format!( + "vector search would filter more than {MAX_LIVE_ROW_IDS} live rows in one \ + ANN segment" + )) + })?; + Ok(()) +} + +/// `row_selections` restricts each source file to the rows a pre-filter allows, +/// keyed by data-file name. A file with **no entry is unrestricted**, an empty +/// entry excludes it, and a non-empty one limits it — see [`FileRowSelection`]. +/// Mirrors Java `rowRangesByFile`. +/// +/// Returns `None` when nothing is restricted, every source file is active, AND no +/// deletion vector is relevant — nothing to mask, so the ANN backend searches +/// unfiltered. Otherwise returns the masked live ids. +/// +/// Java's condition is `allSourcesActive && deletionVectors.isEmpty() && +/// rowRangesByFile.isEmpty()`. The selections half is mirrored exactly (whole-map, +/// not this segment's own files). The deletion-vector half is NOT: ours is +/// segment-local, so a deletion vector on a file this segment does not index leaves +/// it unfiltered where Java would mask it. That predates the bucket-split route and +/// applies to the manifest route too. pub(crate) fn build_live_row_ids( source_files: &[PkVectorSourceFile], active_source_files: &HashSet, deletion_vectors: &HashMap>, - residual_ranges: Option<&HashMap>, + row_selections: Option<&FileRowSelections>, ) -> crate::Result> { let all_active = source_files .iter() @@ -64,13 +104,20 @@ pub(crate) fn build_live_row_ids( let has_relevant_dv = source_files .iter() .any(|f| deletion_vectors.contains_key(f.file_name())); - if residual_ranges.is_none() && all_active && !has_relevant_dv { + // Java's own condition, `rowRangesByFile.isEmpty()`, over the whole bucket-level + // map. Narrowing it to "no entry for one of THIS segment's sources" would leave + // more segments unfiltered, but it also changes which backend entry point they + // take (`search` vs `search_with_filter`), and those can differ in recall. Not + // worth diverging for. + let nothing_restricted = row_selections.is_none_or(FileRowSelections::is_empty); + if nothing_restricted && all_active && !has_relevant_dv { return Ok(None); } let mut live = roaring::RoaringTreemap::new(); let mut deleted = roaring::RoaringTreemap::new(); let mut file_offset: u64 = 0; + let mut budget = MAX_LIVE_ROW_IDS; for source_file in source_files { let row_count = u64::try_from(source_file.row_count()) .map_err(|_| data_invalid("vector source row count must not be negative"))?; @@ -79,40 +126,65 @@ pub(crate) fn build_live_row_ids( .ok_or_else(|| data_invalid("vector source row counts overflow u64"))?; let active = active_source_files.contains(source_file.file_name()); if active && row_count > 0 { - match residual_ranges { - // No residual: the whole active file range is live. + match row_selections.and_then(|selections| selections.get(source_file.file_name())) { + // Unrestricted: the whole active file range is live. This is the + // no-entry case Java spells as `rowRanges == null`. None => { + charge_live_rows(&mut budget, row_count)?; live.insert_range(file_offset..end); } - // Residual present: only allowed physical positions of this file - // become live, mapped into global ordinal space (position + - // file_offset). A missing/empty entry allows no rows. - Some(ranges) => { - if let Some(allowed) = ranges.get(source_file.file_name()) { - // A producer that restricts only some files leaves the rest - // unrestricted, and an adapter has to spell that out as an - // explicit whole-file allow-list. Insert it as one range - // rather than walking every position, which would cost one - // insert per row of the file. `len` plus a maximum of - // `row_count - 1` can only describe the full set, and it - // subsumes the per-position bound check below. - if allowed.len() == row_count && allowed.max() == Some(row_count - 1) { - live.insert_range(file_offset..end); - } else { - for position in allowed.iter() { - if position >= row_count { - return Err(data_invalid(format!( - "residual position {position} is out of range for source file {} ({} rows)", - source_file.file_name(), - row_count - ))); - } - let global = - file_offset.checked_add(position).ok_or_else(|| { - data_invalid("vector residual position overflows u64") - })?; - live.insert(global); + // Restricted to intervals. Added interval-wise, never position-wise: + // these bounds ride in on an engine-supplied split, so walking them + // would be unbounded work driven by untrusted numbers. Mirrors Java + // `live.addRange(range.addOffset(fileOffset))`. + Some(FileRowSelection::Ranges(ranges)) => { + for range in ranges { + // Java checks each range against the SOURCE file's row count. + // On the bucket-split route that count came off the wire as + // well (`PrimaryKeyIndexSourceMeta` is decoded from the + // split's own `GlobalIndexMeta`), so this rejects a range + // that disagrees with its own segment -- it is not a + // resource bound. + let from = u64::try_from(range.from()).map_err(|_| { + data_invalid("vector pre-filter range bound must not be negative") + })?; + let to = u64::try_from(range.to()).map_err(|_| { + data_invalid("vector pre-filter range bound must not be negative") + })?; + if to >= row_count { + return Err(data_invalid(format!( + "pre-filter range [{from}, {to}] is out of range for source file {} ({} rows)", + source_file.file_name(), + row_count + ))); + } + charge_live_rows(&mut budget, to - from + 1)?; + live.insert_range((file_offset + from)..=(file_offset + to)); + } + } + // Restricted to positions a residual predicate left behind. Bounded + // by the rows that read actually returned, so walking them is safe. + Some(FileRowSelection::Positions(allowed)) => { + // `len` plus a maximum of `row_count - 1` can only describe the + // full set; inserting it as one range subsumes the per-position + // bound check below. + if allowed.len() == row_count && allowed.max() == Some(row_count - 1) { + charge_live_rows(&mut budget, row_count)?; + live.insert_range(file_offset..end); + } else { + charge_live_rows(&mut budget, allowed.len())?; + for position in allowed.iter() { + if position >= row_count { + return Err(data_invalid(format!( + "residual position {position} is out of range for source file {} ({} rows)", + source_file.file_name(), + row_count + ))); } + let global = file_offset.checked_add(position).ok_or_else(|| { + data_invalid("vector residual position overflows u64") + })?; + live.insert(global); } } } @@ -121,6 +193,15 @@ pub(crate) fn build_live_row_ids( if active { if let Some(dv) = deletion_vectors.get(source_file.file_name()) { for position in dv.iter() { + // A position past this file's own rows would land inside the NEXT + // source file's ordinal range and delete one of its rows instead. + if position >= row_count { + return Err(data_invalid(format!( + "deleted position {position} is out of range for source file {} ({} rows)", + source_file.file_name(), + row_count + ))); + } let global = file_offset.checked_add(position).ok_or_else(|| { data_invalid("vector source deleted position overflows u64") })?; @@ -144,7 +225,7 @@ pub(crate) fn map_ann_results( source_meta: &PkVectorSourceMeta, active_source_files: &HashSet, deletion_vectors: &HashMap>, - residual_ranges: Option<&HashMap>, + row_selections: Option<&FileRowSelections>, metric: VectorSearchMetric, ) -> crate::Result> { let mut results = Vec::with_capacity(scored.len()); @@ -166,11 +247,13 @@ pub(crate) fn map_ann_results( ))); } } - if let Some(ranges) = residual_ranges { - let allowed = ranges.get(&data_file_name).is_some_and(|r| r.contains(pos)); - if !allowed { + // A file with no entry is unrestricted, so only an entry can reject. + if let Some(selection) = + row_selections.and_then(|selections| selections.get(&data_file_name)) + { + if !selection.contains(pos) { return Err(data_invalid(format!( - "ANN segment returned row position {row_position} in {data_file_name} outside the residual pre-filter" + "ANN segment returned row position {row_position} in {data_file_name} outside the row selection for that file" ))); } } @@ -214,8 +297,8 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { /// Search one ANN segment for a batch of query vectors, returning one /// BEST_FIRST result list per query (outer index aligned to `queries`). The - /// live-row mask (residual ∩ DV) is query-independent, so it is built once and - /// shared across all queries; only the per-query scores differ. Buffered callers + /// live-row mask (selections ∩ DV) is query-independent, so it is built once and + /// then cloned into each query; only the per-query scores differ. Buffered callers /// pass the bytes from `load_segment` by value so they cannot outlive the leaf. #[allow(clippy::too_many_arguments)] fn search_batch( @@ -228,7 +311,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, - residual_ranges: Option<&HashMap>, + row_selections: Option<&FileRowSelections>, ) -> crate::Result>>; #[allow(clippy::too_many_arguments)] @@ -242,7 +325,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, - residual_ranges: Option<&HashMap>, + row_selections: Option<&FileRowSelections>, ) -> crate::Result>> { match segment_source { AnnSegmentSource::Buffered(bytes) => self.search_batch( @@ -254,7 +337,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files, deletion_vectors, search_options, - residual_ranges, + row_selections, ), AnnSegmentSource::Vindex(_) => Err(data_invalid( "ANN searcher does not support a range-backed segment source", @@ -276,7 +359,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, - residual_ranges: Option<&HashMap>, + row_selections: Option<&FileRowSelections>, ) -> crate::Result> { let mut results = self.search_batch( segment, @@ -287,7 +370,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files, deletion_vectors, search_options, - residual_ranges, + row_selections, )?; if results.len() != 1 { return Err(data_invalid(format!( @@ -309,7 +392,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, - residual_ranges: Option<&HashMap>, + row_selections: Option<&FileRowSelections>, ) -> crate::Result> { let mut results = self.search_batch_source( segment, @@ -320,7 +403,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files, deletion_vectors, search_options, - residual_ranges, + row_selections, )?; if results.len() != 1 { return Err(data_invalid(format!( @@ -452,7 +535,7 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, - residual_ranges: Option<&HashMap>, + row_selections: Option<&FileRowSelections>, ) -> crate::Result>> { self.search_batch_source( segment, @@ -463,7 +546,7 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { active_source_files, deletion_vectors, search_options, - residual_ranges, + row_selections, ) } @@ -477,20 +560,23 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, - residual_ranges: Option<&HashMap>, + row_selections: Option<&FileRowSelections>, ) -> crate::Result>> { if limit == 0 { return Err(data_invalid("vector search limit must be positive")); } let source_files = segment.source_meta.source_files(); // The live-row mask depends only on the segment's sources, the active set, - // the deletion vectors, and the residual — none of which vary by query — - // so it is built once and shared across every query's search. + // the deletion vectors, and the selections — none of which vary by query — + // so it is BUILT once, then cloned into each query's search. Handing every + // query one `Arc` instead would save those clones, but it also moves a + // filtered search from Lumina's per-query scalar calls onto its native batch + // call; worth doing, not here. let live = build_live_row_ids( source_files, active_source_files, deletion_vectors, - residual_ranges, + row_selections, )?; let mut searches = Vec::with_capacity(queries.len()); for query in queries { @@ -519,7 +605,7 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { &segment.source_meta, active_source_files, deletion_vectors, - residual_ranges, + row_selections, metric, )? } @@ -564,6 +650,26 @@ mod tests { Arc::new(DeletionVector::from_bitmap(bitmap)) } + /// A residual selection: the physical positions of one file that passed a data + /// predicate. + fn positions(at: &[u64]) -> FileRowSelection { + let mut t = roaring::RoaringTreemap::new(); + for &p in at { + t.insert(p); + } + FileRowSelection::Positions(t) + } + + /// A pre-filter selection in the interval form an engine's split carries. + fn ranges(bounds: &[(i64, i64)]) -> FileRowSelection { + FileRowSelection::Ranges( + bounds + .iter() + .map(|(from, to)| crate::table::RowRange::new(*from, *to)) + .collect(), + ) + } + fn active_set(names: &[&str]) -> HashSet { names.iter().map(|n| (*n).to_string()).collect() } @@ -584,6 +690,50 @@ mod tests { .is_none()); } + #[test] + fn an_empty_selection_map_leaves_the_ann_search_unfiltered() { + // The no-pre-filter shape: a Java split that narrowed nothing. Java returns + // null here (`rowRangesByFile.isEmpty()`), and the ANN backend then searches + // without a filter. Handing it an all-permitting mask instead costs an + // 8-byte id per live row and takes the filtered code path. + let files = [PkVectorSourceFile::new("f0".into(), 3).unwrap()]; + let selections = HashMap::new(); + assert!( + build_live_row_ids( + &files, + &active_set(&["f0"]), + &HashMap::new(), + Some(&selections) + ) + .unwrap() + .is_none(), + "nothing is restricted, so nothing should be masked" + ); + } + + #[test] + fn a_source_file_no_one_restricted_stays_whole() { + // Java reads a missing entry as "every row of this file" and records one only + // for a file its pre-filter narrowed. f1 is narrowed, f0 is not, so f0 must + // stay whole rather than drop out of the search. + let files = vec![ + PkVectorSourceFile::new("f0".into(), 3).unwrap(), + PkVectorSourceFile::new("f1".into(), 2).unwrap(), + ]; + let mut selections = HashMap::new(); + selections.insert("f1".to_string(), positions(&[1])); + let live = build_live_row_ids( + &files, + &active_set(&["f0", "f1"]), + &HashMap::new(), + Some(&selections), + ) + .unwrap() + .unwrap(); + // f0 global 0,1,2 all live; f1 global 3,4 restricted to position 1 -> 4. + assert_eq!(live.iter().collect::>(), vec![0, 1, 2, 4]); + } + #[test] fn test_build_live_row_ids_masks_inactive_source_ordinal_range() { // f0 rows 0..3 (global 0,1,2), f1 rows 0..2 (global 3,4). f1 is inactive, @@ -710,7 +860,8 @@ mod tests { move |_segment: &BucketAnnSegment, _bytes: Bytes, searches: &[VectorSearch]| { let search = &searches[0]; *scorer_limit.lock().unwrap() = search.limit; - *scorer_has_filter.lock().unwrap() = search.include_row_ids.is_some(); + *scorer_has_filter.lock().unwrap() = + search.effective_include_row_ids().is_some(); let mut scores = HashMap::new(); scores.insert(3u64, 0.5f32); // -> (f1, 0) scores.insert(0u64, 0.25f32); // -> (f0, 0), l2 dist 3.0 @@ -836,9 +987,11 @@ mod tests { #[test] fn test_build_live_row_ids_residual_intersects_with_active_and_dv() { // f0 rows 0..3 (global 0,1,2), f1 rows 0..2 (global 3,4). Both active. - // dv on f0 deletes pos1 (global 1). residual allows f0={0,1}, f1 has no - // entry (empty allow). Result: f0 keeps {0} (1 is residual-allowed but - // deleted, 2 not residual-allowed); f1 contributes nothing. + // dv on f0 deletes pos1 (global 1). residual allows f0={0,1}; f1 has no + // entry, which is "unrestricted", so it keeps both its rows. Result: f0 + // keeps {0} (1 is residual-allowed but deleted, 2 not residual-allowed), + // f1 keeps globals 3 and 4. (In production the residual producer registers + // every active file, so an absent one does not arise there.) let files = vec![ PkVectorSourceFile::new("f0".into(), 3).unwrap(), PkVectorSourceFile::new("f1".into(), 2).unwrap(), @@ -846,11 +999,14 @@ mod tests { let mut dvs = HashMap::new(); dvs.insert("f0".to_string(), dv(&[1])); let mut residual = HashMap::new(); - residual.insert("f0".to_string(), treemap(&[0, 1])); + residual.insert( + "f0".to_string(), + FileRowSelection::Positions(treemap(&[0, 1])), + ); let live = build_live_row_ids(&files, &active_set(&["f0", "f1"]), &dvs, Some(&residual)) .unwrap() .unwrap(); - assert_eq!(live.iter().collect::>(), vec![0]); + assert_eq!(live.iter().collect::>(), vec![0, 3, 4]); } #[test] @@ -864,8 +1020,14 @@ mod tests { ]; let active = active_set(&["f0", "f1"]); let mut residual = HashMap::new(); - residual.insert("f0".to_string(), treemap(&[0, 1, 2])); - residual.insert("f1".to_string(), treemap(&[0, 1])); + residual.insert( + "f0".to_string(), + FileRowSelection::Positions(treemap(&[0, 1, 2])), + ); + residual.insert( + "f1".to_string(), + FileRowSelection::Positions(treemap(&[0, 1])), + ); let spelled_out = build_live_row_ids(&files, &active, &HashMap::new(), Some(&residual)) .unwrap() @@ -884,7 +1046,10 @@ mod tests { let mut dvs = HashMap::new(); dvs.insert("f0".to_string(), dv(&[1])); let mut residual = HashMap::new(); - residual.insert("f0".to_string(), treemap(&[0, 1, 2])); + residual.insert( + "f0".to_string(), + FileRowSelection::Positions(treemap(&[0, 1, 2])), + ); let live = build_live_row_ids(&files, &active_set(&["f0"]), &dvs, Some(&residual)) .unwrap() .unwrap(); @@ -900,8 +1065,8 @@ mod tests { PkVectorSourceFile::new("f1".into(), 2).unwrap(), ]; let mut residual = HashMap::new(); - residual.insert("f0".to_string(), treemap(&[2])); - residual.insert("f1".to_string(), treemap(&[1])); + residual.insert("f0".to_string(), FileRowSelection::Positions(treemap(&[2]))); + residual.insert("f1".to_string(), FileRowSelection::Positions(treemap(&[1]))); let live = build_live_row_ids( &files, &active_set(&["f0", "f1"]), @@ -919,7 +1084,10 @@ mod tests { // present, a mask is always required. let files = [PkVectorSourceFile::new("f0".into(), 3).unwrap()]; let mut residual = HashMap::new(); - residual.insert("f0".to_string(), treemap(&[0, 2])); + residual.insert( + "f0".to_string(), + FileRowSelection::Positions(treemap(&[0, 2])), + ); let live = build_live_row_ids( &files, &active_set(&["f0"]), @@ -937,7 +1105,10 @@ mod tests { // naming position 3 is out of range and must fail loud, not be skipped. let files = source_meta(&[("f0", 3)]); let mut residual = HashMap::new(); - residual.insert("f0".to_string(), treemap(&[0, 3])); + residual.insert( + "f0".to_string(), + FileRowSelection::Positions(treemap(&[0, 3])), + ); let err = build_live_row_ids( files.source_files(), &active_set(&["f0"]), @@ -948,13 +1119,35 @@ mod tests { assert!(err.to_string().contains("out of range")); } + #[test] + fn map_ann_results_accepts_a_hit_in_a_file_no_one_restricted() { + // The tri-state on the validation side: the map restricts f0, says nothing + // about f1, and a hit in f1 must be accepted. Reading f1's absence as "no + // rows" would turn every hit in an unrestricted sibling into an error. + let meta = source_meta(&[("f0", 3), ("f1", 5)]); + let mut selections = HashMap::new(); + selections.insert("f0".to_string(), ranges(&[(0, 0)])); + let results = map_ann_results( + &[(3, 0.5)], // ordinal 3 -> (f1, 0) + &meta, + &active_set(&["f0", "f1"]), + &HashMap::new(), + Some(&selections), + VectorSearchMetric::L2, + ) + .expect("a hit in an unrestricted file is allowed"); + assert_eq!(results.len(), 1); + assert_eq!(results[0].data_file_name, "f1"); + assert_eq!(results[0].row_position, 0); + } + #[test] fn test_map_ann_results_rejects_hit_outside_residual_allow_list() { // ordinal 1 -> (f0, 1). Residual allows only {0} in f0, so a hit at position 1 // (e.g. an ANN reader that ignored include_row_ids) must fail loud. let meta = source_meta(&[("f0", 3)]); let mut residual = HashMap::new(); - residual.insert("f0".to_string(), treemap(&[0])); + residual.insert("f0".to_string(), FileRowSelection::Positions(treemap(&[0]))); let err = map_ann_results( &[(1u64, 0.5)], &meta, @@ -964,7 +1157,7 @@ mod tests { VectorSearchMetric::L2, ) .unwrap_err(); - assert!(err.to_string().contains("residual")); + assert!(err.to_string().contains("outside the row selection")); } #[test] @@ -979,8 +1172,7 @@ mod tests { Box::new( move |_segment: &BucketAnnSegment, _bytes: Bytes, searches: &[VectorSearch]| { *scorer_rows.lock().unwrap() = searches[0] - .include_row_ids - .as_ref() + .effective_include_row_ids() .map(|t| t.iter().collect::>()); Ok(vec![None; searches.len()]) }, @@ -988,7 +1180,10 @@ mod tests { ); let segment = BucketAnnSegment::for_test(source_meta(&[("f0", 3)])); let mut residual = HashMap::new(); - residual.insert("f0".to_string(), treemap(&[0, 2])); + residual.insert( + "f0".to_string(), + FileRowSelection::Positions(treemap(&[0, 2])), + ); searcher .search( &segment, @@ -1005,6 +1200,296 @@ mod tests { assert_eq!(seen_rows.lock().unwrap().clone(), Some(vec![0, 2])); } + #[test] + fn the_no_prefilter_shape_reaches_the_backend_with_no_filter_at_all() { + // The backend-facing half of the no-pre-filter case. `include_row_ids` is + // what Lumina turns into a Vec of every live id and a filtered search; + // for a query that filters nothing it must stay unset, exactly as on the + // manifest route. + use std::sync::{Arc, Mutex}; + let seen: Arc>> = Arc::new(Mutex::new(None)); + let scorer_seen = Arc::clone(&seen); + let searcher = vindex_searcher( + "embedding", + Box::new( + move |_segment: &BucketAnnSegment, _bytes: Bytes, searches: &[VectorSearch]| { + *scorer_seen.lock().unwrap() = + Some(searches[0].effective_include_row_ids().is_some()); + Ok(vec![None; searches.len()]) + }, + ), + ); + let segment = BucketAnnSegment::for_test(source_meta(&[("f0", 3)])); + let no_prefilter: FileRowSelections = HashMap::new(); + searcher + .search( + &segment, + Bytes::new(), + &[0.0, 0.0], + VectorSearchMetric::L2, + 2, + &active_set(&["f0"]), + &HashMap::new(), + &HashMap::new(), + Some(&no_prefilter), + ) + .unwrap(); + assert_eq!( + seen.lock().unwrap().clone(), + Some(false), + "a split that narrowed nothing must not take the filtered ANN path" + ); + } + + #[test] + fn a_range_past_the_source_row_count_is_rejected() { + // The range bounds ride in on an engine-supplied split. Java checks each one + // against the SOURCE file's row count, so a range that disagrees with its own + // segment is rejected. On this route that count came off the wire as well, so + // this is a consistency check between the split's own two numbers -- NOT a + // bound on how much the insert can allocate. + let files = [PkVectorSourceFile::new("f0".into(), 3).unwrap()]; + let mut selections = HashMap::new(); + selections.insert("f0".to_string(), ranges(&[(0, 1 << 40)])); + let error = build_live_row_ids( + &files, + &active_set(&["f0"]), + &HashMap::new(), + Some(&selections), + ) + .map(|_| ()) + .expect_err("a range beyond the source file cannot be inserted"); + assert!(error.to_string().contains("out of range"), "{error}"); + } + + #[test] + fn a_deleted_position_past_its_own_source_file_is_rejected() { + // Offsets are cumulative, so a position past f0's rows would land inside f1's + // ordinal range and delete one of ITS rows. + let files = vec![ + PkVectorSourceFile::new("f0".into(), 2).unwrap(), + PkVectorSourceFile::new("f1".into(), 2).unwrap(), + ]; + let mut dvs = HashMap::new(); + dvs.insert("f0".to_string(), dv(&[3])); + let error = build_live_row_ids(&files, &active_set(&["f0", "f1"]), &dvs, None) + .map(|_| ()) + .expect_err("a deleted position must stay inside its own file"); + assert!(error.to_string().contains("out of range"), "{error}"); + } + + #[test] + fn a_live_set_beyond_what_the_backend_can_filter_is_rejected() { + // A mask is necessary (one sibling is restricted), so `huge` is inserted + // wholesale at whatever its wire-supplied row count claims. Without a limit + // on what gets inserted this OOM-kills the process. + let files = vec![ + PkVectorSourceFile::new("huge".into(), i64::MAX).unwrap(), + PkVectorSourceFile::new("small".into(), 1).unwrap(), + ]; + let mut selections = HashMap::new(); + selections.insert("small".to_string(), ranges(&[(0, 0)])); + let error = build_live_row_ids( + &files, + &active_set(&["huge", "small"]), + &HashMap::new(), + Some(&selections), + ) + .map(|_| ()) + .expect_err("more live rows than the backend can be handed"); + assert!(error.to_string().contains("more than"), "{error}"); + } + + #[test] + fn a_deletion_vector_alone_also_reaches_the_live_row_limit() { + // Same exposure with no pre-filter at all: one deletion vector on one of + // this segment's own sources is enough to make a mask necessary. + let files = vec![PkVectorSourceFile::new("huge".into(), i64::MAX).unwrap()]; + let mut dvs = HashMap::new(); + dvs.insert("huge".to_string(), dv(&[0])); + let error = build_live_row_ids(&files, &active_set(&["huge"]), &dvs, None) + .map(|_| ()) + .expect_err("more live rows than the backend can be handed"); + assert!(error.to_string().contains("more than"), "{error}"); + } + + #[test] + fn a_huge_claim_with_nothing_live_is_not_rejected() { + // The limit is on rows actually put in the mask, not on what the metadata + // claims. These three all claim more rows than any backend could filter and + // all leave the mask empty, so none of them may be turned away. + let huge = vec![PkVectorSourceFile::new("huge".into(), i64::MAX).unwrap()]; + + // Explicitly excluded. + let mut excluded = HashMap::new(); + excluded.insert("huge".to_string(), ranges(&[])); + let live = build_live_row_ids( + &huge, + &active_set(&["huge"]), + &HashMap::new(), + Some(&excluded), + ) + .expect("an excluded file puts nothing in the mask") + .expect("a selection was present"); + assert!(live.is_empty()); + + // Inactive: its whole ordinal range is masked out anyway. + let live = build_live_row_ids(&huge, &active_set(&[]), &HashMap::new(), None) + .expect("an inactive source puts nothing in the mask") + .expect("a source was inactive"); + assert!(live.is_empty()); + + // Narrowly restricted: two rows out of an impossible claim. + let mut narrow = HashMap::new(); + narrow.insert("huge".to_string(), ranges(&[(0, 1)])); + let live = build_live_row_ids( + &huge, + &active_set(&["huge"]), + &HashMap::new(), + Some(&narrow), + ) + .expect("a narrow selection puts two rows in the mask") + .expect("a selection was present"); + assert_eq!(live.iter().collect::>(), vec![0, 1]); + } + + #[test] + fn the_live_row_limit_is_javas_dense_filter_limit() { + // Java rejects an include-set above `Integer.MAX_VALUE` before it allocates + // the dense array (`LuminaVectorGlobalIndexReader.toScopedIds`). Exactly that + // many rows is what a backend can still be handed; one more is not. + let at_limit = vec![PkVectorSourceFile::new("f0".into(), i32::MAX as i64).unwrap()]; + let mut dvs = HashMap::new(); + dvs.insert("f0".to_string(), dv(&[0])); + assert!( + build_live_row_ids(&at_limit, &active_set(&["f0"]), &dvs, None).is_ok(), + "exactly the limit is allowed" + ); + + // One row past it, even though the deletion vector would bring the FINAL set + // back under: the charge is on insertion, because that is where the + // allocation happens. + let over = vec![PkVectorSourceFile::new("f0".into(), i32::MAX as i64 + 1).unwrap()]; + assert!( + build_live_row_ids(&over, &active_set(&["f0"]), &dvs, None).is_err(), + "one row past the limit is not" + ); + } + + /// A treemap holding `0..=to`, built as one run so the test itself stays cheap. + fn positions_through(to: u64) -> roaring::RoaringTreemap { + let mut t = roaring::RoaringTreemap::new(); + t.insert_range(0..=to); + t + } + + #[test] + fn an_oversized_range_selection_is_charged() { + // The `Ranges` charge site, distinct from the unrestricted one: the file is + // restricted, so it never reaches the whole-file insert. + let files = vec![PkVectorSourceFile::new("f0".into(), i32::MAX as i64 + 1).unwrap()]; + let mut selections = HashMap::new(); + selections.insert("f0".to_string(), ranges(&[(0, i32::MAX as i64)])); + let error = build_live_row_ids( + &files, + &active_set(&["f0"]), + &HashMap::new(), + Some(&selections), + ) + .map(|_| ()) + .expect_err("a range this wide cannot be filtered"); + assert!(error.to_string().contains("more than"), "{error}"); + } + + #[test] + fn an_oversized_whole_file_position_set_is_charged() { + // The `Positions` whole-file shortcut: `len` equals the row count and the + // maximum is the last row, so it inserts as one range. + let rows = i32::MAX as u64 + 1; + let files = vec![PkVectorSourceFile::new("f0".into(), rows as i64).unwrap()]; + let mut selections = HashMap::new(); + selections.insert( + "f0".to_string(), + FileRowSelection::Positions(positions_through(rows - 1)), + ); + let error = build_live_row_ids( + &files, + &active_set(&["f0"]), + &HashMap::new(), + Some(&selections), + ) + .map(|_| ()) + .expect_err("a whole-file position set this large cannot be filtered"); + assert!(error.to_string().contains("more than"), "{error}"); + } + + #[test] + fn an_oversized_sparse_position_set_is_charged() { + // The per-position `Positions` path: the set is large but is NOT the whole + // file, so the shortcut above does not apply and the loop would walk it. + let rows = i32::MAX as u64 + 5; + let files = vec![PkVectorSourceFile::new("f0".into(), rows as i64).unwrap()]; + let mut selections = HashMap::new(); + selections.insert( + "f0".to_string(), + FileRowSelection::Positions(positions_through(i32::MAX as u64)), + ); + let error = build_live_row_ids( + &files, + &active_set(&["f0"]), + &HashMap::new(), + Some(&selections), + ) + .map(|_| ()) + .expect_err("a position set this large cannot be filtered"); + assert!(error.to_string().contains("more than"), "{error}"); + } + + #[test] + fn an_empty_range_list_excludes_the_source_file() { + // Java's empty `List`: present and permitting nothing, the opposite of + // absent. + let files = vec![ + PkVectorSourceFile::new("f0".into(), 3).unwrap(), + PkVectorSourceFile::new("f1".into(), 2).unwrap(), + ]; + let mut selections = HashMap::new(); + selections.insert("f0".to_string(), ranges(&[])); + let live = build_live_row_ids( + &files, + &active_set(&["f0", "f1"]), + &HashMap::new(), + Some(&selections), + ) + .unwrap() + .unwrap(); + // f0 contributes nothing; f1 is unlisted, so both its rows stay live. + assert_eq!(live.iter().collect::>(), vec![3, 4]); + } + + #[test] + fn range_selections_map_onto_their_source_file_offset() { + // f0 rows global 0,1,2; f1 rows global 3,4. f1 restricted to [1, 1]. This + // pins the offset arithmetic, not the interval-wise insertion -- a + // per-position loop would produce the same bitmap. + let files = vec![ + PkVectorSourceFile::new("f0".into(), 3).unwrap(), + PkVectorSourceFile::new("f1".into(), 2).unwrap(), + ]; + let mut selections = HashMap::new(); + selections.insert("f0".to_string(), ranges(&[(0, 0), (2, 2)])); + selections.insert("f1".to_string(), ranges(&[(1, 1)])); + let live = build_live_row_ids( + &files, + &active_set(&["f0", "f1"]), + &HashMap::new(), + Some(&selections), + ) + .unwrap() + .unwrap(); + assert_eq!(live.iter().collect::>(), vec![0, 2, 4]); + } + #[test] fn test_search_batch_of_one_equals_single_query() { // The single-query `search` wrapper must return exactly what @@ -1061,7 +1546,7 @@ mod tests { #[test] fn test_search_batch_returns_independent_per_query_results() { // Two queries route to different synthetic scores; each result list is - // mapped from that query's own scores, with a shared live-row mask. + // mapped from that query's own scores, over the same live-row mask. let searcher = vindex_searcher( "embedding", Box::new( diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs index bf67e9bb6..7e0a0627b 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -26,6 +26,7 @@ use super::ann::PkVectorAnnSearcher; use super::data_invalid; use super::metric::{java_float_compare, VectorSearchMetric}; use super::result::PkVectorSearchResult; +use super::{FileRowSelection, FileRowSelections}; use crate::deletion_vector::DeletionVector; use crate::spec::PrimaryKeyIndexSourceMeta as PkVectorSourceMeta; use crate::vindex::executor::{ @@ -162,7 +163,7 @@ fn validate_per_query_len( /// closure borrows the allow-list for its lifetime. fn position_excluder( dv: Option>, - residual_allowed: Option<&roaring::RoaringTreemap>, + selection: Option<&FileRowSelection>, ) -> impl Fn(i64) -> bool + Sync + '_ { move |position: i64| -> bool { let dv_deleted = match &dv { @@ -174,12 +175,12 @@ fn position_excluder( if dv_deleted { return true; } - match residual_allowed { - // No residual restriction: the row is allowed. + match selection { + // No entry: the file is unrestricted, so the row is allowed. None => false, - // Residual present: exclude positions outside the allow-list. - Some(allowed) => match u64::try_from(position) { - Ok(p) => !allowed.contains(p), + // Restricted: exclude positions the selection does not list. + Some(selection) => match u64::try_from(position) { + Ok(p) => !selection.contains(p), Err(_) => true, }, } @@ -353,12 +354,11 @@ enum BucketLeaf { /// `ann_searcher` may be `None` only when there are no ANN segments; segments /// present with `None` is an error. /// -/// `residual_ranges` (when `Some`) is a residual-predicate allow-list keyed by -/// data-file name whose value is the set of physical row positions in that file -/// that pass the predicate; only those rows may produce candidates. `None` applies -/// no residual restriction (every row is allowed). A file absent from the map (or -/// with an empty set) has no allowed rows and produces no candidates. Mirrors Java -/// `rowRangesByFile`. +/// `row_selections` is the pre-filter allow-list keyed by data-file name: a file +/// with **no entry is unrestricted**, an empty entry excludes it (the file is +/// skipped without a read), and a non-empty one limits which of its rows may +/// produce candidates. `None` restricts nothing at all. Mirrors Java +/// `rowRangesByFile`; see [`FileRowSelection`]. #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] pub(crate) async fn bucket_search( @@ -381,7 +381,7 @@ pub(crate) async fn bucket_search( exact_limit: usize, search_options: &HashMap, skip_exact_fallback: bool, - residual_ranges: Option<&HashMap>, + row_selections: Option<&FileRowSelections>, concurrency: usize, search_budget: Option, ) -> crate::Result { @@ -516,7 +516,7 @@ pub(crate) async fn bucket_search( let ann_shared = searcher.as_ref().map(|searcher| { ( searcher.clone(), - residual_ranges.map(|r| Arc::new(r.clone())), + row_selections.map(|selections| Arc::new(selections.clone())), Arc::new(active_source_files.clone()), Arc::new(deletion_vectors.clone()), Arc::new(search_options.clone()), @@ -533,15 +533,16 @@ pub(crate) async fn bucket_search( if covered.contains(&file.file_name) { continue; } - let residual_allowed: Option<&roaring::RoaringTreemap> = match residual_ranges { - Some(ranges) => match ranges.get(&file.file_name) { - Some(allowed) if !allowed.is_empty() => Some(allowed), - _ => continue, - }, - None => None, - }; + // No entry means unrestricted; an empty one excludes the file, which is + // skipped without a read. Mirrors Java + // `if (rowRanges != null && rowRanges.isEmpty()) continue;`. + let selection: Option<&FileRowSelection> = + match row_selections.and_then(|selections| selections.get(&file.file_name)) { + Some(selection) if selection.is_excluded() => continue, + other => other, + }; let dv = deletion_vectors.get(&file.file_name).cloned(); - exact_tasks.push((file, Box::new(position_excluder(dv, residual_allowed)))); + exact_tasks.push((file, Box::new(position_excluder(dv, selection)))); } } @@ -658,7 +659,7 @@ pub(crate) async fn bucket_search_batch( exact_limit: usize, search_options: &HashMap, skip_exact_fallback: bool, - residual_ranges: Option<&HashMap>, + row_selections: Option<&FileRowSelections>, concurrency: usize, search_budget: Option, ) -> crate::Result> { @@ -681,7 +682,7 @@ pub(crate) async fn bucket_search_batch( exact_limit, search_options, skip_exact_fallback, - residual_ranges, + row_selections, concurrency, search_budget, ) @@ -806,7 +807,7 @@ pub(crate) async fn bucket_search_batch( Arc::new(queries.iter().map(|q| q.to_vec()).collect()); ( searcher.clone(), - residual_ranges.map(|r| Arc::new(r.clone())), + row_selections.map(|selections| Arc::new(selections.clone())), Arc::new(active_source_files.clone()), Arc::new(deletion_vectors.clone()), Arc::new(search_options.clone()), @@ -821,15 +822,16 @@ pub(crate) async fn bucket_search_batch( if covered.contains(&file.file_name) { continue; } - let residual_allowed: Option<&roaring::RoaringTreemap> = match residual_ranges { - Some(ranges) => match ranges.get(&file.file_name) { - Some(allowed) if !allowed.is_empty() => Some(allowed), - _ => continue, - }, - None => None, - }; + // No entry means unrestricted; an empty one excludes the file, which is + // skipped without a read. Mirrors Java + // `if (rowRanges != null && rowRanges.isEmpty()) continue;`. + let selection: Option<&FileRowSelection> = + match row_selections.and_then(|selections| selections.get(&file.file_name)) { + Some(selection) if selection.is_excluded() => continue, + other => other, + }; let dv = deletion_vectors.get(&file.file_name).cloned(); - exact_tasks.push((file, Box::new(position_excluder(dv, residual_allowed)))); + exact_tasks.push((file, Box::new(position_excluder(dv, selection)))); } } @@ -1075,7 +1077,7 @@ mod tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _residual_ranges: Option<&HashMap>, + _row_selections: Option<&FileRowSelections>, ) -> crate::Result>> { Ok(queries.iter().map(|_| self.result.clone()).collect()) } @@ -1746,6 +1748,11 @@ mod tests { t } + /// A residual selection over one file's physical positions. + fn selected(positions: &[u64]) -> FileRowSelection { + FileRowSelection::Positions(treemap(positions)) + } + #[tokio::test] async fn test_exact_residual_allow_list_restricts_positions() { // No ANN. data-1 has 3 rows: pos0 {1,0} dist 1.0, pos1 {2,0} dist 4.0, @@ -1756,8 +1763,8 @@ mod tests { Some(vec![2.0, 0.0]), Some(vec![3.0, 0.0]), ]); - let mut residual: HashMap = HashMap::new(); - residual.insert("data-1".into(), treemap(&[0, 2])); + let mut residual: FileRowSelections = HashMap::new(); + residual.insert("data-1".into(), selected(&[0, 2])); let out = bucket_search( None, &[], @@ -1798,9 +1805,10 @@ mod tests { } #[tokio::test] - async fn test_exact_residual_file_absent_from_map_is_skipped_without_reading() { - // residual covers only data-1; data-2 has no entry -> no allowed rows, so - // data-2 is skipped entirely (its factory reader is never built). + async fn test_exact_file_absent_from_the_selection_map_is_searched_whole() { + // The selections cover only data-1; data-2 has no entry, which is Java's + // "unrestricted", so data-2 is searched in full rather than skipped. Skipping + // it is the empty-entry case, covered by the test below. let calls = std::sync::Mutex::new(Vec::::new()); let factory = as_search( |file: &BucketActiveFile, @@ -1826,8 +1834,8 @@ mod tests { }) }, ); - let mut residual: HashMap = HashMap::new(); - residual.insert("data-1".into(), treemap(&[0, 1])); + let mut residual: FileRowSelections = HashMap::new(); + residual.insert("data-1".into(), selected(&[0, 1])); let out = bucket_search( None, &[], @@ -1850,9 +1858,12 @@ mod tests { results.extend(out.exact.clone()); results.sort_by(best_first); results.truncate(5); - // Only data-1 rows appear; data-2 was never read. - assert!(results.iter().all(|r| r.data_file_name == "data-1")); - assert_eq!(calls.lock().unwrap().as_slice(), &["data-1".to_string()]); + // Both files were searched: data-1 under its restriction, data-2 unrestricted. + assert_eq!( + calls.lock().unwrap().as_slice(), + &["data-1".to_string(), "data-2".to_string()] + ); + assert!(results.iter().any(|r| r.data_file_name == "data-2")); } #[tokio::test] @@ -1873,8 +1884,8 @@ mod tests { }) }, ); - let mut residual: HashMap = HashMap::new(); - residual.insert("data-1".into(), treemap(&[])); + let mut residual: FileRowSelections = HashMap::new(); + residual.insert("data-1".into(), selected(&[])); let out = bucket_search( None, &[], @@ -1911,8 +1922,8 @@ mod tests { let mut bm = RoaringBitmap::new(); bm.insert(0); // pos0 deleted dvs.insert("data-1".into(), Arc::new(DeletionVector::from_bitmap(bm))); - let mut residual: HashMap = HashMap::new(); - residual.insert("data-1".into(), treemap(&[0, 1, 2])); + let mut residual: FileRowSelections = HashMap::new(); + residual.insert("data-1".into(), selected(&[0, 1, 2])); let out = bucket_search( None, &[], @@ -2546,7 +2557,7 @@ mod tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _residual_ranges: Option<&HashMap>, + _row_selections: Option<&FileRowSelections>, ) -> crate::Result>> { use std::sync::atomic::Ordering::SeqCst; let current = self.inflight.fetch_add(1, SeqCst) + 1; @@ -2842,7 +2853,7 @@ mod tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _residual_ranges: Option<&HashMap>, + _row_selections: Option<&FileRowSelections>, ) -> crate::Result>> { panic!("scorer panic to exercise JoinError mapping"); } @@ -2907,7 +2918,7 @@ mod tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _residual_ranges: Option<&HashMap>, + _row_selections: Option<&FileRowSelections>, ) -> crate::Result>> { let file = segment.source_meta.source_files()[0] .file_name() @@ -3085,7 +3096,7 @@ mod tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _residual_ranges: Option<&HashMap>, + _row_selections: Option<&FileRowSelections>, ) -> crate::Result>> { // Runs on the blocking pool. Announce arrival, then wait (bounded) for the // exact leaf. Both arriving proves overlap; a timeout means no overlap. @@ -3210,7 +3221,7 @@ mod tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _residual_ranges: Option<&HashMap>, + _row_selections: Option<&FileRowSelections>, ) -> crate::Result>> { // The bytes handed to the scorer must be exactly this segment's loaded // bytes (its path), proving load→score threads the right payload. diff --git a/crates/paimon/src/vindex/pkvector/mod.rs b/crates/paimon/src/vindex/pkvector/mod.rs index a1a6446bf..e383a8a69 100644 --- a/crates/paimon/src/vindex/pkvector/mod.rs +++ b/crates/paimon/src/vindex/pkvector/mod.rs @@ -35,3 +35,72 @@ pub(crate) fn data_invalid(message: impl Into) -> crate::Error { source: None, } } + +/// Which physical rows of one data file a bucket search may read. +/// +/// Mirrors Java `rowRangesByFile` (`PkVectorAnnSegmentSearcher.liveRowPositions`, +/// `PrimaryKeyVectorBucketSearch.search`), which is a three-state per file and +/// spells the third state as the absence of a map entry: +/// +/// * **absent from [`FileRowSelections`]** — unrestricted, every row is readable. +/// Java records an entry only for a file its own pre-filter narrowed, so a +/// split that narrowed nothing carries an empty map and restricts nothing. +/// * [`Ranges`](Self::Ranges)/[`Positions`](Self::Positions) **empty** — excluded, +/// no row of the file is readable (Java's empty `List`). +/// * non-empty — restricted to what it lists. +/// +/// The two non-absent variants differ only in where the restriction came from, +/// which decides its shape. A plan built from an engine's bucket split carries +/// interval [`Ranges`] straight off the wire and must never be expanded into +/// positions: the row counts bounding those intervals are untrusted, so +/// materializing one row per allowed position is unbounded work. A residual data +/// predicate produces [`Positions`], whose size is bounded by the rows its own +/// read actually returned. +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum FileRowSelection { + /// Inclusive physical row ranges, sorted and non-overlapping. Empty excludes + /// the file. + Ranges(Vec), + /// Physical row positions. Empty excludes the file. + Positions(roaring::RoaringTreemap), +} + +/// Per-data-file row selections for one split. A file with no entry is +/// unrestricted; see [`FileRowSelection`]. +pub(crate) type FileRowSelections = std::collections::HashMap; + +impl FileRowSelection { + /// Whether the selection permits no row at all, which is how Java's empty + /// `List` reads: the file is skipped rather than searched. + pub(crate) fn is_excluded(&self) -> bool { + match self { + Self::Ranges(ranges) => ranges.is_empty(), + Self::Positions(positions) => positions.is_empty(), + } + } + + /// Whether `position` is permitted. Ranges are binary-searched rather than + /// expanded, mirroring Java `PkVectorAnnSegmentSearcher.contains`. + pub(crate) fn contains(&self, position: u64) -> bool { + match self { + Self::Ranges(ranges) => { + let position = match i64::try_from(position) { + Ok(position) => position, + Err(_) => return false, + }; + ranges + .binary_search_by(|range| { + if position < range.from() { + std::cmp::Ordering::Greater + } else if position > range.to() { + std::cmp::Ordering::Less + } else { + std::cmp::Ordering::Equal + } + }) + .is_ok() + } + Self::Positions(positions) => positions.contains(position), + } + } +} diff --git a/crates/paimon/testdata/pkvector_split/bucket_split_0.bin b/crates/paimon/testdata/pkvector_split/bucket_split_0.bin new file mode 100644 index 000000000..989cb0243 Binary files /dev/null and b/crates/paimon/testdata/pkvector_split/bucket_split_0.bin differ diff --git a/crates/paimon/testdata/pkvector_split/snapshot_id.txt b/crates/paimon/testdata/pkvector_split/snapshot_id.txt new file mode 100644 index 000000000..d8263ee98 --- /dev/null +++ b/crates/paimon/testdata/pkvector_split/snapshot_id.txt @@ -0,0 +1 @@ +2 \ No newline at end of file diff --git a/crates/paimon/testdata/pkvector_split/table/bucket-0/data-44e99441-190a-4e4d-9c08-0402893ea34d-0.parquet b/crates/paimon/testdata/pkvector_split/table/bucket-0/data-44e99441-190a-4e4d-9c08-0402893ea34d-0.parquet new file mode 100644 index 000000000..5d3bae888 Binary files /dev/null and b/crates/paimon/testdata/pkvector_split/table/bucket-0/data-44e99441-190a-4e4d-9c08-0402893ea34d-0.parquet differ diff --git a/crates/paimon/testdata/pkvector_split/table/bucket-0/data-d15b376f-823b-47ff-b13a-97b68d0c0885-0.parquet b/crates/paimon/testdata/pkvector_split/table/bucket-0/data-d15b376f-823b-47ff-b13a-97b68d0c0885-0.parquet new file mode 100644 index 000000000..5d3bae888 Binary files /dev/null and b/crates/paimon/testdata/pkvector_split/table/bucket-0/data-d15b376f-823b-47ff-b13a-97b68d0c0885-0.parquet differ diff --git a/crates/paimon/testdata/pkvector_split/table/index/index-41368f83-71de-4f2a-aad9-4e6b0bb5a97f-0 b/crates/paimon/testdata/pkvector_split/table/index/index-41368f83-71de-4f2a-aad9-4e6b0bb5a97f-0 new file mode 100644 index 000000000..5ebc99c74 Binary files /dev/null and b/crates/paimon/testdata/pkvector_split/table/index/index-41368f83-71de-4f2a-aad9-4e6b0bb5a97f-0 differ diff --git a/crates/paimon/testdata/pkvector_split/table/manifest/index-manifest-211866ae-56c8-4edb-9043-a8e306900588-0 b/crates/paimon/testdata/pkvector_split/table/manifest/index-manifest-211866ae-56c8-4edb-9043-a8e306900588-0 new file mode 100644 index 000000000..a767d40e4 Binary files /dev/null and b/crates/paimon/testdata/pkvector_split/table/manifest/index-manifest-211866ae-56c8-4edb-9043-a8e306900588-0 differ diff --git a/crates/paimon/testdata/pkvector_split/table/manifest/manifest-003ae88e-3c2b-43af-8704-4f5a5e532a0b-0 b/crates/paimon/testdata/pkvector_split/table/manifest/manifest-003ae88e-3c2b-43af-8704-4f5a5e532a0b-0 new file mode 100644 index 000000000..48da11a20 Binary files /dev/null and b/crates/paimon/testdata/pkvector_split/table/manifest/manifest-003ae88e-3c2b-43af-8704-4f5a5e532a0b-0 differ diff --git a/crates/paimon/testdata/pkvector_split/table/manifest/manifest-003ae88e-3c2b-43af-8704-4f5a5e532a0b-1 b/crates/paimon/testdata/pkvector_split/table/manifest/manifest-003ae88e-3c2b-43af-8704-4f5a5e532a0b-1 new file mode 100644 index 000000000..bc847b35c Binary files /dev/null and b/crates/paimon/testdata/pkvector_split/table/manifest/manifest-003ae88e-3c2b-43af-8704-4f5a5e532a0b-1 differ diff --git a/crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-0 b/crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-0 new file mode 100644 index 000000000..c84ff1c69 Binary files /dev/null and b/crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-0 differ diff --git a/crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-1 b/crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-1 new file mode 100644 index 000000000..a06b4bdc5 Binary files /dev/null and b/crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-1 differ diff --git a/crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-2 b/crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-2 new file mode 100644 index 000000000..33f5e82dd Binary files /dev/null and b/crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-2 differ diff --git a/crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-3 b/crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-3 new file mode 100644 index 000000000..77e2a3f37 Binary files /dev/null and b/crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-3 differ diff --git a/crates/paimon/testdata/pkvector_split/table/schema/schema-0 b/crates/paimon/testdata/pkvector_split/table/schema/schema-0 new file mode 100644 index 000000000..6f95dfeb2 --- /dev/null +++ b/crates/paimon/testdata/pkvector_split/table/schema/schema-0 @@ -0,0 +1,30 @@ +{ + "version" : 3, + "id" : 0, + "fields" : [ { + "id" : 0, + "name" : "id", + "type" : "INT NOT NULL" + }, { + "id" : 1, + "name" : "embedding", + "type" : { + "type" : "VECTOR", + "element" : "FLOAT", + "length" : 2 + } + } ], + "highestFieldId" : 1, + "partitionKeys" : [ ], + "primaryKeys" : [ "id" ], + "options" : { + "bucket" : "1", + "merge-engine" : "deduplicate", + "fields.embedding.pk-vector.distance.metric" : "l2", + "fields.embedding.pk-vector.index.type" : "ivf-flat", + "deletion-vectors.enabled" : "true", + "pk-vector.index.columns" : "embedding", + "fields.embedding.pk-vector.index.options.nlist" : "1" + }, + "timeMillis" : 1788269662962 +} \ No newline at end of file diff --git a/crates/paimon/testdata/pkvector_split/table/snapshot/EARLIEST b/crates/paimon/testdata/pkvector_split/table/snapshot/EARLIEST new file mode 100644 index 000000000..56a6051ca --- /dev/null +++ b/crates/paimon/testdata/pkvector_split/table/snapshot/EARLIEST @@ -0,0 +1 @@ +1 \ No newline at end of file diff --git a/crates/paimon/testdata/pkvector_split/table/snapshot/LATEST b/crates/paimon/testdata/pkvector_split/table/snapshot/LATEST new file mode 100644 index 000000000..d8263ee98 --- /dev/null +++ b/crates/paimon/testdata/pkvector_split/table/snapshot/LATEST @@ -0,0 +1 @@ +2 \ No newline at end of file diff --git a/crates/paimon/testdata/pkvector_split/table/snapshot/snapshot-1 b/crates/paimon/testdata/pkvector_split/table/snapshot/snapshot-1 new file mode 100644 index 000000000..05482b3ff --- /dev/null +++ b/crates/paimon/testdata/pkvector_split/table/snapshot/snapshot-1 @@ -0,0 +1,18 @@ +{ + "version" : 3, + "uuid" : "213172c7-d336-4e0f-94c0-37d604cd6d19", + "id" : 1, + "schemaId" : 0, + "baseManifestList" : "manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-0", + "baseManifestListSize" : 1006, + "deltaManifestList" : "manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-1", + "deltaManifestListSize" : 1106, + "commitUser" : "7a084bf2-bcb8-44a9-a28a-9fdfeb5be7e9", + "writerVersion" : "java-2.1-SNAPSHOT-3e510cf1325352003a5f0c1f9a566669eb59085e", + "commitIdentifier" : 9223372036854775807, + "commitKind" : "APPEND", + "timeMillis" : 1788269664787, + "totalRecordCount" : 5, + "deltaRecordCount" : 5, + "nextRowId" : 0 +} \ No newline at end of file diff --git a/crates/paimon/testdata/pkvector_split/table/snapshot/snapshot-2 b/crates/paimon/testdata/pkvector_split/table/snapshot/snapshot-2 new file mode 100644 index 000000000..875091977 --- /dev/null +++ b/crates/paimon/testdata/pkvector_split/table/snapshot/snapshot-2 @@ -0,0 +1,19 @@ +{ + "version" : 3, + "uuid" : "cf376efd-d09c-446d-a3b5-76bcf4cd4efe", + "id" : 2, + "schemaId" : 0, + "baseManifestList" : "manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-2", + "baseManifestListSize" : 1106, + "deltaManifestList" : "manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-3", + "deltaManifestListSize" : 1113, + "indexManifest" : "index-manifest-211866ae-56c8-4edb-9043-a8e306900588-0", + "commitUser" : "7a084bf2-bcb8-44a9-a28a-9fdfeb5be7e9", + "writerVersion" : "java-2.1-SNAPSHOT-3e510cf1325352003a5f0c1f9a566669eb59085e", + "commitIdentifier" : 9223372036854775807, + "commitKind" : "COMPACT", + "timeMillis" : 1788269664884, + "totalRecordCount" : 5, + "deltaRecordCount" : 0, + "nextRowId" : 0 +} \ No newline at end of file diff --git a/crates/paimon/tests/pk_vector_bucket_split_read_test.rs b/crates/paimon/tests/pk_vector_bucket_split_read_test.rs new file mode 100644 index 000000000..66cf334f8 --- /dev/null +++ b/crates/paimon/tests/pk_vector_bucket_split_read_test.rs @@ -0,0 +1,371 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Read a primary-key vector table over a bucket split that JAVA PLANNED, the +//! shape an engine uses when planning runs in Paimon Java and execution is +//! shipped elsewhere. +//! +//! Both halves of the fixture come from ONE Java run, which is what makes this a +//! cross-language test rather than a round-trip of our own bytes: the table +//! directory was written by Java's `ivf-flat` indexer, and the split bytes are +//! what Java's `PrimaryKeyVectorScan` planned over that same table and +//! serialized through `BucketVectorSearchSplit.serialize`. The split names its +//! data and index files by their generated UUIDs, so bytes from a different run +//! would reference files that do not exist. +//! +//! Provenance of `testdata/pkvector_split` (opaque binary, regenerate rather +//! than hand-edit): +//! * Source: Apache Paimon Java master `3e510cf132`. +//! * Generator: `PkVectorSplitFixtureGenerator` (module `paimon-vector`). +//! * Command: `mvn -pl paimon-vector test -Dtest=PkVectorSplitFixtureGenerator \ +//! -Dgen.pkvector.split.fixture=true -Dgen.pkvector.split.out=`. +//! * Config: primary key `id`, vector column `embedding`, `ivf-flat`, +//! `nlist = 1` (exact, deterministic single inverted list), `deduplicate` +//! merge engine, deletion vectors enabled, one bucket. Compacted, since the +//! ANN segment is only built at `level > 0`. +//! * Rows: `id == row position`, vectors `[0,0] [1,0] [2,0] [3,0] [4,0]`. +//! * Query `[0, 0]`, squared-L2 distances `[0, 1, 4, 9, 16]`; top-3 -> ids +//! `[0, 1, 2]`, scores `1/(1+d) = [1.0, 0.5, 0.2]`. +//! +//! The split embeds its bucket directory as an ABSOLUTE path, because that is +//! what Java serializes and what a real engine ships. The fixture is therefore +//! staged into a temp dir and that one path rewritten, below. + +// Gated off Windows for the whole file: the fixture is opened via a `file://` URL +// built from a tempdir path, which the fs lister cannot strip a Windows prefix +// from. Matches how `pk_vector_java_fixture_test`, `pk_vector_baseline_test` and +// `rest_catalog_test` gate their `file://` tempdir tests. +#![cfg(not(target_os = "windows"))] + +use std::path::Path; + +use arrow_array::{Array, Float32Array, Int32Array, RecordBatch}; +use futures::TryStreamExt; +use paimon::catalog::Identifier; +use paimon::io::{FileIO, FileIOBuilder}; +use paimon::table::{SchemaManager, Table}; + +const FIXTURE: &str = "testdata/pkvector_split"; +const TABLE_DIR: &str = "table"; +const VECTOR_COLUMN: &str = "embedding"; + +/// The data file the fixture's ANN segment indexes, named by the split itself. +const DATA_FILE: &str = "data-d15b376f-823b-47ff-b13a-97b68d0c0885-0.parquet"; + +/// The bucket path the generator baked into the split bytes, as a `writeUTF` +/// string: a 2-byte big-endian length followed by the bytes. +const GENERATED_BUCKET_PATH: &str = "/tmp/pkvfixture/warehouse/default.db/pk_vector_split/bucket-0"; + +/// Stage the committed fixture into a private temp root and rewrite the one +/// absolute path the split carries, so the split points at the staged table. +/// Returns the temp dir (kept alive by the caller), the opened table, and the +/// split bytes. +async fn open_bucket_split_fixture() -> (tempfile::TempDir, Table, Vec>) { + let manifest_dir = env!("CARGO_MANIFEST_DIR"); + let src = Path::new(manifest_dir).join(FIXTURE); + let tmp = tempfile::tempdir().expect("create temp dir"); + let dst = tmp.path().join(TABLE_DIR); + copy_dir(&src.join(TABLE_DIR), &dst); + + let staged_bucket_path = format!("{}/bucket-0", dst.display()); + let mut splits = Vec::new(); + for entry in std::fs::read_dir(&src).expect("read fixture dir") { + let path = entry.expect("fixture entry").path(); + if path.extension().is_some_and(|e| e == "bin") { + let bytes = std::fs::read(&path).expect("read split bytes"); + splits.push(rewrite_bucket_path( + &bytes, + GENERATED_BUCKET_PATH, + &staged_bucket_path, + )); + } + } + assert!(!splits.is_empty(), "fixture carries no split bytes"); + + let location = format!("file://{}", dst.display()); + let file_io: FileIO = FileIOBuilder::new("file").build().expect("build fs FileIO"); + let schema = SchemaManager::new(file_io.clone(), location.clone()) + .latest() + .await + .expect("failed to list schemas") + .expect("fixture table has no schema"); + let table = Table::new( + file_io, + Identifier::new("default", "pk_vector_split"), + location, + (*schema).clone(), + None, + ); + (tmp, table, splits) +} + +/// Replace one `writeUTF`-encoded string in place, rewriting its 2-byte +/// big-endian length prefix. Both strings are ASCII, so byte length is the +/// encoded length. +fn rewrite_bucket_path(bytes: &[u8], from: &str, to: &str) -> Vec { + let mut needle = (from.len() as u16).to_be_bytes().to_vec(); + needle.extend_from_slice(from.as_bytes()); + let at = bytes + .windows(needle.len()) + .position(|w| w == needle) + .unwrap_or_else(|| { + panic!("split bytes do not carry the generated bucket path '{from}'; regenerate the fixture and update GENERATED_BUCKET_PATH") + }); + + let mut out = Vec::with_capacity(bytes.len() + to.len()); + out.extend_from_slice(&bytes[..at]); + out.extend_from_slice(&(to.len() as u16).to_be_bytes()); + out.extend_from_slice(to.as_bytes()); + out.extend_from_slice(&bytes[at + needle.len()..]); + out +} + +fn copy_dir(src: &Path, dst: &Path) { + std::fs::create_dir_all(dst).unwrap(); + for entry in std::fs::read_dir(src).unwrap() { + let entry = entry.unwrap(); + let from = entry.path(); + let to = dst.join(entry.file_name()); + if from.is_dir() { + copy_dir(&from, &to); + } else { + std::fs::copy(&from, &to).unwrap(); + } + } +} + +fn batch_i32(batches: &[RecordBatch], column: &str) -> Vec { + let mut out = Vec::new(); + for batch in batches { + let idx = batch + .schema() + .index_of(column) + .unwrap_or_else(|_| panic!("column '{column}' missing from output")); + let array = batch + .column(idx) + .as_any() + .downcast_ref::() + .unwrap_or_else(|| panic!("column '{column}' is not int32")); + out.extend((0..array.len()).map(|i| array.value(i))); + } + out +} + +fn batch_f32(batches: &[RecordBatch], column: &str) -> Vec { + let mut out = Vec::new(); + for batch in batches { + let idx = batch + .schema() + .index_of(column) + .unwrap_or_else(|_| panic!("column '{column}' missing from output")); + let array = batch + .column(idx) + .as_any() + .downcast_ref::() + .unwrap_or_else(|| panic!("column '{column}' is not float32")); + out.extend((0..array.len()).map(|i| array.value(i))); + } + out +} + +async fn read_over_splits(table: &Table, splits: &[Vec], limit: usize) -> Vec { + let refs: Vec<&[u8]> = splits.iter().map(Vec::as_slice).collect(); + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vector(vec![0.0, 0.0]) + .with_limit(limit) + .with_projection(&["id"]); + builder + .execute_read_for_bucket_splits(&refs) + .await + .expect("bucket-split read over the Java fixture failed") + .try_collect::>() + .await + .expect("collecting read batches failed") +} + +/// The read is driven entirely by the Java-planned split: no index manifest is +/// consulted, and the rows come back best-first with their scores. +#[tokio::test] +async fn reads_java_planned_bucket_split() { + let (_tmp, table, splits) = open_bucket_split_fixture().await; + let batches = read_over_splits(&table, &splits, 3).await; + + assert_eq!( + batch_i32(&batches, "id"), + vec![0, 1, 2], + "rows must be best-first: squared-L2 from [0,0] is id*id" + ); + + let scores = batch_f32(&batches, "__paimon_search_score"); + for (got, want) in scores.iter().zip(&[1.0f32, 0.5, 0.2]) { + assert!( + (got - want).abs() < 1e-4, + "score diverges: got {got}, want {want}" + ); + } +} + +/// The split route and the manifest route are two ways to reach the same plan +/// over the same snapshot, so on a table whose splits cover every bucket they +/// must agree exactly. Necessary but NOT sufficient on its own -- the two routes +/// agreeing is also what a read that quietly ignored the split and re-planned +/// from the manifest would produce. `restricts_the_read_to_the_splits_row_ranges` +/// is the test that separates them. +#[tokio::test] +async fn agrees_with_the_manifest_route() { + let (_tmp, table, splits) = open_bucket_split_fixture().await; + let from_splits = read_over_splits(&table, &splits, 3).await; + + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vector(vec![0.0, 0.0]) + .with_limit(3) + .with_projection(&["id"]); + let from_manifest = builder + .execute_read() + .await + .expect("manifest-route read failed") + .try_collect::>() + .await + .expect("collecting manifest-route batches failed"); + + assert_eq!( + batch_i32(&from_splits, "id"), + batch_i32(&from_manifest, "id"), + "the split route must return the manifest route's rows" + ); + assert_eq!( + batch_f32(&from_splits, "__paimon_search_score"), + batch_f32(&from_manifest, "__paimon_search_score"), + "the split route must return the manifest route's scores" + ); +} + +/// A limit below the number of matching rows is applied to the search, not to +/// the output alone. +#[tokio::test] +async fn honors_a_narrower_limit() { + let (_tmp, table, splits) = open_bucket_split_fixture().await; + assert_eq!( + batch_i32(&read_over_splits(&table, &splits, 1).await, "id"), + vec![0] + ); +} + +/// No splits cannot pin a snapshot, so it is rejected rather than answered as an +/// empty read -- which would look identical to a query that matched nothing. +#[tokio::test] +async fn rejects_an_empty_split_list() { + let (_tmp, table, _splits) = open_bucket_split_fixture().await; + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vector(vec![0.0, 0.0]) + .with_limit(3); + let error = match builder.execute_read_for_bucket_splits(&[]).await { + Ok(_) => panic!("an empty split list must be rejected"), + Err(e) => e, + }; + assert!( + error.to_string().contains("at least one split"), + "unexpected error: {error}" + ); +} + +/// The bytes come from outside the process, so a corrupt buffer must fail as +/// invalid data rather than as an internal fault. +#[tokio::test] +async fn rejects_corrupt_split_bytes() { + let (_tmp, table, splits) = open_bucket_split_fixture().await; + let mut corrupt = splits[0].clone(); + corrupt[0] ^= 0xFF; // break the PKVSPLIT magic + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vector(vec![0.0, 0.0]) + .with_limit(3); + assert!(builder + .execute_read_for_bucket_splits(&[&corrupt]) + .await + .is_err()); +} + +/// The split's per-file row ranges are the read's authority, and nothing else +/// carries them: the index manifest places no positional restriction. Narrowing +/// the range to the first two rows must drop `id = 2` from a top-3 -- a read that +/// re-planned from the manifest, or that treated the ranges as advisory, would +/// still return three rows. +/// +/// This is what makes the suite able to tell the two routes apart, so it is also +/// the test that fails first if `plan_for_bucket_vector_splits` stops being +/// reached. +#[tokio::test] +async fn restricts_the_read_to_the_splits_row_ranges() { + let (_tmp, table, splits) = open_bucket_split_fixture().await; + let narrowed: Vec> = splits + .iter() + .map(|bytes| with_row_range(bytes, DATA_FILE, 1)) + .collect(); + + // Limit ABOVE the number of allowed rows on purpose: with a limit of 3 the + // answer would be the same either way, and the test would pass without the + // range being applied at all. + let ids = batch_i32(&read_over_splits(&table, &narrowed, 5).await, "id"); + assert_eq!( + ids, + vec![0, 1], + "only the rows the split allows may produce candidates" + ); + + // The unrestricted split over the same query returns every row, so the + // difference above is the range and nothing else. + assert_eq!( + batch_i32(&read_over_splits(&table, &splits, 5).await, "id"), + vec![0, 1, 2, 3, 4] + ); +} + +/// Append a row-range entry restricting `file` to `0..=to`. +/// +/// The fixture's own `rangeFileCount` is ZERO, which is not an oversight: Java +/// only records a range for a file its pre-filter narrowed, and this query has no +/// pre-filter. That asymmetry is exactly what the read side has to normalize -- +/// an omitted file means the whole file, while an explicitly listed one means +/// only what it lists -- so the two cases are worth driving separately, and this +/// builds the listed one. The section is last in the byte form and empty here, so +/// appending is rewriting its count. +fn with_row_range(bytes: &[u8], file: &str, to: i64) -> Vec { + let head = bytes.len() - 4; + assert_eq!( + i32::from_be_bytes(bytes[head..].try_into().unwrap()), + 0, + "fixture split was expected to carry no row ranges; regenerate and revisit" + ); + + let mut out = Vec::with_capacity(bytes.len() + file.len() + 24); + out.extend_from_slice(&bytes[..head]); + out.extend_from_slice(&1i32.to_be_bytes()); // rangeFileCount + out.extend_from_slice(&(file.len() as u16).to_be_bytes()); + out.extend_from_slice(file.as_bytes()); + out.extend_from_slice(&1i32.to_be_bytes()); // rangeCount + out.extend_from_slice(&0i64.to_be_bytes()); // from + out.extend_from_slice(&to.to_be_bytes()); + out +}