JingsongLi commented on code in PR #771:
URL: https://github.com/apache/paimon-rust/pull/771#discussion_r3911862455
##########
crates/paimon/src/table/vector_search_builder.rs:
##########
@@ -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<ArrowRecordBatchStream> {
+ // 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::<crate::Result<Vec<_>>>()?;
+
+ // 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)?;
Review Comment:
[P1] Bound or preserve decoded row ranges before bitmap expansion
This new entry point forwards engine-supplied DataFileMeta.row_count values
into plan_for_bucket_vector_splits without a resource bound. When a file omits
its range entry (which is valid in the Java encoding), planning synthesizes [0,
row_count - 1], and the search path materializes that range into a
RoaringTreemap. Both roaring 0.11.4 and the current 0.11.5 implementation
iterate every high-u32 shard in insert_range and construct a full u32 bitmap
for each intermediate shard; row_count = 2^40 already creates 256 full bitmaps
(about 16.8 million containers), while a value near i64::MAX can hang or OOM
the worker before any storage I/O. The later i32 row-position validation runs
only after this expansion. Please reject unreasonable row counts/endpoints and,
preferably, preserve unrestricted or interval-form selections through the ANN
layer instead of eagerly materializing every allowed position.
##########
crates/paimon/src/table/vector_search_builder.rs:
##########
@@ -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<ArrowRecordBatchStream> {
+ // 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::<crate::Result<Vec<_>>>()?;
+
+ // 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,
+ )
Review Comment:
[P2] Keep unrestricted files off the filtered ANN path
A Java split with rangeFileCount == 0 is the normal no-prefilter form
(including the committed fixture), but the planner normalizes every omitted
file into an explicit full-file range. This call then converts those ranges
into RoaringTreemaps and passes Some(filter) to ANN even though every row is
allowed. Lumina subsequently collects the bitmap into a Vec<u64> and invokes
search_with_filter, adding O(live_rows) setup and at least 8 * live_rows bytes
for the ID vector per segment search. Please preserve an unrestricted sentinel
(with a per-file unrestricted/restricted/excluded state for mixed splits), and
add a backend-facing test asserting that the no-prefilter fixture uses the
unfiltered ANN path.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
