From be9cf53f8f5748f5703c6befa947d135f9c61b49 Mon Sep 17 00:00:00 2001 From: chaoyang Date: Sat, 12 Sep 2026 12:23:38 +0800 Subject: [PATCH 1/2] [python] Batch vector writes when building global indexes --- .../dev/benchmark_vindex_batch_write.py | 268 ++++++++++++++++++ .../globalindex/create_global_index.py | 38 ++- .../vindex/vindex_vector_index_writer.py | 66 +++++ .../pypaimon/tests/vindex_batch_write_test.py | 195 +++++++++++++ 4 files changed, 560 insertions(+), 7 deletions(-) create mode 100644 paimon-python/dev/benchmark_vindex_batch_write.py create mode 100644 paimon-python/pypaimon/tests/vindex_batch_write_test.py diff --git a/paimon-python/dev/benchmark_vindex_batch_write.py b/paimon-python/dev/benchmark_vindex_batch_write.py new file mode 100644 index 000000000000..8d782f595b0a --- /dev/null +++ b/paimon-python/dev/benchmark_vindex_batch_write.py @@ -0,0 +1,268 @@ +#!/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 vector batch conversion/writing and native builds in fresh processes. + +Run from paimon-python with the project dependencies installed:: + + PYTHONPATH=. python dev/benchmark_vindex_batch_write.py prepare \ + --warehouse /tmp/paimon-batch-bench --rows 65536 --dimension 256 + PYTHONPATH=. python dev/benchmark_vindex_batch_write.py run \ + --warehouse /tmp/paimon-batch-bench --mode batch --batch-size 1024 + +Repeat run in separate processes for baseline, scalar-batches, convert-only, batch. +All modes retain whole-shard Arrow reading and the original training sampler. +scalar-batches converts one batch at a time to Python lists; convert-only uses +columnar conversion but writes each row separately; batch is the production path. +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.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 scalar_build(per_batch): + # Both reference modes retain the whole Arrow table through native finish, + # just like master. Only the lifetime of Python row lists differs. + 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): + table = table_read.to_arrow([split]) + if table is None or table.num_rows == 0: + continue + writer = self._create_generic_index_writer(index_path, field) + try: + batches = table.to_batches(max_chunksize=build_module.ADD_BATCH_SIZE) \ + if per_batch else [table] + for batch in batches: + for value, row_id in build_module._extract_index_rows( + batch, self._index_columns[0], SpecialFields.ROW_ID.name, + row_range): + writer.write(value, row_id - row_range.from_) + adds = build_module._to_index_manifest_entries( + self._table, split.partition, row_range, field.id, + self._index_type, writer.finish()) + finally: + writer.close() + if adds: + messages.append(CommitMessage( + partition=tuple(split.partition.values), bucket=0, + new_files=[], index_adds=adds)) + return messages + return build + + +class RowWrites: + """Ablate batching of file writes while preserving columnar conversion. + + These are buffered Python file writes, not a count of OS write syscalls. + """ + + def __init__(self, stream, row_bytes): + self.stream = stream + self.row_bytes = row_bytes + + def __getattr__(self, name): + return getattr(self.stream, name) + + def write(self, data): + data = memoryview(data).cast('B') + written = 0 + for start in range(0, len(data), self.row_bytes): + written += self.stream.write(data[start:start + self.row_bytes]) + return written + + +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: + stack.enter_context(patch.object(build_module, 'ADD_BATCH_SIZE', args.batch_size)) + if args.mode in ('baseline', 'scalar-batches'): + stack.enter_context(patch.object( + GlobalIndexBuilder, '_build_generic_index', + scalar_build(args.mode == 'scalar-batches'))) + if args.mode == 'convert-only': + original_ensure = VindexVectorIndexWriter._ensure_temp_files + + def ensure(writer): + original_ensure(writer) + if not isinstance(writer._row_id_temp, RowWrites): + writer._row_id_temp = RowWrites(writer._row_id_temp, 8) + writer._vector_temp = RowWrites(writer._vector_temp, dimension * 4) + + stack.enter_context(patch.object(VindexVectorIndexWriter, '_ensure_temp_files', ensure)) + 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', 'scalar-batches', 'convert-only', 'batch'], + default='batch') + 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..b2e125e42159 100644 --- a/paimon-python/pypaimon/globalindex/create_global_index.py +++ b/paimon-python/pypaimon/globalindex/create_global_index.py @@ -21,6 +21,7 @@ from typing import Dict, List, Optional, Sequence, Union import pyarrow as pa +import pyarrow.compute as pc from pypaimon.common.options.core_options import CoreOptions from pypaimon.common.options.options import Options @@ -51,6 +52,7 @@ VINDEX_IDENTIFIERS, ) from pypaimon.globalindex.vindex.vindex_vector_index_writer import ( + ADD_BATCH_SIZE, VindexVectorIndexWriter, ) from pypaimon.index.index_file_meta import IndexFileMeta @@ -302,13 +304,20 @@ def _build_generic_index( writer = self._create_generic_index_writer(index_path, index_field) 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_) + if self._index_type in VINDEX_IDENTIFIERS: + if table.column(SpecialFields.ROW_ID.name).null_count: + raise ValueError("Cannot build global index because _ROW_ID is null.") + for batch in table.to_batches(max_chunksize=ADD_BATCH_SIZE): + _write_vector_batch( + writer, batch, self._index_columns[0], index_range) + else: + 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_) index_adds = _to_index_manifest_entries( self._table, @@ -444,6 +453,21 @@ def compare(left, right): return sorted(rows, key=cmp_to_key(compare)) +def _write_vector_batch(writer, batch, index_column, row_range): + row_ids = batch.column(SpecialFields.ROW_ID.name) + if row_ids.null_count: + raise ValueError("Cannot build global index because _ROW_ID is null.") + vectors = batch.column(index_column) + selected = pc.and_( + pc.greater_equal(row_ids, row_range.from_), + pc.less_equal(row_ids, row_range.to), + ) + if not pc.all(selected).as_py(): + row_ids = pc.filter(row_ids, selected) + vectors = pc.filter(vectors, selected) + writer.write_batch(vectors, pc.subtract(row_ids, row_range.from_)) + + def _extract_index_rows( table: pa.Table, index_column: str, diff --git a/paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py b/paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py index aa56a8ab5975..87c334fa1b2f 100644 --- a/paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py +++ b/paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py @@ -86,6 +86,52 @@ def write(self, vector, relative_row_id: int) -> None: self._vector_temp.write(array("f", materialized).tobytes()) self._vector_count += 1 + def write_batch(self, vectors, relative_row_ids) -> None: + """Write Arrow arrays without materializing valid float32 vectors as lists. + + Unsupported layouts and invalid vectors use the scalar path so that + validation errors and the order of successfully written rows match write(). + """ + if self._closed: + raise RuntimeError("VindexVectorIndexWriter is already closed.") + if len(vectors) != len(relative_row_ids): + raise ValueError("Vector and row ID batch lengths differ.") + if relative_row_ids.null_count: + raise ValueError("Cannot build global index because _ROW_ID is null.") + if len(vectors) == 0: + return + + import numpy as np + import pyarrow as pa + import pyarrow.compute as pc + + valid_vectors, valid_ids = vectors, relative_row_ids + if vectors.null_count: + valid = pc.is_valid(vectors) + valid_vectors = pc.filter(vectors, valid) + valid_ids = pc.filter(relative_row_ids, valid) + if len(valid_vectors) == 0: + self._row_count += len(vectors) + return + + values = _float32_batch_values(np, pa, valid_vectors, self._dimension) + if (values is not None and values.null_count == 0 + and valid_ids.type == pa.int64()): + data = values.to_numpy(zero_copy_only=True) + if np.isfinite(data).all(): + ids = np.ascontiguousarray( + valid_ids.to_numpy(zero_copy_only=False), dtype=np.int64) + data = np.ascontiguousarray(data, dtype=np.float32) + self._row_count += len(vectors) + self._ensure_temp_files() + self._row_id_temp.write(memoryview(ids).cast('B')) + self._vector_temp.write(memoryview(data).cast('B')) + self._vector_count += len(valid_vectors) + return + + for vector, row_id in zip(vectors.to_pylist(), relative_row_ids.to_pylist()): + self.write(vector, row_id) + def finish(self) -> List[ResultEntry]: if self._closed: raise RuntimeError("VindexVectorIndexWriter is already closed.") @@ -355,6 +401,26 @@ def _sample_training_vectors(np, vectors, sample_ratio: float): return np.ascontiguousarray(vectors[indexes]) +def _float32_batch_values(np, pa, vectors, dimension): + vector_type = vectors.type + if not ( + pa.types.is_list(vector_type) + or pa.types.is_large_list(vector_type) + or pa.types.is_fixed_size_list(vector_type) + ) or vector_type.value_type != pa.float32(): + return None + + if pa.types.is_fixed_size_list(vector_type): + if vector_type.list_size != dimension: + return None + return vectors.values.slice(vectors.offset * dimension, len(vectors) * dimension) + + offsets = vectors.offsets.to_numpy(zero_copy_only=True) + if not np.all(np.diff(offsets) == dimension): + return None + return vectors.values.slice(int(offsets[0]), int(offsets[-1] - offsets[0])) + + def _materialize_vector( value, dimension: int, relative_row_id: int ) -> List[float]: diff --git a/paimon-python/pypaimon/tests/vindex_batch_write_test.py b/paimon-python/pypaimon/tests/vindex_batch_write_test.py new file mode 100644 index 000000000000..72da3bb74531 --- /dev/null +++ b/paimon-python/pypaimon/tests/vindex_batch_write_test.py @@ -0,0 +1,195 @@ +# 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 os +import unittest +from unittest.mock import Mock, patch + +import numpy as np +import pyarrow as pa + +from pypaimon.globalindex.create_global_index import GlobalIndexBuilder, _write_vector_batch +from pypaimon.globalindex.vindex.vindex_vector_index_writer import VindexVectorIndexWriter +from pypaimon.schema.data_types import ArrayType, AtomicType +from pypaimon.utils.range import Range + + +class VindexBatchWriteTest(unittest.TestCase): + + def _writer(self): + writer = VindexVectorIndexWriter( + Mock(), '/unused', ArrayType(True, AtomicType('FLOAT')), + 'ivf-flat', {'ivf-flat.dimension': '2'}, 'embedding') + self.addCleanup(writer.close) + return writer + + def _contents(self, writer): + writer._close_temp_files() + contents = [] + for path in (writer._row_id_temp_path, writer._vector_temp_path): + if path is None: + contents.append(b'') + else: + with open(path, 'rb') as stream: + contents.append(stream.read()) + return writer._row_count, writer._vector_count, contents + + def _assert_parity(self, vectors, row_ids, fast=True): + scalar, batch = self._writer(), self._writer() + for vector, row_id in zip(vectors.to_pylist(), row_ids.to_pylist()): + scalar.write(vector, row_id) + with patch.object(batch, 'write', wraps=batch.write) as write: + batch.write_batch(vectors, row_ids) + if fast: + write.assert_not_called() + self.assertEqual(self._contents(scalar), self._contents(batch)) + + def test_sliced_list_large_list_and_fixed_size_list(self): + for array_type in (pa.list_(pa.float32()), pa.large_list(pa.float32()), + pa.list_(pa.float32(), 2)): + with self.subTest(array_type=array_type): + vectors = pa.array([ + [99, 99], [1, 2], None, [-0.0, 1e-40], [3, 4], [88, 88], + ], type=array_type).slice(1, 4) + row_ids = pa.array([77, 0, 3, 8, 2 ** 62, 99], type=pa.int64()).slice(1, 4) + self._assert_parity(vectors, row_ids) + + def test_null_parent_ignores_invalid_child_values(self): + values = pa.array([1, 2, float('nan'), None, 3, 4], type=pa.float32()) + vectors = pa.Array.from_buffers(pa.list_(pa.float32()), 3, [ + pa.py_buffer(b'\x05'), + pa.py_buffer(np.array([0, 2, 4, 6], dtype=np.int32)), + ], children=[values]) + self._assert_parity(vectors, pa.array([0, 1, 2], type=pa.int64())) + + def test_empty_and_all_null_batches_do_not_create_files(self): + for data in ([], [None, None]): + with self.subTest(data=data): + writer = self._writer() + writer.write_batch(pa.array(data, type=pa.list_(pa.float32())), + pa.array(range(len(data)), type=pa.int64())) + self.assertEqual((len(data), 0, [b'', b'']), self._contents(writer)) + + def test_multiple_batches_and_scalar_writes_can_be_interleaved(self): + writer = self._writer() + writer.write([1, 2], 0) + writer.write_batch(pa.array([[3, 4], None], type=pa.list_(pa.float32())), + pa.array([2, 3], type=pa.int64())) + writer.write([5, 6], 4) + writer.write_batch(pa.array([[7, 8]], type=pa.list_(pa.float32(), 2)), + pa.array([6], type=pa.int64())) + count, valid_count, contents = self._contents(writer) + self.assertEqual((5, 4), (count, valid_count)) + self.assertEqual([0, 2, 4, 6], np.frombuffer(contents[0], dtype=np.int64).tolist()) + self.assertEqual(list(range(1, 9)), np.frombuffer(contents[1], dtype=np.float32).tolist()) + + def test_invalid_vectors_preserve_scalar_error_and_written_prefix(self): + cases = [ + [[1, 2], [3]], + [[1, 2], [None, 4]], + [[1, 2], [float('nan'), 4]], + [[1, 2], [3, float('inf')]], + [[1, 2], [3, float('-inf')]], + [[float('nan'), 2], [3]], + [[1, 2], [float('nan'), None]], + [[1, 2], [None, float('nan')]], + ] + for data in cases: + with self.subTest(data=data): + vectors = pa.array(data, type=pa.list_(pa.float32())) + ids = pa.array([5, 9], type=pa.int64()) + scalar, batch = self._writer(), self._writer() + with self.assertRaises(ValueError) as old_error: + for vector, row_id in zip(vectors.to_pylist(), ids.to_pylist()): + scalar.write(vector, row_id) + with self.assertRaises(ValueError) as new_error: + batch.write_batch(vectors, ids) + self.assertEqual(str(old_error.exception), str(new_error.exception)) + self.assertEqual(self._contents(scalar), self._contents(batch)) + + def test_float64_and_non_int64_ids_fall_back_to_scalar(self): + self._assert_parity( + pa.array([[1.1, 2.2], None], type=pa.list_(pa.float64())), + pa.array([0, 1], type=pa.int64()), fast=False) + self._assert_parity( + pa.array([[1, 2], [3, 4]], type=pa.list_(pa.float32())), + pa.array([0, 1], type=pa.int32()), fast=False) + + def test_batch_length_and_null_row_id_validation(self): + vectors = pa.array([[1, 2]], type=pa.list_(pa.float32())) + writer = self._writer() + with self.assertRaisesRegex(ValueError, 'batch lengths differ'): + writer.write_batch(vectors, pa.array([], type=pa.int64())) + with self.assertRaisesRegex(ValueError, '_ROW_ID is null'): + writer.write_batch(vectors, pa.array([None], type=pa.int64())) + self.assertEqual((0, 0, [b'', b'']), self._contents(writer)) + + def test_close_removes_batch_files_and_rejects_further_writes(self): + writer = self._writer() + vectors = pa.array([[1, 2]], type=pa.list_(pa.float32())) + ids = pa.array([0], type=pa.int64()) + writer.write_batch(vectors, ids) + paths = writer._row_id_temp_path, writer._vector_temp_path + writer.close() + self.assertTrue(all(not os.path.exists(path) for path in paths)) + with self.assertRaisesRegex(RuntimeError, 'already closed'): + writer.write_batch(vectors, ids) + + def test_builder_filters_ranges_before_vector_validation(self): + vectors = pa.array([[float('nan'), 0], [1, 2], None, [3, 4], [5]], + type=pa.list_(pa.float32())) + ids = pa.array([9, 10, 11, 19, 20], type=pa.int64()) + batch = pa.RecordBatch.from_arrays([vectors, ids], ['embedding', '_ROW_ID']) + writer = self._writer() + _write_vector_batch(writer, batch, 'embedding', Range(10, 19)) + count, valid_count, contents = self._contents(writer) + self.assertEqual((3, 2), (count, valid_count)) + self.assertEqual([0, 9], np.frombuffer(contents[0], dtype=np.int64).tolist()) + self.assertEqual([1, 2, 3, 4], np.frombuffer(contents[1], dtype=np.float32).tolist()) + batch = pa.RecordBatch.from_arrays([ + pa.array([[1, 2]], type=pa.list_(pa.float32())), + pa.array([None], type=pa.int64()), + ], ['embedding', '_ROW_ID']) + with self.assertRaisesRegex(ValueError, '_ROW_ID is null'): + _write_vector_batch(self._writer(), batch, 'embedding', Range(10, 19)) + + def test_null_row_ids_are_rejected_before_writing_any_batch(self): + builder = object.__new__(GlobalIndexBuilder) + builder._core_options = Mock() + builder._core_options.global_index_row_count_per_shard.return_value = 10 + builder._index_type = 'ivf-flat' + builder._index_columns = ['embedding'] + writer = Mock() + builder._create_generic_index_writer = Mock(return_value=writer) + read = Mock() + read.to_arrow.return_value = pa.table({ + 'embedding': pa.array([[1], [3, 4]], type=pa.list_(pa.float32())), + '_ROW_ID': pa.array([0, None], type=pa.int64()), + }) + module = 'pypaimon.globalindex.create_global_index' + with patch(module + '._split_by_global_index_shard', return_value=[ + (Mock(), Range(0, 9)), + ]), patch(module + '.ADD_BATCH_SIZE', 1): + with self.assertRaisesRegex(ValueError, '_ROW_ID is null'): + builder._build_generic_index([], [], Mock(), read, '/unused') + writer.write_batch.assert_not_called() + writer.finish.assert_not_called() + writer.close.assert_called_once() + + +if __name__ == '__main__': + unittest.main() From 10c9664c9852c758c1b66a6f8e445270c147373a 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_vindex_batch_write.py | 268 ------------------ 1 file changed, 268 deletions(-) delete mode 100644 paimon-python/dev/benchmark_vindex_batch_write.py diff --git a/paimon-python/dev/benchmark_vindex_batch_write.py b/paimon-python/dev/benchmark_vindex_batch_write.py deleted file mode 100644 index 8d782f595b0a..000000000000 --- a/paimon-python/dev/benchmark_vindex_batch_write.py +++ /dev/null @@ -1,268 +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 vector batch conversion/writing and native builds in fresh processes. - -Run from paimon-python with the project dependencies installed:: - - PYTHONPATH=. python dev/benchmark_vindex_batch_write.py prepare \ - --warehouse /tmp/paimon-batch-bench --rows 65536 --dimension 256 - PYTHONPATH=. python dev/benchmark_vindex_batch_write.py run \ - --warehouse /tmp/paimon-batch-bench --mode batch --batch-size 1024 - -Repeat run in separate processes for baseline, scalar-batches, convert-only, batch. -All modes retain whole-shard Arrow reading and the original training sampler. -scalar-batches converts one batch at a time to Python lists; convert-only uses -columnar conversion but writes each row separately; batch is the production path. -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.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 scalar_build(per_batch): - # Both reference modes retain the whole Arrow table through native finish, - # just like master. Only the lifetime of Python row lists differs. - 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): - table = table_read.to_arrow([split]) - if table is None or table.num_rows == 0: - continue - writer = self._create_generic_index_writer(index_path, field) - try: - batches = table.to_batches(max_chunksize=build_module.ADD_BATCH_SIZE) \ - if per_batch else [table] - for batch in batches: - for value, row_id in build_module._extract_index_rows( - batch, self._index_columns[0], SpecialFields.ROW_ID.name, - row_range): - writer.write(value, row_id - row_range.from_) - adds = build_module._to_index_manifest_entries( - self._table, split.partition, row_range, field.id, - self._index_type, writer.finish()) - finally: - writer.close() - if adds: - messages.append(CommitMessage( - partition=tuple(split.partition.values), bucket=0, - new_files=[], index_adds=adds)) - return messages - return build - - -class RowWrites: - """Ablate batching of file writes while preserving columnar conversion. - - These are buffered Python file writes, not a count of OS write syscalls. - """ - - def __init__(self, stream, row_bytes): - self.stream = stream - self.row_bytes = row_bytes - - def __getattr__(self, name): - return getattr(self.stream, name) - - def write(self, data): - data = memoryview(data).cast('B') - written = 0 - for start in range(0, len(data), self.row_bytes): - written += self.stream.write(data[start:start + self.row_bytes]) - return written - - -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: - stack.enter_context(patch.object(build_module, 'ADD_BATCH_SIZE', args.batch_size)) - if args.mode in ('baseline', 'scalar-batches'): - stack.enter_context(patch.object( - GlobalIndexBuilder, '_build_generic_index', - scalar_build(args.mode == 'scalar-batches'))) - if args.mode == 'convert-only': - original_ensure = VindexVectorIndexWriter._ensure_temp_files - - def ensure(writer): - original_ensure(writer) - if not isinstance(writer._row_id_temp, RowWrites): - writer._row_id_temp = RowWrites(writer._row_id_temp, 8) - writer._vector_temp = RowWrites(writer._vector_temp, dimension * 4) - - stack.enter_context(patch.object(VindexVectorIndexWriter, '_ensure_temp_files', ensure)) - 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', 'scalar-batches', 'convert-only', 'batch'], - default='batch') - 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()