From 32f6b5b2cfad037e0448911b854a26f19b64c42c Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Tue, 1 Sep 2026 21:58:15 +0800 Subject: [PATCH 1/2] feat(table): execute a primary-key vector search over engine-planned splits A planner running in Paimon Java enumerates one BucketVectorSearchSplit per bucket -- a bucket is never divided, because the ANN current-segment decision needs the bucket's whole active file set -- and ships its bytes to a worker outside the JVM. Decoding those bytes, resolving their index files and planning from them already landed; this is the part that runs the search and returns rows. `VectorSearchBuilder::execute_read_for_bucket_splits` takes the splits as 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. Everything after planning is the ordinary primary-key vector read, so search, optional refine, Top-K and materialization are shared with `execute_read` rather than reimplemented, and the output is the same -- 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. Rejected rather than worked around: a data-evolution column, which plans through the global index instead, so answering it would silently use a different plan than the caller supplied; and an empty split list, which pins no snapshot and would otherwise be indistinguishable from a query that matched nothing. `paimon_vector_search_builder_execute_read_for_bucket_splits` exposes it over the C ABI as a new symbol, with the array-of-buffers marshalling the existing terminal does not need. A null array, a zero count or a null entry is the caller's input error and is reported as one, rather than reaching the decoder as corrupt data. The end-to-end test is driven by a fixture where BOTH halves come from one Java run: the table directory written by the ivf-flat indexer, and the split bytes its own planner produced over that same table. A split names its files by generated UUID, so bytes from a different run would reference files that do not exist. Java records a row range only for a file its pre-filter narrowed, and the fixture's query has no pre-filter, so its rangeFileCount is zero and the committed split exercises the omitted-file normalization. The listed case is built on top of it, and is what distinguishes this route from the manifest one: with the manifest plan substituted, that test is the one that fails. --- bindings/c/src/tests.rs | 83 ++++ bindings/c/src/vector_search.rs | 118 ++++++ crates/paimon/src/table/pk_vector_scan.rs | 4 - .../paimon/src/table/vector_search_builder.rs | 138 +++++++ .../pkvector_split/bucket_split_0.bin | Bin 0 -> 900 bytes .../testdata/pkvector_split/snapshot_id.txt | 1 + ...9441-190a-4e4d-9c08-0402893ea34d-0.parquet | Bin 0 -> 1303 bytes ...376f-823b-47ff-b13a-97b68d0c0885-0.parquet | Bin 0 -> 1303 bytes ...dex-41368f83-71de-4f2a-aad9-4e6b0bb5a97f-0 | Bin 0 -> 145 bytes ...est-211866ae-56c8-4edb-9043-a8e306900588-0 | Bin 0 -> 1426 bytes ...est-003ae88e-3c2b-43af-8704-4f5a5e532a0b-0 | Bin 0 -> 2187 bytes ...est-003ae88e-3c2b-43af-8704-4f5a5e532a0b-1 | Bin 0 -> 2235 bytes ...ist-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-0 | Bin 0 -> 1006 bytes ...ist-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-1 | Bin 0 -> 1106 bytes ...ist-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-2 | Bin 0 -> 1106 bytes ...ist-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-3 | Bin 0 -> 1113 bytes .../pkvector_split/table/schema/schema-0 | 30 ++ .../pkvector_split/table/snapshot/EARLIEST | 1 + .../pkvector_split/table/snapshot/LATEST | 1 + .../pkvector_split/table/snapshot/snapshot-1 | 18 + .../pkvector_split/table/snapshot/snapshot-2 | 19 + .../tests/pk_vector_bucket_split_read_test.rs | 371 ++++++++++++++++++ 22 files changed, 780 insertions(+), 4 deletions(-) create mode 100644 crates/paimon/testdata/pkvector_split/bucket_split_0.bin create mode 100644 crates/paimon/testdata/pkvector_split/snapshot_id.txt create mode 100644 crates/paimon/testdata/pkvector_split/table/bucket-0/data-44e99441-190a-4e4d-9c08-0402893ea34d-0.parquet create mode 100644 crates/paimon/testdata/pkvector_split/table/bucket-0/data-d15b376f-823b-47ff-b13a-97b68d0c0885-0.parquet create mode 100644 crates/paimon/testdata/pkvector_split/table/index/index-41368f83-71de-4f2a-aad9-4e6b0bb5a97f-0 create mode 100644 crates/paimon/testdata/pkvector_split/table/manifest/index-manifest-211866ae-56c8-4edb-9043-a8e306900588-0 create mode 100644 crates/paimon/testdata/pkvector_split/table/manifest/manifest-003ae88e-3c2b-43af-8704-4f5a5e532a0b-0 create mode 100644 crates/paimon/testdata/pkvector_split/table/manifest/manifest-003ae88e-3c2b-43af-8704-4f5a5e532a0b-1 create mode 100644 crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-0 create mode 100644 crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-1 create mode 100644 crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-2 create mode 100644 crates/paimon/testdata/pkvector_split/table/manifest/manifest-list-1a76c8fa-98c6-4afc-b9c2-81e4b8a7269f-3 create mode 100644 crates/paimon/testdata/pkvector_split/table/schema/schema-0 create mode 100644 crates/paimon/testdata/pkvector_split/table/snapshot/EARLIEST create mode 100644 crates/paimon/testdata/pkvector_split/table/snapshot/LATEST create mode 100644 crates/paimon/testdata/pkvector_split/table/snapshot/snapshot-1 create mode 100644 crates/paimon/testdata/pkvector_split/table/snapshot/snapshot-2 create mode 100644 crates/paimon/tests/pk_vector_bucket_split_read_test.rs 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/table/pk_vector_scan.rs b/crates/paimon/src/table/pk_vector_scan.rs index 21d70d1dc..56f608953 100644 --- a/crates/paimon/src/table/pk_vector_scan.rs +++ b/crates/paimon/src/table/pk_vector_scan.rs @@ -414,9 +414,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 +448,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, diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 447d69953..8dd039b03 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, }; @@ -416,6 +417,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 +905,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( 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 0000000000000000000000000000000000000000..989cb024362cc5f9a22251450739998f9d945543 GIT binary patch literal 900 zcmcIiJx;?g6m}3uAa!6s0v5yp>L#hvBpr}CQo&FO#7MDgTSP51q)v&1SURzC1CGEQ zxB(LrVnJeH;0Qb?b%+owNId29_x_)M(%U;a=N3Og3G(#QoA z>jRUc#={WAETAdl2ukMAn^Ce^rB z!SP}Y6a0ZWRTOi##Yo^w95xBVULYlAE}^&`FhAtZzzM?^b{a8FE>fXPm*2XV{_khd zh5Fpva}I&^K6?O=vaT@8we(l}mCiD)Q2x8WnH%wZ-mI7kUih5C&`;C>)B(GITYwhx zgoX9B-hi3`qN&7kL{)UoMTR|>_@N954qPs5iK5a(UW-5gC6Wr6zZV|QADUG4@IqScC$&^Uc_Z5JM-qf@B7|+Ga2h2 zH-JC{lJIeJWvfjC1(!0&ghZa5{7r_|i2Zo*m*Vi^U^sA%eq%5Q)d%bMlmlTe2(g40 z#R%aq9(EO19udiqAq3mo8&6mQD4>Hw7P9?n;0!eQymV)&MSzQlx~6t?^EwLtOE|Oh z{qpYhD^Xy=a~@dc{#nisj0M?ralQf+PiB|~C3D&=ndM5_TD2@j zxT0$6q)s%H944|_*9Wa!;S;>9%L%9WPA52t*sX9@SQ9+RsFNTSeS;5eQ~CaNQDmXpX22{>Q(*Al4}2FqNp$qxPJqAbgZ^})rGQLA6=Yjn!)joTB% z5Aa7G{@$Z+F#W67)@>!5b-OE~?$_PfIJ1P`7p%(?3=m#rT+wCz+!w$rH$0o2yR~3e zPL$?Ci{{*Xai$zHx{ThgXeg3fC{7nbr=s2Yo<1X*?}>L8x=;5;drYIxh{ug=-poyz a`Fz~SnNz(+r#h(-()^O2c!S6AFYzx*|H;n) literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..5d3bae888f18a2f862eccbe3b267a9fb824bda10 GIT binary patch literal 1303 zcmbVMO=uHA6rM~b+3nh5i+Mv97Re<`Dq6FxX(|Lwt0}eBQmZKlg_xhM2AZ_CX^*l{ z=&3?4UPPe>5iK5a(UW-5gC6Wr6zZV|QADUG4@IqScC$&^Uc_Z5JM-qf@B7|+Ga2h2 zH-JC{lJIeJWvfjC1(!0&ghZa5{7r_|i2Zo*m*Vi^U^sA%eq%5Q)d%bMlmlTe2(g40 z#R%aq9(EO19udiqAq3mo8&6mQD4>Hw7P9?n;0!eQymV)&MSzQlx~6t?^EwLtOE|Oh z{qpYhD^Xy=a~@dc{#nisj0M?ralQf+PiB|~C3D&=ndM5_TD2@j zxT0$6q)s%H944|_*9Wa!;S;>9%L%9WPA52t*sX9@SQ9+RsFNTSeS;5eQ~CaNQDmXpX22{>Q(*Al4}2FqNp$qxPJqAbgZ^})rGQLA6=Yjn!)joTB% z5Aa7G{@$Z+F#W67)@>!5b-OE~?$_PfIJ1P`7p%(?3=m#rT+wCz+!w$rH$0o2yR~3e zPL$?Ci{{*Xai$zHx{ThgXeg3fC{7nbr=s2Yo<1X*?}>L8x=;5;drYIxh{ug=-poyz a`Fz~SnNz(+r#h(-()^O2c!S6AFYzx*|H;n) literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..5ebc99c744ede9119a115162111642b8f05a64b4 GIT binary patch literal 145 zcmeZq3-e@TU|?VZVh{}itPl#$hjMXI4j_pLs4Pe?Og~5n%w_}vs63Eru!pi8AhZLN HZUEB&2d)Gr literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..a767d40e4cb9cbd7bcbfb0e3463d9670bbe00433 GIT binary patch literal 1426 zcmeZI%3@>@ODrqO*DFrWNXij}OQt6?U^hq(p?d;0qUC82g@ z=9MVb>L3(&d-}QHmJM(W3h@j z%u7cx#?8~mHQvwBmk3qCo>54ua`Mqk3i1z+clHnULvj;_8W&d|SCGeofkEUP;vWxdL7&!|dRJtPP1&kyENyl!Uu0AgDo_I11Fl>Sx<6-KEPf<$XD8(0M5EuBmhB(5# z1CK{YzQm_H*grJL*)^Ud?O0MlY^?~m1-)i%Dav{$Bl&5efFe4-&;!6%*OLiZ*pI8=lbpI(~|mKb`%KsYklJ`c!28f{M6L@CRz=$)AB})_h}G({$#io0}zOY1&1Tv?Q4@M`5CZ z2+G*Sg9t+q96JvNBJMB{Jb3WnNpT1&GR5O?=;oKCO|oWfCwTGJ_Iuvvd*A1Kp0~-N zD-BzA3)l2~)DVXyRaiYR2F%fd1Ui5&iFrK(B0v+=z;~k{h}MsJJr@+JRmXwSEItIteOneH`43VVJVd*YeGeI%p0pbLOS3m8=XqI5t5e$u* zlCh4WlZ(Y6Tb+bWvKYxbD5g5Biy)zUT1Ok6FQGRW**f7VEa8y>!_;Lhx2|a$#SjBN zl#;K&(8)eNmgFekYynJyp!hhv=r#_5`XsRadJx3GGc3k>@@z@~-r{u&H_*cBE`^aC z^oXl8SRX}kB$DwqoH%^J;1m2m_E&czcuwIUV7dS}8kej(gC!{@(1zm_Nws|xY>xK!8 zCw##VFp@9@N$mX*<(9&=ihbEYDI`7F%M7k&)#R_mj@}c}6k5~07~373t+YUE zF?10$!`dqQd{yqPzQlbrfgb7K59aZ@0r?B*oLB_LMq4Xg+-3j&iusHr?!1odsAU8_ z6FRl>;Cs`%(;$4Pr!6Pni#Lup9-VFf9qN7_4o!CLX`49q_1ueRH|Ku6Tn@iqTD?2` zbM4H>%liiA9``OiRpJjD*Jf_dF3*g!ptWah{=-}C>dD3Stxc2JvwQC^oLHQ`^k$&h zm|qZ^Z&06J#qYd+MBYk$esXvkKcX~j4;^fs`q8p7{%wBS`OwskFD?7ci7;N-(jUq~ B(L(?L literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..bc847b35c1a398291df37ae1d0c7bf72a5cb918b GIT binary patch literal 2235 zcmdT_Z%7ki9NyGWNDM>#RFX$f*qYIYC?)dTUh~-P&Tcm~WgX}4oF}(G-cHG6Wb`Q_ z5c?KJLYNf#Q6=z$9}+7<#G){gf{0Ko$nVYGyW8&0o!boh(zk8z^LyUk?|Gl+H?phO zUR22rNKt1z*u#g}WYrKN55{<;4w3j^RN@c^5@ExZ@9`KLRDO&~Atwtz_wddbD}NBc$T?j3 zPe=%|BCYL9SqSrSnGMHIhXqj-;5M@O%CzC!00%T+tYJ=HmNI9iHBA%%%I|gqw>mHb zHzsJ;rO{cm;OFW@+c~74av57kFIN=KRcy)zOd)Q_o@UTB%O-y=cKDuvBH)^CLFs1e zY^epBi=hdle6*^v$ye##@=M%CVepau{a^;K9+1C~PK$+QY`C=oi3;`qHyGvk;Nct4 zjuDQPI2-rr!TfS7n>=^BtiJkP!9dZMf|E0~zirJI57-J`A3Huic_H*+=hF4B&n8A6 zFU}{Q57j3Nx7Is6z0y?eHgSSmCGE36N+x#7UOqfC{q*wU`-Y{B;K9wYcOcI|xW z*;6%y0{A?~w74tpErpB49Wwc7)IM8$XSn>;3;VCqqKUq>`)_8z^y-tU$-Zx!zR1SM xe&Xlm-D_iW+ZAQD;gu;{Vy=3@c5~$D%2vtxEquY=U>iU5z5IwgQs~Yc;SWj}?Gyk2 literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..c84ff1c696462782e355326fb7ab956879738397 GIT binary patch literal 1006 zcmbVLJ5Iwe7zT;C0|HhgOKerd+$}UlMC5^*q#&xuxQW{aB~Gzjq*CP^a08CQEjR!d zV8t$_2~EYLTO5B6f8TGr{oC!G1ANbNjYlJr;kU*Ni0PPs7GNaGSqz$h;;gpjV;rfU zIU8zNevgPY#z~e_Z7wz>lrSuaq7M?1#vHVIvo9sog&ogy9l5#iN~i!uvjm))7DSG2 z!xF*;OX#qU^33ZBD$Q3gPGBS5Fra~yrdR5%MHzh+Un3@YfxTx(-^aOtMPtx_@ODrqO*DFrWNX<=r#aFFVQdy9yWTjM;nw(#hqNJmgmzayeFD^(-1_|aD zrRyaE*%_&N1&Nut`FVO^!_rgpQi~ExQbF3&GE;L>ij}OQt6?U^hq(p?d;0qUC82g@ z=9MVb>L3)jdHT4<`#Ji$B9s)D6lLb6W2y@Fj6zbClaHpxFVr{Q(Z$8pB_3pEFv2{n zDqUQCTtkRZ8{ile;u!+;LVR$DV+g`j0)ZDFkEsz$@c4TA#fLfigd#gPsj?)s7{&L# zjuAwt@(cCxiFfu7^+R%Hv{GJaPL2{VyrIFLSX7i)2@IOdlGI#KOhM(9z+n_y8>^#~ zlA4xSnp2`=1=3j?TZVXi*pSxc@~pshjv;mAP(_bnt$(L)r?yQFE3trY?1y5#(n)MTNm;*!L?l*FPG zhEtaITGoGGJM_GGc<|I^r6Q(9j2c_@|1x;EGBG5z@QW|lYWbhp#?315)ThwY9Q76l zIq$nst2A?Cua~}_HZ6+d#?*zDQ#dYK*_$oV-&360RlRIVMH;)xP97G1B@QNr0}qTD IM9}R302^;^i2wiq literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..33f5e82dd13a7d0be600b277f74d790f4efb2208 GIT binary patch literal 1106 zcmeZI%3@>@ODrqO*DFrWNX<=r#aFFVQdy9yWTjM;nw(#hqNJmgmzayeFD^(-1_|aD zrRyaE*%_&N1&Nut`FVO^!_rgpQi~ExQbF3&GE;L>ij}OQt6?U^hq(p?d;0qUC82g@ z=9MVb>L3)jdHT4<`#Ji$B9s)D6lLb6W2y@Fj6zbClaHpxFVr{Q(Z$8pB_3pEFv2{n zDqUQCTtkRZ8{ile;u!+;LVR$DV+g`j0)ZDFkEsz$@c4TA#fLfigd#gPsj?)s7{&L# zjuAwt@(cCxiFfu7^+R%Hv{GJaPL2{VyrIFLSX7i)2@IOdlGI#KOhM(9z+n_y8>^#~ zlA4xSnp2`=1=3j?TZVXi*pSxc@~pshjv;mAP(_bnt$(L)r?yQFE3trY?1y5#(n)MTNm;*!L?l*FPG z2CoS$Nr7E{+J;wLj&GRl?!~l-QDdwAUj`3XCWfRIe(?ocE&nsyxLF0B`V^X)qu$~m z=Y2P7m1b`2_0reVrbThwn7Yt%3dcn&d$T3_dx}%Ls+UcvNMl#o$-~00#KFXH;DIrN H2)Z2r?9*vo literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..77e2a3f37e6ff3be212956297e82d57f3d57b886 GIT binary patch literal 1113 zcmbVLJ#W)M74`<)=anJMG{qf08?IC<(StIw4X#!WfXNb?H6m=1! ze#Qb+M=2y5&-oO5@iJqR29)b#+L%I=WNBkAHlZnHz^R~z5siZ!b&t+|E0G}{I7Z(t zHy2(}%282Y0bUsmKGyoPz=WoA$xxBvj^fBy4H|`@-KHqur?b#3Mu#lyoR{NMh8ih3HO0%f~?R2Q3 w`UhjBr9!B5UNl=uD{MoXw$%e@j+N)%ku2{^`;sh`_x!pqRu`A5w6|{c6K~>QX8-^I literal 0 HcmV?d00001 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 +} From 57ad142a83c63a6ed96f1ff2eedbf0f365099a54 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Wed, 2 Sep 2026 19:04:39 +0800 Subject: [PATCH 2/2] fix(table): keep a bucket split's unrestricted files off the filtered path Review of #771 asked for two things on the new bucket-split entry point: reject unreasonable row counts/endpoints, and preferably preserve unrestricted or interval-form selections through the ANN layer rather than eagerly materializing every allowed position. Both problems came from one place: planning normalized a file the split left unlisted into an explicit `[0, row_count - 1]`, because the search kernel read a missing entry as "no rows allowed" -- the opposite of what Java means by it. Java's `rowRangesByFile` is three states, and spells the third as an absent entry: absent is unrestricted, an empty list excludes the file, a non-empty one restricts it (`PkVectorAnnSegmentSearcher.liveRowPositions`, `.contains`, `PrimaryKeyVectorBucketSearch.search`). It never expands a range into positions. Mirror that. `FileRowSelection` is the per-file state, absence carries the unrestricted case, and ranges stay ranges from planning through the ANN mask, the ANN result check and the exact fallback. * Planning no longer sizes anything from a wire row count: `whole_file_range` and `positions_in_ranges` are gone, and with them the only callers that turned one into a `RoaringTreemap`. * A split that narrowed nothing (`rangeFileCount == 0`, the committed fixture and the shape an engine ships by default) now reaches the ANN backend with no filter at all, instead of an all-permitting mask that Lumina turns into one `u64` per live row and a filtered search. * The residual predicate's map keeps its own, opposite contract -- it is exhaustive over the files a search can read from, because the kernel reads a file's absence as unrestricted and would otherwise drop the predicate for it. On the endpoints, three checks, none of them a number invented here: * A listed range is checked against its source file's row count, as Java checks it. Both numbers come off the same split, so this catches a range that disagrees with its own segment; it is not a resource bound. * Row counts are checked non-negative at decode for EVERY data file, not only for the ones the message goes on to list ranges for. That check used to ride on the normalization this commit removes. * `MAX_LIVE_ROW_IDS` bounds what goes into one segment's mask, charged before each of the four insertions. Removing the planning-side expansion is not enough on its own: the mask spans the segment's SOURCE row counts, which arrive on the wire too (`PrimaryKeyIndexSourceMeta` is decoded from the split's own `GlobalIndexMeta`, and only its sign is checked), so once anything makes a mask necessary -- a restriction on one file of this bucket, or a deletion vector on one of this segment's own sources -- an unrestricted file is inserted wholesale at whatever size that count claims. The added tests OOM-kill the process, or run for minutes, without it. What that limit is and is not. It stops a forged count from making mask construction unbounded, and it is charged on insertion so nothing is allocated for a span that could not be searched anyway. It is NOT a memory guarantee: at `i32::MAX` a Lumina filter is still ~16 GiB once densified. The value is Java's for the same quantity -- `LuminaVectorGlobalIndexReader.toScopedIds` refuses an include-set above `Integer.MAX_VALUE` before allocating the dense array, and `to_scoped_ids` now mirrors that on our side, which had no such guard -- so it accepts exactly what Java accepts rather than being tuned here. Two consequences are deliberate: it is applied backend-neutrally, so a vindex segment above the limit is refused even though vindex serializes the bitmap instead of densifying it; and charging on insertion is marginally stricter than Java's post-`andNot` check, so a segment of `i32::MAX + 1` rows whose deletion vector would bring the final set back to `i32::MAX` is refused. Both are pinned by tests. Neither is reachable by a real bucket, which is a compaction unit whose ANN segment is built in memory. Charging insertions rather than the claimed span is what keeps this from rejecting valid data, and there is a test for that: a segment CLAIMING an impossible row count is accepted when its live set is empty (all sources inactive, an explicitly empty selection) or narrow. Two bounds on the CLAIM were tried first and are wrong: * A cap on the total claimed source span turns away exactly those cases. * Bounding the claimed span by the index file's byte length is unsound, not merely weak. `PkVectorAnnSegmentFile` advances the ordinal space by each source's FULL physical row count while writing a vector only for rows that are non-null and not excluded, so a legitimate million-row source with mostly null vectors has a small payload. That bound would reject valid data. Adding `to_scoped_ids` also put the batch search's steps in Java's order: `index.size()` decides first, then the dense filter is built, and the metric is read last (`LuminaVectorGlobalIndexReader`, both paths). An empty index therefore reports no hits rather than an oversized-filter or malformed-metric error. An include-set that permits nothing still short-circuits ahead of all of it, with no native call at all. Unrelated to either ask, but in the same function: a deletion vector's positions were added to the mask without checking them against their own source file, so a position past that file's rows landed inside the NEXT file's ordinal range and deleted one of ITS rows. Deliberately unchanged, being pre-existing and reachable without this entry point: caller-controlled Top-K heap sizes, `DeletionVector::iter` materializing its positions, the `deleted` bitmap and the residual's own position sets (both sized by data actually on disk), the scalar Lumina path's own step order, and the segment-local deletion-vector condition in the early return (Java's is whole-map, so any bucket deletion vector puts every segment on the filtered path there). Java parity note: the early "no mask needed" return keeps Java's whole-map `rowRangesByFile.isEmpty()` rather than asking whether this segment's own sources are restricted. The narrower test would leave more segments unfiltered in a mixed split, but it also changes which backend entry point they take, and `search` and `search_with_filter` can differ in recall. Sharing one `Arc` for the mask across a batch was tried for the same reason and reverted: it moves a filtered search from per-query scalar calls onto the native batch call. Tests: 2602 `--lib` (2677 with `fulltext`), all `-p paimon` integration tests. Each of the four charge sites, the dense-filter guard, the batch step order, and the ANN, merge and planning assertions were mutation-checked individually against the code they cover. --- crates/paimon/src/lumina/reader.rs | 102 ++- .../src/table/pk_vector_bucket_split.rs | 57 +- .../src/table/pk_vector_orchestrator.rs | 51 +- crates/paimon/src/table/pk_vector_scan.rs | 217 +++--- .../paimon/src/table/vector_search_builder.rs | 276 +++++--- crates/paimon/src/vindex/pkvector/ann.rs | 647 +++++++++++++++--- crates/paimon/src/vindex/pkvector/bucket.rs | 117 ++-- crates/paimon/src/vindex/pkvector/mod.rs | 69 ++ 8 files changed, 1199 insertions(+), 337 deletions(-) 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 56f608953..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>>>, } @@ -573,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| { @@ -593,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, @@ -1264,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); @@ -1329,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() @@ -1369,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] @@ -1470,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] @@ -1538,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 8dd039b03..93767012f 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -45,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; @@ -61,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}; @@ -1052,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: /// -/// 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. +/// * 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". +/// +/// 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 { @@ -1076,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(), )) } } @@ -1369,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(), @@ -1447,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?; @@ -2633,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 @@ -2665,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?; @@ -8240,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] @@ -8428,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] @@ -8446,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) @@ -8455,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), + } + } +}