Skip to content

[core] Rerank primary-key vector candidates with exact distances#550

Merged
JingsongLi merged 2 commits into
apache:mainfrom
JunRuiLee:feat/pk-vector-rerank-reconcile
Jul 20, 2026
Merged

[core] Rerank primary-key vector candidates with exact distances#550
JingsongLi merged 2 commits into
apache:mainfrom
JunRuiLee:feat/pk-vector-rerank-reconcile

Conversation

@JunRuiLee

@JunRuiLee JunRuiLee commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

Purpose

Part of #514. Adds exact rerank to primary-key vector search, mirroring Java PrimaryKeyVectorRead's refine/rerank. On a PK-vector table, execute_read() materializes rows in file-local physical coordinates; this lets a configured refine factor over-fetch approximate (ANN) candidates and reorder them by exact distance for higher recall, while leaving the no-refine path byte-identical.

Brief change log

  • Dual-limit bucket search: bucket_search returns separate indexed/exact candidate lists, each bounded independently — the ANN/indexed channel over-fetches (indexed_limit = limit × refine_factor) while the exact-fallback channel stays bounded by the final limit. PkVectorOrchestrator::search_candidates takes the dual limits and returns an OrchestratorSearchResult; merge_candidates runs a global Top-K over the union of the two channels.
  • Rerank: when a positive refine factor is configured, rerank_indexed_positional rereads the over-fetched ANN candidates' vectors from their data files by file-local position (reusing the position read path), recomputes exact distance, and re-runs the global Top-K. Exact-fallback candidates are already exact and are not reranked.
  • Refine-factor resolution reuses the existing option-key precedence (refine_factor / refine-factor / rerank_factor / rerank-factor, with fields.<col>. → index-type → ivf. → bare prefix order); factor 0 (unset) leaves indexed_limit == limit, so output is byte-identical to the no-rerank path.
  • Fail-loud validation: append_batch_vectors rejects null child elements in vector arrays (never a defaulted 0.0), build_indexed_splits validates split_index is in range before indexing, and the rerank kernel fails loud on ownership/dedup/NULL/dimension/file-not-in-plan violations — never a panic or a silently corrupted result.

Tests

  • Dual-limit bucket behavior; merge_candidates union Top-K; rerank kernel (ownership order, dedup, NULL element, dimension mismatch, file-not-in-plan); end-to-end refine integration matching an exact brute-force ground truth.
  • Regression invariant: refine_factor == 0 preserves the exact pre-rerank output (baseline + Java-fixture read-back tests unchanged).
  • Fail-loud coverage: null vector element, out-of-range split_index.
  • cargo test -p paimon green; cargo build -p paimon-datafusion clean (cross-crate Send gate); cargo clippy -D warnings and cargo fmt --check clean.

API and Format

No on-disk format change, no new result columns. refine_factor == 0 (the default) is byte-identical to the current output.

Documentation

Code comments only; no user-facing docs change.

@JunRuiLee
JunRuiLee force-pushed the feat/pk-vector-rerank-reconcile branch from 75fb950 to e2377cd Compare July 20, 2026 06:53
@JunRuiLee
JunRuiLee marked this pull request as ready for review July 20, 2026 06:53

@leaves12138 leaves12138 left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The rerank implementation and focused tests look solid overall, but refine-factor validation is currently data-dependent: an empty PK-vector plan returns before parsing the option. Please address the inline finding so invalid query/table configuration fails consistently.

// separately from table options) so a broad query key cannot be overridden
// by a more specific table key: query options take precedence as a whole.
// `search_options` above is the merged view used only to drive the ANN read.
let refine_factor = configured_refine_factor(

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

configured_refine_factor is reached only after PkVectorScan::plan() and the plan.splits.is_empty() early return above. As a result, an empty PK-vector table accepts refine_factor=0 or refine_factor=abc and returns a normal empty reader, while the identical query starts failing as soon as the table has a searchable split. Java resolves this option in the PrimaryKeyVectorRead constructor before planning, and configuration validity should not depend on whether data currently exists. Please resolve the factor/search limit before the empty-plan return (ideally before planning/index I/O) and add an empty-plan regression test.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 1ab7337. configured_refine_factor / indexed_search_limit are now resolved before PkVectorScan::plan() and the empty-plan early return, so an invalid factor fails loud regardless of whether the table currently has searchable data — matching Java, which resolves this in the PrimaryKeyVectorRead constructor before planning. Added an empty-table regression test covering both invalid forms you named (0 → "must be positive", and a non-integer → "must be an integer") and both option sources (query and table options, since the fix moved resolution of both ahead of planning).

Reconcile the primary-key vector exact rerank feature onto the physical-coordinate
read path. On a PK-vector table, execute_read() materializes rows in file-local
physical coordinates. This adds:

- Dual-limit bucket search: ANN/indexed candidates can over-fetch (indexed_limit =
  limit × refine_factor) while exact-fallback remains bounded by the final limit.
- When a positive refine factor is configured, rerank the over-fetched ANN candidates
  by rereading their vectors from data files, recomputing exact distance, and re-running
  the global Top-K merge.
- Fail-loud validation: null vector elements and out-of-range split indices return
  DataInvalid errors instead of panicking or silently corrupting results.

Changes:
- bucket_search returns separate indexed/exact lists, each bounded independently
- PkVectorOrchestrator.search_candidates takes dual limits, returns OrchestratorSearchResult
- merge_candidates performs global Top-K over the union of indexed and exact channels
- rerank_indexed_positional reads vectors by file-local position, validates ownership order
- append_batch_vectors rejects null child elements in vector arrays
- build_indexed_splits validates split_index is in range before indexing

Tests added: dual-limit bucket behavior, merge_candidates union Top-K, rerank kernel
(ownership/dedup/NULL/dimension/file-not-in-plan), e2e refine integration, baseline
regression (refine=0 preserves exact output), null element fail-loud, out-of-range
split_index fail-loud.
The refine factor was resolved only after PkVectorScan::plan() and the
empty-plan early return, so an empty primary-key vector table accepted an
invalid or zero refine factor and returned an empty reader, while the same
query started failing once the table had a searchable split. Configuration
validity must not depend on whether data currently exists. Resolve the refine
factor (and its search limit) before planning so an invalid factor fails loud
regardless of table state. Add an empty-table regression test.
@JunRuiLee
JunRuiLee force-pushed the feat/pk-vector-rerank-reconcile branch from 1ab7337 to 51c5c42 Compare July 20, 2026 10:06

@leaves12138 leaves12138 left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-reviewed at 51c5c42. The refine factor and derived indexed search limit are now resolved before planning and the empty-plan return, so invalid query or table configuration fails consistently; the new empty-table regression covers the reported cases. The dual-limit search and exact positional rerank remain aligned with the Java flow. Verified locally with git diff --check, cargo fmt --check, cargo test -p paimon --lib (1705 passed, 1 ignored), cargo test -p paimon --test pk_vector_baseline_test (8 passed), and cargo check -p paimon-datafusion.

@JingsongLi
JingsongLi merged commit 4211dda into apache:main Jul 20, 2026
12 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants