diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index fe9d96536719..3e8615033342 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -18,7 +18,7 @@ """Vector search read to read index files.""" from abc import ABC, abstractmethod -from concurrent.futures import wait +from concurrent.futures import ThreadPoolExecutor, wait from pypaimon.globalindex.batch_vector_search import BatchVectorSearch from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta @@ -320,6 +320,10 @@ def _read_raw_candidate_vectors(self, raw_row_ranges, candidates, return raw_vectors def _read_raw_arrow(self, raw_row_ranges, include_filter, snapshot=None): + reader, splits = self._plan_raw_read(raw_row_ranges, include_filter, snapshot) + return reader.to_arrow(splits) + + def _plan_raw_read(self, raw_row_ranges, include_filter, snapshot=None): read_table = global_index_live_row_filter.table_at_snapshot( self._table, snapshot) read_builder = read_table.new_read_builder() @@ -331,7 +335,7 @@ def _read_raw_arrow(self, raw_row_ranges, include_filter, snapshot=None): read_builder = read_builder.with_projection( self._raw_search_projection(include_filter)) plan = read_builder.new_scan().with_row_ranges(raw_row_ranges).plan() - return read_builder.new_read().to_arrow(plan.splits()) + return read_builder.new_read(), plan.splits() def _score_raw_vectors(self, candidates, raw_vectors, query_vector, metric, top_k): top_k_heap = [] @@ -596,13 +600,62 @@ 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) - 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)) - return results + 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) + ] + + 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] + 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] + + table_read, splits = self._plan_raw_read(raw_row_ranges, True, snapshot) + metric = _raw_search_metric( + self._table, self._vector_column, self._options, 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 + + heaps = [[] for _ in self._query_vectors] + 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: + for batch in batch_reader: + row_ids = batch.column(SpecialFields.ROW_ID.name).to_pylist() + vectors = batch.column(self._vector_column.name).to_pylist() + for row_id, stored in zip(row_ids, vectors): + if stored is None: + continue + stored_vector = _to_vector_list(stored) + for query, heap in zip(self._query_vectors, heaps): + _check_vector_dimension(query, stored_vector) + _offer_score( + heap, self._limit, row_id, + _compute_score(query, stored_vector, metric)) + del batch, row_ids, vectors + return [_scored_result(heap) for heap in heaps] def _create_vector_reader(index_type, file_io, index_path, index_io_meta_list, options=None): diff --git a/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py b/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py new file mode 100644 index 000000000000..ad84a853160c --- /dev/null +++ b/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py @@ -0,0 +1,289 @@ +# 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. + +import unittest +import threading +from unittest.mock import patch + +import pyarrow as pa + +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 + + +def _scores(result): + getter = result.score_getter() + return {row_id: getter(row_id) for row_id in result.results()} + + +class BatchVectorRawScanTest(BatchModeMixin, DataEvolutionTestBase, unittest.TestCase): + + pa_schema = pa.schema([ + ('id', pa.int32()), ('embedding', pa.list_(pa.float32())), ('pt', pa.int32()), + ]) + table_options = { + '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): + return pa.table({'id': list(range(len(vectors))), 'embedding': vectors, + 'pt': [partition] * len(vectors)}, schema=self.pa_schema) + + def _reader(self, table, queries, metric='l2', **kwargs): + return BatchVectorSearchReadImpl( + table, limit=2, vector_column=table.field_dict['embedding'], + query_vectors=queries, options={'ivf-flat.metric': metric}, **kwargs) + + def test_shared_scan_matches_individual_queries_for_all_metrics(self): + table = self._create_table() + self._write_arrow(table, self._data([ + [1, 0], [0, 1], None, [0, 0], [1, 0], [-1, 2], [2, -1], + ])) + queries = [[1, 0], [0, 1], [0, 0], [1, 0]] + ranges = [Range(0, 4), Range(3, 6)] + for metric in ('l2', 'cosine', 'inner_product'): + with self.subTest(metric=metric): + 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() + 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): + 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]) + + def test_scoring_finishes_each_batch_before_reading_the_next(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_score + + 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_scores = 0 + try: + for batch in source: + batch_sizes.append(batch.num_rows) + expected_scores += batch.num_rows * 2 + yield batch + self.assertEqual(expected_scores, 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_score', wraps=_compute_score) as score: + reader._read_raw_batch_search([Range(0, 3)], None, 'ivf-flat') + self.assertEqual(8, 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) + + def test_public_batch_search_uses_planned_snapshot(self): + table = self._create_table() + self._write_arrow(table, self._data([[0, 0]])) + builder = table.new_batch_vector_search_builder().with_vector_column( + 'embedding').with_query_vectors([[2, 0], [0, 0]]).with_limit(1) + plan = builder.new_vector_search_scan().scan() + 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())) + current = builder.execute_batch_local() + self.assertEqual([1.0], list(_scores(current[0]).values())) + self.assertNotEqual(list(old[0].results()), list(current[0].results())) + + 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) + + 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)) + + +if __name__ == '__main__': + unittest.main() diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py b/paimon-python/pypaimon/tests/vector_search_filter_test.py index 4d0c39e29eae..c43804e6f0da 100644 --- a/paimon-python/pypaimon/tests/vector_search_filter_test.py +++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py @@ -201,6 +201,9 @@ def plan(self): return _Plan(self._row_ids) class _Read: + def _resolve_parallelism(self, runtime, num_splits): + return 1 + def to_arrow(self, splits): row_ids = list(splits) return pa.table({ @@ -209,6 +212,11 @@ def to_arrow(self, splits): "_ROW_ID": pa.array(row_ids, type=pa.int64()), }) + def _new_arrow_batch_reader(self, splits): + table = self.to_arrow(splits) + batches = (batch for batch in table.to_batches(max_chunksize=2)) + return pa.RecordBatchReader.from_batches(table.schema, batches), batches + class _Builder: def with_partition_filter(self, predicate): calls["partition_filter"] = predicate @@ -2910,13 +2918,12 @@ def close(self_inner): reader = BatchVectorSearchReadImpl( table, limit=5, vector_column=embedding_field, query_vectors=[[1.0], [2.0]], filter_=None) - with mock.patch.object( - reader, "_read_raw_search", - return_value=DictBasedScoredIndexResult({8: 0.9})) as raw_read: - results = reader.read_batch([split, raw]) + raw_calls = _install_raw_vector_read_builder( + table, "embedding", {8: [1.5]}) + 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_calls["raw_read_count"]) self.assertEqual([1, 8], sorted(list(results[0].results()))) self.assertEqual([2, 8], sorted(list(results[1].results())))