diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index 46c45f01f3d0..5fe186fed22b 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -18,8 +18,9 @@ """Vector search read to read index files.""" from abc import ABC, abstractmethod -from concurrent.futures import wait +from concurrent.futures import ThreadPoolExecutor, as_completed +from pypaimon.common.options.core_options import CoreOptions from pypaimon.globalindex.batch_vector_search import BatchVectorSearch from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta from pypaimon.globalindex.global_index_result import GlobalIndexResult @@ -87,6 +88,19 @@ def __init__( self._partition_filter = partition_filter self._options = dict(options or {}) + @property + def _index_thread_num(self): + _opts = self._table.options + _get = getattr(_opts, 'global_index_thread_num', None) + value = _get() if _get else None + 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 + ) + return value + def _pre_filters(self, splits, snapshot=None): # type: (list) -> List[RoaringBitmap64] """Evaluate live-row/scalar filters and return one bitmap per index split.""" @@ -259,6 +273,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): @@ -270,25 +328,12 @@ def _read_raw_search(self, raw_row_ranges, pre_filter, query_vector, if table is None or table.num_rows == 0: return DictBasedScoredIndexResult({}) - top_k_heap = [] metric = _raw_search_metric( self._table, self._vector_column, self._options, index_type) - row_ids = table.column(SpecialFields.ROW_ID.name).to_pylist() - vectors = table.column(self._vector_column.name).to_pylist() - for row_id, stored in zip(row_ids, vectors): - if score_candidates is not None and row_id not in score_candidates: - continue - if stored is None: - continue - stored_vector = _to_vector_list(stored) - _check_vector_dimension(query_vector, stored_vector) - _offer_score( - top_k_heap, - self._limit, - row_id, - _compute_score(query_vector, stored_vector, metric), - ) - return _scored_result(top_k_heap) + + return _raw_search_from_arrow( + table, self._vector_column.name, query_vector, metric, self._limit, + score_candidates) def _read_raw_vectors(self, candidates, include_filter=True, snapshot=None): return self._read_raw_candidate_vectors( @@ -330,19 +375,30 @@ def _read_raw_arrow(self, raw_row_ranges, include_filter, snapshot=None): return read_builder.new_read().to_arrow(plan.splits()) def _score_raw_vectors(self, candidates, raw_vectors, query_vector, metric, top_k): - top_k_heap = [] + import numpy as np + + row_ids = [] + vectors = [] for row_id in candidates: stored_vector = raw_vectors.get(row_id) if stored_vector is None: continue - _check_vector_dimension(query_vector, stored_vector) - _offer_score( - top_k_heap, - top_k, - row_id, - _compute_score(query_vector, stored_vector, metric), - ) - return _scored_result(top_k_heap) + row_ids.append(row_id) + vectors.append(stored_vector) + + 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): @@ -502,27 +558,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) - futures = [ - self._eval( - 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] - ) - for i, split in enumerate(splits) - ] - - wait(futures) - merged_scores = {} - for future in 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) + 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) indexed = DictBasedScoredIndexResult(merged_scores).top_k(search_limit) return self._maybe_rerank_indexed_result( @@ -552,22 +621,34 @@ 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) - futures = [ - self._eval_batch( + + 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, - None if not pre_filters else pre_filters[i], - ) - for i, split in enumerate(index_splits) - ] - - wait(futures) + 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 future in futures: - split_results = future.result() + for split_results in split_results_list: for i in range(n): split_result = split_results[i] if split_result is None: @@ -584,18 +665,36 @@ 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_batch_raw_search( + raw_ranges, raw_pre_filter, self._query_vectors, raw_index_type, + snapshot=snapshot) + results = [] for i in range(n): - raw = self._read_raw_search( - raw_ranges, raw_pre_filter, self._query_vectors[i], raw_index_type, - snapshot=snapshot) - results.append(indexed_results[i].or_(raw).top_k(self._limit)) + results.append(indexed_results[i].or_(raw_results[i]).top_k(self._limit)) return results + 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.""" @@ -824,3 +923,238 @@ def _compute_score(query, stored, metric): if metric == "inner_product": return sum(float(q) * float(s) for q, s in zip(query, stored)) raise ValueError("Unknown vector search metric: %s" % metric) + + +def _raw_search_vectorized(row_ids, vectors, query_vector, metric, limit, + score_candidates=None): + """Vectorized raw search using numpy for batch distance computation.""" + import numpy as np + + # Filter by score_candidates and null vectors. + if score_candidates is not None: + candidate_set = set(score_candidates) + filtered = [(rid, vec) for rid, vec in zip(row_ids, vectors) + if rid in candidate_set and vec is not None] + else: + filtered = [(rid, vec) for rid, vec in zip(row_ids, vectors) + if vec is not None] + + if not filtered: + return DictBasedScoredIndexResult({}) + + filtered_ids, filtered_vecs = zip(*filtered) + row_id_array = np.array(filtered_ids, dtype=np.int64) + stored_matrix = np.array( + [_to_vector_list(v) for v in filtered_vecs], dtype=np.float32) + query_np = np.array( + _to_vector_list(query_vector) if not isinstance(query_vector, np.ndarray) + else query_vector, dtype=np.float32) + + return _numpy_topk(row_id_array, stored_matrix, query_np, metric, limit) + + +def _raw_search_from_arrow(arrow_table, vector_column_name, query_vector, + metric, limit, score_candidates=None): + """Vectorized raw search directly from Arrow table (avoids Python list intermediary).""" + import numpy as np + import pyarrow.compute as pc + + row_ids_col = arrow_table.column(SpecialFields.ROW_ID.name) + vectors_col = arrow_table.column(vector_column_name) + + # Filter out null vectors at the Arrow level before conversion. + valid_mask = pc.is_valid(vectors_col) + if not pc.all(valid_mask).as_py(): + arrow_table = arrow_table.filter(valid_mask) + row_ids_col = arrow_table.column(SpecialFields.ROW_ID.name) + vectors_col = arrow_table.column(vector_column_name) + + # Try fast path: fixed-size list → direct numpy reshape. + row_id_array = row_ids_col.to_numpy() + try: + # ChunkedArray has no .values; combine to a single array first. + if hasattr(vectors_col, 'combine_chunks'): + vectors_arr = vectors_col.combine_chunks() + else: + vectors_arr = vectors_col + flat = vectors_arr.values + dim = vectors_arr.type.list_size + if dim is not None and flat is not None: + stored_matrix = flat.to_numpy(zero_copy_only=False).reshape(-1, dim).astype( + np.float32) + else: + stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float32) + except (AttributeError, TypeError, ValueError): + stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float32) + + query_np = np.asarray(query_vector, dtype=np.float32) + + if stored_matrix.shape[1] != query_np.shape[0]: + raise ValueError( + "Query vector dimension mismatch: expected %d, got %d" + % (stored_matrix.shape[1], query_np.shape[0])) + + # Handle null vectors and score_candidates filtering. + if score_candidates is not None: + candidate_set = set(score_candidates) + mask = np.array([rid in candidate_set for rid in row_id_array], dtype=bool) + # Also mask null vectors (check for any NaN row). + null_mask = ~np.isnan(stored_matrix).any(axis=1) + mask = mask & null_mask + row_id_array = row_id_array[mask] + stored_matrix = stored_matrix[mask] + else: + null_mask = ~np.isnan(stored_matrix).any(axis=1) + if not null_mask.all(): + row_id_array = row_id_array[null_mask] + stored_matrix = stored_matrix[null_mask] + + if len(row_id_array) == 0: + return DictBasedScoredIndexResult({}) + + return _numpy_topk(row_id_array, stored_matrix, query_np, metric, limit) + + +def _numpy_topk(row_id_array, stored_matrix, query_np, metric, limit): + """Core numpy distance computation + topK selection.""" + import numpy as np + + if metric == "l2": + diffs = stored_matrix - query_np + dists = np.sum(diffs * diffs, axis=1) + scores = 1.0 / (1.0 + dists) + elif metric == "cosine": + dots = stored_matrix @ query_np + norms = np.linalg.norm(stored_matrix, axis=1) * np.linalg.norm(query_np) + norms = np.where(norms == 0, 1.0, norms) + scores = dots / norms + elif metric == "inner_product": + scores = stored_matrix @ query_np + else: + raise ValueError("Unknown vector search metric: %s" % metric) + + n = len(scores) + if n <= limit: + top_indices = np.argsort(-scores) + else: + top_indices = np.argpartition(-scores, limit)[:limit] + top_indices = top_indices[np.argsort(-scores[top_indices])] + + return DictBasedScoredIndexResult( + {int(row_id_array[i]): float(scores[i]) for i in top_indices} + ) + + +def _raw_batch_search_from_arrow(arrow_table, vector_column_name, query_vectors, + metric, limit, score_candidates=None): + """Batch raw search: multiple queries against the same Arrow table in one SGEMM call.""" + import numpy as np + import pyarrow.compute as pc + + row_ids_col = arrow_table.column(SpecialFields.ROW_ID.name) + vectors_col = arrow_table.column(vector_column_name) + + valid_mask = pc.is_valid(vectors_col) + if not pc.all(valid_mask).as_py(): + arrow_table = arrow_table.filter(valid_mask) + row_ids_col = arrow_table.column(SpecialFields.ROW_ID.name) + vectors_col = arrow_table.column(vector_column_name) + + row_id_array = row_ids_col.to_numpy() + try: + if hasattr(vectors_col, 'combine_chunks'): + vectors_arr = vectors_col.combine_chunks() + else: + vectors_arr = vectors_col + flat = vectors_arr.values + dim = vectors_arr.type.list_size + if dim is not None and flat is not None: + stored_matrix = flat.to_numpy(zero_copy_only=False).reshape(-1, dim).astype( + np.float32) + else: + stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float32) + except (AttributeError, TypeError, ValueError): + stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float32) + + query_matrix = np.array( + [q if isinstance(q, np.ndarray) else list(q) for q in query_vectors], + dtype=np.float32) + + if stored_matrix.shape[1] != query_matrix.shape[1]: + raise ValueError( + "Query vector dimension mismatch: expected %d, got %d" + % (stored_matrix.shape[1], query_matrix.shape[1])) + + if score_candidates is not None: + candidate_set = set(score_candidates) + mask = np.array([rid in candidate_set for rid in row_id_array], dtype=bool) + null_mask = ~np.isnan(stored_matrix).any(axis=1) + mask = mask & null_mask + row_id_array = row_id_array[mask] + stored_matrix = stored_matrix[mask] + else: + null_mask = ~np.isnan(stored_matrix).any(axis=1) + if not null_mask.all(): + row_id_array = row_id_array[null_mask] + stored_matrix = stored_matrix[null_mask] + + if len(row_id_array) == 0: + return [DictBasedScoredIndexResult({}) for _ in range(len(query_vectors))] + + return _numpy_batch_topk(row_id_array, stored_matrix, query_matrix, metric, limit) + + +def _numpy_batch_topk(row_id_array, stored_matrix, query_matrix, metric, limit): + """Batch 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] + + # 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) + + results = [] + 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: + 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 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/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() diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py b/paimon-python/pypaimon/tests/vector_search_filter_test.py index e3bbccc911b5..68614405fc20 100644 --- a/paimon-python/pypaimon/tests/vector_search_filter_test.py +++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py @@ -2910,13 +2910,14 @@ def close(self_inner): reader = BatchVectorSearchReadImpl( table, limit=5, vector_column=embedding_field, query_vectors=[[1.0], [2.0]], filter_=None) + raw_result = DictBasedScoredIndexResult({8: 0.9}) with mock.patch.object( - reader, "_read_raw_search", - return_value=DictBasedScoredIndexResult({8: 0.9})) as raw_read: + 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(2, raw_read.call_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()))) @@ -3566,5 +3567,110 @@ def tearDown(self): mock.patch.stopall() +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()