JunRuiLee commented on code in PR #771:
URL: https://github.com/apache/paimon-rust/pull/771#discussion_r3920819665
##########
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:
Fixed in 57ad142.
Both comments trace to one line: planning normalized an unlisted file into
`[0, row_count - 1]`, because our kernel read a missing entry as "no rows" —
the opposite of Java. Now mirrored: `whole_file_range` and
`positions_in_ranges` are gone, and ranges stay ranges through the mask, the
result check and the exact fallback.
Endpoints are rejected in three places: a listed range against its source
file's row count (as Java checks it), non-negative row counts at decode for
*every* data file, and `MAX_LIVE_ROW_IDS` charged before each insertion into
the mask.
That last one turned out to be necessary: removing the planning-side
expansion alone is not enough, because the mask spans the segment's *source*
row counts, which also come off the wire. One restricted sibling file or one
deletion vector is enough to insert an unrestricted file wholesale — the added
tests OOM-kill the process without the bound. The limit is Java's own for this
quantity (`LuminaVectorGlobalIndexReader.toScopedIds` refuses above
`Integer.MAX_VALUE`); our dense conversion had no such guard, so
`to_scoped_ids` mirrors it.
One thing worth flagging: bounding the claimed span by the index file's byte
length looks reasonable but is unsound — `PkVectorAnnSegmentFile` advances the
ordinal space by each source's full physical row count while writing vectors
only for non-null, non-excluded rows, so a legitimate mostly-null source has a
small payload.
Unrelated bug found in the same function: deletion-vector positions were
added without checking them against their own source file, so a position past
that file's rows deleted a row in the *next* file.
--
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]
