[python] Optimize vector raw search with numpy vectorization and ThreadPoolExecutor

[839224346]

Purpose

Optimize pypaimon's vector raw search path by replacing the pure-Python loop-based distance computation with numpy vectorized operations, and improve index split concurrency using ThreadPoolExecutor.

Key changes:

1. Numpy vectorized distance computation — Replace per-row Python loop (_compute_score + heap) with batch matrix operations (_raw_search_from_arrow + _numpy_topk), leveraging Arrow's zero-copy buffer for direct numpy matrix construction.

2. O(n) top-K selection — Use np.argpartition instead of a heap-based approach, reducing top-K selection from O(n·log k) to O(n).

3. ThreadPoolExecutor for index splits — Replace the old wait(futures) pattern with ThreadPoolExecutor + as_completed, adding synchronous _eval_sync / _eval_batch_sync methods that properly manage reader lifecycle with try/finally. Single-split case avoids thread pool overhead entirely.

4. Standalone benchmark script — Added benchmark_vector_search_standalone.py for reproducible performance validation without requiring a running Paimon table.

Why:

The raw search path is a fallback for data that hasn't yet been indexed by Faiss/HNSW (e.g., newly written data before compaction). Previously, this path used a pure-Python loop iterating row by row — acceptable for small datasets but extremely slow at scale.

The ThreadPoolExecutor change also fixes a subtle issue: the old _eval() returned futures with reader references via callbacks, but the reader lifecycle wasn't guaranteed in error paths. The new _eval_sync uses explicit try/finally.

Tests

# Standalone benchmark (no Paimon table required)
python3 pypaimon/tests/benchmark_vector_search_standalone.py --num-rows 100000 --dim 768

# Existing test suite
python -m pytest pypaimon/tests/ -q -k "vector"

# Static analysis
python -m pyflakes pypaimon/table/source/vector_search_read.py
git diff --check

Performance results (768 dimensions, top-100, isolated processes):

Environment: 16-core CPU, 32GB RAM, numpy 2.x + OpenBLAS 0.3.34 (Haswell, 64-bit int)

Each path runs in its own process to avoid memory contention — matches production behavior.

100K rows:

Metric	Python loop	Numpy (fast path)	Speedup
L2	4,449 ms	695 ms	6.4×
Cosine	6,214 ms	607 ms	10.2×
Inner Product	4,908 ms	400 ms	12.3×

500K rows:

Metric	Python loop	Numpy (fast path)	Speedup
L2	23,997 ms	3,318 ms	7.2×
Cosine	32,129 ms	2,229 ms	14.4×
Inner Product	21,696 ms	1,862 ms	11.7×

Top-K correctness: 100% overlap between Python loop and numpy path across all metrics.

[JingsongLi]

Two correctness issues are noted inline.

[JingsongLi]

Have we measured the crossover point on small raw tails using the actual _raw_search_from_arrow path? For a single query, stored_matrix @ query_np is GEMV rather than SGEMM, and Arrow-to-NumPy materialization (astype copies by default), BLAS dispatch/thread startup, temporary arrays, and the full stored-norm scan for cosine can dominate when rows * dim is small. Raw fallback often represents only the newly written unindexed tail. The current fast benchmark starts from a pre-built NumPy matrix, so it excludes these costs. Please add small-N benchmarks (for example 1/8/32/128/512/2K rows across representative dimensions) and consider a measured hybrid threshold; batching queries would also let us reuse stored norms and use true SGEMM.

[839224346]

Thanks for the thorough review! I've added end-to-end benchmarks (Arrow table construction → result) and addressed both points:

Small-N crossover benchmark (benchmark_small_n_crossover.py)

Tested N=1/8/32/128/512/2048 × dim=128/768, timing from Arrow table construction:

