Skip to content

[WIP][spark] Distribute lateral vector search across index splits - #8879

Open
shyjsarah wants to merge 4 commits into
apache:masterfrom
shyjsarah:codex/vector-multi-partition-20260727
Open

[WIP][spark] Distribute lateral vector search across index splits#8879
shyjsarah wants to merge 4 commits into
apache:masterfrom
shyjsarah:codex/vector-multi-partition-20260727

Conversation

@shyjsarah

Copy link
Copy Markdown
Contributor

Purpose

Distribute Spark lateral vector_search over compatible vector-index split groups so multi-partition searches can use executor parallelism instead of searching every index split serially in one task.

Changes

  • group vector-index splits by estimated search cost
  • execute local Top-K search independently for each split group
  • merge and deduplicate partial results into a deterministic global Top-K per query
  • preserve the outer row when a split group returns no result
  • pin the snapshot before materializing distributed results
  • fall back to the existing local path for unsupported plans, including deletion vectors, raw-data fallback, and refine/rerank
  • add an opt-in Spark connector option; the feature remains disabled by default:
spark.paimon.vector-search.lateral-join.distributed.enabled=true

The maximum number of split groups can be controlled with:

spark.paimon.vector-search.lateral-join.distributed.max-split-groups=16

This does not change IVF-PQ nprobe and does not introduce a cross-partition index.

Tests

  • VectorSearchBuilderTest
  • LateralVectorSearchExecutionTest
  • TableValuedFunctionsTest

@shyjsarah

Copy link
Copy Markdown
Contributor Author

Added a follow-up commit to balance the distributed result-merge stage.

In the 100k-query / multi-partition validation, the merge stage had 16 tasks but several tasks received almost no work. The default hash partitioning of QueryBatchId can collide for the first batches from different input partitions.

This change introduces a dedicated partitioner based on (input partition + batch ordinal) % merge parallelism. It keeps every identical query batch on the same reducer while spreading batches deterministically. The logic is independent of the table, filter, partition count, and configured parallelism.

Validation:

  • LateralVectorSearchExecutionTest: passed on Scala 2.12 and 2.13
  • TableValuedFunctionsTest: 23/23 passed
  • Spotless and diff checks passed

Next I will validate the packaged change with the multi-partition workload.

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.

1 participant