diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index 0df1fec9839b..76838014619f 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 from contextlib import closing +from concurrent.futures import ThreadPoolExecutor from threading import Lock from pypaimon.common.options.core_options import CoreOptions @@ -124,6 +124,19 @@ 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) + 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.""" @@ -324,6 +337,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 +392,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 +465,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): @@ -713,15 +775,13 @@ 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) - # 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: + 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]) @@ -734,49 +794,33 @@ def _read_batch(self, splits, snapshot): indexed_results = self._maybe_rerank_indexed_results( indexed_results, index_type, self._query_vectors, snapshot) - # Each query: merge indexed results with the raw (brute-force) fallback. + # Batch raw search: read Arrow table once, compute all queries in one SGEMM. raw_pre_filter = self._raw_pre_filter(raw_splits, snapshot) raw_ranges = _raw_row_ranges(raw_splits) raw_index_type = _raw_search_index_type(raw_splits) raw_results = self._read_raw_batch_search( - raw_ranges, raw_pre_filter, raw_index_type, snapshot) - return [ - indexed.or_(raw).top_k(self._limit) - for indexed, raw in zip(indexed_results, raw_results) - ] + raw_ranges, raw_pre_filter, raw_index_type, + snapshot=snapshot) + + results = [] + for i in range(n): + results.append(indexed_results[i].or_(raw_results[i]).top_k(self._limit)) + return results def _read_raw_batch_search(self, raw_row_ranges, pre_filter, index_type=None, snapshot=None): - """Scan raw rows once, keeping a separate top-k heap for each query.""" - heaps = [[] for _ in self._query_vectors] + from pypaimon.read.table_read import _ClosableArrowBatchReader + raw_row_ranges = _filtered_raw_row_ranges(raw_row_ranges, pre_filter) - if not raw_row_ranges or not heaps: - return [_scored_result(heap) for heap in heaps] + n = len(self._query_vectors) + if not raw_row_ranges: + return [DictBasedScoredIndexResult({}) for _ in range(n)] - 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 - + table_read, splits = self._plan_raw_read(raw_row_ranges, True, snapshot) reader, batches = table_read._new_arrow_batch_reader(splits) - # Close the underlying iterator as well if scoring fails mid-batch. + # Score each bounded batch and merge only per-query Top-K candidates, so + # peak memory tracks the configured read batch, not the whole raw tail. with _ClosableArrowBatchReader(reader, batches) as batch_reader: return self._score_raw_batch_queries(batch_reader, metric) @@ -1200,3 +1244,183 @@ def _compute_score(query, stored, metric): if metric == "inner_product": return sum(float(q) * float(s) for q, s in zip(query, stored)) raise ValueError("Unknown vector search metric: %s" % metric) + + +def _raw_search_vectorized(row_ids, vectors, query_vector, metric, limit, + score_candidates=None): + """Vectorized raw search using numpy for batch distance computation.""" + import numpy as np + + # Filter by score_candidates and null vectors. + if score_candidates is not None: + candidate_set = set(score_candidates) + filtered = [(rid, vec) for rid, vec in zip(row_ids, vectors) + if rid in candidate_set and vec is not None] + else: + filtered = [(rid, vec) for rid, vec in zip(row_ids, vectors) + if vec is not None] + + if not filtered: + return DictBasedScoredIndexResult({}) + + filtered_ids, filtered_vecs = zip(*filtered) + row_id_array = np.array(filtered_ids, dtype=np.int64) + stored_matrix = np.array( + [_to_vector_list(v) for v in filtered_vecs], dtype=np.float32) + query_np = np.array( + _to_vector_list(query_vector) if not isinstance(query_vector, np.ndarray) + else query_vector, dtype=np.float32) + + return _numpy_topk(row_id_array, stored_matrix, query_np, metric, limit) + + +def _raw_search_from_arrow(arrow_table, vector_column_name, query_vector, + metric, limit, score_candidates=None): + """Vectorized raw search directly from Arrow table (avoids Python list intermediary).""" + import numpy as np + import pyarrow as pa + import pyarrow.compute as pc + + row_ids_col = arrow_table.column(SpecialFields.ROW_ID.name) + vectors_col = arrow_table.column(vector_column_name) + + # Filter out null vectors at the Arrow level before conversion. + valid_mask = pc.is_valid(vectors_col) + if not pc.all(valid_mask).as_py(): + arrow_table = arrow_table.filter(valid_mask) + row_ids_col = arrow_table.column(SpecialFields.ROW_ID.name) + vectors_col = arrow_table.column(vector_column_name) + + # Filter by score_candidates at Arrow level to avoid ragged-array issues + # when non-candidate rows have mismatched dimensions. + if score_candidates is not None: + candidate_set = set(score_candidates) + candidate_mask = pc.is_in(row_ids_col, pa.array( + list(candidate_set), type=row_ids_col.type)) + arrow_table = arrow_table.filter(candidate_mask) + row_ids_col = arrow_table.column(SpecialFields.ROW_ID.name) + vectors_col = arrow_table.column(vector_column_name) + + if arrow_table.num_rows == 0: + return DictBasedScoredIndexResult({}) + + # Try fast path: fixed-size list → direct numpy reshape. + row_id_array = row_ids_col.to_numpy() + try: + # ChunkedArray has no .values; combine to a single array first. + if hasattr(vectors_col, 'combine_chunks'): + vectors_arr = vectors_col.combine_chunks() + else: + vectors_arr = vectors_col + flat = vectors_arr.values + dim = vectors_arr.type.list_size + if dim is not None and flat is not None: + stored_matrix = flat.to_numpy(zero_copy_only=False).reshape(-1, dim).astype( + np.float64) + else: + stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float64) + except (AttributeError, TypeError, ValueError): + stored_matrix = np.array(vectors_col.to_pylist(), dtype=np.float64) + + query_np = np.asarray(query_vector, dtype=np.float64) + + if stored_matrix.ndim < 2 or stored_matrix.shape[0] == 0: + return DictBasedScoredIndexResult({}) + + if stored_matrix.shape[1] != query_np.shape[0]: + raise ValueError( + "Query vector dimension mismatch: expected %d, got %d" + % (stored_matrix.shape[1], query_np.shape[0])) + + # Filter NaN vectors (from variable-length lists that got padded). + null_mask = ~np.isnan(stored_matrix).any(axis=1) + if not null_mask.all(): + row_id_array = row_id_array[null_mask] + stored_matrix = stored_matrix[null_mask] + + if len(row_id_array) == 0: + return DictBasedScoredIndexResult({}) + + return _numpy_topk(row_id_array, stored_matrix, query_np, metric, limit) + + +def _numpy_topk(row_id_array, stored_matrix, query_np, metric, limit): + """Core numpy distance computation + topK selection with row tiling.""" + import numpy as np + + ROW_TILE = 65536 + n_rows = stored_matrix.shape[0] + + if n_rows <= ROW_TILE: + scores = _compute_scores_single(stored_matrix, query_np, metric) + top_indices = _topk_indices(scores, row_id_array, limit) + return DictBasedScoredIndexResult( + {int(row_id_array[i]): float(scores[i]) for i in top_indices} + ) + + best_ids = np.empty(0, dtype=row_id_array.dtype) + # Accumulate in float64. Downcasting to float32 would collapse finite but + # large/small legal scores into Infinity or zero across the tile merge, + # selecting the wrong row for inputs with more than ROW_TILE rows. + best_scores = np.empty(0, dtype=np.float64) + + 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, np.asarray(tile_scores, dtype=np.float64)[tile_top]]) + + final_top = _topk_indices(best_scores, best_ids, limit) + return DictBasedScoredIndexResult( + {int(best_ids[i]): float(best_scores[i]) for i in final_top} + ) + + +def _compute_scores_single(stored_chunk, query_np, metric): + import numpy as np + + if metric == "l2": + diffs = stored_chunk - query_np + dists = np.sum(diffs * diffs, axis=1) + return 1.0 / (1.0 + dists) + elif metric == "cosine": + dots = stored_chunk @ query_np + norms = np.linalg.norm(stored_chunk, axis=1) * np.linalg.norm(query_np) + norms = np.where(norms == 0, 1.0, norms) + return dots / norms + elif metric == "inner_product": + return stored_chunk @ query_np + else: + raise ValueError("Unknown vector search metric: %s" % metric) + + +def _topk_indices(scores, row_id_array, limit): + """Select top-limit indices by (highest score, smallest row_id) tie-break.""" + import numpy as np + + n = len(scores) + if n <= limit: + return np.lexsort((row_id_array, -scores)) + + part_idx = np.argpartition(-scores, limit)[:limit] + kth_score = np.min(scores[part_idx]) + + above_mask = scores[part_idx] > kth_score + above = part_idx[above_mask] + + all_tie_idx = np.where(scores == kth_score)[0] + n_ties_needed = limit - len(above) + + if len(all_tie_idx) <= n_ties_needed: + result = np.concatenate([above, all_tie_idx]) + else: + tie_order = np.argsort(row_id_array[all_tie_idx]) + result = np.concatenate([above, all_tie_idx[tie_order[:n_ties_needed]]]) + + final_order = np.lexsort((row_id_array[result], -scores[result])) + return result[final_order] 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/vector_search_filter_test.py b/paimon-python/pypaimon/tests/vector_search_filter_test.py index 837e792ea7fa..6438a27203ae 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_raw_batch_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,222 @@ def test_reader_lifetime_for_single_and_batch_search(self): self.assertTrue(stream.closed) +class RawBatchStreamingBoundTest(unittest.TestCase): + """Batch raw scoring must stay bounded, releasing each batch as it reads.""" + + def test_batches_are_released_before_later_batches_are_read(self): + import weakref + + import numpy as np + import pyarrow as pa + + from pypaimon.table.source.vector_search_read import BatchVectorSearchReadImpl + + column = _field(1, "embedding", "FLOAT") + reader = BatchVectorSearchReadImpl( + _StubTable([column], []), limit=2, vector_column=column, + query_vectors=[[1.0, 0.0], [0.0, 1.0]]) + + # One batch per read; the previous batch must be unreferenced before the + # next is produced, so peak memory tracks a single batch, not the tail. + live = [] + n_batches = 4 + + def batches(): + for i in range(n_batches): + for ref in live: + self.assertIsNone( + ref(), "a previous batch was still alive when the next " + "batch was read") + base = i * 2 + batch = pa.record_batch({ + "_ROW_ID": pa.array([base, base + 1], type=pa.int64()), + "embedding": pa.array( + [[float(base), 0.0], [0.0, float(base + 1)]], + type=pa.list_(pa.float32())), + }) + live.append(weakref.ref(batch)) + yield batch + del batch + + results = reader._score_raw_batch_queries(batches(), "l2") + self.assertEqual(len(results), 2) + for result in results: + self.assertLessEqual(len(result.results()), 2) + # Every batch was observed released once the generator advanced past it. + for ref in live: + self.assertIsNone(ref()) + + +class RawBatchSearchFromArrowTest(unittest.TestCase): + """Unit tests for the numpy top-K selection shared by raw search paths.""" + + 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, + ) + + 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) + + 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_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_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]) + + def test_numpy_topk_cross_tile_scores_stay_float64(self): + """Row-tile merge must keep float64 scores, not collapse to inf/zero. + + With more than ROW_TILE (65,536) rows, downcasting accumulated scores to + float32 turns finite legal FLOAT inputs into Infinity (inner_product) or + zero (L2), selecting the wrong row. The tiled path must match the + single-tile result on the same data. + """ + import numpy as np + from pypaimon.table.source.vector_search_read import _numpy_topk + + n_rows = 65537 + row_ids = np.arange(n_rows, dtype=np.int64) + # Callers hand _numpy_topk a float64 matrix (the Arrow path upcasts); + # the regression is the float32 merge accumulator, not the compute step. + stored = np.full((n_rows, 1), 1e23, dtype=np.float64) + stored[n_rows - 1] = 2e23 + + for metric, q in (("inner_product", 1e20), ("l2", 3e23)): + query = np.array([q], dtype=np.float64) + full = _numpy_topk(row_ids, stored, query, metric, 1) + self.assertEqual( + [n_rows - 1], list(full.results()), + "tiled %s selected the wrong row" % metric) + getter = full.score_getter() + score = getter(n_rows - 1) + self.assertTrue(np.isfinite(score) and score != 0.0, + "tiled %s score collapsed to %r" % (metric, score)) + + def test_public_batch_matches_base_on_cross_tile_extremes(self): + """execute_batch_local must match the exact-base result past a tile edge. + + Reproduces the P2 overflow/underflow case end to end: the streaming batch + path and the single-query base path must select the same row and finite + score on 65,537 one-dimensional ARRAY rows. + """ + import numpy as np + import pyarrow as pa + from pypaimon.table.source.vector_search_read import ( + BatchVectorSearchReadImpl, + _raw_search_from_arrow, + ) + + n_rows = 65537 + vectors = [[1e23]] * (n_rows - 1) + [[2e23]] + table = pa.table({ + "_ROW_ID": pa.array(list(range(n_rows)), type=pa.int64()), + "vector": pa.array(vectors, type=pa.list_(pa.float32())), + }) + column = _field(1, "vector", "FLOAT") + + for metric, q in (("inner_product", 1e20), ("l2", 3e23)): + query = [np.float32(q)] + base = _raw_search_from_arrow(table, "vector", np.float32(query), + metric, 1) + reader = BatchVectorSearchReadImpl( + _StubTable([column], []), limit=1, vector_column=column, + query_vectors=[query]) + batch = reader._score_raw_batch_queries( + table.to_batches(max_chunksize=16384), metric)[0] + self.assertEqual(list(base.results()), list(batch.results()), + "batch and base diverged for metric=%s" % metric) + base_getter = base.score_getter() + batch_getter = batch.score_getter() + for row_id in base.results(): + self.assertTrue(np.isfinite(base_getter(row_id))) + self.assertAlmostEqual(base_getter(row_id), batch_getter(row_id)) + + if __name__ == "__main__": unittest.main()