From 2f72668ec9d0562385aaeb486f2539eb5e0daa0e Mon Sep 17 00:00:00 2001 From: chaoyang Date: Sat, 12 Sep 2026 12:04:43 +0800 Subject: [PATCH 1/2] [python] Stream batches when building generic global indexes --- .../dev/benchmark_global_index_streaming.py | 247 ++++++++++++++++++ .../globalindex/create_global_index.py | 41 ++- .../pypaimon/tests/global_index_build_test.py | 163 +++++++++++- 3 files changed, 434 insertions(+), 17 deletions(-) create mode 100644 paimon-python/dev/benchmark_global_index_streaming.py diff --git a/paimon-python/dev/benchmark_global_index_streaming.py b/paimon-python/dev/benchmark_global_index_streaming.py new file mode 100644 index 000000000000..c0125a725567 --- /dev/null +++ b/paimon-python/dev/benchmark_global_index_streaming.py @@ -0,0 +1,247 @@ +#!/usr/bin/env python +# 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. + +"""Measure global-index ingestion and native builds in fresh processes. + +Run from paimon-python with the project dependencies installed:: + + PYTHONPATH=. python dev/benchmark_global_index_streaming.py prepare \ + --warehouse /tmp/paimon-stream-bench --rows 65536 --dimension 256 + PYTHONPATH=. python dev/benchmark_global_index_streaming.py run \ + --warehouse /tmp/paimon-stream-bench --mode stream --batch-size 1024 + +Repeat run in separate processes for baseline, python-only, arrow-only, stream. +Use --native to include IVF-Flat training, serialization and query validation +(requires paimon-vindex==0.4.0). Runs build but do not commit index manifests. +Generated index files are removed after validation. Use a dedicated warehouse. +Do not compare dataset creation RSS to run RSS. Filesystem cache is uncontrolled. +""" + +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.globalindex.create_global_index import GlobalIndexBuilder +from pypaimon.globalindex.vindex.vindex_vector_index_writer import VindexVectorIndexWriter +from pypaimon.read.table_read import _ClosableArrowBatchReader +from pypaimon.table.special_fields import SpecialFields +from pypaimon.write.commit_message import CommitMessage + +build_module = importlib.import_module('pypaimon.globalindex.create_global_index') + + +def peak_rss_mib(): + rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return rss / (1024 * 1024 if sys.platform == 'darwin' else 1024) + + +def digest_file(path): + digest = hashlib.sha256() + with open(path, 'rb') as stream: + for block in iter(lambda: stream.read(1024 * 1024), b''): + digest.update(block) + return digest.hexdigest() + + +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', + }), 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) + data = rng.standard_normal((count, args.dimension)).astype(np.float32) + array = pa.ListArray.from_arrays( + np.arange(count + 1, dtype=np.int32) * args.dimension, + pa.array(data.reshape(-1))) + writer.write_arrow(pa.Table.from_arrays([array], 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 ablation_build(mode): + # The baseline preserves the original builder's whole-table lifetime. + # All variants use the same row extraction, writer, shard plan and finish. + def build(self, splits, ranges, field, table_read, index_path): + messages = [] + shard_size = self._core_options.global_index_row_count_per_shard() + for split, row_range in build_module._split_by_global_index_shard( + splits, shard_size, ranges): + writer = None + try: + if mode == 'arrow-only': + reader, batches = table_read._new_arrow_batch_reader([split]) + with _ClosableArrowBatchReader(reader, batches) as reader: + rows = [] + for batch in reader: + rows.extend(build_module._extract_index_rows( + batch, self._index_columns[0], + SpecialFields.ROW_ID.name, row_range)) + del batch + chunks = [rows] + else: + table = table_read.to_arrow([split]) + if table is None or table.num_rows == 0: + continue + batches = table.to_batches() if mode == 'python-only' else [table] + chunks = (build_module._extract_index_rows( + batch, self._index_columns[0], SpecialFields.ROW_ID.name, + row_range) for batch in batches) + writer = self._create_generic_index_writer(index_path, field) + for rows in chunks: + for value, row_id in rows: + writer.write(value, row_id - row_range.from_) + del rows, chunks + adds = build_module._to_index_manifest_entries( + self._table, split.partition, row_range, field.id, + self._index_type, writer.finish()) + finally: + if writer is not None: + writer.close() + if adds: + messages.append(CommitMessage( + partition=tuple(split.partition.values), bucket=0, + new_files=[], index_adds=adds)) + return messages + return build + + +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)}) + dimension = metadata['dimension'] + result = dict(metadata, mode=args.mode, batch_size=args.batch_size, + native=args.native, python=platform.python_version(), + pyarrow=pa.__version__, numpy=np.__version__, platform=platform.platform()) + if args.native: + # Import before timing so native initialization does not skew modes. + from paimon_vindex import VectorIndexReader, SearchParams + from pypaimon.globalindex.vindex.vindex_vector_global_index_reader import PaimonVindexInput + original_finish = VindexVectorIndexWriter.finish + hash_seconds = [0.0] + hashes = [] + + def finish(writer): + assert writer._train_sample_ratio == 0.25 + result['train_sample_ratio'] = writer._train_sample_ratio + writer._close_temp_files() + result['ingest_seconds'] = time.perf_counter() - started + result['ingest_peak_rss_mib'] = peak_rss_mib() + before_hash = time.perf_counter() + hashes.append({ + 'row_ids': digest_file(writer._row_id_temp_path), + 'vectors': digest_file(writer._vector_temp_path), + 'vector_count': writer._vector_count, + }) + hash_seconds[0] += time.perf_counter() - before_hash + return original_finish(writer) if args.native else [] + + builder = GlobalIndexBuilder(table, 'embedding', 'ivf-flat', options={ + 'global-index.row-count-per-shard': str(metadata['rows'] + 1), + 'ivf-flat.dimension': str(dimension), 'ivf-flat.nlist': '16', + 'ivf-flat.train.sample-ratio': '0.25', 'ivf-flat.distance.metric': 'l2', + }) + result['before_build_peak_rss_mib'] = peak_rss_mib() + started = time.perf_counter() + with ExitStack() as stack: + if args.mode != 'stream': + stack.enter_context(patch.object( + GlobalIndexBuilder, '_build_generic_index', ablation_build(args.mode))) + stack.enter_context(patch.object(VindexVectorIndexWriter, 'finish', finish)) + messages = builder.build() + result['build_seconds'] = time.perf_counter() - started - hash_seconds[0] + result['build_peak_rss_mib'] = peak_rss_mib() + result['input_hashes'] = hashes + assert len(hashes) == 1, 'Benchmark requires one index shard' + assert hashes[0]['vector_count'] == metadata['rows'] + entries = [entry for msg in messages for entry in msg.index_adds] + paths = [table.path_factory().global_index_path_factory().to_path( + entry.index_file.file_name) for entry in entries] + try: + if args.native: + assert len(paths) == 1 + queries = np.random.RandomState(42).standard_normal((16, dimension)).astype(np.float32) + scores = {} + with table.file_io.new_input_stream(paths[0]) as stream: + reader = VectorIndexReader(PaimonVindexInput(stream)) + try: + for nprobe in (4, 16): + scores[str(nprobe)] = [] + for query in queries: + ids, distances = reader.search( + query, SearchParams.ivf(top_k=10, nprobe=nprobe)) + scores[str(nprobe)].append([ids.tolist(), distances.tolist()]) + finally: + reader.close() + result['queries'] = scores + result['query_sha256'] = hashlib.sha256(json.dumps(scores).encode()).hexdigest() + finally: + for path in paths: + table.file_io.delete(path) + output = json.dumps(result, 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('--rows', type=int, default=65536) + parser.add_argument('--dimension', type=int, default=256) + parser.add_argument('--batch-size', type=int, default=1024) + parser.add_argument('--mode', choices=['baseline', 'python-only', 'arrow-only', 'stream'], + default='stream') + parser.add_argument('--native', action='store_true') + parser.add_argument('--output') + args = parser.parse_args() + if min(args.rows, args.dimension, args.batch_size) <= 0: + parser.error('rows, dimension and batch-size must be positive') + (prepare if args.action == 'prepare' else run)(args) + + +if __name__ == '__main__': + main() diff --git a/paimon-python/pypaimon/globalindex/create_global_index.py b/paimon-python/pypaimon/globalindex/create_global_index.py index bd9d34a05a7a..e9199e0868ff 100644 --- a/paimon-python/pypaimon/globalindex/create_global_index.py +++ b/paimon-python/pypaimon/globalindex/create_global_index.py @@ -286,6 +286,8 @@ def _create_sorted_index_writer(self, index_path: str, key_serializer): def _build_generic_index( self, splits, unindexed_ranges, index_field, table_read, index_path: str ) -> List[CommitMessage]: + from pypaimon.read.table_read import _ClosableArrowBatchReader + rows_per_shard = self._core_options.global_index_row_count_per_shard() if rows_per_shard <= 0: raise ValueError( @@ -296,19 +298,29 @@ def _build_generic_index( for index_split, index_range in _split_by_global_index_shard( splits, rows_per_shard, unindexed_ranges ): - table = table_read.to_arrow([index_split]) - if table is None or table.num_rows == 0: - continue - - writer = self._create_generic_index_writer(index_path, index_field) + writer = None try: - for value, row_id in _extract_index_rows( - table, - self._index_columns[0], - SpecialFields.ROW_ID.name, - index_range, - ): - writer.write(value, row_id - index_range.from_) + reader, batches = table_read._new_arrow_batch_reader([index_split]) + # Close the Python iterator explicitly on failure as well as + # the Arrow reader, which may retain a suspended generator. + with _ClosableArrowBatchReader(reader, batches) as batch_reader: + for batch in batch_reader: + if batch.num_rows == 0: + continue + if writer is None: + writer = self._create_generic_index_writer( + index_path, index_field) + for value, row_id in _extract_index_rows( + batch, + self._index_columns[0], + SpecialFields.ROW_ID.name, + index_range, + ): + writer.write(value, row_id - index_range.from_) + del batch + + if writer is None: + continue index_adds = _to_index_manifest_entries( self._table, @@ -319,7 +331,8 @@ def _build_generic_index( writer.finish(), ) finally: - writer.close() + if writer is not None: + writer.close() if index_adds: messages.append( CommitMessage( @@ -445,7 +458,7 @@ def compare(left, right): def _extract_index_rows( - table: pa.Table, + table: Union[pa.Table, pa.RecordBatch], index_column: str, row_id_column: str, row_range: Optional[Range] = None, diff --git a/paimon-python/pypaimon/tests/global_index_build_test.py b/paimon-python/pypaimon/tests/global_index_build_test.py index 186270fcb79a..3812a1ed8408 100644 --- a/paimon-python/pypaimon/tests/global_index_build_test.py +++ b/paimon-python/pypaimon/tests/global_index_build_test.py @@ -22,6 +22,7 @@ import struct import sys import types +from unittest.mock import Mock, patch import pyarrow as pa @@ -556,7 +557,8 @@ def test_create_vindex_global_index_from_python(self): ('id', pa.int32()), ('embedding', pa.list_(pa.float32())), ]) - table = self._create_table(pa_schema=schema, options=self.table_options) + table = self._create_table(pa_schema=schema, options=dict( + self.table_options, **{'read.batch-size': '1'})) vectors = pa.array( [[1.0, 0.0], [0.0, 1.0], None], type=pa.list_(pa.float32()), @@ -610,12 +612,56 @@ def test_create_vindex_global_index_from_python(self): table.path_factory().global_index_path_factory().to_path( entry.index_file.file_name))) + def test_create_vindex_streaming_failure_cleans_resources(self): + from pypaimon.read.table_read import TableRead + + schema = pa.schema([('embedding', pa.list_(pa.float32()))]) + table = self._create_table(pa_schema=schema, options=dict( + self.table_options, **{'read.batch-size': '1'})) + self._write_arrow(table, pa.table( + {'embedding': [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]]}, schema=schema)) + snapshot_id = table.snapshot_manager().get_latest_snapshot().id + original_write = VindexVectorIndexWriter.write + original_batches = TableRead._arrow_batch_generator + temp_paths = [] + closed = [] + generators = [] + + def failing_write(writer, vector, row_id): + original_write(writer, vector, row_id) + if row_id == 1: + temp_paths.extend([writer._row_id_temp_path, writer._vector_temp_path]) + raise RuntimeError('injected write failure') + + def tracked_batches(reader, *args): + def generate(): + try: + yield from original_batches(reader, *args) + finally: + closed.append(True) + generator = generate() + generators.append(generator) + return generator + + with patch.object(VindexVectorIndexWriter, 'write', failing_write), \ + patch.object(TableRead, '_arrow_batch_generator', tracked_batches): + with self.assertRaisesRegex(RuntimeError, 'injected write failure'): + table.create_global_index('embedding', index_type='ivf-flat', options={ + 'ivf-flat.dimension': '2', + }) + + self.assertEqual([True], closed) + self.assertEqual(2, len(temp_paths)) + self.assertTrue(all(not os.path.exists(path) for path in temp_paths)) + self.assertEqual(snapshot_id, table.snapshot_manager().get_latest_snapshot().id) + def test_create_vindex_global_index_respects_row_count_per_shard(self): schema = pa.schema([ ('id', pa.int32()), ('embedding', pa.list_(pa.float32())), ]) - table = self._create_table(pa_schema=schema, options=self.table_options) + table = self._create_table(pa_schema=schema, options=dict( + self.table_options, **{'read.batch-size': '1'})) vectors = pa.array( [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5], [0.2, 0.8], [0.9, 0.1]], type=pa.list_(pa.float32()), @@ -774,7 +820,8 @@ def test_create_native_fulltext_global_index_from_python(self): ('id', pa.int32()), ('content', pa.string()), ]) - table = self._create_table(pa_schema=schema, options=self.table_options) + table = self._create_table(pa_schema=schema, options=dict( + self.table_options, **{'read.batch-size': '1'})) self._write_arrow(table, pa.table( { 'id': [1, 2, 3], @@ -1181,5 +1228,115 @@ def test_java_scalar_key_serializers_round_trip(self): self.assertEqual(value, actual) +class GenericIndexStreamingTest(unittest.TestCase): + + schema = pa.schema([ + ('embedding', pa.list_(pa.float32())), + ('_ROW_ID', pa.int64()), + ]) + + def setUp(self): + self.builder = object.__new__(GlobalIndexBuilder) + self.builder._table = Mock() + self.builder._core_options = Mock() + self.builder._core_options.global_index_row_count_per_shard.return_value = 10 + self.builder._index_columns = ['embedding'] + self.builder._index_type = 'ivf-flat' + self.writer = Mock() + self.writer.finish.return_value = [] + self.builder._create_generic_index_writer = Mock(return_value=self.writer) + self.read = Mock() + self.read.to_arrow.side_effect = AssertionError('Must not materialize a shard') + self.events = [] + + def _batch(self, values, row_ids): + return pa.RecordBatch.from_arrays([ + pa.array(values, type=self.schema.field(0).type), + pa.array(row_ids, type=pa.int64()), + ], schema=self.schema) + + def _build(self, batches): + def generate(): + try: + for batch in batches: + self.events.append('read') + yield batch + finally: + self.events.append('reader closed') + + # Retain the generator: cleanup must be explicit, not depend on GC. + self.generator = generate() + reader = pa.RecordBatchReader.from_batches(self.schema, self.generator) + self.read._new_arrow_batch_reader.return_value = reader, self.generator + module = 'pypaimon.globalindex.create_global_index' + with patch(module + '._split_by_global_index_shard', return_value=[ + (_FakeSplit([]), Range(10, 19)), + ]), patch(module + '._to_index_manifest_entries', return_value=[]): + return self.builder._build_generic_index( + [], [], Mock(), self.read, '/unused') + + def test_batches_are_written_before_reading_the_next_batch(self): + written = [] + + def write(value, row_id): + self.events.append('write') + written.append((value, row_id)) + + def finish(): + self.assertEqual('reader closed', self.events[-1]) + return [] + + self.writer.write.side_effect = write + self.writer.finish.side_effect = finish + self._build([ + self._batch([], []), + self._batch([[9.0], [10.0], None], [9, 10, 11]), + self._batch([[19.0], [20.0]], [19, 20]), + ]) + self.assertEqual([([10.0], 0), (None, 1), ([19.0], 9)], written) + self.assertEqual([ + 'read', 'read', 'write', 'write', 'read', 'write', 'reader closed', + ], self.events) + self.writer.finish.assert_called_once() + self.writer.close.assert_called_once() + self.read.to_arrow.assert_not_called() + + def test_empty_input_does_not_create_a_writer(self): + for batches in ([], [self._batch([], [])]): + with self.subTest(batches=len(batches)): + self.assertEqual([], self._build(batches)) + self.builder._create_generic_index_writer.assert_not_called() + self.assertEqual('reader closed', self.events[-1]) + + def test_failures_close_reader_and_writer(self): + for failure in ('create', 'read', 'write', 'finish', 'null_row_id'): + with self.subTest(failure=failure): + self.setUp() + error = RuntimeError('injected failure') + if failure == 'create': + self.builder._create_generic_index_writer.side_effect = error + elif failure in ('write', 'finish'): + getattr(self.writer, failure).side_effect = error + + def batches(): + yield self._batch([[10.0]], [10]) + if failure == 'read': + raise error + yield self._batch([[11.0]], [ + None if failure == 'null_row_id' else 11]) + + exception = ValueError if failure == 'null_row_id' else RuntimeError + message = '_ROW_ID is null' if failure == 'null_row_id' else 'injected failure' + with self.assertRaisesRegex(exception, message): + self._build(batches()) + self.assertEqual('reader closed', self.events[-1]) + if failure == 'create': + self.writer.close.assert_not_called() + else: + self.writer.close.assert_called_once() + if failure != 'finish': + self.writer.finish.assert_not_called() + + if __name__ == "__main__": unittest.main() From 267eab2900d6d39fc69556d0725cab0456af9ba7 Mon Sep 17 00:00:00 2001 From: chaoyang Date: Sun, 13 Sep 2026 09:14:41 +0800 Subject: [PATCH 2/2] [python] Remove benchmark artifacts from vector optimization --- .../dev/benchmark_global_index_streaming.py | 247 ------------------ 1 file changed, 247 deletions(-) delete mode 100644 paimon-python/dev/benchmark_global_index_streaming.py diff --git a/paimon-python/dev/benchmark_global_index_streaming.py b/paimon-python/dev/benchmark_global_index_streaming.py deleted file mode 100644 index c0125a725567..000000000000 --- a/paimon-python/dev/benchmark_global_index_streaming.py +++ /dev/null @@ -1,247 +0,0 @@ -#!/usr/bin/env python -# 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. - -"""Measure global-index ingestion and native builds in fresh processes. - -Run from paimon-python with the project dependencies installed:: - - PYTHONPATH=. python dev/benchmark_global_index_streaming.py prepare \ - --warehouse /tmp/paimon-stream-bench --rows 65536 --dimension 256 - PYTHONPATH=. python dev/benchmark_global_index_streaming.py run \ - --warehouse /tmp/paimon-stream-bench --mode stream --batch-size 1024 - -Repeat run in separate processes for baseline, python-only, arrow-only, stream. -Use --native to include IVF-Flat training, serialization and query validation -(requires paimon-vindex==0.4.0). Runs build but do not commit index manifests. -Generated index files are removed after validation. Use a dedicated warehouse. -Do not compare dataset creation RSS to run RSS. Filesystem cache is uncontrolled. -""" - -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.globalindex.create_global_index import GlobalIndexBuilder -from pypaimon.globalindex.vindex.vindex_vector_index_writer import VindexVectorIndexWriter -from pypaimon.read.table_read import _ClosableArrowBatchReader -from pypaimon.table.special_fields import SpecialFields -from pypaimon.write.commit_message import CommitMessage - -build_module = importlib.import_module('pypaimon.globalindex.create_global_index') - - -def peak_rss_mib(): - rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss - return rss / (1024 * 1024 if sys.platform == 'darwin' else 1024) - - -def digest_file(path): - digest = hashlib.sha256() - with open(path, 'rb') as stream: - for block in iter(lambda: stream.read(1024 * 1024), b''): - digest.update(block) - return digest.hexdigest() - - -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', - }), 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) - data = rng.standard_normal((count, args.dimension)).astype(np.float32) - array = pa.ListArray.from_arrays( - np.arange(count + 1, dtype=np.int32) * args.dimension, - pa.array(data.reshape(-1))) - writer.write_arrow(pa.Table.from_arrays([array], 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 ablation_build(mode): - # The baseline preserves the original builder's whole-table lifetime. - # All variants use the same row extraction, writer, shard plan and finish. - def build(self, splits, ranges, field, table_read, index_path): - messages = [] - shard_size = self._core_options.global_index_row_count_per_shard() - for split, row_range in build_module._split_by_global_index_shard( - splits, shard_size, ranges): - writer = None - try: - if mode == 'arrow-only': - reader, batches = table_read._new_arrow_batch_reader([split]) - with _ClosableArrowBatchReader(reader, batches) as reader: - rows = [] - for batch in reader: - rows.extend(build_module._extract_index_rows( - batch, self._index_columns[0], - SpecialFields.ROW_ID.name, row_range)) - del batch - chunks = [rows] - else: - table = table_read.to_arrow([split]) - if table is None or table.num_rows == 0: - continue - batches = table.to_batches() if mode == 'python-only' else [table] - chunks = (build_module._extract_index_rows( - batch, self._index_columns[0], SpecialFields.ROW_ID.name, - row_range) for batch in batches) - writer = self._create_generic_index_writer(index_path, field) - for rows in chunks: - for value, row_id in rows: - writer.write(value, row_id - row_range.from_) - del rows, chunks - adds = build_module._to_index_manifest_entries( - self._table, split.partition, row_range, field.id, - self._index_type, writer.finish()) - finally: - if writer is not None: - writer.close() - if adds: - messages.append(CommitMessage( - partition=tuple(split.partition.values), bucket=0, - new_files=[], index_adds=adds)) - return messages - return build - - -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)}) - dimension = metadata['dimension'] - result = dict(metadata, mode=args.mode, batch_size=args.batch_size, - native=args.native, python=platform.python_version(), - pyarrow=pa.__version__, numpy=np.__version__, platform=platform.platform()) - if args.native: - # Import before timing so native initialization does not skew modes. - from paimon_vindex import VectorIndexReader, SearchParams - from pypaimon.globalindex.vindex.vindex_vector_global_index_reader import PaimonVindexInput - original_finish = VindexVectorIndexWriter.finish - hash_seconds = [0.0] - hashes = [] - - def finish(writer): - assert writer._train_sample_ratio == 0.25 - result['train_sample_ratio'] = writer._train_sample_ratio - writer._close_temp_files() - result['ingest_seconds'] = time.perf_counter() - started - result['ingest_peak_rss_mib'] = peak_rss_mib() - before_hash = time.perf_counter() - hashes.append({ - 'row_ids': digest_file(writer._row_id_temp_path), - 'vectors': digest_file(writer._vector_temp_path), - 'vector_count': writer._vector_count, - }) - hash_seconds[0] += time.perf_counter() - before_hash - return original_finish(writer) if args.native else [] - - builder = GlobalIndexBuilder(table, 'embedding', 'ivf-flat', options={ - 'global-index.row-count-per-shard': str(metadata['rows'] + 1), - 'ivf-flat.dimension': str(dimension), 'ivf-flat.nlist': '16', - 'ivf-flat.train.sample-ratio': '0.25', 'ivf-flat.distance.metric': 'l2', - }) - result['before_build_peak_rss_mib'] = peak_rss_mib() - started = time.perf_counter() - with ExitStack() as stack: - if args.mode != 'stream': - stack.enter_context(patch.object( - GlobalIndexBuilder, '_build_generic_index', ablation_build(args.mode))) - stack.enter_context(patch.object(VindexVectorIndexWriter, 'finish', finish)) - messages = builder.build() - result['build_seconds'] = time.perf_counter() - started - hash_seconds[0] - result['build_peak_rss_mib'] = peak_rss_mib() - result['input_hashes'] = hashes - assert len(hashes) == 1, 'Benchmark requires one index shard' - assert hashes[0]['vector_count'] == metadata['rows'] - entries = [entry for msg in messages for entry in msg.index_adds] - paths = [table.path_factory().global_index_path_factory().to_path( - entry.index_file.file_name) for entry in entries] - try: - if args.native: - assert len(paths) == 1 - queries = np.random.RandomState(42).standard_normal((16, dimension)).astype(np.float32) - scores = {} - with table.file_io.new_input_stream(paths[0]) as stream: - reader = VectorIndexReader(PaimonVindexInput(stream)) - try: - for nprobe in (4, 16): - scores[str(nprobe)] = [] - for query in queries: - ids, distances = reader.search( - query, SearchParams.ivf(top_k=10, nprobe=nprobe)) - scores[str(nprobe)].append([ids.tolist(), distances.tolist()]) - finally: - reader.close() - result['queries'] = scores - result['query_sha256'] = hashlib.sha256(json.dumps(scores).encode()).hexdigest() - finally: - for path in paths: - table.file_io.delete(path) - output = json.dumps(result, 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('--rows', type=int, default=65536) - parser.add_argument('--dimension', type=int, default=256) - parser.add_argument('--batch-size', type=int, default=1024) - parser.add_argument('--mode', choices=['baseline', 'python-only', 'arrow-only', 'stream'], - default='stream') - parser.add_argument('--native', action='store_true') - parser.add_argument('--output') - args = parser.parse_args() - if min(args.rows, args.dimension, args.batch_size) <= 0: - parser.error('rows, dimension and batch-size must be positive') - (prepare if args.action == 'prepare' else run)(args) - - -if __name__ == '__main__': - main()