Path	N=1	N=8	N=128	N=2048
FixedSizeList (real format)	numpy 1.35x faster	9.2x	111x	582x
Variable-length list (to_pylist fallback)	scalar 1.29x faster	numpy 1.43x	1.69x	1.52x

The bottleneck is Python's per-element loop (dim multiplications per row), not BLAS startup. Since Paimon vector columns are FixedSizeList, numpy wins even at N=1 — no hybrid threshold needed.

Batch query SGEMM optimization

Added _raw_batch_search_from_arrow that reads the Arrow table once and computes stored_matrix @ query_matrix.T in a single SGEMM call, reusing stored norms for cosine:

Queries	Loop (μs)	Batch SGEMM (μs)	Speedup
1	297	306	~1x (no regression)
4	1,211	403	3x
8	2,674	513	5.2x
32	12,179	3,069	4x

BatchVectorSearchReadImpl._read_batch now calls the batch path instead of the per-query loop.

[JingsongLi]

[P1] Could we avoid the norm-expansion formula here? In float32 it suffers catastrophic cancellation and can change the Top-K result. For example, with query [100000] and stored vectors [99906] and [99904], direct subtraction gives distances 8836 and 9216 and correctly selects the first row, while this expression produces 10240 and 8192 and selects the farther row. The single-query and previous paths are correct. Please use a numerically stable distance calculation, such as direct differences, possibly in bounded tiles.

[839224346]

Fixed. The L2 branch now uses direct subtraction (sum((a-b)²)) instead of the norm-expansion formula.

Cosine and inner_product still use SGEMM — they don't suffer from this issue because their computation structure has no large-magnitude subtraction:

• inner_product: pure dot product (a·b), only accumulation, no cancellation risk.
• cosine: (a·b) / (||a|| × ||b||) is a ratio where numerator and denominator are well-matched in scale — there's no "big minus big yields small" pattern that destroys significant digits.

The L2 problem is specific to ||a||² + ||b||² - 2a·b where the first two terms and the third are close in magnitude, causing catastrophic cancellation in float32.

[JingsongLi]

[P1] Could we tile this computation instead of materializing the complete rows × queries score space? In the L2 branch, dots, dists, and all_scores simultaneously retain full matrices, while the batch API does not bound the query count. At 1,000,000 rows × 128 queries, each float32 matrix is about 512 MB, so these intermediates alone exceed 1.5 GB before Arrow and stored-vector buffers. The previous per-query path had bounded peak memory. Please process rows or queries in bounded blocks and merge each query’s Top-K incrementally.

[839224346]

Fixed.

[JingsongLi]

[P2] Using or here makes the validation below ineffective for zero: a configured global-index.thread-num=0 becomes the default 32 before value < 1 is checked. This differs from the Java contract, which rejects non-positive values, and unexpectedly enables up to 32 concurrent readers. Please apply the default only when the configured value is None, then reject every value below 1.

[839224346]

Fixed.

[JingsongLi]

[P2] Please handle the empty table immediately after filtering parent nulls, before accessing stored_matrix.shape[1]. For a nullable ARRAY<FLOAT> column whose scanned values are all null, the filter leaves zero rows and the ListArray fallback builds np.array([], dtype=np.float32) with shape (0,), so this line raises IndexError: tuple index out of range. The previous loop skipped those null vectors and returned an empty result. The batch helper has the same issue at its dimension check. Please return the appropriate empty result after filtering when no rows remain, and add all-null ListArray coverage.

[JingsongLi]

[P2] Please preserve the existing smaller-row-ID tie-break at the partial Top-K boundary. np.argpartition partitions only by score, so equal-score rows are selected arbitrarily. With row IDs [4, 3, 2, 1], identical vectors, and limit=2, both new NumPy paths return [3, 4], while _offer_score and ScoredGlobalIndexResult.top_k define [1, 2]. A later .top_k() cannot repair this because this helper has already discarded all but k rows. Please include row ID as the secondary selection key; the batch Top-K branches need the same treatment.