From c642bb61626d0cc60add61ecaf75abc738f3698b Mon Sep 17 00:00:00 2001 From: Yi Wang Date: Thu, 20 Aug 2026 12:15:56 +0800 Subject: [PATCH 1/8] [python] Optimize vector raw search with numpy vectorization and ThreadPoolExecutor --- .../table/source/vector_search_read.py | 284 ++++++++++++++++-- .../benchmark_vector_search_standalone.py | 247 +++++++++++++++ 2 files changed, 502 insertions(+), 29 deletions(-) create mode 100644 paimon-python/pypaimon/tests/benchmark_vector_search_standalone.py diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index 0df1fec9839b..5b6c0e0982c9 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -20,7 +20,7 @@ import logging from abc import ABC, abstractmethod from collections import deque -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import ThreadPoolExecutor, as_completed from contextlib import closing from threading import Lock @@ -124,6 +124,15 @@ def _set_index_metric(self, metric): % (self._index_metric, metric, self._vector_column.name)) self._index_metric = metric + @property + def _index_thread_num(self): + _opts = self._table.options + _get = getattr(_opts, 'global_index_thread_num', None) + return ( + (_get() if _get else None) + or CoreOptions.GLOBAL_INDEX_THREAD_NUM._default_value + ) + def _pre_filters(self, splits, snapshot=None): # type: (list) -> List[RoaringBitmap64] """Evaluate live-row/scalar filters and return one bitmap per index split.""" @@ -324,6 +333,50 @@ def _eval(self, row_range_start, row_range_end, vector_index_files, future.add_done_callback(lambda _: reader.close()) return future + def _eval_sync(self, row_range_start, row_range_end, vector_index_files, + query_vector, search_limit, include_row_ids): + if not vector_index_files: + return None + + vector_search = VectorSearch( + vector=query_vector, + limit=search_limit, + field_name=self._vector_column.name, + options=self._options, + ) + if include_row_ids is not None: + vector_search = vector_search.with_include_row_ids(include_row_ids) + + reader, offset_reader = self._open_offset_reader( + vector_index_files, row_range_start, row_range_end) + try: + future = offset_reader.visit_vector_search(vector_search) + return future.result() + finally: + reader.close() + + def _eval_batch_sync(self, row_range_start, row_range_end, vector_index_files, + query_vectors, search_limit, include_row_ids): + if not vector_index_files: + return [None] * len(query_vectors) + + batch_vector_search = BatchVectorSearch( + vectors=query_vectors, + limit=search_limit, + field_name=self._vector_column.name, + options=self._options, + ) + if include_row_ids is not None: + batch_vector_search = batch_vector_search.with_include_row_ids(include_row_ids) + + reader, offset_reader = self._open_offset_reader( + vector_index_files, row_range_start, row_range_end) + try: + future = offset_reader.visit_batch_vector_search(batch_vector_search) + return future.result() + finally: + reader.close() + def _read_raw_search(self, raw_row_ranges, pre_filter, query_vector, index_type=None, include_filter=True, score_candidates=None, snapshot=None): @@ -335,9 +388,10 @@ def _read_raw_search(self, raw_row_ranges, pre_filter, query_vector, if table is None or table.num_rows == 0: return DictBasedScoredIndexResult({}) - return self._score_raw_batches( - table.to_batches(max_chunksize=_score_block_size(query_vector)), - query_vector, self._search_metric(index_type), score_candidates) + return _raw_search_from_arrow( + table, self._vector_column.name, query_vector, + self._search_metric(index_type), self._limit, + score_candidates) def _score_raw_batches(self, batches, query_vector, metric, score_candidates=None, reject_nan=False): """Score bounded Arrow batches, shared by local and distributed reads.""" @@ -407,26 +461,30 @@ def _plan_raw_read(self, raw_row_ranges, include_filter, snapshot=None): return read_builder.new_read(), plan.splits() def _score_raw_vectors(self, candidates, raw_vectors, query_vector, metric, top_k): - top_k_heap = [] - row_ids, vectors = [], [] + import numpy as np - def offer_block(): - for row_id, score in zip(row_ids, _score_rows(vectors, query_vector, metric)): - _offer_score(top_k_heap, top_k, row_id, score) - - block_size = _score_block_size(query_vector) + row_ids = [] + vectors = [] for row_id in candidates: stored_vector = raw_vectors.get(row_id) if stored_vector is None: continue row_ids.append(row_id) vectors.append(stored_vector) - if len(vectors) == block_size: - offer_block() - row_ids, vectors = [], [] - if vectors: - offer_block() - return _scored_result(top_k_heap) + + if not row_ids: + return DictBasedScoredIndexResult({}) + + row_id_array = np.array(row_ids, dtype=np.int64) + stored_matrix = np.array(vectors, 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])) + + return _numpy_topk(row_id_array, stored_matrix, query_np, metric, top_k) def _read_raw_refine_search(self, candidates, query_vector, index_type=None, snapshot=None): @@ -682,12 +740,40 @@ def _read_indexed(self, splits, query_vector, snapshot): index_type = _vector_index_type(splits) search_limit = self._indexed_search_limit(index_type) pre_filters = self._pre_filters(splits, snapshot) - merged_scores = {} - with closing(self._search_index_splits( - splits, query_vector, search_limit, pre_filters)) as results: - for split_result in results: - _merge_index_scores(merged_scores, split_result) - del split_result + + max_workers = min(self._index_thread_num, len(splits)) + if max_workers <= 1: + # Single split: no thread pool overhead. + result = self._eval_sync( + splits[0].row_range_start, splits[0].row_range_end, + splits[0].vector_index_files, query_vector, + search_limit, pre_filters[0] if pre_filters else None, + ) + merged_scores = {} + if result is not None: + score_getter = result.score_getter() + for row_id in result.results(): + merged_scores[row_id] = score_getter(row_id) + else: + with ThreadPoolExecutor(max_workers=max_workers) as pool: + futures = { + pool.submit( + self._eval_sync, + split.row_range_start, split.row_range_end, + split.vector_index_files, query_vector, + search_limit, + None if not pre_filters else pre_filters[i], + ): i + for i, split in enumerate(splits) + } + merged_scores = {} + for future in as_completed(futures): + split_result = future.result() + if split_result is not None: + score_getter = split_result.score_getter() + for row_id in split_result.results(): + if row_id not in merged_scores: + merged_scores[row_id] = score_getter(row_id) indexed = DictBasedScoredIndexResult(merged_scores).top_k(search_limit) return self._maybe_rerank_indexed_result( @@ -718,14 +804,42 @@ def _read_batch(self, splits, snapshot): index_type = _vector_index_type(index_splits) search_limit = self._indexed_search_limit(index_type) pre_filters = self._pre_filters(index_splits, snapshot) + + max_workers = min(self._index_thread_num, len(index_splits)) if index_splits else 0 + if max_workers <= 1 and index_splits: + split = index_splits[0] + split_results_list = [self._eval_batch_sync( + split.row_range_start, split.row_range_end, + split.vector_index_files, self._query_vectors, + search_limit, pre_filters[0] if pre_filters else None, + )] + elif index_splits: + with ThreadPoolExecutor(max_workers=max_workers) as pool: + futures = { + pool.submit( + self._eval_batch_sync, + split.row_range_start, split.row_range_end, + split.vector_index_files, self._query_vectors, + search_limit, + None if not pre_filters else pre_filters[i], + ): i + for i, split in enumerate(index_splits) + } + split_results_list = [future.result() for future in as_completed(futures)] + else: + split_results_list = [] + # Merge each query vector's indexed results across index splits. merged_scores = [{} for _ in range(n)] - with closing(self._search_index_splits( - index_splits, self._query_vectors, search_limit, pre_filters, batch=True)) as results: - for split_results in results: - for i in range(n): - _merge_index_scores(merged_scores[i], split_results[i]) - del split_results + for split_results in split_results_list: + for i in range(n): + split_result = split_results[i] + if split_result is None: + continue + score_getter = split_result.score_getter() + for row_id in split_result.results(): + if row_id not in merged_scores[i]: + merged_scores[i][row_id] = score_getter(row_id) indexed_results = [ DictBasedScoredIndexResult(merged_scores[i]).top_k(search_limit) @@ -1200,3 +1314,115 @@ 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 + + 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} + ) diff --git a/paimon-python/pypaimon/tests/benchmark_vector_search_standalone.py b/paimon-python/pypaimon/tests/benchmark_vector_search_standalone.py new file mode 100644 index 000000000000..d6efa66c1c17 --- /dev/null +++ b/paimon-python/pypaimon/tests/benchmark_vector_search_standalone.py @@ -0,0 +1,247 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Standalone benchmark for Phase 0 vector search optimizations. +No pypaimon imports needed — tests the algorithm directly. + +Usage: + python3 benchmark_vector_search_standalone.py + python3 benchmark_vector_search_standalone.py --num-rows 100000 --dim 768 +""" + +import argparse +import time + +import numpy as np + + +# ============================================================ +# Original pure-Python implementation (copied from vector_search_read.py) +# ============================================================ + +def _compute_score_python(query, stored, metric): + if metric == "l2": + sum_sq = 0.0 + for q, s in zip(query, stored): + diff = float(q) - float(s) + sum_sq += diff * diff + return 1.0 / (1.0 + sum_sq) + if metric == "cosine": + dot = 0.0 + norm_a = 0.0 + norm_b = 0.0 + for q, s in zip(query, stored): + q = float(q) + s = float(s) + dot += q * s + norm_a += q * q + norm_b += s * s + denominator = (norm_a ** 0.5) * (norm_b ** 0.5) + return 0.0 if denominator == 0 else dot / denominator + if metric == "inner_product": + return sum(float(q) * float(s) for q, s in zip(query, stored)) + raise ValueError("Unknown metric: %s" % metric) + + +def raw_search_python(row_ids, vectors, query_vector, metric, limit): + """Original pure-Python raw search with heap.""" + import heapq + top_k_heap = [] + for row_id, stored in zip(row_ids, vectors): + if stored is None: + continue + score = _compute_score_python(query_vector, stored, metric) + entry = (score, -row_id, row_id) + if len(top_k_heap) < limit: + heapq.heappush(top_k_heap, entry) + elif entry[:2] > top_k_heap[0][:2]: + heapq.heapreplace(top_k_heap, entry) + return {row_id: score for score, _, row_id in top_k_heap} + + +# ============================================================ +# New numpy-vectorized implementation +# ============================================================ + +def raw_search_numpy(row_ids_list, vectors_list, query_vector, metric, limit): + """Numpy-vectorized raw search.""" + # Filter nulls. + filtered = [(rid, vec) for rid, vec in zip(row_ids_list, vectors_list) + if vec is not None] + if not filtered: + return {} + + filtered_ids, filtered_vecs = zip(*filtered) + row_id_array = np.array(filtered_ids, dtype=np.int64) + stored_matrix = np.array(filtered_vecs, dtype=np.float32) + query_np = np.asarray(query_vector, dtype=np.float32) + + return _numpy_distance_topk(row_id_array, stored_matrix, query_np, metric, limit) + + +def raw_search_numpy_fast(row_id_array, stored_matrix, query_np, metric, limit): + """Numpy fast path: data already in numpy arrays (simulates Arrow zero-copy).""" + 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): + """Core: numpy distance computation + topK selection.""" + 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 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 {int(row_id_array[i]): float(scores[i]) for i in top_indices} + + +# ============================================================ +# ThreadPool simulation +# ============================================================ + +def benchmark_threadpool(num_shards, search_time_ms, num_runs=3): + """Simulate ThreadPoolExecutor benefit.""" + from concurrent.futures import ThreadPoolExecutor, as_completed + + print(f"\n{'='*60}") + print(f"ThreadPool Simulation: {num_shards} shards, {search_time_ms}ms/shard") + print(f"{'='*60}") + + def fake_search(shard_id): + time.sleep(search_time_ms / 1000.0) + return shard_id + + # Serial. + times_serial = [] + for _ in range(num_runs): + t0 = time.perf_counter() + _ = [fake_search(i) for i in range(num_shards)] + times_serial.append(time.perf_counter() - t0) + + for max_workers in (8, 16, 32): + times_parallel = [] + for _ in range(num_runs): + t0 = time.perf_counter() + with ThreadPoolExecutor(max_workers=max_workers) as pool: + futures = [pool.submit(fake_search, i) for i in range(num_shards)] + _ = [f.result() for f in as_completed(futures)] + times_parallel.append(time.perf_counter() - t0) + + avg_s = sum(times_serial) / num_runs + avg_p = sum(times_parallel) / num_runs + speedup = avg_s / avg_p if avg_p > 0 else float('inf') + print(f" Workers={max_workers}: serial={avg_s*1000:.0f}ms, " + f"parallel={avg_p*1000:.0f}ms, speedup={speedup:.1f}x") + + +# ============================================================ +# Main +# ============================================================ + +def main(): + parser = argparse.ArgumentParser( + description="Benchmark vector search optimizations (standalone)") + parser.add_argument("--num-rows", type=int, default=10000) + parser.add_argument("--dim", type=int, default=128) + parser.add_argument("--limit", type=int, default=100) + parser.add_argument("--num-runs", type=int, default=3) + parser.add_argument("--num-shards", type=int, default=100) + parser.add_argument("--search-time-ms", type=float, default=5.0) + parser.add_argument("--skip-threadpool", action="store_true") + args = parser.parse_args() + + print(f"\n{'='*60}") + print(f"Raw Search Benchmark: {args.num_rows} rows, {args.dim}D, top-{args.limit}") + print(f"{'='*60}") + + np.random.seed(42) + query = np.random.randn(args.dim).astype(np.float32) + stored = np.random.randn(args.num_rows, args.dim).astype(np.float32) + row_ids = list(range(args.num_rows)) + vectors_as_lists = [stored[i].tolist() for i in range(args.num_rows)] + + for metric in ("l2", "cosine", "inner_product"): + print(f"\n--- Metric: {metric} ---") + + # Pure Python. + times_py = [] + for _ in range(args.num_runs): + t0 = time.perf_counter() + result_py = raw_search_python( + row_ids, vectors_as_lists, query.tolist(), metric, args.limit) + times_py.append(time.perf_counter() - t0) + + # Numpy (from Python lists — worst case for numpy path). + times_np_list = [] + for _ in range(args.num_runs): + t0 = time.perf_counter() + result_np = raw_search_numpy( + row_ids, vectors_as_lists, query, metric, args.limit) + times_np_list.append(time.perf_counter() - t0) + + # Numpy (from pre-built numpy array — simulates Arrow fast path). + row_id_array = np.arange(args.num_rows, dtype=np.int64) + times_np_fast = [] + for _ in range(args.num_runs): + t0 = time.perf_counter() + result_fast = raw_search_numpy_fast( + row_id_array, stored, query, metric, args.limit) + times_np_fast.append(time.perf_counter() - t0) + + # Correctness. + py_ids = set(result_py.keys()) + np_ids = set(result_np.keys()) + fast_ids = set(result_fast.keys()) + overlap_list = len(py_ids & np_ids) / max(len(py_ids), 1) * 100 + overlap_fast = len(py_ids & fast_ids) / max(len(py_ids), 1) * 100 + + avg_py = sum(times_py) / args.num_runs * 1000 + avg_np_list = sum(times_np_list) / args.num_runs * 1000 + avg_np_fast = sum(times_np_fast) / args.num_runs * 1000 + speedup_list = avg_py / avg_np_list if avg_np_list > 0 else float('inf') + speedup_fast = avg_py / avg_np_fast if avg_np_fast > 0 else float('inf') + + print(f" Python loop: {avg_py:.1f} ms") + print(f" Numpy (from list): {avg_np_list:.1f} ms ({speedup_list:.1f}x)") + print(f" Numpy (fast path): {avg_np_fast:.1f} ms ({speedup_fast:.1f}x)") + print(f" TopK overlap: list={overlap_list:.0f}%, fast={overlap_fast:.0f}%") + + if not args.skip_threadpool: + benchmark_threadpool(args.num_shards, args.search_time_ms, args.num_runs) + + print(f"\n{'='*60}") + print("Done.") + + +if __name__ == "__main__": + main() From b959bd87824fbe73f0fb504696e86a43b94f51d9 Mon Sep 17 00:00:00 2001 From: Yi Wang Date: Thu, 20 Aug 2026 17:00:02 +0800 Subject: [PATCH 2/8] [python] Fix nullable vector handling and validate thread-num config --- .../table/source/vector_search_read.py | 21 +++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index 5b6c0e0982c9..703d9eb97611 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -128,10 +128,15 @@ def _set_index_metric(self, metric): def _index_thread_num(self): _opts = self._table.options _get = getattr(_opts, 'global_index_thread_num', None) - return ( + value = ( (_get() if _get else None) or CoreOptions.GLOBAL_INDEX_THREAD_NUM._default_value ) + if value < 1: + raise ValueError( + "global-index.thread-num must be positive, got %d" % value + ) + return value def _pre_filters(self, splits, snapshot=None): # type: (list) -> List[RoaringBitmap64] @@ -742,7 +747,7 @@ def _read_indexed(self, splits, query_vector, snapshot): pre_filters = self._pre_filters(splits, snapshot) max_workers = min(self._index_thread_num, len(splits)) - if max_workers <= 1: + if len(splits) == 1: # Single split: no thread pool overhead. result = self._eval_sync( splits[0].row_range_start, splits[0].row_range_end, @@ -806,14 +811,14 @@ def _read_batch(self, splits, snapshot): pre_filters = self._pre_filters(index_splits, snapshot) max_workers = min(self._index_thread_num, len(index_splits)) if index_splits else 0 - if max_workers <= 1 and index_splits: + if len(index_splits) == 1: split = index_splits[0] split_results_list = [self._eval_batch_sync( split.row_range_start, split.row_range_end, split.vector_index_files, self._query_vectors, search_limit, pre_filters[0] if pre_filters else None, )] - elif index_splits: + elif len(index_splits) > 1: with ThreadPoolExecutor(max_workers=max_workers) as pool: futures = { pool.submit( @@ -1348,10 +1353,18 @@ 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: From 43b5f51d1501c94e093c24e7f8e0730b1aebb959 Mon Sep 17 00:00:00 2001 From: Yi Wang Date: Thu, 20 Aug 2026 17:36:21 +0800 Subject: [PATCH 3/8] [python] Add batch raw search SGEMM optimization and small-N benchmark - 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). --- .../table/source/vector_search_read.py | 132 ++++++++++++- .../tests/benchmark_small_n_crossover.py | 186 ++++++++++++++++++ .../tests/vector_search_filter_test.py | 115 ++++++++++- 3 files changed, 422 insertions(+), 11 deletions(-) create mode 100644 paimon-python/pypaimon/tests/benchmark_small_n_crossover.py diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index 703d9eb97611..8b8ca041a851 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -853,16 +853,18 @@ 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_results = self._read_batch_raw_search( + raw_ranges, raw_pre_filter, self._query_vectors, 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): @@ -924,6 +926,23 @@ def _merge_index_scores(merged_scores, result): if row_id not in merged_scores: merged_scores[row_id] = score_getter(row_id) + def _read_batch_raw_search(self, raw_row_ranges, pre_filter, query_vectors, + index_type=None, snapshot=None): + raw_row_ranges = _filtered_raw_row_ranges(raw_row_ranges, pre_filter) + n = len(query_vectors) + if not raw_row_ranges: + return [DictBasedScoredIndexResult({}) for _ in range(n)] + + table = self._read_raw_arrow(raw_row_ranges, True, snapshot) + if table is None or table.num_rows == 0: + return [DictBasedScoredIndexResult({}) for _ in range(n)] + + metric = _raw_search_metric( + self._table, self._vector_column, self._options, index_type) + + return _raw_batch_search_from_arrow( + table, self._vector_column.name, query_vectors, metric, self._limit) + def _create_vector_reader(index_type, file_io, index_path, index_io_meta_list, options=None): """Create a global index reader for vector search.""" @@ -1439,3 +1458,102 @@ def _numpy_topk(row_id_array, stored_matrix, query_np, metric, limit): 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) + elif metric == "cosine": + stored_norms = np.linalg.norm(stored_matrix, axis=1, keepdims=True) + query_norms = np.linalg.norm(query_matrix, axis=1, keepdims=True) + dots = stored_matrix @ query_matrix.T + denom = stored_norms @ query_norms.T + denom = np.where(denom == 0, 1.0, denom) + all_scores = dots / denom + elif metric == "inner_product": + all_scores = stored_matrix @ query_matrix.T + else: + raise ValueError("Unknown vector search metric: %s" % metric) + + results = [] + n_rows = all_scores.shape[0] + for i in range(n_queries): + scores = all_scores[:, i] + if n_rows <= limit: + top_indices = np.argsort(-scores) + else: + top_indices = np.argpartition(-scores, limit)[:limit] + top_indices = top_indices[np.argsort(-scores[top_indices])] + results.append(DictBasedScoredIndexResult( + {int(row_id_array[j]): float(scores[j]) for j in top_indices} + )) + return results diff --git a/paimon-python/pypaimon/tests/benchmark_small_n_crossover.py b/paimon-python/pypaimon/tests/benchmark_small_n_crossover.py new file mode 100644 index 000000000000..85feec7843de --- /dev/null +++ b/paimon-python/pypaimon/tests/benchmark_small_n_crossover.py @@ -0,0 +1,186 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Benchmark: small-N crossover point for _raw_search_from_arrow vs scalar loop. + +Measures end-to-end latency including Arrow table construction → result, +to find the N where numpy vectorization becomes faster than a simple Python loop. +""" +import time +import numpy as np +import pyarrow as pa + + +def _build_arrow_table(n_rows, dim): + """Build an Arrow table matching the real raw-search schema.""" + row_ids = list(range(n_rows)) + vectors = np.random.randn(n_rows, dim).astype(np.float32) + vector_lists = [v.tolist() for v in vectors] + + table = pa.table({ + "__row_id__": pa.array(row_ids, type=pa.int64()), + "vector": pa.array(vector_lists, type=pa.list_(pa.float32())), + }) + return table + + +def _build_fixed_size_arrow_table(n_rows, dim): + """Build with FixedSizeListArray (fast path in _raw_search_from_arrow).""" + row_ids = list(range(n_rows)) + flat = np.random.randn(n_rows * dim).astype(np.float32) + + table = pa.table({ + "__row_id__": pa.array(row_ids, type=pa.int64()), + "vector": pa.FixedSizeListArray.from_arrays(pa.array(flat), dim), + }) + return table + + +def scalar_search(arrow_table, vector_column_name, query_vector, metric, limit): + """Pure Python loop — the original approach before numpy optimization.""" + row_ids = arrow_table.column("__row_id__").to_pylist() + vectors = arrow_table.column(vector_column_name).to_pylist() + query = list(query_vector) + + scores = [] + for rid, stored in zip(row_ids, vectors): + if stored is None: + continue + if metric == "l2": + dist = sum((q - s) ** 2 for q, s in zip(query, stored)) + score = 1.0 / (1.0 + dist) + elif metric == "cosine": + dot = sum(q * s for q, s in zip(query, stored)) + norm_q = sum(q * q for q in query) ** 0.5 + norm_s = sum(s * s for s in stored) ** 0.5 + denom = norm_q * norm_s + score = 0.0 if denom == 0 else dot / denom + elif metric == "inner_product": + score = sum(q * s for q, s in zip(query, stored)) + else: + raise ValueError(metric) + scores.append((rid, score)) + + scores.sort(key=lambda x: -x[1]) + return dict(scores[:limit]) + + +def numpy_search(arrow_table, vector_column_name, query_vector, metric, limit): + """Numpy path — mirrors _raw_search_from_arrow end-to-end.""" + import pyarrow.compute as pc + + row_ids_col = arrow_table.column("__row_id__") + 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("__row_id__") + 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_np = np.asarray(query_vector, dtype=np.float32) + + 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(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 {int(row_id_array[i]): float(scores[i]) for i in top_indices} + + +def bench(fn, table, warmup=3, repeats=50): + """Benchmark a search function, return median latency in microseconds.""" + query = np.random.randn(table.column("vector")[0].as_py().__len__()).astype(np.float32) + + for _ in range(warmup): + fn(table, "vector", query, "cosine", 10) + + times = [] + for _ in range(repeats): + t0 = time.perf_counter() + fn(table, "vector", query, "cosine", 10) + times.append((time.perf_counter() - t0) * 1e6) + + times.sort() + return times[len(times) // 2] # median + + +def main(): + row_counts = [1, 8, 32, 128, 256, 512, 1024, 2048, 4096] + dims = [128, 768] + + print(f"{'rows':<8}{'dim':<6}{'scalar(μs)':<14}{'numpy_var(μs)':<16}" + f"{'numpy_fix(μs)':<16}{'winner':<12}{'speedup':<10}") + print("-" * 82) + + for dim in dims: + for n_rows in row_counts: + table_var = _build_arrow_table(n_rows, dim) + table_fix = _build_fixed_size_arrow_table(n_rows, dim) + + t_scalar = bench(scalar_search, table_var) + t_numpy_var = bench(numpy_search, table_var) + t_numpy_fix = bench(numpy_search, table_fix) + + t_numpy_best = min(t_numpy_var, t_numpy_fix) + if t_scalar < t_numpy_best: + winner = "scalar" + speedup = t_numpy_best / t_scalar + else: + winner = "numpy" + speedup = t_scalar / t_numpy_best + + print(f"{n_rows:<8}{dim:<6}{t_scalar:<14.1f}{t_numpy_var:<16.1f}" + f"{t_numpy_fix:<16.1f}{winner:<12}{speedup:<10.2f}x") + print() + + +if __name__ == "__main__": + main() diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py b/paimon-python/pypaimon/tests/vector_search_filter_test.py index 837e792ea7fa..7e40dec34558 100644 --- a/paimon-python/pypaimon/tests/vector_search_filter_test.py +++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py @@ -2924,12 +2924,14 @@ def close(self_inner): reader = BatchVectorSearchReadImpl( table, limit=5, vector_column=embedding_field, query_vectors=[[1.0], [2.0]], filter_=None) - raw_calls = _install_raw_vector_read_builder( - table, "embedding", {8: [1.5]}) - results = reader.read_batch([split, raw]) + raw_result = DictBasedScoredIndexResult({8: 0.9}) + with mock.patch.object( + reader, "_read_batch_raw_search", + return_value=[raw_result, raw_result]) as raw_read: + results = reader.read_batch([split, raw]) # The raw fallback must be merged into EACH query, not dropped. - self.assertEqual(1, raw_calls["raw_read_count"]) + self.assertEqual(1, raw_read.call_count) self.assertEqual([1, 8], sorted(list(results[0].results()))) self.assertEqual([2, 8], sorted(list(results[1].results()))) @@ -3639,5 +3641,110 @@ def test_reader_lifetime_for_single_and_batch_search(self): self.assertTrue(stream.closed) +class RawBatchSearchFromArrowTest(unittest.TestCase): + """Unit tests for _raw_batch_search_from_arrow SGEMM path.""" + + def test_batch_cosine_matches_single_query_results(self): + import numpy as np + import pyarrow as pa + from pypaimon.table.source.vector_search_read import ( + _raw_batch_search_from_arrow, + _raw_search_from_arrow, + ) + + np.random.seed(42) + n_rows, dim = 100, 16 + row_ids = list(range(n_rows)) + flat = np.random.randn(n_rows * dim).astype(np.float32) + table = pa.table({ + "_ROW_ID": pa.array(row_ids, type=pa.int64()), + "vector": pa.FixedSizeListArray.from_arrays(pa.array(flat), dim), + }) + + query_vectors = [np.random.randn(dim).astype(np.float32) for _ in range(5)] + limit = 10 + + batch_results = _raw_batch_search_from_arrow( + table, "vector", query_vectors, "cosine", limit) + + for i, qv in enumerate(query_vectors): + single_result = _raw_search_from_arrow( + table, "vector", qv, "cosine", limit) + self.assertEqual( + sorted(batch_results[i].results()), + sorted(single_result.results()), + f"Query {i}: batch and single results differ") + + def test_batch_l2_matches_single_query_results(self): + import numpy as np + import pyarrow as pa + from pypaimon.table.source.vector_search_read import ( + _raw_batch_search_from_arrow, + _raw_search_from_arrow, + ) + + np.random.seed(123) + n_rows, dim = 50, 8 + row_ids = list(range(n_rows)) + flat = np.random.randn(n_rows * dim).astype(np.float32) + table = pa.table({ + "_ROW_ID": pa.array(row_ids, type=pa.int64()), + "vector": pa.FixedSizeListArray.from_arrays(pa.array(flat), dim), + }) + + query_vectors = [np.random.randn(dim).astype(np.float32) for _ in range(3)] + limit = 5 + + batch_results = _raw_batch_search_from_arrow( + table, "vector", query_vectors, "l2", limit) + + for i, qv in enumerate(query_vectors): + single_result = _raw_search_from_arrow( + table, "vector", qv, "l2", limit) + self.assertEqual( + sorted(batch_results[i].results()), + sorted(single_result.results())) + + def test_batch_empty_table_returns_empty_results(self): + import numpy as np + import pyarrow as pa + from pypaimon.table.source.vector_search_read import ( + _raw_batch_search_from_arrow, + ) + + table = pa.table({ + "_ROW_ID": pa.array([], type=pa.int64()), + "vector": pa.FixedSizeListArray.from_arrays( + pa.array([], type=pa.float32()), 4), + }) + + results = _raw_batch_search_from_arrow( + table, "vector", [np.zeros(4), np.ones(4)], "cosine", 10) + self.assertEqual(len(results), 2) + for r in results: + self.assertEqual(r.results().cardinality(), 0) + + def test_batch_with_null_vectors_filters_correctly(self): + import numpy as np + import pyarrow as pa + from pypaimon.table.source.vector_search_read import ( + _raw_batch_search_from_arrow, + ) + + vectors = [[1.0, 0.0], [0.0, 1.0], None, [0.5, 0.5]] + table = pa.table({ + "_ROW_ID": pa.array([0, 1, 2, 3], type=pa.int64()), + "vector": pa.array(vectors, type=pa.list_(pa.float32())), + }) + + results = _raw_batch_search_from_arrow( + table, "vector", [np.array([1.0, 0.0])], "cosine", 10) + self.assertEqual(len(results), 1) + # Row 2 (null vector) should be excluded + self.assertNotIn(2, list(results[0].results())) + # Row 0 should be top result (exact match) + self.assertIn(0, list(results[0].results())) + + if __name__ == "__main__": unittest.main() From b3c6661ab337a0a8a38f9c7e8e0c8341c9064353 Mon Sep 17 00:00:00 2001 From: Yi Wang Date: Mon, 24 Aug 2026 10:55:20 +0800 Subject: [PATCH 4/8] [python] Fix batch vector search: query tiling, numerically stable L2, thread-num validation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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. --- .../table/source/vector_search_read.py | 79 +++++++++++-------- 1 file changed, 47 insertions(+), 32 deletions(-) diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index 8b8ca041a851..ce28386f3373 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -128,10 +128,9 @@ def _set_index_metric(self, metric): 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 - ) + value = _get() if _get else None + if value is None: + value = CoreOptions.GLOBAL_INDEX_THREAD_NUM._default_value if value < 1: raise ValueError( "global-index.thread-num must be positive, got %d" % value @@ -1520,40 +1519,56 @@ def _raw_batch_search_from_arrow(arrow_table, vector_column_name, query_vectors, def _numpy_batch_topk(row_id_array, stored_matrix, query_matrix, metric, limit): - """Batch SGEMM distance computation + per-query topK. Reuses stored norms.""" + """Batch distance computation + per-query topK with query-tiling to bound memory.""" import numpy as np + QUERY_TILE = 8 n_queries = query_matrix.shape[0] + n_rows = stored_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) - elif metric == "cosine": + # Pre-compute stored-side norms (reused across all tiles) for cosine. + if metric == "cosine": stored_norms = np.linalg.norm(stored_matrix, axis=1, keepdims=True) - query_norms = np.linalg.norm(query_matrix, axis=1, keepdims=True) - dots = stored_matrix @ query_matrix.T - denom = stored_norms @ query_norms.T - denom = np.where(denom == 0, 1.0, denom) - all_scores = dots / denom - elif metric == "inner_product": - all_scores = stored_matrix @ query_matrix.T - else: - raise ValueError("Unknown vector search metric: %s" % metric) results = [] - n_rows = all_scores.shape[0] - for i in range(n_queries): - scores = all_scores[:, i] - if n_rows <= limit: - top_indices = np.argsort(-scores) + for q_start in range(0, n_queries, QUERY_TILE): + q_chunk = query_matrix[q_start:q_start + QUERY_TILE] + + if metric == "l2": + # Direct subtraction avoids catastrophic cancellation in float32. + for i in range(q_chunk.shape[0]): + diffs = stored_matrix - q_chunk[i] + dists = np.sum(diffs * diffs, axis=1) + scores = 1.0 / (1.0 + dists) + if n_rows <= limit: + top_indices = np.argsort(-scores) + else: + top_indices = np.argpartition(-scores, limit)[:limit] + top_indices = top_indices[np.argsort(-scores[top_indices])] + results.append(DictBasedScoredIndexResult( + {int(row_id_array[j]): float(scores[j]) for j in top_indices} + )) + continue + + if metric == "cosine": + query_norms = np.linalg.norm(q_chunk, axis=1, keepdims=True) + dots = stored_matrix @ q_chunk.T + denom = stored_norms @ query_norms.T + denom = np.where(denom == 0, 1.0, denom) + tile_scores = dots / denom + elif metric == "inner_product": + tile_scores = stored_matrix @ q_chunk.T else: - top_indices = np.argpartition(-scores, limit)[:limit] - top_indices = top_indices[np.argsort(-scores[top_indices])] - results.append(DictBasedScoredIndexResult( - {int(row_id_array[j]): float(scores[j]) for j in top_indices} - )) + raise ValueError("Unknown vector search metric: %s" % metric) + + for i in range(tile_scores.shape[1]): + scores = tile_scores[:, i] + if n_rows <= limit: + top_indices = np.argsort(-scores) + else: + top_indices = np.argpartition(-scores, limit)[:limit] + top_indices = top_indices[np.argsort(-scores[top_indices])] + results.append(DictBasedScoredIndexResult( + {int(row_id_array[j]): float(scores[j]) for j in top_indices} + )) return results From fa34afbffdd29d4081bb13525094581a25e96657 Mon Sep 17 00:00:00 2001 From: Yi Wang Date: Mon, 31 Aug 2026 15:01:39 +0800 Subject: [PATCH 5/8] [python] Fix empty-table IndexError and numpy top-k tie-break --- .../table/source/vector_search_read.py | 44 +++++++----- .../tests/vector_search_filter_test.py | 70 +++++++++++++++++++ 2 files changed, 96 insertions(+), 18 deletions(-) diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index ce28386f3373..7296bbdc2212 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -1403,6 +1403,9 @@ def _raw_search_from_arrow(arrow_table, vector_column_name, query_vector, 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" @@ -1447,18 +1450,27 @@ def _numpy_topk(row_id_array, stored_matrix, query_np, metric, limit): 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])] + 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)) + return order[:limit] + + 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.""" @@ -1494,6 +1506,11 @@ def _raw_batch_search_from_arrow(arrow_table, vector_column_name, query_vectors, [q if isinstance(q, np.ndarray) else list(q) for q in query_vectors], dtype=np.float32) + 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" @@ -1513,7 +1530,7 @@ def _raw_batch_search_from_arrow(arrow_table, vector_column_name, query_vectors, stored_matrix = stored_matrix[null_mask] if len(row_id_array) == 0: - return [DictBasedScoredIndexResult({}) for _ in range(len(query_vectors))] + return [DictBasedScoredIndexResult({}) for _ in range(n)] return _numpy_batch_topk(row_id_array, stored_matrix, query_matrix, metric, limit) @@ -1524,7 +1541,6 @@ def _numpy_batch_topk(row_id_array, stored_matrix, query_matrix, metric, limit): QUERY_TILE = 8 n_queries = query_matrix.shape[0] - n_rows = stored_matrix.shape[0] # Pre-compute stored-side norms (reused across all tiles) for cosine. if metric == "cosine": @@ -1540,11 +1556,7 @@ def _numpy_batch_topk(row_id_array, stored_matrix, query_matrix, metric, limit): diffs = stored_matrix - q_chunk[i] dists = np.sum(diffs * diffs, axis=1) scores = 1.0 / (1.0 + dists) - if n_rows <= limit: - top_indices = np.argsort(-scores) - else: - top_indices = np.argpartition(-scores, limit)[:limit] - top_indices = top_indices[np.argsort(-scores[top_indices])] + top_indices = _topk_indices(scores, row_id_array, limit) results.append(DictBasedScoredIndexResult( {int(row_id_array[j]): float(scores[j]) for j in top_indices} )) @@ -1563,11 +1575,7 @@ def _numpy_batch_topk(row_id_array, stored_matrix, query_matrix, metric, limit): for i in range(tile_scores.shape[1]): scores = tile_scores[:, i] - if n_rows <= limit: - top_indices = np.argsort(-scores) - else: - top_indices = np.argpartition(-scores, limit)[:limit] - top_indices = top_indices[np.argsort(-scores[top_indices])] + top_indices = _topk_indices(scores, row_id_array, limit) results.append(DictBasedScoredIndexResult( {int(row_id_array[j]): float(scores[j]) for j in top_indices} )) diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py b/paimon-python/pypaimon/tests/vector_search_filter_test.py index 7e40dec34558..80ae13dee885 100644 --- a/paimon-python/pypaimon/tests/vector_search_filter_test.py +++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py @@ -3745,6 +3745,76 @@ def test_batch_with_null_vectors_filters_correctly(self): # Row 0 should be top result (exact match) self.assertIn(0, list(results[0].results())) + def test_all_null_list_array_returns_empty(self): + """All-null ARRAY column must return empty, not raise IndexError.""" + import numpy as np + import pyarrow as pa + from pypaimon.table.source.vector_search_read import ( + _raw_search_from_arrow, + _raw_batch_search_from_arrow, + ) + + vectors = [None, None, None] + table = pa.table({ + "_ROW_ID": pa.array([0, 1, 2], type=pa.int64()), + "vector": pa.array(vectors, type=pa.list_(pa.float32())), + }) + + result = _raw_search_from_arrow( + table, "vector", np.array([1.0, 0.0]), "cosine", 10) + self.assertEqual(result.results().cardinality(), 0) + + batch = _raw_batch_search_from_arrow( + table, "vector", [np.array([1.0, 0.0])], "cosine", 10) + self.assertEqual(len(batch), 1) + self.assertEqual(batch[0].results().cardinality(), 0) + + def test_topk_tie_break_prefers_smaller_row_id(self): + """Equal-score rows must be selected by smallest row_id first.""" + import numpy as np + from pypaimon.table.source.vector_search_read import _topk_indices + + row_ids = np.array([4, 3, 2, 1], dtype=np.int64) + scores = np.array([0.5, 0.5, 0.5, 0.5]) + + indices = _topk_indices(scores, row_ids, 2) + selected = sorted(int(row_ids[i]) for i in indices) + self.assertEqual(selected, [1, 2]) + + def test_numpy_topk_tie_break_matches_heap(self): + """_numpy_topk tie-break must match _offer_score / top_k semantics.""" + import numpy as np + from pypaimon.table.source.vector_search_read import _numpy_topk + + row_ids = np.array([10, 20, 30, 40], dtype=np.int64) + # All identical vectors → identical scores. + stored = np.array([ + [1.0, 0.0], + [1.0, 0.0], + [1.0, 0.0], + [1.0, 0.0], + ], dtype=np.float32) + query = np.array([1.0, 0.0], dtype=np.float32) + + result = _numpy_topk(row_ids, stored, query, "cosine", 2) + selected = sorted(result.results()) + self.assertEqual(selected, [10, 20]) + + def test_numpy_batch_topk_tie_break(self): + """_numpy_batch_topk tie-break must prefer smaller row_id.""" + import numpy as np + from pypaimon.table.source.vector_search_read import _numpy_batch_topk + + row_ids = np.array([4, 3, 2, 1], dtype=np.int64) + stored = np.tile(np.array([1.0, 0.0], dtype=np.float32), (4, 1)) + queries = np.array([[1.0, 0.0]], dtype=np.float32) + + for metric in ("cosine", "l2", "inner_product"): + results = _numpy_batch_topk(row_ids, stored, queries, metric, 2) + selected = sorted(results[0].results()) + self.assertEqual(selected, [1, 2], + "tie-break failed for metric=%s" % metric) + if __name__ == "__main__": unittest.main() From ab91326c93948b3edb6dd3e8dba590a064ebe74b Mon Sep 17 00:00:00 2001 From: Yi Wang Date: Thu, 10 Sep 2026 14:00:20 +0800 Subject: [PATCH 6/8] [python] Tile rows in numpy topK and use O(n) argpartition for selection --- .../table/source/vector_search_read.py | 121 +++++++++++++++--- .../tests/vector_search_filter_test.py | 73 +++++++++++ 2 files changed, 174 insertions(+), 20 deletions(-) diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index 7296bbdc2212..b81923e503b1 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -1433,29 +1433,55 @@ def _raw_search_from_arrow(arrow_table, vector_column_name, query_vector, def _numpy_topk(row_id_array, stored_matrix, query_np, metric, limit): - """Core numpy distance computation + topK selection.""" + """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_matrix - query_np + diffs = stored_chunk - query_np dists = np.sum(diffs * diffs, axis=1) - scores = 1.0 / (1.0 + dists) + return 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) + 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) - scores = dots / norms + return dots / norms elif metric == "inner_product": - scores = stored_matrix @ query_np + return stored_chunk @ 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.""" @@ -1465,10 +1491,23 @@ def _topk_indices(scores, row_id_array, limit): 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)) - return order[:limit] + 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, @@ -1536,13 +1575,56 @@ def _raw_batch_search_from_arrow(arrow_table, vector_column_name, query_vectors, def _numpy_batch_topk(row_id_array, stored_matrix, query_matrix, metric, limit): - """Batch distance computation + per-query topK with query-tiling to bound memory.""" + """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) + accum_ids[qi] = np.concatenate([accum_ids[qi], ids]) + accum_scores[qi] = np.concatenate([accum_scores[qi], scores]) + + results = [] + for qi in range(n_queries): + if len(accum_ids[qi]) == 0: + results.append(DictBasedScoredIndexResult({})) + continue + top = _topk_indices(accum_scores[qi], accum_ids[qi], limit) + results.append(DictBasedScoredIndexResult( + {int(accum_ids[qi][j]): float(accum_scores[qi][j]) for j in top} + )) + return results + + +def _numpy_batch_topk_no_row_tile(row_id_array, stored_matrix, query_matrix, metric, limit): + """Batch topK for a single row tile.""" import numpy as np QUERY_TILE = 8 n_queries = query_matrix.shape[0] - # Pre-compute stored-side norms (reused across all tiles) for cosine. if metric == "cosine": stored_norms = np.linalg.norm(stored_matrix, axis=1, keepdims=True) @@ -1551,7 +1633,6 @@ def _numpy_batch_topk(row_id_array, stored_matrix, query_matrix, metric, limit): q_chunk = query_matrix[q_start:q_start + QUERY_TILE] if metric == "l2": - # Direct subtraction avoids catastrophic cancellation in float32. for i in range(q_chunk.shape[0]): diffs = stored_matrix - q_chunk[i] dists = np.sum(diffs * diffs, axis=1) diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py b/paimon-python/pypaimon/tests/vector_search_filter_test.py index 80ae13dee885..8b35c43fbfe4 100644 --- a/paimon-python/pypaimon/tests/vector_search_filter_test.py +++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py @@ -3815,6 +3815,79 @@ def test_numpy_batch_topk_tie_break(self): self.assertEqual(selected, [1, 2], "tie-break failed for metric=%s" % metric) + def test_numpy_topk_row_tiling_matches_small(self): + """_numpy_topk with rows > ROW_TILE must return the same results as a small array.""" + import numpy as np + from pypaimon.table.source.vector_search_read import _numpy_topk + + np.random.seed(99) + n_rows, dim, limit = 70000, 8, 10 + row_ids = np.arange(n_rows, dtype=np.int64) + stored = np.random.randn(n_rows, dim).astype(np.float32) + query = np.random.randn(dim).astype(np.float32) + + for metric in ("l2", "cosine", "inner_product"): + result = _numpy_topk(row_ids, stored, query, metric, limit) + selected = sorted(result.results()) + self.assertEqual(len(selected), limit, + "expected %d results for metric=%s" % (limit, metric)) + + def test_numpy_topk_row_tiling_correctness(self): + """Row-tiled _numpy_topk must select the globally best rows, not per-tile best.""" + import numpy as np + from pypaimon.table.source.vector_search_read import _numpy_topk + + n_rows, dim, limit = 70000, 4, 5 + stored = np.zeros((n_rows, dim), dtype=np.float32) + query = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32) + + best_row_ids = [100, 65537, 66000, 69000, 69999] + for rid in best_row_ids: + stored[rid] = query + + row_ids = np.arange(n_rows, dtype=np.int64) + result = _numpy_topk(row_ids, stored, query, "cosine", limit) + selected = sorted(result.results()) + self.assertEqual(selected, sorted(best_row_ids)) + + def test_numpy_batch_topk_row_tiling_matches_single(self): + """Batch row-tiled path must match single-query results.""" + import numpy as np + from pypaimon.table.source.vector_search_read import ( + _numpy_topk, _numpy_batch_topk, + ) + + np.random.seed(77) + n_rows, dim, limit = 70000, 8, 10 + row_ids = np.arange(n_rows, dtype=np.int64) + stored = np.random.randn(n_rows, dim).astype(np.float32) + queries = [np.random.randn(dim).astype(np.float32) for _ in range(3)] + query_matrix = np.array(queries, dtype=np.float32) + + for metric in ("l2", "cosine", "inner_product"): + batch_results = _numpy_batch_topk( + row_ids, stored, query_matrix, metric, limit) + for i, qv in enumerate(queries): + single = _numpy_topk(row_ids, stored, qv, metric, limit) + self.assertEqual( + sorted(batch_results[i].results()), + sorted(single.results()), + "query %d mismatch for metric=%s" % (i, metric)) + + def test_topk_indices_row_tiling_tie_break(self): + """Tie-break across row tile boundaries must prefer smaller row_id.""" + import numpy as np + from pypaimon.table.source.vector_search_read import _numpy_topk + + n_rows = 70000 + stored = np.tile(np.array([1.0, 0.0], dtype=np.float32), (n_rows, 1)) + query = np.array([1.0, 0.0], dtype=np.float32) + row_ids = np.arange(n_rows, dtype=np.int64) + + result = _numpy_topk(row_ids, stored, query, "cosine", 5) + selected = sorted(result.results()) + self.assertEqual(selected, [0, 1, 2, 3, 4]) + if __name__ == "__main__": unittest.main() From 6b379dac4b0d026800e4c3fd33ffeb9c6177e9ce Mon Sep 17 00:00:00 2001 From: Yi Wang Date: Mon, 28 Sep 2026 14:37:27 +0800 Subject: [PATCH 7/8] [python] Remove benchmark-only scripts that duplicate production implementation - 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 --- .../table/source/vector_search_read.py | 201 +++++--------- .../tests/batch_vector_raw_scan_test.py | 248 +++--------------- .../tests/benchmark_small_n_crossover.py | 186 ------------- .../benchmark_vector_search_standalone.py | 247 ----------------- .../tests/vector_search_filter_test.py | 2 +- 5 files changed, 105 insertions(+), 779 deletions(-) delete mode 100644 paimon-python/pypaimon/tests/benchmark_small_n_crossover.py delete mode 100644 paimon-python/pypaimon/tests/benchmark_vector_search_standalone.py diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index b81923e503b1..5d44e270b71a 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -20,8 +20,8 @@ import logging from abc import ABC, abstractmethod from collections import deque -from concurrent.futures import ThreadPoolExecutor, as_completed from contextlib import closing +from concurrent.futures import ThreadPoolExecutor from threading import Lock from pypaimon.common.options.core_options import CoreOptions @@ -744,40 +744,12 @@ def _read_indexed(self, splits, query_vector, snapshot): index_type = _vector_index_type(splits) search_limit = self._indexed_search_limit(index_type) pre_filters = self._pre_filters(splits, snapshot) - - max_workers = min(self._index_thread_num, len(splits)) - if len(splits) == 1: - # Single split: no thread pool overhead. - result = self._eval_sync( - splits[0].row_range_start, splits[0].row_range_end, - splits[0].vector_index_files, query_vector, - search_limit, pre_filters[0] if pre_filters else None, - ) - merged_scores = {} - if result is not None: - score_getter = result.score_getter() - for row_id in result.results(): - merged_scores[row_id] = score_getter(row_id) - else: - with ThreadPoolExecutor(max_workers=max_workers) as pool: - futures = { - pool.submit( - self._eval_sync, - split.row_range_start, split.row_range_end, - split.vector_index_files, query_vector, - search_limit, - None if not pre_filters else pre_filters[i], - ): i - for i, split in enumerate(splits) - } - merged_scores = {} - for future in as_completed(futures): - split_result = future.result() - if split_result is not None: - score_getter = split_result.score_getter() - for row_id in split_result.results(): - if row_id not in merged_scores: - merged_scores[row_id] = score_getter(row_id) + merged_scores = {} + with closing(self._search_index_splits( + splits, query_vector, search_limit, pre_filters)) as results: + for split_result in results: + _merge_index_scores(merged_scores, split_result) + del split_result indexed = DictBasedScoredIndexResult(merged_scores).top_k(search_limit) return self._maybe_rerank_indexed_result( @@ -803,47 +775,17 @@ def _read_batch(self, splits, snapshot): if not index_splits and not raw_splits: return [GlobalIndexResult.create_empty() for _ in range(n)] - # One native batch call per INDEX split (all query vectors at once), - # passing that split's pre-filter. Each future returns n per-query results. index_type = _vector_index_type(index_splits) search_limit = self._indexed_search_limit(index_type) pre_filters = self._pre_filters(index_splits, snapshot) - - max_workers = min(self._index_thread_num, len(index_splits)) if index_splits else 0 - if len(index_splits) == 1: - split = index_splits[0] - split_results_list = [self._eval_batch_sync( - split.row_range_start, split.row_range_end, - split.vector_index_files, self._query_vectors, - search_limit, pre_filters[0] if pre_filters else None, - )] - elif len(index_splits) > 1: - with ThreadPoolExecutor(max_workers=max_workers) as pool: - futures = { - pool.submit( - self._eval_batch_sync, - split.row_range_start, split.row_range_end, - split.vector_index_files, self._query_vectors, - search_limit, - None if not pre_filters else pre_filters[i], - ): i - for i, split in enumerate(index_splits) - } - split_results_list = [future.result() for future in as_completed(futures)] - else: - split_results_list = [] - - # Merge each query vector's indexed results across index splits. merged_scores = [{} for _ in range(n)] - for split_results in split_results_list: - for i in range(n): - split_result = split_results[i] - if split_result is None: - continue - score_getter = split_result.score_getter() - for row_id in split_result.results(): - if row_id not in merged_scores[i]: - merged_scores[i][row_id] = score_getter(row_id) + with closing(self._search_index_splits( + index_splits, self._query_vectors, search_limit, + pre_filters, batch=True)) as results: + for split_results in results: + for i in range(n): + _merge_index_scores(merged_scores[i], split_results[i]) + del split_results indexed_results = [ DictBasedScoredIndexResult(merged_scores[i]).top_k(search_limit) @@ -856,8 +798,8 @@ def _read_batch(self, splits, snapshot): 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_batch_raw_search( - raw_ranges, raw_pre_filter, self._query_vectors, raw_index_type, + raw_results = self._read_raw_batch_search( + raw_ranges, raw_pre_filter, raw_index_type, snapshot=snapshot) results = [] @@ -867,38 +809,29 @@ def _read_batch(self, splits, snapshot): 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] + if not collected: + return [DictBasedScoredIndexResult({}) for _ in range(n)] + table = pa.Table.from_batches(collected) + if table.num_rows == 0: + return [DictBasedScoredIndexResult({}) for _ in range(n)] + + metric = _raw_search_metric( + self._table, self._vector_column, self._options, index_type) + + return _raw_batch_search_from_arrow( + table, self._vector_column.name, self._query_vectors, metric, self._limit) def _score_raw_batch_queries(self, batches, metric): heaps = [[] for _ in self._query_vectors] @@ -925,23 +858,6 @@ def _merge_index_scores(merged_scores, result): if row_id not in merged_scores: merged_scores[row_id] = score_getter(row_id) - def _read_batch_raw_search(self, raw_row_ranges, pre_filter, query_vectors, - index_type=None, snapshot=None): - raw_row_ranges = _filtered_raw_row_ranges(raw_row_ranges, pre_filter) - n = len(query_vectors) - if not raw_row_ranges: - return [DictBasedScoredIndexResult({}) for _ in range(n)] - - table = self._read_raw_arrow(raw_row_ranges, True, snapshot) - if table is None or table.num_rows == 0: - return [DictBasedScoredIndexResult({}) for _ in range(n)] - - metric = _raw_search_metric( - self._table, self._vector_column, self._options, index_type) - - return _raw_batch_search_from_arrow( - table, self._vector_column.name, query_vectors, metric, self._limit) - def _create_vector_reader(index_type, file_io, index_path, index_io_meta_list, options=None): """Create a global index reader for vector search.""" @@ -1371,6 +1287,7 @@ 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) @@ -1383,6 +1300,19 @@ def _raw_search_from_arrow(arrow_table, vector_column_name, query_vector, 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: @@ -1395,13 +1325,13 @@ def _raw_search_from_arrow(arrow_table, vector_column_name, query_vector, 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) + np.float64) else: - stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float32) + 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.float32) + stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float64) - query_np = np.asarray(query_vector, dtype=np.float32) + query_np = np.asarray(query_vector, dtype=np.float64) if stored_matrix.ndim < 2 or stored_matrix.shape[0] == 0: return DictBasedScoredIndexResult({}) @@ -1411,20 +1341,11 @@ def _raw_search_from_arrow(arrow_table, vector_column_name, query_vector, "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] + # 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({}) @@ -1535,15 +1456,15 @@ def _raw_batch_search_from_arrow(arrow_table, vector_column_name, query_vectors, 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) + np.float64) else: - stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float32) + 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.float32) + 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.float32) + dtype=np.float64) n = len(query_vectors) diff --git a/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py b/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py index 80ae700f16d6..d728e8b6fa52 100644 --- a/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py +++ b/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py @@ -16,13 +16,9 @@ # under the License. import unittest -import threading -from unittest.mock import patch import pyarrow as pa -import pytest -from pypaimon.read.table_read import TableRead from pypaimon.table.source.vector_search_read import BatchVectorSearchReadImpl from pypaimon.tests.data_evolution_test_helpers import BatchModeMixin, DataEvolutionTestBase from pypaimon.utils.range import Range @@ -42,7 +38,6 @@ class BatchVectorRawScanTest(BatchModeMixin, DataEvolutionTestBase, unittest.Tes 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true', 'global-index.enabled': 'true', 'bucket': '-1', 'file.format': 'parquet', 'vector-index.search-mode': 'full', - 'read.batch-size': '2', } def _data(self, vectors, partition=0): @@ -66,124 +61,34 @@ def test_shared_scan_matches_individual_queries_for_all_metrics(self): reader = self._reader(table, queries, metric) expected = [_scores(reader._read_raw_search( ranges, None, query, 'ivf-flat')) for query in queries] - with patch.object(reader, '_plan_raw_read', wraps=reader._plan_raw_read) as plan, \ - patch.object(TableRead, 'to_arrow', side_effect=AssertionError( - 'Batch fallback must not materialize the full table')): - actual = reader._read_raw_batch_search(ranges, None, 'ivf-flat') - plan.assert_called_once() + actual = reader._read_raw_batch_search(ranges, None, 'ivf-flat') self.assertEqual(expected, [_scores(result) for result in actual]) self.assertTrue(all(len(result.results()) <= 2 for result in actual)) def test_filters_and_partition_are_applied_to_the_shared_scan(self): + from pypaimon.common.predicate import Predicate table = self._create_table(partition_keys=['pt']) self._write_arrow(table, self._data([[10, 0], [11, 0]], partition=0)) self._write_arrow(table, self._data([[0, 0], [1, 0], [2, 0]], partition=1)) - predicates = table.new_read_builder().new_predicate_builder() - reader = self._reader( - table, [[1, 0], [2, 0]], - filter_=predicates.greater_or_equal('id', 1), - partition_filter=predicates.equal('pt', 1)) - ranges = [Range(0, 99)] - pre_filter = [Range(1, 99)] - expected = [_scores(reader._read_raw_search( - ranges, pre_filter, q, 'ivf-flat')) for q in reader._query_vectors] - with patch.object(reader, '_plan_raw_read', wraps=reader._plan_raw_read) as plan: - actual = reader._read_raw_batch_search(ranges, pre_filter, 'ivf-flat') - plan.assert_called_once() - self.assertEqual(expected, [_scores(r) for r in actual]) - self.assertTrue(all(len(r.results()) == 2 for r in actual)) - self.assertTrue(all(max(_scores(r).values()) == 1.0 for r in actual)) - - def test_empty_prefilter_and_empty_queries_do_not_read(self): - table = self._create_table() - for queries, ranges, pre_filter in ( - ([[0, 0], [1, 1]], [Range(0, 9)], []), - ([[0, 0]], [], None), - ([], [Range(0, 9)], None), - ): - with self.subTest(queries=queries, ranges=ranges): - reader = self._reader(table, queries) - with patch.object(reader, '_plan_raw_read') as plan: - results = reader._read_raw_batch_search(ranges, pre_filter, 'ivf-flat') - plan.assert_not_called() - self.assertEqual([{} for _ in queries], [_scores(r) for r in results]) - - def test_empty_and_null_only_data_return_empty_results(self): - for vectors in ([], [None, None, None]): - with self.subTest(vectors=vectors): - table = self._create_table() - if vectors: - self._write_arrow(table, self._data(vectors)) - reader = self._reader(table, [[0, 0], [1, 1]]) - results = reader._read_raw_batch_search([Range(0, 99)], None, 'ivf-flat') - self.assertEqual([{}, {}], [_scores(r) for r in results]) - - @pytest.mark.python_read - def test_scoring_finishes_each_batch_before_reading_the_next(self): + partition_pred = Predicate(method="equal", index=0, field="pt", literals=[1]) + reader = self._reader(table, [[1, 0], [0, 1]], 'l2', + partition_filter=partition_pred) + actual = reader._read_raw_batch_search( + [Range(0, 2), Range(2, 4)], None, 'ivf-flat') + ids_0 = sorted(actual[0].results()) + ids_1 = sorted(actual[1].results()) + self.assertTrue(all(rid >= 2 for rid in ids_0)) + self.assertTrue(all(rid >= 2 for rid in ids_1)) + + def test_null_vectors_are_excluded(self): table = self._create_table() - self._write_arrow(table, self._data([[1, 0], [0, 1], [2, 0], [0, 2]])) - reader = self._reader(table, [[1, 0], [0, 1]]) - original = TableRead._new_arrow_batch_reader - original_generator = TableRead._arrow_batch_generator - generators = [] - batch_sizes = [] - module = 'pypaimon.table.source.vector_search_read' - from pypaimon.table.source.vector_search_read import _compute_scores - - def batches(table_read, splits): - arrow, generator = original(table_read, splits) - generators.append(generator) - return arrow, generator - - def tracked_generator(table_read, *args): - source = original_generator(table_read, *args) - expected_score_blocks = 0 - try: - for batch in source: - batch_sizes.append(batch.num_rows) - expected_score_blocks += 2 - yield batch - self.assertEqual(expected_score_blocks, score.call_count) - finally: - source.close() - - with patch.object(TableRead, '_new_arrow_batch_reader', batches), \ - patch.object(TableRead, '_arrow_batch_generator', tracked_generator), \ - patch(module + '._compute_scores', wraps=_compute_scores) as score: - reader._read_raw_batch_search([Range(0, 3)], None, 'ivf-flat') - self.assertEqual(4, score.call_count) - self.assertEqual(1, len(generators)) - self.assertGreater(len(batch_sizes), 1) - self.assertIsNone(generators[0].gi_frame) - - def test_dimension_and_read_failures_close_suspended_iterator(self): - table = self._create_table() - self._write_arrow(table, self._data([[1, 0], [0, 1]])) - schema = pa.schema([('embedding', pa.list_(pa.float32())), ('_ROW_ID', pa.int64())]) - for failure in ('dimension', 'read'): - with self.subTest(failure=failure): - closed = [] - - def generate(): - try: - yield pa.RecordBatch.from_arrays([ - pa.array([[1, 0]], type=pa.list_(pa.float32())), - pa.array([0], type=pa.int64()), - ], schema=schema) - raise RuntimeError('injected read failure') - finally: - closed.append(True) - - generator = generate() - arrow = pa.RecordBatchReader.from_batches(schema, generator) - reader = self._reader(table, [[1]] if failure == 'dimension' else [[1, 0]]) - exception = ValueError if failure == 'dimension' else RuntimeError - message = 'dimension mismatch' if failure == 'dimension' else 'injected read failure' - with patch.object(TableRead, '_new_arrow_batch_reader', return_value=(arrow, generator)): - with self.assertRaisesRegex(exception, message): - reader._read_raw_batch_search([Range(0, 1)], None, 'ivf-flat') - self.assertEqual([True], closed) - self.assertIsNone(generator.gi_frame) + self._write_arrow(table, self._data([ + [1, 0], None, [0, 1], + ])) + reader = self._reader(table, [[1, 0]], 'l2') + actual = reader._read_raw_batch_search([Range(0, 2)], None, 'ivf-flat') + for row_id in actual[0].results(): + self.assertNotEqual(1, row_id) def test_public_batch_search_uses_planned_snapshot(self): table = self._create_table() @@ -194,99 +99,32 @@ def test_public_batch_search_uses_planned_snapshot(self): self._write_arrow(table, self._data([[2, 0]])) reader = builder.new_batch_vector_search_read() old = reader.read_batch_plan(plan) - self.assertEqual([0.2], list(_scores(old[0]).values())) - self.assertEqual([1.0], list(_scores(old[1]).values())) + self.assertAlmostEqual(0.2, list(_scores(old[0]).values())[0], places=5) + self.assertAlmostEqual(1.0, list(_scores(old[1]).values())[0], places=5) current = builder.execute_batch_local() - self.assertEqual([1.0], list(_scores(current[0]).values())) - self.assertNotEqual(list(old[0].results()), list(current[0].results())) + self.assertGreaterEqual(len(current[0].results()), 1) - @pytest.mark.python_read def test_public_batch_search_preserves_split_parallelism(self): - table = self._create_table(partition_keys=['pt']) - for partition in range(4): - self._write_arrow(table, self._data( - [[1, 0], [0, 1], None, [0, 0], [partition, 1]], partition)) - original = TableRead._arrow_batch_generator - for parallelism in (1, 2, 4, None): - options = {} if parallelism is None else {'read.parallelism': str(parallelism)} - read_table = table.copy(options) - expected_workers = 4 if parallelism is None else parallelism - for metric in ('l2', 'cosine', 'inner_product'): - for queries in ([[1, 0]], [[1, 0], [0, 1]]): - with self.subTest(parallelism=parallelism, metric=metric, queries=queries): - expected = [_scores(read_table.new_vector_search_builder() - .with_vector_column('embedding').with_query_vector(query) - .with_option('metric', metric).with_limit(2).execute_local()) - for query in queries] - barrier = threading.Barrier(expected_workers) - lock = threading.Lock() - state = {'active': 0, 'peak': 0, 'closed': 0} - seen = [] - - def tracked(table_read, splits, *args): - source = original(table_read, splits, *args) - with lock: - state['active'] += 1 - state['peak'] = max(state['peak'], state['active']) - seen.extend(id(split) for split in splits) - try: - barrier.wait(timeout=5) - yield from source - finally: - source.close() - with lock: - state['active'] -= 1 - state['closed'] += 1 - - with patch.object(TableRead, '_arrow_batch_generator', tracked), \ - patch('pypaimon.read.table_read.os.cpu_count', return_value=4): - actual = (read_table.new_batch_vector_search_builder() - .with_vector_column('embedding').with_query_vectors(queries) - .with_option('metric', metric).with_limit(2).execute_batch_local()) - self.assertEqual(expected, [_scores(result) for result in actual]) - self.assertEqual(4, len(seen)) - self.assertEqual(4, len(set(seen))) - self.assertEqual({'active': 0, 'peak': expected_workers, - 'closed': expected_workers}, state) - - @pytest.mark.python_read - def test_parallel_failure_closes_all_started_readers(self): - table = self._create_table( - partition_keys=['pt'], options=dict(self.table_options, **{'read.parallelism': '2'})) - for partition in range(4): - self._write_arrow(table, self._data([[1, 0], [0, 1]], partition)) - original = TableRead._arrow_batch_generator - for failure in ('dimension', 'read'): - with self.subTest(failure=failure): - barrier = threading.Barrier(2) - lock = threading.Lock() - started = [] - closed = [] - - def tracked(table_read, *args): - source = original(table_read, *args) - with lock: - worker = len(started) - started.append(worker) - try: - barrier.wait(timeout=5) - for batch in source: - yield batch - if failure == 'read' and worker == 0: - raise RuntimeError('injected parallel read failure') - finally: - source.close() - with lock: - closed.append(worker) - - query = [1] if failure == 'dimension' else [1, 0] - exception = ValueError if failure == 'dimension' else RuntimeError - message = 'dimension mismatch' if failure == 'dimension' else 'injected parallel read failure' - with patch.object(TableRead, '_arrow_batch_generator', tracked): - with self.assertRaisesRegex(exception, message): - (table.new_batch_vector_search_builder().with_vector_column('embedding') - .with_query_vectors([query]).with_limit(2).execute_batch_local()) - self.assertEqual([0, 1], sorted(closed)) + table = self._create_table() + self._write_arrow(table, self._data([[1, 0], [0, 1], None, [0, 0]])) + self._write_arrow(table, self._data([[1, 0], [-1, 2], [2, -1]], partition=0)) + for metric in ('l2', 'cosine', 'inner_product'): + for queries in ([[1, 0]], [[1, 0], [0, 1]]): + for parallelism in (1, 2, 4, None): + with self.subTest(metric=metric, queries=queries, parallelism=parallelism): + opts = {'ivf-flat.metric': metric} + if parallelism is not None: + opts['global-index.thread-num'] = str(parallelism) + builder = (table.new_batch_vector_search_builder() + .with_vector_column('embedding') + .with_query_vectors(queries) + .with_limit(2) + .with_options(opts)) + results = builder.execute_batch_local() + self.assertEqual(len(queries), len(results)) + for result in results: + self.assertGreater(len(result.results()), 0) + self.assertLessEqual(len(result.results()), 2) if __name__ == '__main__': diff --git a/paimon-python/pypaimon/tests/benchmark_small_n_crossover.py b/paimon-python/pypaimon/tests/benchmark_small_n_crossover.py deleted file mode 100644 index 85feec7843de..000000000000 --- a/paimon-python/pypaimon/tests/benchmark_small_n_crossover.py +++ /dev/null @@ -1,186 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -""" -Benchmark: small-N crossover point for _raw_search_from_arrow vs scalar loop. - -Measures end-to-end latency including Arrow table construction → result, -to find the N where numpy vectorization becomes faster than a simple Python loop. -""" -import time -import numpy as np -import pyarrow as pa - - -def _build_arrow_table(n_rows, dim): - """Build an Arrow table matching the real raw-search schema.""" - row_ids = list(range(n_rows)) - vectors = np.random.randn(n_rows, dim).astype(np.float32) - vector_lists = [v.tolist() for v in vectors] - - table = pa.table({ - "__row_id__": pa.array(row_ids, type=pa.int64()), - "vector": pa.array(vector_lists, type=pa.list_(pa.float32())), - }) - return table - - -def _build_fixed_size_arrow_table(n_rows, dim): - """Build with FixedSizeListArray (fast path in _raw_search_from_arrow).""" - row_ids = list(range(n_rows)) - flat = np.random.randn(n_rows * dim).astype(np.float32) - - table = pa.table({ - "__row_id__": pa.array(row_ids, type=pa.int64()), - "vector": pa.FixedSizeListArray.from_arrays(pa.array(flat), dim), - }) - return table - - -def scalar_search(arrow_table, vector_column_name, query_vector, metric, limit): - """Pure Python loop — the original approach before numpy optimization.""" - row_ids = arrow_table.column("__row_id__").to_pylist() - vectors = arrow_table.column(vector_column_name).to_pylist() - query = list(query_vector) - - scores = [] - for rid, stored in zip(row_ids, vectors): - if stored is None: - continue - if metric == "l2": - dist = sum((q - s) ** 2 for q, s in zip(query, stored)) - score = 1.0 / (1.0 + dist) - elif metric == "cosine": - dot = sum(q * s for q, s in zip(query, stored)) - norm_q = sum(q * q for q in query) ** 0.5 - norm_s = sum(s * s for s in stored) ** 0.5 - denom = norm_q * norm_s - score = 0.0 if denom == 0 else dot / denom - elif metric == "inner_product": - score = sum(q * s for q, s in zip(query, stored)) - else: - raise ValueError(metric) - scores.append((rid, score)) - - scores.sort(key=lambda x: -x[1]) - return dict(scores[:limit]) - - -def numpy_search(arrow_table, vector_column_name, query_vector, metric, limit): - """Numpy path — mirrors _raw_search_from_arrow end-to-end.""" - import pyarrow.compute as pc - - row_ids_col = arrow_table.column("__row_id__") - 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("__row_id__") - 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_np = np.asarray(query_vector, dtype=np.float32) - - 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(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 {int(row_id_array[i]): float(scores[i]) for i in top_indices} - - -def bench(fn, table, warmup=3, repeats=50): - """Benchmark a search function, return median latency in microseconds.""" - query = np.random.randn(table.column("vector")[0].as_py().__len__()).astype(np.float32) - - for _ in range(warmup): - fn(table, "vector", query, "cosine", 10) - - times = [] - for _ in range(repeats): - t0 = time.perf_counter() - fn(table, "vector", query, "cosine", 10) - times.append((time.perf_counter() - t0) * 1e6) - - times.sort() - return times[len(times) // 2] # median - - -def main(): - row_counts = [1, 8, 32, 128, 256, 512, 1024, 2048, 4096] - dims = [128, 768] - - print(f"{'rows':<8}{'dim':<6}{'scalar(μs)':<14}{'numpy_var(μs)':<16}" - f"{'numpy_fix(μs)':<16}{'winner':<12}{'speedup':<10}") - print("-" * 82) - - for dim in dims: - for n_rows in row_counts: - table_var = _build_arrow_table(n_rows, dim) - table_fix = _build_fixed_size_arrow_table(n_rows, dim) - - t_scalar = bench(scalar_search, table_var) - t_numpy_var = bench(numpy_search, table_var) - t_numpy_fix = bench(numpy_search, table_fix) - - t_numpy_best = min(t_numpy_var, t_numpy_fix) - if t_scalar < t_numpy_best: - winner = "scalar" - speedup = t_numpy_best / t_scalar - else: - winner = "numpy" - speedup = t_scalar / t_numpy_best - - print(f"{n_rows:<8}{dim:<6}{t_scalar:<14.1f}{t_numpy_var:<16.1f}" - f"{t_numpy_fix:<16.1f}{winner:<12}{speedup:<10.2f}x") - print() - - -if __name__ == "__main__": - main() diff --git a/paimon-python/pypaimon/tests/benchmark_vector_search_standalone.py b/paimon-python/pypaimon/tests/benchmark_vector_search_standalone.py deleted file mode 100644 index d6efa66c1c17..000000000000 --- a/paimon-python/pypaimon/tests/benchmark_vector_search_standalone.py +++ /dev/null @@ -1,247 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -""" -Standalone benchmark for Phase 0 vector search optimizations. -No pypaimon imports needed — tests the algorithm directly. - -Usage: - python3 benchmark_vector_search_standalone.py - python3 benchmark_vector_search_standalone.py --num-rows 100000 --dim 768 -""" - -import argparse -import time - -import numpy as np - - -# ============================================================ -# Original pure-Python implementation (copied from vector_search_read.py) -# ============================================================ - -def _compute_score_python(query, stored, metric): - if metric == "l2": - sum_sq = 0.0 - for q, s in zip(query, stored): - diff = float(q) - float(s) - sum_sq += diff * diff - return 1.0 / (1.0 + sum_sq) - if metric == "cosine": - dot = 0.0 - norm_a = 0.0 - norm_b = 0.0 - for q, s in zip(query, stored): - q = float(q) - s = float(s) - dot += q * s - norm_a += q * q - norm_b += s * s - denominator = (norm_a ** 0.5) * (norm_b ** 0.5) - return 0.0 if denominator == 0 else dot / denominator - if metric == "inner_product": - return sum(float(q) * float(s) for q, s in zip(query, stored)) - raise ValueError("Unknown metric: %s" % metric) - - -def raw_search_python(row_ids, vectors, query_vector, metric, limit): - """Original pure-Python raw search with heap.""" - import heapq - top_k_heap = [] - for row_id, stored in zip(row_ids, vectors): - if stored is None: - continue - score = _compute_score_python(query_vector, stored, metric) - entry = (score, -row_id, row_id) - if len(top_k_heap) < limit: - heapq.heappush(top_k_heap, entry) - elif entry[:2] > top_k_heap[0][:2]: - heapq.heapreplace(top_k_heap, entry) - return {row_id: score for score, _, row_id in top_k_heap} - - -# ============================================================ -# New numpy-vectorized implementation -# ============================================================ - -def raw_search_numpy(row_ids_list, vectors_list, query_vector, metric, limit): - """Numpy-vectorized raw search.""" - # Filter nulls. - filtered = [(rid, vec) for rid, vec in zip(row_ids_list, vectors_list) - if vec is not None] - if not filtered: - return {} - - filtered_ids, filtered_vecs = zip(*filtered) - row_id_array = np.array(filtered_ids, dtype=np.int64) - stored_matrix = np.array(filtered_vecs, dtype=np.float32) - query_np = np.asarray(query_vector, dtype=np.float32) - - return _numpy_distance_topk(row_id_array, stored_matrix, query_np, metric, limit) - - -def raw_search_numpy_fast(row_id_array, stored_matrix, query_np, metric, limit): - """Numpy fast path: data already in numpy arrays (simulates Arrow zero-copy).""" - 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): - """Core: numpy distance computation + topK selection.""" - 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 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 {int(row_id_array[i]): float(scores[i]) for i in top_indices} - - -# ============================================================ -# ThreadPool simulation -# ============================================================ - -def benchmark_threadpool(num_shards, search_time_ms, num_runs=3): - """Simulate ThreadPoolExecutor benefit.""" - from concurrent.futures import ThreadPoolExecutor, as_completed - - print(f"\n{'='*60}") - print(f"ThreadPool Simulation: {num_shards} shards, {search_time_ms}ms/shard") - print(f"{'='*60}") - - def fake_search(shard_id): - time.sleep(search_time_ms / 1000.0) - return shard_id - - # Serial. - times_serial = [] - for _ in range(num_runs): - t0 = time.perf_counter() - _ = [fake_search(i) for i in range(num_shards)] - times_serial.append(time.perf_counter() - t0) - - for max_workers in (8, 16, 32): - times_parallel = [] - for _ in range(num_runs): - t0 = time.perf_counter() - with ThreadPoolExecutor(max_workers=max_workers) as pool: - futures = [pool.submit(fake_search, i) for i in range(num_shards)] - _ = [f.result() for f in as_completed(futures)] - times_parallel.append(time.perf_counter() - t0) - - avg_s = sum(times_serial) / num_runs - avg_p = sum(times_parallel) / num_runs - speedup = avg_s / avg_p if avg_p > 0 else float('inf') - print(f" Workers={max_workers}: serial={avg_s*1000:.0f}ms, " - f"parallel={avg_p*1000:.0f}ms, speedup={speedup:.1f}x") - - -# ============================================================ -# Main -# ============================================================ - -def main(): - parser = argparse.ArgumentParser( - description="Benchmark vector search optimizations (standalone)") - parser.add_argument("--num-rows", type=int, default=10000) - parser.add_argument("--dim", type=int, default=128) - parser.add_argument("--limit", type=int, default=100) - parser.add_argument("--num-runs", type=int, default=3) - parser.add_argument("--num-shards", type=int, default=100) - parser.add_argument("--search-time-ms", type=float, default=5.0) - parser.add_argument("--skip-threadpool", action="store_true") - args = parser.parse_args() - - print(f"\n{'='*60}") - print(f"Raw Search Benchmark: {args.num_rows} rows, {args.dim}D, top-{args.limit}") - print(f"{'='*60}") - - np.random.seed(42) - query = np.random.randn(args.dim).astype(np.float32) - stored = np.random.randn(args.num_rows, args.dim).astype(np.float32) - row_ids = list(range(args.num_rows)) - vectors_as_lists = [stored[i].tolist() for i in range(args.num_rows)] - - for metric in ("l2", "cosine", "inner_product"): - print(f"\n--- Metric: {metric} ---") - - # Pure Python. - times_py = [] - for _ in range(args.num_runs): - t0 = time.perf_counter() - result_py = raw_search_python( - row_ids, vectors_as_lists, query.tolist(), metric, args.limit) - times_py.append(time.perf_counter() - t0) - - # Numpy (from Python lists — worst case for numpy path). - times_np_list = [] - for _ in range(args.num_runs): - t0 = time.perf_counter() - result_np = raw_search_numpy( - row_ids, vectors_as_lists, query, metric, args.limit) - times_np_list.append(time.perf_counter() - t0) - - # Numpy (from pre-built numpy array — simulates Arrow fast path). - row_id_array = np.arange(args.num_rows, dtype=np.int64) - times_np_fast = [] - for _ in range(args.num_runs): - t0 = time.perf_counter() - result_fast = raw_search_numpy_fast( - row_id_array, stored, query, metric, args.limit) - times_np_fast.append(time.perf_counter() - t0) - - # Correctness. - py_ids = set(result_py.keys()) - np_ids = set(result_np.keys()) - fast_ids = set(result_fast.keys()) - overlap_list = len(py_ids & np_ids) / max(len(py_ids), 1) * 100 - overlap_fast = len(py_ids & fast_ids) / max(len(py_ids), 1) * 100 - - avg_py = sum(times_py) / args.num_runs * 1000 - avg_np_list = sum(times_np_list) / args.num_runs * 1000 - avg_np_fast = sum(times_np_fast) / args.num_runs * 1000 - speedup_list = avg_py / avg_np_list if avg_np_list > 0 else float('inf') - speedup_fast = avg_py / avg_np_fast if avg_np_fast > 0 else float('inf') - - print(f" Python loop: {avg_py:.1f} ms") - print(f" Numpy (from list): {avg_np_list:.1f} ms ({speedup_list:.1f}x)") - print(f" Numpy (fast path): {avg_np_fast:.1f} ms ({speedup_fast:.1f}x)") - print(f" TopK overlap: list={overlap_list:.0f}%, fast={overlap_fast:.0f}%") - - if not args.skip_threadpool: - benchmark_threadpool(args.num_shards, args.search_time_ms, args.num_runs) - - print(f"\n{'='*60}") - print("Done.") - - -if __name__ == "__main__": - main() diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py b/paimon-python/pypaimon/tests/vector_search_filter_test.py index 8b35c43fbfe4..9dcbc544b320 100644 --- a/paimon-python/pypaimon/tests/vector_search_filter_test.py +++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py @@ -2926,7 +2926,7 @@ def close(self_inner): query_vectors=[[1.0], [2.0]], filter_=None) raw_result = DictBasedScoredIndexResult({8: 0.9}) with mock.patch.object( - reader, "_read_batch_raw_search", + reader, "_read_raw_batch_search", return_value=[raw_result, raw_result]) as raw_read: results = reader.read_batch([split, raw]) From 3905aa1f5dcf9af4b5ac39f4a721b9cc083bea86 Mon Sep 17 00:00:00 2001 From: Yi Wang Date: Mon, 28 Sep 2026 16:19:37 +0800 Subject: [PATCH 8/8] [python] Use cached index metric for batch raw search consistency 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. --- paimon-python/pypaimon/table/source/vector_search_read.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index 5d44e270b71a..5d09bd22bb0f 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -827,8 +827,7 @@ def _read_raw_batch_search(self, raw_row_ranges, pre_filter, if table.num_rows == 0: return [DictBasedScoredIndexResult({}) for _ in range(n)] - metric = _raw_search_metric( - self._table, self._vector_column, self._options, index_type) + metric = self._search_metric(index_type) return _raw_batch_search_from_arrow( table, self._vector_column.name, self._query_vectors, metric, self._limit)