JingsongLi commented on code in PR #9315:
URL: https://github.com/apache/paimon/pull/9315#discussion_r3932206252
##########
paimon-python/pypaimon/table/source/vector_search_read.py:
##########
@@ -824,3 +923,246 @@ 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.ndim < 2 or stored_matrix.shape[0] == 0:
+ return DictBasedScoredIndexResult({})
+
+ 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
Review Comment:
[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.
##########
paimon-python/pypaimon/table/source/vector_search_read.py:
##########
@@ -824,3 +923,246 @@ 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.ndim < 2 or stored_matrix.shape[0] == 0:
+ return DictBasedScoredIndexResult({})
+
+ 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)
+ norms = np.where(norms == 0, 1.0, norms)
+ scores = dots / norms
+ elif metric == "inner_product":
+ scores = stored_matrix @ query_np
+ else:
+ raise ValueError("Unknown vector search metric: %s" % metric)
+
+ top_indices = _topk_indices(scores, row_id_array, limit)
+
+ return DictBasedScoredIndexResult(
+ {int(row_id_array[i]): float(scores[i]) for i in top_indices}
+ )
+
+
+def _topk_indices(scores, row_id_array, limit):
+ """Select top-limit indices by (highest score, smallest row_id)
tie-break."""
+ import numpy as np
+
+ n = len(scores)
+ if n <= limit:
+ return np.lexsort((row_id_array, -scores))
+
+ # 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))
Review Comment:
[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.
--
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]
