Skip to content

Execute primary-key vector search over Java-planned bucket splits (engine-distributed read) #755

Description

@JunRuiLee

Search before asking

  • I searched in the issues and found nothing similar.

Motivation

paimon-rust can already run a primary-key (bucket-local ANN) vector search end to end on its own — that was #514, now closed. What it cannot do is act as the execution kernel for an engine that plans the query somewhere else.

That is the shape needed to bring PK-vector search to external Paimon tables in a distributed engine. Apache Doris is being wired to read Paimon through the paimon-rust C FFI (apache/doris#65883 adds a PaimonRustReader on the BE side). For a vector search over an external Paimon table, planning happens in Paimon Java on the Doris FE, and each unit of work is shipped to a BE, which calls into paimon-rust.

Java already produces the right unit of work. PrimaryKeyVectorScan aggregates a whole bucket into a single BucketVectorSearchSplit (BucketAccumulator.build()), so a bucket is never split across tasks and the ANN current-segment decision stays correct. apache/paimon#9386 then gave that class a public, versioned byte form (PKVSPLIT v1), so a non-Java engine can receive one.

What is missing on the Rust side is everything after decoding those bytes:

  • plan from the split — its payload files, its per-file row ranges, and the snapshot it pins — instead of re-enumerating the index manifest for the whole table;
  • run search, optional refine, Top-K and materialization for that split;
  • expose it over the C FFI so a BE can call it.

Worth doing inside paimon-rust rather than in each engine: the semantics here are Paimon's, not the engine's — exact fallback over the data files no ANN segment covers, deletion-vector filtering of stale candidates, a deterministic global Top-K, and materialization by physical row position. Reimplementing that per engine is how it drifts.

Solution

Four steps. The first is merged; the second is open.

1. Decode the BucketVectorSearchSplit byte form — #746 (merged).
BucketVectorSearchSplit::deserialize for the PKVSPLIT v1 envelope (i64 magic + i32 version + embedded DataSplit.serialize + IndexFileMetaSerializer.serializeList + per-file row ranges), with golden fixtures produced by Java.

One thing still to settle, at or before step 4: deserialize_binary_array_str validates only that each variable-length region is within bounds — it does not check element ordering, aliasing, or trailing bytes, so n elements may all point at the same large body and each get cloned, giving an output bound around len²/8. It is reachable from DataFileMeta row decoding in crates/paimon/src/spec/data_file.rs, therefore from the embedded DataSplit. #746 (merged) enforces the writer's real invariants for the row-array variant (each element body starts exactly at the previous padded end; the array ends exactly at the last element's padded end); the same rule should be applied to the string variant in its own PR, before a C entry point starts accepting arbitrary split bytes.

2. Resolve index files by external path and bucket layout — #752 (open).
PkVectorScan currently builds every index path as <table>/index/<file> and ignores both external_path and bucket-local placement (index-file-in-data-file-dir). Java can write either. Split-driven execution has to land after this, or it will fail to find index files that Java legitimately placed elsewhere.

3. Plan from a decoded bucket split (not yet opened).

Two parts, best reviewed as two commits in one PR.

(a) Separate planning from searching. plan_and_search_pk_candidates_batch currently plans inside itself — it calls PkVectorScan::new(..).plan(), which reads the snapshot's index manifest and enumerates the whole table's ANN payloads. A bucket split already carries its plan, so that entry point cannot be reused as-is: it would re-read the manifest, ignore the split's payload list, and drop the per-file row ranges. Split it into query/parameter resolution and search-over-a-supplied-plan, keeping the existing function as a thin wrapper over both so the current path is provably unchanged.

Expose both the raw indexed/exact candidate layer and the merged/reranked layer in this step, even though only the merged layer has a caller here. A candidate-only phase (see "Anything else?") needs the raw layer, and extracting only the merged layer means reopening this refactor later.

(b) Build the plan from the split. PkVectorScan::plan_for_bucket_vector_splits: use the split's payload files and rowRangesByFile instead of reading the index manifest; require that all supplied splits pin the same snapshot; apply partition pruning; and keep PkVectorScanPlan.snapshot_id populated, since it stays authoritative even when pruning leaves zero searchable splits.

One asymmetry to handle explicitly: Java only inserts a rowRangesByFile entry for IndexedSplit files, while the Rust kernel treats a missing key as "no rows allowed". Files omitted from the map must be normalized to unrestricted full-file ranges, or valid splits will silently return nothing.

4. Execute a bucket split end to end, plus the C entry point (not yet opened).
VectorSearchBuilder::execute_read_for_bucket_splits and paimon_vector_search_builder_execute_read_for_bucket_splits: split bytes in; search, optional refine, local Top-K and row materialization in one call; Arrow record batches with __paimon_search_score out. Query parameters and projection reuse the existing with_* builder methods. Covered by an end-to-end test driven by Java-produced split bytes.

After step 4 an engine can distribute one call per bucket and merge the per-bucket Top-K itself.

Anything else?

A staged form is deliberately left out of the five steps above. Java + Spark does something stronger: search tasks return lightweight candidates, the driver merges them globally, reranks only the survivors, and emits IndexedSplits that executors read back. With the per-bucket form above, refine and full-row reads happen within a bucket, so a multi-bucket query can rerank and materialize rows that the global Top-K then discards. Adding candidate-only search, global merge/rerank, and deferred materialization is the natural follow-up — but the candidate and materialize wire formats should not be frozen before there is a consumer for them, so it is better proposed separately.

Related gaps, independent of this lane. DataSplit decoding accepts only versions 8 and 9 while Java reads 1..9 (the current Java writer always emits 9). DataSplit.total_buckets is stored as i32 with an absent value silently read as 1, while Java models it as nullable — normal planning always sets it and PK-vector tables require fixed or postpone bucketing, so this is a wire-conformance gap rather than a blocker for the steps above.

No new Java API is required. BatchVectorSearchBuilder (partition filter, filter, limit, vector column, vectors, newVectorScan()), VectorScan.scan(), PrimaryKeyVectorScan.Plan.splits()/snapshotId() and BucketVectorSearchSplit.serialize(DataOutputView) are all public as of apache/paimon#9386.

Predecessor: #514 (primary-key vector read, closed). Prerequisite already merged: #745 (accept DataSplit version 9).

Willingness to contribute

  • I'm willing to submit a PR!

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions