From ed70c289815b0997e48db7cb08b3c327ea402d14 Mon Sep 17 00:00:00 2001 From: chaoyang Date: Sat, 12 Sep 2026 14:37:47 +0800 Subject: [PATCH 1/3] [python] Share raw vector scans across batch queries --- .../dev/benchmark_batch_vector_raw_scan.py | 199 +++++++++++++++++ .../table/source/vector_search_read.py | 50 ++++- .../tests/batch_vector_raw_scan_test.py | 202 ++++++++++++++++++ .../tests/vector_search_filter_test.py | 14 +- 4 files changed, 452 insertions(+), 13 deletions(-) create mode 100644 paimon-python/dev/benchmark_batch_vector_raw_scan.py create mode 100644 paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py diff --git a/paimon-python/dev/benchmark_batch_vector_raw_scan.py b/paimon-python/dev/benchmark_batch_vector_raw_scan.py new file mode 100644 index 000000000000..f156293fd4ac --- /dev/null +++ b/paimon-python/dev/benchmark_batch_vector_raw_scan.py @@ -0,0 +1,199 @@ +# 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. + +"""Compare repeated raw scans with eager and streaming shared scans. + +Run from paimon-python with PYTHONPATH=. and project dependencies installed:: + + python dev/benchmark_batch_vector_raw_scan.py prepare --warehouse /tmp/raw-scan-bench + python dev/benchmark_batch_vector_raw_scan.py run --warehouse /tmp/raw-scan-bench \ + --mode shared-stream --queries 8 --output /tmp/shared.json + +Run each mode in a fresh process. Preparation is excluded from measurements. +The table has no vector indexes and uses vector-index.search-mode=full. All +variants keep the existing scalar distance calculation and top-k rules. +Filesystem cache is not controlled. Reported row counts measure rows delivered +by Arrow readers, not physical disk reads. Read parallelism is fixed to one. +""" + +import argparse +from contextlib import ExitStack +import hashlib +import importlib +import json +import os +import platform +import resource +import sys +import time +from unittest.mock import patch + +import numpy as np +import pyarrow as pa + +from pypaimon import CatalogFactory, Schema +from pypaimon.read.table_read import TableRead +from pypaimon.table.source.vector_search_read import BatchVectorSearchReadImpl +from pypaimon.table.special_fields import SpecialFields + +search_module = importlib.import_module('pypaimon.table.source.vector_search_read') + + +def peak_rss_mib(): + rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return rss / (1024 * 1024 if sys.platform == 'darwin' else 1024) + + +def prepare(args): + catalog = CatalogFactory.create({'warehouse': args.warehouse}) + catalog.create_database('default', True) + schema = pa.schema([('embedding', pa.list_(pa.float32()))]) + catalog.create_table('default.vectors', Schema.from_pyarrow_schema(schema, options={ + 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true', + 'global-index.enabled': 'true', 'bucket': '-1', 'file.format': 'parquet', + 'vector-index.search-mode': 'full', 'read.parallelism': '1', + }), False) + table = catalog.get_table('default.vectors') + wb = table.new_batch_write_builder() + writer, commit = wb.new_write(), wb.new_commit() + rng = np.random.RandomState(20260912) + try: + for start in range(0, args.rows, 4096): + count = min(4096, args.rows - start) + values = rng.standard_normal((count, args.dimension)).astype(np.float32) + vectors = pa.ListArray.from_arrays( + np.arange(count + 1, dtype=np.int32) * args.dimension, + pa.array(values.reshape(-1))) + writer.write_arrow(pa.Table.from_arrays([vectors], schema=schema)) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + with open(os.path.join(args.warehouse, 'benchmark.json'), 'w') as stream: + json.dump({'rows': args.rows, 'dimension': args.dimension}, stream) + + +def repeated(self, ranges, pre_filter, index_type=None, snapshot=None): + return [self._read_raw_search(ranges, pre_filter, query, index_type, snapshot=snapshot) + for query in self._query_vectors] + + +def shared_table(self, ranges, pre_filter, index_type=None, snapshot=None): + heaps = [[] for _ in self._query_vectors] + ranges = search_module._filtered_raw_row_ranges(ranges, pre_filter) + if not ranges or not heaps: + return [search_module._scored_result(heap) for heap in heaps] + table = self._read_raw_arrow(ranges, True, snapshot) + if table is None or table.num_rows == 0: + return [search_module._scored_result(heap) for heap in heaps] + metric = search_module._raw_search_metric( + self._table, self._vector_column, self._options, index_type) + ids = table.column(SpecialFields.ROW_ID.name).to_pylist() + vectors = table.column(self._vector_column.name).to_pylist() + for row_id, stored in zip(ids, vectors): + if stored is None: + continue + vector = search_module._to_vector_list(stored) + for query, heap in zip(self._query_vectors, heaps): + search_module._check_vector_dimension(query, vector) + search_module._offer_score(heap, self._limit, row_id, + search_module._compute_score(query, vector, metric)) + return [search_module._scored_result(heap) for heap in heaps] + + +def run(args): + with open(os.path.join(args.warehouse, 'benchmark.json')) as stream: + metadata = json.load(stream) + table = CatalogFactory.create({'warehouse': args.warehouse}).get_table('default.vectors') + table = table.copy({'read.batch-size': str(args.batch_size), 'read.parallelism': '1'}) + queries = np.random.RandomState(42).standard_normal( + (args.queries, metadata['dimension'])).astype(np.float32).tolist() + builder = table.new_batch_vector_search_builder().with_vector_column( + 'embedding').with_query_vectors(queries).with_limit(args.top_k).with_option('metric', 'l2') + stats = dict(metadata, mode=args.mode, queries=args.queries, top_k=args.top_k, + batch_size=args.batch_size, raw_plans=0, delivered_rows=0, + delivered_batches=0, source_passes=0, + before_peak_rss_mib=peak_rss_mib(), python=platform.python_version(), + pyarrow=pa.__version__, numpy=np.__version__, platform=platform.platform()) + original_plan = BatchVectorSearchReadImpl._plan_raw_read + original_generator = TableRead._arrow_batch_generator + + def plan(reader, *a, **kw): + stats['raw_plans'] += 1 + return original_plan(reader, *a, **kw) + + def batches(reader, *a, **kw): + stats['source_passes'] += 1 + source = original_generator(reader, *a, **kw) + try: + for batch in source: + stats['delivered_rows'] += batch.num_rows + stats['delivered_batches'] += 1 + yield batch + finally: + source.close() + + with ExitStack() as stack: + stack.enter_context(patch.object(BatchVectorSearchReadImpl, '_plan_raw_read', plan)) + stack.enter_context(patch.object(TableRead, '_arrow_batch_generator', batches)) + if args.mode != 'shared-stream': + stack.enter_context(patch.object( + BatchVectorSearchReadImpl, '_read_raw_batch_search', + repeated if args.mode == 'repeated' else shared_table)) + start = time.perf_counter() + results = builder.execute_batch_local() + stats['seconds'] = time.perf_counter() - start + stats['peak_rss_mib'] = peak_rss_mib() + expected_passes = args.queries if args.mode == 'repeated' else 1 + assert stats['raw_plans'] == expected_passes, stats + assert stats['source_passes'] == expected_passes, stats + assert stats['delivered_rows'] == metadata['rows'] * expected_passes, stats + scores = [] + for result in results: + getter = result.score_getter() + scores.append([[int(row_id), float(getter(row_id))] for row_id in sorted(result.results())]) + assert len(scores) == args.queries + assert all(len(score) == min(args.top_k, metadata['rows']) for score in scores) + stats['results'] = scores + stats['result_sha256'] = hashlib.sha256(json.dumps(scores).encode()).hexdigest() + output = json.dumps(stats, sort_keys=True) + if args.output: + with open(args.output, 'w') as stream: + stream.write(output + '\n') + print(output) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('action', choices=['prepare', 'run']) + parser.add_argument('--warehouse', required=True) + parser.add_argument('--mode', choices=['repeated', 'shared-table', 'shared-stream'], + default='shared-stream') + parser.add_argument('--rows', type=int, default=16384) + parser.add_argument('--dimension', type=int, default=128) + parser.add_argument('--queries', type=int, default=8) + parser.add_argument('--top-k', type=int, default=10) + parser.add_argument('--batch-size', type=int, default=1024) + parser.add_argument('--output') + args = parser.parse_args() + if min(args.rows, args.dimension, args.queries, args.top_k, args.batch_size) <= 0: + parser.error('Numeric arguments must be positive') + (prepare if args.action == 'prepare' else run)(args) + + +if __name__ == '__main__': + main() diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index fe9d96536719..b0a5c01fea93 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -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,43 @@ 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.""" + from pypaimon.read.table_read import _ClosableArrowBatchReader + + 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) + 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: + metric = _raw_search_metric( + self._table, self._vector_column, self._options, index_type) + 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..dec90577f4a3 --- /dev/null +++ b/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py @@ -0,0 +1,202 @@ +# 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 +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())) + + +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..6ff3fd6e431b 100644 --- a/paimon-python/pypaimon/tests/vector_search_filter_test.py +++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py @@ -209,6 +209,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 +2915,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()))) From 9a23e2d277336ea9d40f7a2839e8c17c3ba2b40e Mon Sep 17 00:00:00 2001 From: chaoyang Date: Sun, 13 Sep 2026 09:14:41 +0800 Subject: [PATCH 2/3] [python] Preserve split concurrency in shared vector scans --- .../table/source/vector_search_read.py | 29 +++++-- .../tests/batch_vector_raw_scan_test.py | 87 +++++++++++++++++++ .../tests/vector_search_filter_test.py | 3 + 3 files changed, 114 insertions(+), 5 deletions(-) diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index b0a5c01fea93..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 @@ -610,19 +610,38 @@ def _read_batch(self, splits, snapshot): def _read_raw_batch_search(self, raw_row_ranges, pre_filter, index_type=None, snapshot=None): """Scan raw rows once, keeping a separate top-k heap for each query.""" - from pypaimon.read.table_read import _ClosableArrowBatchReader - 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: - metric = _raw_search_metric( - self._table, self._vector_column, self._options, index_type) for batch in batch_reader: row_ids = batch.column(SpecialFields.ROW_ID.name).to_pylist() vectors = batch.column(self._vector_column.name).to_pylist() 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 dec90577f4a3..ad84a853160c 100644 --- a/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py +++ b/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py @@ -16,6 +16,7 @@ # under the License. import unittest +import threading from unittest.mock import patch import pyarrow as pa @@ -197,6 +198,92 @@ def test_public_batch_search_uses_planned_snapshot(self): 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 6ff3fd6e431b..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({ From abfb2bde6bc94360e68370d82f194d0e9609cb47 Mon Sep 17 00:00:00 2001 From: chaoyang Date: Sun, 13 Sep 2026 09:26:41 +0800 Subject: [PATCH 3/3] [python] Remove shared vector scan benchmark --- .../dev/benchmark_batch_vector_raw_scan.py | 199 ------------------ 1 file changed, 199 deletions(-) delete mode 100644 paimon-python/dev/benchmark_batch_vector_raw_scan.py diff --git a/paimon-python/dev/benchmark_batch_vector_raw_scan.py b/paimon-python/dev/benchmark_batch_vector_raw_scan.py deleted file mode 100644 index f156293fd4ac..000000000000 --- a/paimon-python/dev/benchmark_batch_vector_raw_scan.py +++ /dev/null @@ -1,199 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -"""Compare repeated raw scans with eager and streaming shared scans. - -Run from paimon-python with PYTHONPATH=. and project dependencies installed:: - - python dev/benchmark_batch_vector_raw_scan.py prepare --warehouse /tmp/raw-scan-bench - python dev/benchmark_batch_vector_raw_scan.py run --warehouse /tmp/raw-scan-bench \ - --mode shared-stream --queries 8 --output /tmp/shared.json - -Run each mode in a fresh process. Preparation is excluded from measurements. -The table has no vector indexes and uses vector-index.search-mode=full. All -variants keep the existing scalar distance calculation and top-k rules. -Filesystem cache is not controlled. Reported row counts measure rows delivered -by Arrow readers, not physical disk reads. Read parallelism is fixed to one. -""" - -import argparse -from contextlib import ExitStack -import hashlib -import importlib -import json -import os -import platform -import resource -import sys -import time -from unittest.mock import patch - -import numpy as np -import pyarrow as pa - -from pypaimon import CatalogFactory, Schema -from pypaimon.read.table_read import TableRead -from pypaimon.table.source.vector_search_read import BatchVectorSearchReadImpl -from pypaimon.table.special_fields import SpecialFields - -search_module = importlib.import_module('pypaimon.table.source.vector_search_read') - - -def peak_rss_mib(): - rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss - return rss / (1024 * 1024 if sys.platform == 'darwin' else 1024) - - -def prepare(args): - catalog = CatalogFactory.create({'warehouse': args.warehouse}) - catalog.create_database('default', True) - schema = pa.schema([('embedding', pa.list_(pa.float32()))]) - catalog.create_table('default.vectors', Schema.from_pyarrow_schema(schema, options={ - 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true', - 'global-index.enabled': 'true', 'bucket': '-1', 'file.format': 'parquet', - 'vector-index.search-mode': 'full', 'read.parallelism': '1', - }), False) - table = catalog.get_table('default.vectors') - wb = table.new_batch_write_builder() - writer, commit = wb.new_write(), wb.new_commit() - rng = np.random.RandomState(20260912) - try: - for start in range(0, args.rows, 4096): - count = min(4096, args.rows - start) - values = rng.standard_normal((count, args.dimension)).astype(np.float32) - vectors = pa.ListArray.from_arrays( - np.arange(count + 1, dtype=np.int32) * args.dimension, - pa.array(values.reshape(-1))) - writer.write_arrow(pa.Table.from_arrays([vectors], schema=schema)) - commit.commit(writer.prepare_commit()) - finally: - writer.close() - commit.close() - with open(os.path.join(args.warehouse, 'benchmark.json'), 'w') as stream: - json.dump({'rows': args.rows, 'dimension': args.dimension}, stream) - - -def repeated(self, ranges, pre_filter, index_type=None, snapshot=None): - return [self._read_raw_search(ranges, pre_filter, query, index_type, snapshot=snapshot) - for query in self._query_vectors] - - -def shared_table(self, ranges, pre_filter, index_type=None, snapshot=None): - heaps = [[] for _ in self._query_vectors] - ranges = search_module._filtered_raw_row_ranges(ranges, pre_filter) - if not ranges or not heaps: - return [search_module._scored_result(heap) for heap in heaps] - table = self._read_raw_arrow(ranges, True, snapshot) - if table is None or table.num_rows == 0: - return [search_module._scored_result(heap) for heap in heaps] - metric = search_module._raw_search_metric( - self._table, self._vector_column, self._options, index_type) - ids = table.column(SpecialFields.ROW_ID.name).to_pylist() - vectors = table.column(self._vector_column.name).to_pylist() - for row_id, stored in zip(ids, vectors): - if stored is None: - continue - vector = search_module._to_vector_list(stored) - for query, heap in zip(self._query_vectors, heaps): - search_module._check_vector_dimension(query, vector) - search_module._offer_score(heap, self._limit, row_id, - search_module._compute_score(query, vector, metric)) - return [search_module._scored_result(heap) for heap in heaps] - - -def run(args): - with open(os.path.join(args.warehouse, 'benchmark.json')) as stream: - metadata = json.load(stream) - table = CatalogFactory.create({'warehouse': args.warehouse}).get_table('default.vectors') - table = table.copy({'read.batch-size': str(args.batch_size), 'read.parallelism': '1'}) - queries = np.random.RandomState(42).standard_normal( - (args.queries, metadata['dimension'])).astype(np.float32).tolist() - builder = table.new_batch_vector_search_builder().with_vector_column( - 'embedding').with_query_vectors(queries).with_limit(args.top_k).with_option('metric', 'l2') - stats = dict(metadata, mode=args.mode, queries=args.queries, top_k=args.top_k, - batch_size=args.batch_size, raw_plans=0, delivered_rows=0, - delivered_batches=0, source_passes=0, - before_peak_rss_mib=peak_rss_mib(), python=platform.python_version(), - pyarrow=pa.__version__, numpy=np.__version__, platform=platform.platform()) - original_plan = BatchVectorSearchReadImpl._plan_raw_read - original_generator = TableRead._arrow_batch_generator - - def plan(reader, *a, **kw): - stats['raw_plans'] += 1 - return original_plan(reader, *a, **kw) - - def batches(reader, *a, **kw): - stats['source_passes'] += 1 - source = original_generator(reader, *a, **kw) - try: - for batch in source: - stats['delivered_rows'] += batch.num_rows - stats['delivered_batches'] += 1 - yield batch - finally: - source.close() - - with ExitStack() as stack: - stack.enter_context(patch.object(BatchVectorSearchReadImpl, '_plan_raw_read', plan)) - stack.enter_context(patch.object(TableRead, '_arrow_batch_generator', batches)) - if args.mode != 'shared-stream': - stack.enter_context(patch.object( - BatchVectorSearchReadImpl, '_read_raw_batch_search', - repeated if args.mode == 'repeated' else shared_table)) - start = time.perf_counter() - results = builder.execute_batch_local() - stats['seconds'] = time.perf_counter() - start - stats['peak_rss_mib'] = peak_rss_mib() - expected_passes = args.queries if args.mode == 'repeated' else 1 - assert stats['raw_plans'] == expected_passes, stats - assert stats['source_passes'] == expected_passes, stats - assert stats['delivered_rows'] == metadata['rows'] * expected_passes, stats - scores = [] - for result in results: - getter = result.score_getter() - scores.append([[int(row_id), float(getter(row_id))] for row_id in sorted(result.results())]) - assert len(scores) == args.queries - assert all(len(score) == min(args.top_k, metadata['rows']) for score in scores) - stats['results'] = scores - stats['result_sha256'] = hashlib.sha256(json.dumps(scores).encode()).hexdigest() - output = json.dumps(stats, sort_keys=True) - if args.output: - with open(args.output, 'w') as stream: - stream.write(output + '\n') - print(output) - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument('action', choices=['prepare', 'run']) - parser.add_argument('--warehouse', required=True) - parser.add_argument('--mode', choices=['repeated', 'shared-table', 'shared-stream'], - default='shared-stream') - parser.add_argument('--rows', type=int, default=16384) - parser.add_argument('--dimension', type=int, default=128) - parser.add_argument('--queries', type=int, default=8) - parser.add_argument('--top-k', type=int, default=10) - parser.add_argument('--batch-size', type=int, default=1024) - parser.add_argument('--output') - args = parser.parse_args() - if min(args.rows, args.dimension, args.queries, args.top_k, args.batch_size) <= 0: - parser.error('Numeric arguments must be positive') - (prepare if args.action == 'prepare' else run)(args) - - -if __name__ == '__main__': - main()