Skip to content

[python] Share raw vector scans across batch queries - #9753

Merged
JingsongLi merged 3 commits into
apache:masterfrom
TheR1sing3un:codex/batch-vector-shared-raw-scan
Sep 13, 2026
Merged

[python] Share raw vector scans across batch queries#9753
JingsongLi merged 3 commits into
apache:masterfrom
TheR1sing3un:codex/batch-vector-shared-raw-scan

Conversation

@TheR1sing3un

@TheR1sing3un TheR1sing3un commented Sep 12, 2026

Copy link
Copy Markdown
Member

Purpose

Batch vector search reads and converts the same uncovered rows once per query. Share a snapshot-pinned raw scan across the query batch, consume Arrow batches incrementally, and maintain a separate Top-K heap for each query.

Preserve the existing split-read parallelism, including the read.parallelism table option and automatic worker count. Each worker streams its assigned splits, scores all queries, and retains bounded Top-K state. Merge the worker results with the same score and row-ID tie-breaking rules. Close source iterators on success or failure and wait for running workers before propagating an error.

Scoring buffers scale with the worker count: at most one Arrow batch and its converted vectors plus Q Top-K heaps per worker. There is no full-shard table or Q-by-N score matrix in this fallback. Storage decoder buffers and scan metadata remain separate costs.

Tests

  • python -m pytest pypaimon/tests/batch_vector_raw_scan_test.py pypaimon/tests/vector_search_filter_test.py -q: 85 passed.
  • Real four-partition Parquet regression through execute_batch_local(): one/two queries, L2/cosine/inner product, parallelism 1/2/4/automatic, identical scores and IDs versus individual searches, and one scan of every split. Synchronization barriers verify concurrency without timing thresholds.
  • Parallel dimension/read failures close all started source iterators before returning. Existing coverage includes scalar/partition/prefilters, snapshot pinning, nulls, ties, empty input, mixed indexed/raw results and incremental batch consumption.
  • Repository-configured flake8, license-header checks and git diff --check passed.

Benchmark and ablation

The experiment calls the public execute_batch_local() API on real local Paimon/Parquet tables with no vector index and vector-index.search-mode=full.

  • repeated: reproduces the previous per-query raw scan loop.
  • shared-table: scans once but materializes the full table/Python lists, isolating the shared-scan benefit.
  • shared-stream: this change, sharing the scan and consuming one batch at a time.

All three use identical scalar distance and top-k helpers. Float32 vectors, 128 dimensions, L2, top-k 10, batch size 1,024, read parallelism 1. Each configuration runs three times in fresh processes, sequentially with shuffled configuration order. Timing includes planning, reading, conversion, scoring and merging; excludes data generation/process startup. RSS is total process peak including imports and Arrow buffers. Environment: macOS 26.4.1 arm64, Python 3.9.6, PyArrow 19.0.1, NumPy 2.0.2; OMP/OPENBLAS thread counts 1. Filesystem cache is not flushed.

Rows Queries Mode Median seconds (min–max) Median peak RSS, MiB Raw scans
16,384 1 repeated 0.410 (0.409–0.427) 255.7 1
16,384 1 shared-table 0.414 (0.403–0.433) 255.6 1
16,384 1 shared-stream 0.395 (0.390–0.406) 170.5 1
16,384 8 repeated 3.199 (3.124–3.253) 256.1 8
16,384 8 shared-table 1.369 (1.369–1.413) 255.8 1
16,384 8 shared-stream 1.357 (1.355–1.377) 170.6 1
16,384 32 repeated 12.599 (12.472–12.713) 256.5 32
16,384 32 shared-table 4.635 (4.624–4.637) 256.7 1
16,384 32 shared-stream 4.646 (4.530–4.677) 171.0 1
65,536 8 repeated 12.589 (12.282–12.817) 581.0 8
65,536 8 shared-table 5.521 (5.428–5.785) 564.8 1
65,536 8 shared-stream 5.357 (5.353–5.375) 192.9 1

The 8-query cases improve 2.35–2.36x; 32 queries improve 2.71x. The shared-table ablation shows that shared scanning accounts for the timing improvement, while streaming lowers peak RSS. For 65,536 rows and 8 queries, peak RSS falls from 581.0 to 192.9 MiB (66.8%). Single-query timing is similar.

Batch-size sensitivity at 16,384 rows / 8 queries:

Batch rows Median seconds Median peak RSS, MiB
256 1.395 164.2
1024 1.357 170.6
4096 1.386 196.8

All 42 runs produce exactly matching row IDs and floating-point scores within each dataset/query-count group, including across batch sizes. Instrumentation verifies Q raw plans/reader passes and N×Q delivered rows for repeated, versus one plan/pass and N delivered rows for both shared modes. These are logical reader counters, not physical disk reads.

These measurements fix read parallelism to 1, isolating scan sharing and streaming memory. The revised implementation also preserves split concurrency; that validation is reported separately below.

Split concurrency experiment

A local four-partition Parquet table with 20 two-dimensional rows, L2, Top-K 2 and batch size 2. Inject 150 ms at each split-reader open. Timing covers the public batch API after data preparation and a warm-up query. Each configuration runs three times in fresh Python 3.9 processes on macOS arm64; values below are medians. This isolates split-open latency and is not a live object-store throughput measurement.

Mode Read parallelism One query (ms) Two queries (ms) Split opens (one/two queries) Peak concurrent opens
Original repeated scans 4 162.79 327.89 4 / 8 4
Shared serial scan 4 configured 639.54 638.22 4 / 4 1
Shared parallel scan 2 321.54 322.09 4 / 4 2
Shared parallel scan 4 165.18 161.38 4 / 4 4

All 24 runs return exactly identical IDs and scores for the same query set. Single-query latency returns to the original parallel-read level; two queries reuse the same four split reads.

@TheR1sing3un
TheR1sing3un marked this pull request as ready for review September 12, 2026 06:48

@JingsongLi JingsongLi left a comment

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.

Reviewed c718d0b. Requirement fit: supported; one performance regression needs attention.

Sharing the raw scan across queries removes real repeated work, and the focused tests passed (82 tests). The new streaming path, however, serializes split I/O that the previous path ran concurrently; the existing benchmark fixes read.parallelism to 1 and therefore cannot reveal this regression.

Comment thread paimon-python/pypaimon/table/source/vector_search_read.py
@TheR1sing3un
TheR1sing3un force-pushed the codex/batch-vector-shared-raw-scan branch from c718d0b to 9a23e2d Compare September 13, 2026 01:15
@JingsongLi
JingsongLi merged commit 4641a45 into apache:master Sep 13, 2026
9 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.

2 participants