JingsongLi commented on code in PR #9315:
URL: https://github.com/apache/paimon/pull/9315#discussion_r3828939441
##########
paimon-python/pypaimon/table/source/vector_search_read.py:
##########
@@ -824,3 +924,222 @@ 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)
+ 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)
+
+ n = len(scores)
+ if n <= limit:
+ top_indices = np.argsort(-scores)
+ else:
+ top_indices = np.argpartition(-scores, limit)[:limit]
+ top_indices = top_indices[np.argsort(-scores[top_indices])]
+
+ return DictBasedScoredIndexResult(
+ {int(row_id_array[i]): float(scores[i]) for i in top_indices}
+ )
+
+
+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.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_matrix = np.array(
+ [q if isinstance(q, np.ndarray) else list(q) for q in query_vectors],
+ dtype=np.float32)
+
+ 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(len(query_vectors))]
+
+ 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 SGEMM distance computation + per-query topK. Reuses stored
norms."""
+ import numpy as np
+
+ n_queries = query_matrix.shape[0]
+
+ if metric == "l2":
+ 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
Review Comment:
[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.
##########
paimon-python/pypaimon/table/source/vector_search_read.py:
##########
@@ -824,3 +924,222 @@ 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)
+ 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)
+
+ n = len(scores)
+ if n <= limit:
+ top_indices = np.argsort(-scores)
+ else:
+ top_indices = np.argpartition(-scores, limit)[:limit]
+ top_indices = top_indices[np.argsort(-scores[top_indices])]
+
+ return DictBasedScoredIndexResult(
+ {int(row_id_array[i]): float(scores[i]) for i in top_indices}
+ )
+
+
+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.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_matrix = np.array(
+ [q if isinstance(q, np.ndarray) else list(q) for q in query_vectors],
+ dtype=np.float32)
+
+ 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(len(query_vectors))]
+
+ 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 SGEMM distance computation + per-query topK. Reuses stored
norms."""
+ import numpy as np
+
+ n_queries = query_matrix.shape[0]
+
+ if metric == "l2":
+ 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
+ np.maximum(dists, 0, out=dists)
+ all_scores = 1.0 / (1.0 + dists)
Review Comment:
[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.
##########
paimon-python/pypaimon/table/source/vector_search_read.py:
##########
@@ -87,6 +88,20 @@ def __init__(
self._partition_filter = partition_filter
self._options = dict(options or {})
+ @property
+ def _index_thread_num(self):
+ _opts = self._table.options
+ _get = getattr(_opts, 'global_index_thread_num', None)
+ value = (
+ (_get() if _get else None)
+ or CoreOptions.GLOBAL_INDEX_THREAD_NUM._default_value
Review Comment:
[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.
--
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]
