[python] Optimize vector raw search with numpy vectorization and ThreadPoolExecutor by vibex-wang · Pull Request #9315 · apache/paimon

vibex-wang · 2026-08-20T07:31:50Z

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:

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.
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).
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.
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 · 2026-08-20T09:05:14Z

+        scores = 1.0 / (1.0 + dists)
+    elif metric == "cosine":
+        dots = stored_matrix @ query_np
+        norms = np.linalg.norm(stored_matrix, axis=1) * np.linalg.norm(query_np)


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.

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 · 2026-08-21T09:07:20Z

+        stored_sq = np.sum(stored_matrix * stored_matrix, axis=1, keepdims=True)
+        query_sq = np.sum(query_matrix * query_matrix, axis=1, keepdims=True)
+        dots = stored_matrix @ query_matrix.T
+        dists = stored_sq + query_sq.T - 2 * dots


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

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 · 2026-08-21T09:07:23Z

+        dots = stored_matrix @ query_matrix.T
+        dists = stored_sq + query_sq.T - 2 * dots
+        np.maximum(dists, 0, out=dists)
+        all_scores = 1.0 / (1.0 + dists)


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

JingsongLi · 2026-08-21T09:07:25Z

+        _get = getattr(_opts, 'global_index_thread_num', None)
+        value = (
+            (_get() if _get else None)
+            or CoreOptions.GLOBAL_INDEX_THREAD_NUM._default_value


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

JingsongLi · 2026-08-27T08:31:25Z

+
+    query_np = np.asarray(query_vector, dtype=np.float32)
+
+    if stored_matrix.shape[1] != query_np.shape[0]:


[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 · 2026-08-27T08:31:28Z

+    if n <= limit:
+        top_indices = np.argsort(-scores)
+    else:
+        top_indices = np.argpartition(-scores, limit)[:limit]


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

JingsongLi · 2026-09-04T08:21:09Z

+
+    # Full lexsort is O(n log n) but guarantees correct tie-break at the
+    # partition boundary where argpartition alone would pick arbitrarily.
+    order = np.lexsort((row_id_array, -scores))


[P1] This full lexsort restores the exact tie-break, but it also turns every raw Top-K back into O(n log n) and discards the PR’s O(n) selection optimization. For the advertised 500K-row/top-100 case, every query now sorts all 500K scores; the benchmark numbers were recorded before this change, so they no longer validate the current path. Please keep argpartition for the strict-above-boundary set and resolve only the kth-score ties by smallest row ID (or benchmark and justify another bounded selection), then rerun the published benchmark.

JingsongLi · 2026-09-04T08:21:09Z

+    import numpy as np
+
+    if metric == "l2":
+        diffs = stored_matrix - query_np


[P1] The numerically stable L2 implementation still materializes an unbounded rows × dimension temporary. At 1M rows × 768 float32 dimensions, stored_matrix is about 3.1 GB and diffs allocates another 3.1 GB (plus dists/scores and Arrow buffers); the batch branch repeats the same allocation for every query. This can OOM exactly on the large raw tails this PR targets, whereas the previous row loop had bounded auxiliary memory. Please tile rows as well as queries and merge each tile’s Top-K incrementally; the same fix is needed in _numpy_batch_topk.

JingsongLi · 2026-09-10T08:18:05Z

+    return _numpy_distance_topk(row_id_array, stored_matrix, query_np, metric, limit)
+
+
+def _numpy_distance_topk(row_id_array, stored_matrix, query_np, metric, limit):


Please remove the benchmark-only code from this PR, including both benchmark_vector_search_standalone.py and benchmark_small_n_crossover.py. These scripts duplicate the production vector-search implementation, are not regression tests, and have already diverged from the current row-tiling and tie-break logic. Keeping copied benchmark implementations here creates misleading validation and an ongoing maintenance burden; performance measurements can remain in the PR description instead.

Just pushed, please take another look.

…adPoolExecutor

- Add _raw_batch_search_from_arrow: computes all query vectors against the same stored matrix in one SGEMM call (stored_matrix @ query_matrix.T), reusing stored norms for cosine metric. - Modify BatchVectorSearchReadImpl._read_batch to read the raw Arrow table once and call the batch path instead of per-query loop. - Add small-N crossover benchmark proving numpy is faster even at N=1 for FixedSizeList (the real storage format). - Benchmark shows 3-5x speedup for multi-query batch raw search (4 queries: 3x, 8 queries: 5.2x).

…, thread-num validation - Tile _numpy_batch_topk by QUERY_TILE=8 to bound peak memory from O(N*Q) to O(N*tile_size), addressing unbounded query count. - Replace norm-expansion L2 (||a||²+||b||²-2a·b) with direct subtraction to avoid float32 catastrophic cancellation. - Fix thread-num config: use 'is None' instead of 'or' so that configured zero/negative values are properly rejected.

…ementation - Delete benchmark_small_n_crossover.py and benchmark_vector_search_standalone.py - Unify _read_raw_batch_search method name (was _read_batch_raw_search) - Delegate _read_indexed/_read_batch to _search_index_splits for Ray compatibility - Use float64 in numpy scoring paths to match streaming/Ray precision - Filter score_candidates at Arrow level before numpy conversion - Restore _score_raw_batch_queries for Ray worker usage - Adapt batch_vector_raw_scan_test to numpy implementation

Batch raw search used _raw_search_metric() directly, bypassing the _index_metric cache populated during the earlier index-split phase. This could cause metric divergence when the index file's actual metric differs from the table configuration. Switch to self._search_metric(index_type) to match the single-query _read_raw_search path and upstream's original implementation.

vibex-wang force-pushed the feature/vector-search-numpy-optimization branch 2 times, most recently from fd33d47 to 4997653 Compare August 20, 2026 07:48

JingsongLi reviewed Aug 20, 2026

View reviewed changes

Comment thread paimon-python/pypaimon/table/source/vector_search_read.py Outdated

Comment thread paimon-python/pypaimon/table/source/vector_search_read.py Outdated

JingsongLi reviewed Aug 20, 2026

View reviewed changes

vibex-wang force-pushed the feature/vector-search-numpy-optimization branch 2 times, most recently from ee5ee3a to a160bdd Compare August 20, 2026 11:05

JingsongLi reviewed Aug 21, 2026

View reviewed changes

JingsongLi reviewed Aug 27, 2026

View reviewed changes

vibex-wang force-pushed the feature/vector-search-numpy-optimization branch from 79c9607 to 6a2ddb9 Compare August 31, 2026 08:36

JingsongLi reviewed Sep 4, 2026

View reviewed changes

vibex-wang force-pushed the feature/vector-search-numpy-optimization branch 5 times, most recently from a517994 to 81aa708 Compare September 10, 2026 07:38

JingsongLi reviewed Sep 10, 2026

View reviewed changes

vibex-wang force-pushed the feature/vector-search-numpy-optimization branch 2 times, most recently from 9f78dce to 7e71708 Compare September 28, 2026 03:35

vibex-wang added 6 commits September 28, 2026 12:29

[python] Optimize vector raw search with numpy vectorization and Thre…

c642bb6

…adPoolExecutor

[python] Fix nullable vector handling and validate thread-num config

b959bd8

[python] Fix empty-table IndexError and numpy top-k tie-break

fa34afb

[python] Tile rows in numpy topK and use O(n) argpartition for selection

ab91326

vibex-wang force-pushed the feature/vector-search-numpy-optimization branch 2 times, most recently from 8ab6c03 to fb2e903 Compare September 28, 2026 05:56

vibex-wang force-pushed the feature/vector-search-numpy-optimization branch 2 times, most recently from 66d97d2 to 58086d0 Compare September 28, 2026 06:37

vibex-wang force-pushed the feature/vector-search-numpy-optimization branch from 58086d0 to 6b379da Compare September 28, 2026 07:28

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

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


		query_np = np.asarray(query_vector, dtype=np.float32)

		if stored_matrix.shape[1] != query_np.shape[0]:

		return _numpy_distance_topk(row_id_array, stored_matrix, query_np, metric, limit)


		def _numpy_distance_topk(row_id_array, stored_matrix, query_np, metric, limit):

Conversation

vibex-wang commented Aug 20, 2026

Purpose

Tests

Uh oh!

JingsongLi left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants