Skip to content

[python] Support batch vector search with Ray - #9950

Open
TheR1sing3un wants to merge 2 commits into
apache:masterfrom
TheR1sing3un:codex/ray-batch-vector-search
Open

TheR1sing3un wants to merge 2 commits into
apache:masterfrom
TheR1sing3un:codex/ray-batch-vector-search

Conversation

@TheR1sing3un

@TheR1sing3un TheR1sing3un commented Sep 18, 2026

Copy link
Copy Markdown
Member

Purpose

Add search_vectors(...).to_arrow(execution="ray", concurrency=..., ray_remote_args=...) for data-evolution tables. Each Ray task handles the whole query batch for one split: index workers reuse an open shard across bounded query blocks, and raw workers stream each complete data split once with an independent top-k per query.

Reuse the existing global candidate selection, driver-side batch refinement, and shared final row lookup. All queries and workers use one pinned snapshot, preserving input order, persisted index metrics, filters, deletions, and deterministic duplicate precedence. Batch execution uses the existing scoring behavior and bounded task scheduling and cancellation. Local execution remains the default.

Tests

  • Python 3.11 / Ray 2.54.0: 149 passed across the new batch Ray suite, existing single-vector Ray suite, batch raw scan, streaming refinement, and shared lookup suites.
  • Ray 2.44.0 / NumPy 1.24.3 / PyArrow 18.1.0: 75 passed across both Ray vector search suites, including real workers and native vector indexes.
  • Cover query blocking and reader cleanup, one task per data split, one final lookup, independent global candidate sets, reversed completion order, empty and historical snapshots, concurrent commits, filters, and invalid inputs. Add raw batch coverage to the CI Ray compatibility matrix.
  • Flake8, Python 3.6 syntax parsing, and git diff --check passed. Upstream Rust planner validation is delegated to CI.


def _score_refine_splits(self, table_read, splits, queries_by_row, query_vectors, metric):
def _score_refine_splits(self, table_read, splits, queries_by_row, query_vectors, metric,
reject_nan=False):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Do not handle nan, let is go.

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.

2 participants