839224346 commented on code in PR #9315:
URL: https://github.com/apache/paimon/pull/9315#discussion_r3820819485
##########
paimon-python/pypaimon/table/source/vector_search_read.py:
##########
@@ -824,3 +906,123 @@ def _compute_score(query, stored, metric):
if metric == "inner_product":
return sum(float(q) * float(s) for q, s in zip(query, stored))
raise ValueError("Unknown vector search metric: %s" % metric)
+
+
+def _raw_search_vectorized(row_ids, vectors, query_vector, metric, limit,
+ score_candidates=None):
+ """Vectorized raw search using numpy for batch distance computation."""
+ import numpy as np
+
+ # Filter by score_candidates and null vectors.
+ if score_candidates is not None:
+ candidate_set = set(score_candidates)
+ filtered = [(rid, vec) for rid, vec in zip(row_ids, vectors)
+ if rid in candidate_set and vec is not None]
+ else:
+ filtered = [(rid, vec) for rid, vec in zip(row_ids, vectors)
+ if vec is not None]
+
+ if not filtered:
+ return DictBasedScoredIndexResult({})
+
+ filtered_ids, filtered_vecs = zip(*filtered)
+ row_id_array = np.array(filtered_ids, dtype=np.int64)
+ stored_matrix = np.array(
+ [_to_vector_list(v) for v in filtered_vecs], dtype=np.float32)
+ query_np = np.array(
+ _to_vector_list(query_vector) if not isinstance(query_vector,
np.ndarray)
+ else query_vector, dtype=np.float32)
+
+ return _numpy_topk(row_id_array, stored_matrix, query_np, metric, limit)
+
+
+def _raw_search_from_arrow(arrow_table, vector_column_name, query_vector,
+ metric, limit, score_candidates=None):
+ """Vectorized raw search directly from Arrow table (avoids Python list
intermediary)."""
+ import numpy as np
+ import pyarrow.compute as pc
+
+ row_ids_col = arrow_table.column(SpecialFields.ROW_ID.name)
+ vectors_col = arrow_table.column(vector_column_name)
+
+ # Filter out null vectors at the Arrow level before conversion.
+ valid_mask = pc.is_valid(vectors_col)
+ if not pc.all(valid_mask).as_py():
+ arrow_table = arrow_table.filter(valid_mask)
+ row_ids_col = arrow_table.column(SpecialFields.ROW_ID.name)
+ vectors_col = arrow_table.column(vector_column_name)
+
+ # Try fast path: fixed-size list → direct numpy reshape.
+ row_id_array = row_ids_col.to_numpy()
+ try:
+ # ChunkedArray has no .values; combine to a single array first.
+ if hasattr(vectors_col, 'combine_chunks'):
+ vectors_arr = vectors_col.combine_chunks()
+ else:
+ vectors_arr = vectors_col
+ flat = vectors_arr.values
+ dim = vectors_arr.type.list_size
+ if dim is not None and flat is not None:
+ stored_matrix = flat.to_numpy(zero_copy_only=False).reshape(-1,
dim).astype(
+ np.float32)
+ else:
+ stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float32)
+ except (AttributeError, TypeError, ValueError):
+ stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float32)
+
+ query_np = np.asarray(query_vector, dtype=np.float32)
+
+ if stored_matrix.shape[1] != query_np.shape[0]:
+ raise ValueError(
+ "Query vector dimension mismatch: expected %d, got %d"
+ % (stored_matrix.shape[1], query_np.shape[0]))
+
+ # Handle null vectors and score_candidates filtering.
+ if score_candidates is not None:
+ candidate_set = set(score_candidates)
+ mask = np.array([rid in candidate_set for rid in row_id_array],
dtype=bool)
+ # Also mask null vectors (check for any NaN row).
+ null_mask = ~np.isnan(stored_matrix).any(axis=1)
+ mask = mask & null_mask
+ row_id_array = row_id_array[mask]
+ stored_matrix = stored_matrix[mask]
+ else:
+ null_mask = ~np.isnan(stored_matrix).any(axis=1)
+ if not null_mask.all():
+ row_id_array = row_id_array[null_mask]
+ stored_matrix = stored_matrix[null_mask]
+
+ if len(row_id_array) == 0:
+ return DictBasedScoredIndexResult({})
+
+ return _numpy_topk(row_id_array, stored_matrix, query_np, metric, limit)
+
+
+def _numpy_topk(row_id_array, stored_matrix, query_np, metric, limit):
+ """Core numpy distance computation + topK selection."""
+ import numpy as np
+
+ if metric == "l2":
+ diffs = stored_matrix - query_np
+ dists = np.sum(diffs * diffs, axis=1)
+ 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)
Review Comment:
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
