JingsongLi commented on code in PR #9315:
URL: https://github.com/apache/paimon/pull/9315#discussion_r4183340731
##########
paimon-python/pypaimon/table/source/vector_search_read.py:
##########
@@ -734,51 +794,43 @@ def _read_batch(self, splits, snapshot):
indexed_results = self._maybe_rerank_indexed_results(
indexed_results, index_type, self._query_vectors, snapshot)
- # Each query: merge indexed results with the raw (brute-force)
fallback.
+ # Batch raw search: read Arrow table once, compute all queries in one
SGEMM.
raw_pre_filter = self._raw_pre_filter(raw_splits, snapshot)
raw_ranges = _raw_row_ranges(raw_splits)
raw_index_type = _raw_search_index_type(raw_splits)
raw_results = self._read_raw_batch_search(
- raw_ranges, raw_pre_filter, raw_index_type, snapshot)
- return [
- indexed.or_(raw).top_k(self._limit)
- for indexed, raw in zip(indexed_results, raw_results)
- ]
+ raw_ranges, raw_pre_filter, raw_index_type,
+ snapshot=snapshot)
+
+ results = []
+ for i in range(n):
+
results.append(indexed_results[i].or_(raw_results[i]).top_k(self._limit))
+ return results
def _read_raw_batch_search(self, raw_row_ranges, pre_filter,
index_type=None, snapshot=None):
- """Scan raw rows once, keeping a separate top-k heap for each query."""
- heaps = [[] for _ in self._query_vectors]
+ from pypaimon.read.table_read import _ClosableArrowBatchReader
+
raw_row_ranges = _filtered_raw_row_ranges(raw_row_ranges, pre_filter)
- if not raw_row_ranges or not heaps:
- return [_scored_result(heap) for heap in heaps]
+ n = len(self._query_vectors)
+ if not raw_row_ranges:
+ return [DictBasedScoredIndexResult({}) for _ in range(n)]
+ import pyarrow as pa
table_read, splits = self._plan_raw_read(raw_row_ranges, True,
snapshot)
- metric = self._search_metric(index_type)
- workers = min(len(splits), table_read._resolve_parallelism(None,
len(splits)))
- if workers <= 1:
- return self._score_raw_splits(table_read, splits, metric)
-
- # Keep only one streaming reader and Q top-k heaps per worker, even
- # when the plan contains many splits. Each split is scanned once.
- with ThreadPoolExecutor(max_workers=workers) as executor:
- futures = [executor.submit(
- self._score_raw_splits, table_read, splits[i::workers], metric)
- for i in range(workers)]
- for future in futures:
- for heap, result in zip(heaps, future.result()):
- score_getter = result.score_getter()
- for row_id in result.results():
- _offer_score(heap, self._limit, row_id,
score_getter(row_id))
- return [_scored_result(heap) for heap in heaps]
-
- def _score_raw_splits(self, table_read, splits, metric):
- from pypaimon.read.table_read import _ClosableArrowBatchReader
-
reader, batches = table_read._new_arrow_batch_reader(splits)
- # Close the underlying iterator as well if scoring fails mid-batch.
with _ClosableArrowBatchReader(reader, batches) as batch_reader:
- return self._score_raw_batch_queries(batch_reader, metric)
+ collected = [b for b in batch_reader]
Review Comment:
[P1] Keep raw batch scoring bounded before matrix conversion
This replaces the base streaming raw search with collected=[b for b in
batch_reader], then creates a full Arrow table and a full float64 stored
matrix. Row tiling happens only afterwards, so it does not bound these
allocations. ARRAY<FLOAT> additionally takes the regular ListArray fallback
through to_pylist(), expanding every component into a Python float before
building that matrix. I measured the real public execute_batch_local on
committed Parquet tables in separate processes (768 dimensions, two queries,
limit=10, read.batch-size=256, read.parallelism=1, full raw mode): 4,096
rows/12 MiB payload grows search peak by 197.5 MiB here versus 29.9 MiB with
the exact base reader; 16,384 rows/48 MiB grows it by 725.3 MiB versus 42.1
MiB. Both return identical Top-K IDs. This scales with the complete unindexed
tail despite the configured read batch and small result limit, so large newly
written tails can exhaust memory that the original bounded path handled. Retain
vectorized scoring
and Top-K selection per bounded Arrow/matrix block, merging only each query's
bounded best candidates; use bounded flatten/offset conversion for regular
ListArray instead of materializing all vectors. Restore tests asserting that
batches are released before later batches are read.
##########
paimon-python/pypaimon/table/source/vector_search_read.py:
##########
@@ -1200,3 +1252,332 @@ 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 as pa
+ 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)
+
+ # Filter by score_candidates at Arrow level to avoid ragged-array issues
+ # when non-candidate rows have mismatched dimensions.
+ if score_candidates is not None:
+ candidate_set = set(score_candidates)
+ candidate_mask = pc.is_in(row_ids_col, pa.array(
+ list(candidate_set), type=row_ids_col.type))
+ arrow_table = arrow_table.filter(candidate_mask)
+ row_ids_col = arrow_table.column(SpecialFields.ROW_ID.name)
+ vectors_col = arrow_table.column(vector_column_name)
+
+ if arrow_table.num_rows == 0:
+ return DictBasedScoredIndexResult({})
+
+ # 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.float64)
+ else:
+ stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float64)
+ except (AttributeError, TypeError, ValueError):
+ stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float64)
+
+ query_np = np.asarray(query_vector, dtype=np.float64)
+
+ 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]))
+
+ # Filter NaN vectors (from variable-length lists that got padded).
+ 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 with row tiling."""
+ import numpy as np
+
+ ROW_TILE = 65536
+ n_rows = stored_matrix.shape[0]
+
+ if n_rows <= ROW_TILE:
+ scores = _compute_scores_single(stored_matrix, query_np, 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}
+ )
+
+ best_ids = np.empty(0, dtype=row_id_array.dtype)
+ best_scores = np.empty(0, dtype=np.float32)
+
+ for r_start in range(0, n_rows, ROW_TILE):
+ r_end = min(r_start + ROW_TILE, n_rows)
+ tile_scores = _compute_scores_single(
+ stored_matrix[r_start:r_end], query_np, metric)
+ tile_rids = row_id_array[r_start:r_end]
+ tile_top = _topk_indices(tile_scores, tile_rids, limit)
+
+ best_ids = np.concatenate([best_ids, tile_rids[tile_top]])
+ best_scores = np.concatenate([best_scores, tile_scores[tile_top]])
+
+ final_top = _topk_indices(best_scores, best_ids, limit)
+ return DictBasedScoredIndexResult(
+ {int(best_ids[i]): float(best_scores[i]) for i in final_top}
+ )
+
+
+def _compute_scores_single(stored_chunk, query_np, metric):
+ import numpy as np
+
+ if metric == "l2":
+ diffs = stored_chunk - query_np
+ dists = np.sum(diffs * diffs, axis=1)
+ return 1.0 / (1.0 + dists)
+ elif metric == "cosine":
+ dots = stored_chunk @ query_np
+ norms = np.linalg.norm(stored_chunk, axis=1) * np.linalg.norm(query_np)
+ norms = np.where(norms == 0, 1.0, norms)
+ return dots / norms
+ elif metric == "inner_product":
+ return stored_chunk @ query_np
+ else:
+ raise ValueError("Unknown vector search metric: %s" % metric)
+
+
+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))
+
+ part_idx = np.argpartition(-scores, limit)[:limit]
+ kth_score = np.min(scores[part_idx])
+
+ above_mask = scores[part_idx] > kth_score
+ above = part_idx[above_mask]
+
+ all_tie_idx = np.where(scores == kth_score)[0]
+ n_ties_needed = limit - len(above)
+
+ if len(all_tie_idx) <= n_ties_needed:
+ result = np.concatenate([above, all_tie_idx])
+ else:
+ tie_order = np.argsort(row_id_array[all_tie_idx])
+ result = np.concatenate([above,
all_tie_idx[tie_order[:n_ties_needed]]])
+
+ final_order = np.lexsort((row_id_array[result], -scores[result]))
+ return result[final_order]
+
+
+def _raw_batch_search_from_arrow(arrow_table, vector_column_name,
query_vectors,
+ metric, limit, score_candidates=None):
+ """Batch raw search: multiple queries against the same Arrow table in one
SGEMM call."""
+ 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)
+
+ 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)
+
+ row_id_array = row_ids_col.to_numpy()
+ try:
+ 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.float64)
+ else:
+ stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float64)
+ except (AttributeError, TypeError, ValueError):
+ stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float64)
+
+ query_matrix = np.array(
+ [q if isinstance(q, np.ndarray) else list(q) for q in query_vectors],
+ dtype=np.float64)
+
+ n = len(query_vectors)
+
+ if stored_matrix.ndim < 2 or stored_matrix.shape[0] == 0:
+ return [DictBasedScoredIndexResult({}) for _ in range(n)]
+
+ if stored_matrix.shape[1] != query_matrix.shape[1]:
+ raise ValueError(
+ "Query vector dimension mismatch: expected %d, got %d"
+ % (stored_matrix.shape[1], query_matrix.shape[1]))
+
+ 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)
+ 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({}) for _ in range(n)]
+
+ return _numpy_batch_topk(row_id_array, stored_matrix, query_matrix,
metric, limit)
+
+
+def _numpy_batch_topk(row_id_array, stored_matrix, query_matrix, metric,
limit):
+ """Batch distance computation + per-query topK with row and query
tiling."""
+ import numpy as np
+
+ ROW_TILE = 65536
+ n_queries = query_matrix.shape[0]
+ n_rows = stored_matrix.shape[0]
+
+ if n_rows <= ROW_TILE:
+ return _numpy_batch_topk_no_row_tile(
+ row_id_array, stored_matrix, query_matrix, metric, limit)
+
+ accum_ids = [np.empty(0, dtype=row_id_array.dtype) for _ in
range(n_queries)]
+ accum_scores = [np.empty(0, dtype=np.float32) for _ in range(n_queries)]
+
+ for r_start in range(0, n_rows, ROW_TILE):
+ r_end = min(r_start + ROW_TILE, n_rows)
+
+ tile_results = _numpy_batch_topk_no_row_tile(
+ row_id_array[r_start:r_end], stored_matrix[r_start:r_end],
+ query_matrix, metric, limit)
+
+ for qi, res in enumerate(tile_results):
+ if res.results().cardinality() == 0:
+ continue
+ score_getter = res.score_getter()
+ rid_list = list(res.results())
+ ids = np.array(rid_list, dtype=row_id_array.dtype)
+ scores = np.array([score_getter(r) for r in rid_list],
dtype=np.float32)
Review Comment:
[P2] Preserve float64 scores when merging row tiles
The per-tile scorer produces float64 scores, but this array converts them to
float32 before the global Top-K merge. With more than 65,536 rows, finite legal
FLOAT inputs can therefore collapse into Infinity or zero and select the wrong
row. I reproduced the complete public scalar and batch paths on one committed
Parquet table with 65,537 one-dimensional ARRAY<FLOAT> rows: the first 65,536
vectors are 1e23 and the last is 2e23. For inner_product with
query=float32(1e20), the current scalar and exact-base batch select row 65536
with a finite score around 2e43; current execute_batch_local selects row 0 with
Infinity. For L2 with query=float32(3e23), scalar/base batch again select 65536
(score around 1e-46), while current batch selects row 0 with zero. All input
components are finite float32, and this is unrelated to signed zero. Keep
accumulated scores in float64 and add cross-tile overflow/underflow regressions
which compare the public batch result to scalar/base behavior.
--
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]
