[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.