From 50d1ec2bb919ec703e80f2613bbb6b9bf91325ee Mon Sep 17 00:00:00 2001 From: chaoyang Date: Sat, 12 Sep 2026 16:30:28 +0800 Subject: [PATCH 1/2] [python] Stream training samples into native vector index trainers --- paimon-python/README.md | 21 +++ .../benchmark/vindex_training_bench.py | 156 ++++++++++++++++++ .../vindex/vindex_vector_index_writer.py | 53 +++--- .../pypaimon/tests/global_index_build_test.py | 30 +++- .../pypaimon/tests/vindex_training_test.py | 115 +++++++++++++ 5 files changed, 352 insertions(+), 23 deletions(-) create mode 100644 paimon-python/pypaimon/benchmark/vindex_training_bench.py create mode 100644 paimon-python/pypaimon/tests/vindex_training_test.py diff --git a/paimon-python/README.md b/paimon-python/README.md index f864ff265bf0..89f6c817a363 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -286,3 +286,24 @@ unsupported platform such as Windows), `pypaimon` automatically falls back to the `pyarrow` (`libhdfs`/JVM) path and logs a warning. Disable the fallback with `hdfs.client.fallback-to-pyarrow=false` if you want hard failures instead. + + +# Native vector index training + +The native vector index writer submits training vectors in bounded batches. +`.train.sample-ratio` (or its field-level override) still selects +the same evenly spaced non-null vectors in the same order. Native training +receives the final corpus size for automatic IVF sizing. This bounds Python +training buffers; native training and index construction have their own +memory requirements. + +With `pypaimon[vindex]` installed, compare complete writer builds and training +buffer strategies using: + +```shell +python -m pypaimon.benchmark.vindex_training_bench --output /tmp/training.json +``` + +Each variant runs in a fresh process and reports peak RSS, ingestion time, +finish time, and an index-file checksum. The benchmark checks identical +index bytes for each sampling ratio across all variants. diff --git a/paimon-python/pypaimon/benchmark/vindex_training_bench.py b/paimon-python/pypaimon/benchmark/vindex_training_bench.py new file mode 100644 index 000000000000..0107fe5159c8 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/vindex_training_bench.py @@ -0,0 +1,156 @@ +# 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. + +"""Benchmark complete vector writer builds in isolated processes. + +python -m pypaimon.benchmark.vindex_training_bench --output /tmp/training.json +Requires pypaimon[vindex]. Compares full-file one-shot training, bounded reads +into a complete sample matrix, and streaming native training. Source reads +and ingestion are bounded and identical in all variants. Reports process +peak RSS, ingestion/finish time and an index digest for result equivalence. +""" + +import argparse +import hashlib +import json +import math +import os +import platform +import resource +import subprocess +import sys +import tempfile +import time +from importlib.metadata import version + +import numpy as np + +from pypaimon.filesystem.local_file_io import LocalFileIO +from pypaimon.globalindex.vindex.vindex_vector_index_writer import ( + VindexVectorIndexWriter, _iter_training_batches, +) +from pypaimon.schema.data_types import ArrayType, AtomicType + + +def one_shot(writer, np, trainer_type, bounded): + count = max(1, math.ceil(writer._vector_count * writer._train_sample_ratio)) + if bounded: + sample = np.empty((count, writer._dimension), dtype=np.float32) + position = 0 + with open(writer._vector_temp_path, "rb") as stream: + for batch in _iter_training_batches( + np, stream, writer._vector_count, writer._dimension, + writer._train_sample_ratio, + ): + sample[position:position + len(batch)] = batch + position += len(batch) + else: + vectors = np.fromfile(writer._vector_temp_path, dtype=np.float32).reshape( + writer._vector_count, writer._dimension) + sample = (vectors if count == len(vectors) else + vectors[np.arange(count) * len(vectors) // count]) + return trainer_type.train(writer._training_options(), sample) + + +def worker(args): + with tempfile.TemporaryDirectory(prefix="paimon-stream-train-") as directory: + options = {args.index_type + ".dimension": str(args.dimension), + args.index_type + ".train.sample-ratio": str(args.ratio)} + if args.index_type != "diskann": + options[args.index_type + ".nlist"] = "64" + writer = VindexVectorIndexWriter( + LocalFileIO(), directory, ArrayType(True, AtomicType("FLOAT")), + args.index_type, options, "embedding") + if args.mode != "streaming": + writer._train = lambda np, trainer: one_shot( + writer, np, trainer, args.mode == "sample-matrix") + start = time.perf_counter() + try: + row_id = 0 + with open(args.source, "rb") as stream: + while True: + batch = np.fromfile(stream, dtype=np.float32, count=10000 * args.dimension) + if not batch.size: + break + for vector in batch.reshape(-1, args.dimension): + writer.write(vector, row_id) + row_id += 1 + ingestion_s = time.perf_counter() - start + start_finish = time.perf_counter() + writer.finish() + finish_s = time.perf_counter() - start_finish + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + peak_mib = peak / (1024 ** 2 if sys.platform == "darwin" else 1024) + digest = hashlib.sha256() + with open(writer._file_path(), "rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + finally: + writer.close() + return {"mode": args.mode, "ratio": args.ratio, "ingestion_s": ingestion_s, + "finish_s": finish_s, "total_s": ingestion_s + finish_s, + "peak_rss_mib": peak_mib, "sha256": digest.hexdigest()} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, default=200000) + parser.add_argument("--dimension", type=int, default=128) + parser.add_argument("--index-type", default="ivf-flat") + parser.add_argument("--ratios", type=float, nargs="+", default=[1.0, 0.1]) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--output") + parser.add_argument("--source", help=argparse.SUPPRESS) + parser.add_argument("--mode", help=argparse.SUPPRESS) + parser.add_argument("--ratio", type=float, help=argparse.SUPPRESS) + args = parser.parse_args() + if args.mode: + print(json.dumps(worker(args))) + return + if not args.output: + parser.error("--output is required") + records = [] + with tempfile.TemporaryDirectory(prefix="paimon-train-source-") as directory: + source = os.path.join(directory, "vectors") + rng = np.random.default_rng(42) + with open(source, "wb") as stream: + for start in range(0, args.rows, 10000): + rng.standard_normal((min(10000, args.rows - start), args.dimension)).astype( + np.float32).tofile(stream) + for ratio in args.ratios: + expected = None + for _ in range(args.repeats): + for mode in ("baseline", "sample-matrix", "streaming"): + process = subprocess.run( + [sys.executable, "-m", "pypaimon.benchmark.vindex_training_bench", + "--source", source, "--mode", mode, "--ratio", str(ratio), + "--dimension", str(args.dimension), "--index-type", args.index_type], + check=True, capture_output=True, text=True) + record = json.loads(process.stdout) + if expected is None: + expected = record["sha256"] + assert record["sha256"] == expected, "Index bytes changed" + records.append(record) + print(json.dumps(record), flush=True) + with open(args.output, "w") as output: + json.dump({"platform": platform.platform(), "python": platform.python_version(), + "paimon_vindex": version("paimon-vindex"), "parameters": vars(args), + "records": records}, output, indent=2) + + +if __name__ == "__main__": + main() 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..346ab4061ac4 100644 --- a/paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py +++ b/paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py @@ -106,18 +106,8 @@ def finish(self) -> List[ResultEntry]: self._close_temp_files() self._file_io.check_or_mkdirs(self._index_path) - vectors = np.fromfile( - self._vector_temp_path, - dtype=np.float32, - count=self._vector_count * self._dimension, - ).reshape(self._vector_count, self._dimension) - training_vectors = _sample_training_vectors( - np, vectors, self._train_sample_ratio) - training = VectorIndexTrainer.train( - self._training_options(), training_vectors) + training = self._train(np, VectorIndexTrainer) try: - del training_vectors - del vectors with VectorIndexWriter(training) as writer: self._add_vectors_in_batches(np, writer) with self._file_io.new_output_stream(file_path) as output_stream: @@ -132,6 +122,16 @@ def finish(self) -> List[ResultEntry]: return [ResultEntry(self.file_name, self._row_count, b"{}")] + def _train(self, np, trainer_type): + with open(self._vector_temp_path, "rb") as vector_file: + with trainer_type.create(self._training_options()) as trainer: + for batch in _iter_training_batches( + np, vector_file, self._vector_count, self._dimension, + self._train_sample_ratio, batch_size=ADD_BATCH_SIZE, + ): + trainer.add_training_vectors(batch) + return trainer.finish_training() + def _file_path(self) -> str: return "%s/%s" % (self._index_path, self.file_name) @@ -343,16 +343,31 @@ def _is_float_type(data_type: DataType) -> bool: ) -def _sample_training_vectors(np, vectors, sample_ratio: float): - vector_count = vectors.shape[0] +def _iter_training_batches( + np, vector_file, vector_count: int, dimension: int, sample_ratio: float, + batch_size: int = ADD_BATCH_SIZE, +): + """Yield the existing evenly spaced sample using bounded reads and buffers.""" train_count = max(1, min(vector_count, int(math.ceil( vector_count * sample_ratio)))) - if train_count == vector_count: - return vectors - indexes = ( - np.arange(train_count, dtype=np.int64) * vector_count // train_count - ) - return np.ascontiguousarray(vectors[indexes]) + position = 0 + item_size = np.dtype(np.float32).itemsize + while position < train_count: + start = position * vector_count // train_count + end = min(start + batch_size, vector_count) + # First sample position whose source row is at or beyond this block. + next_position = min(train_count, (end * train_count + vector_count - 1) // vector_count) + vector_file.seek(start * dimension * item_size) + vectors = np.fromfile( + vector_file, dtype=np.float32, count=(end - start) * dimension, + ).reshape(end - start, dimension) + if train_count == vector_count: + yield vectors + else: + indexes = np.arange(position, next_position, dtype=np.int64) + indexes = indexes * vector_count // train_count - start + yield np.ascontiguousarray(vectors[indexes]) + position = next_position def _materialize_vector( diff --git a/paimon-python/pypaimon/tests/global_index_build_test.py b/paimon-python/pypaimon/tests/global_index_build_test.py index 186270fcb79a..f7c7563a2209 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 +import tempfile import pyarrow as pa @@ -41,7 +42,7 @@ ) from pypaimon.globalindex.vindex.vindex_vector_index_writer import ( VindexVectorIndexWriter, - _sample_training_vectors, + _iter_training_batches, native_options, train_sample_ratio, ) @@ -93,9 +94,26 @@ def close(self): class _FakeVectorIndexTrainer: + def __init__(self, options): + self.options = options + self.batches = [] + @classmethod - def train(cls, options, data): - return _FakeVectorIndexTraining(options, data) + def create(cls, options): + return cls(options) + + def add_training_vectors(self, data): + self.batches.append(data.copy()) + + def finish_training(self): + import numpy as np + return _FakeVectorIndexTraining(self.options, np.concatenate(self.batches)) + + def __enter__(self): + return self + + def __exit__(self, *args): + pass class _FakeVectorIndexWriter: @@ -1028,7 +1046,11 @@ def test_vindex_training_sample_ratio(self): import numpy as np vectors = np.arange(20, dtype=np.float32).reshape(10, 2) - sampled = _sample_training_vectors(np, vectors, 0.4) + with tempfile.TemporaryFile() as vector_file: + vectors.tofile(vector_file) + vector_file.flush() + sampled = np.concatenate(list(_iter_training_batches( + np, vector_file, 10, 2, 0.4, batch_size=3))) self.assertEqual( [[0.0, 1.0], [4.0, 5.0], [10.0, 11.0], [14.0, 15.0]], sampled.tolist(), diff --git a/paimon-python/pypaimon/tests/vindex_training_test.py b/paimon-python/pypaimon/tests/vindex_training_test.py new file mode 100644 index 000000000000..a238e3029660 --- /dev/null +++ b/paimon-python/pypaimon/tests/vindex_training_test.py @@ -0,0 +1,115 @@ +# 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 math +import os +import tempfile +import unittest +from unittest import mock + +import numpy as np + +from pypaimon.filesystem.local_file_io import LocalFileIO +from pypaimon.globalindex.vindex.vindex_vector_index_writer import ( + VindexVectorIndexWriter, _iter_training_batches, +) +from pypaimon.schema.data_types import ArrayType, AtomicType + + +class VindexTrainingTest(unittest.TestCase): + + def test_batches_preserve_sample_positions_and_bound_reads(self): + vectors = np.arange(10003 * 3, dtype=np.float32).reshape(-1, 3) + with tempfile.TemporaryFile() as stream: + vectors.tofile(stream) + stream.flush() + for ratio in (1.0, 0.999, 0.37, 0.01, 1e-8): + for batch_size in (1, 17, 1000): + with self.subTest(ratio=ratio, batch_size=batch_size): + with mock.patch.object(np, "fromfile", wraps=np.fromfile) as read: + batches = list(_iter_training_batches( + np, stream, len(vectors), 3, ratio, batch_size)) + count = max(1, math.ceil(len(vectors) * ratio)) + indexes = np.arange(count) * len(vectors) // count + np.testing.assert_array_equal(vectors[indexes], np.concatenate(batches)) + self.assertTrue(all(b.flags.c_contiguous for b in batches)) + self.assertTrue(all(len(b) <= batch_size for b in batches)) + self.assertTrue(all(c[1]["count"] <= batch_size * 3 + for c in read.call_args_list)) + + def test_training_failure_closes_trainer_and_removes_temp_files(self): + with tempfile.TemporaryDirectory() as directory: + writer = self._writer(directory, {}) + writer.write([1.0] * 8, 0) + paths = [writer._vector_temp_path, writer._row_id_temp_path] + for phase in ("add_training_vectors", "finish_training"): + trainer = mock.MagicMock() + trainer.__enter__.return_value = trainer + getattr(trainer, phase).side_effect = RuntimeError("training failed") + module = mock.Mock() + module.VectorIndexTrainer.create.return_value = trainer + with mock.patch.dict("sys.modules", {"paimon_vindex": module}): + with self.assertRaisesRegex(RuntimeError, "training failed"): + writer.finish() + trainer.__exit__.assert_called_once() + self.assertTrue(all(not os.path.exists(path) for path in paths)) + self.assertFalse(os.path.exists(writer._file_path())) + writer = self._writer(directory, {}) + writer.write([1.0] * 8, 0) + paths = [writer._vector_temp_path, writer._row_id_temp_path] + writer.close() + + def test_native_streamed_build_matches_one_shot(self): + try: + from paimon_vindex import VectorIndexTrainer, VectorIndexWriter + except ImportError: + self.skipTest("paimon-vindex is not installed") + vectors = np.random.default_rng(42).standard_normal((2049, 8)).astype(np.float32) + with tempfile.TemporaryDirectory() as directory: + for ratio in (1.0, 0.37): + for index_type in ("ivf-flat", "ivf-pq", "ivf-sq", "ivf-rq", "diskann"): + with self.subTest(ratio=ratio, index_type=index_type): + options = {index_type + ".train.sample-ratio": str(ratio)} + if index_type != "diskann": + options[index_type + ".nlist"] = "16" + writer = self._writer(directory, options, index_type) + writer.write(None, 0) + for i, vector in enumerate(vectors): + writer.write(vector, i + 1) + reference_path = os.path.join(directory, "reference") + count = math.ceil(len(vectors) * ratio) + sample = vectors[np.arange(count) * len(vectors) // count] + with VectorIndexTrainer.train(writer._training_options(), sample) as training: + with VectorIndexWriter(training) as native: + native.add_vectors(np.arange(1, len(vectors) + 1), vectors) + with open(reference_path, "wb") as output: + native.write(output) + with mock.patch( + "pypaimon.globalindex.vindex.vindex_vector_index_writer.ADD_BATCH_SIZE", 127 + ): + result = writer.finish() + self.assertEqual(1, len(result)) + with open(reference_path, "rb") as reference, open(writer._file_path(), "rb") as actual: + self.assertEqual(reference.read(), actual.read()) + + @staticmethod + def _writer(directory, options, index_type="ivf-flat"): + options = dict(options) + options[index_type + ".dimension"] = "8" + return VindexVectorIndexWriter( + LocalFileIO(), directory, ArrayType(True, AtomicType("FLOAT")), + index_type, options, "embedding") From 65568590e57be2f302f4fc6e1c60b6d48b14daa4 Mon Sep 17 00:00:00 2001 From: chaoyang Date: Sun, 13 Sep 2026 09:14:42 +0800 Subject: [PATCH 2/2] [python] Remove benchmark artifacts from vector optimization --- paimon-python/README.md | 11 -- .../benchmark/vindex_training_bench.py | 156 ------------------ 2 files changed, 167 deletions(-) delete mode 100644 paimon-python/pypaimon/benchmark/vindex_training_bench.py diff --git a/paimon-python/README.md b/paimon-python/README.md index 89f6c817a363..e4e22b264b18 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -296,14 +296,3 @@ the same evenly spaced non-null vectors in the same order. Native training receives the final corpus size for automatic IVF sizing. This bounds Python training buffers; native training and index construction have their own memory requirements. - -With `pypaimon[vindex]` installed, compare complete writer builds and training -buffer strategies using: - -```shell -python -m pypaimon.benchmark.vindex_training_bench --output /tmp/training.json -``` - -Each variant runs in a fresh process and reports peak RSS, ingestion time, -finish time, and an index-file checksum. The benchmark checks identical -index bytes for each sampling ratio across all variants. diff --git a/paimon-python/pypaimon/benchmark/vindex_training_bench.py b/paimon-python/pypaimon/benchmark/vindex_training_bench.py deleted file mode 100644 index 0107fe5159c8..000000000000 --- a/paimon-python/pypaimon/benchmark/vindex_training_bench.py +++ /dev/null @@ -1,156 +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. - -"""Benchmark complete vector writer builds in isolated processes. - -python -m pypaimon.benchmark.vindex_training_bench --output /tmp/training.json -Requires pypaimon[vindex]. Compares full-file one-shot training, bounded reads -into a complete sample matrix, and streaming native training. Source reads -and ingestion are bounded and identical in all variants. Reports process -peak RSS, ingestion/finish time and an index digest for result equivalence. -""" - -import argparse -import hashlib -import json -import math -import os -import platform -import resource -import subprocess -import sys -import tempfile -import time -from importlib.metadata import version - -import numpy as np - -from pypaimon.filesystem.local_file_io import LocalFileIO -from pypaimon.globalindex.vindex.vindex_vector_index_writer import ( - VindexVectorIndexWriter, _iter_training_batches, -) -from pypaimon.schema.data_types import ArrayType, AtomicType - - -def one_shot(writer, np, trainer_type, bounded): - count = max(1, math.ceil(writer._vector_count * writer._train_sample_ratio)) - if bounded: - sample = np.empty((count, writer._dimension), dtype=np.float32) - position = 0 - with open(writer._vector_temp_path, "rb") as stream: - for batch in _iter_training_batches( - np, stream, writer._vector_count, writer._dimension, - writer._train_sample_ratio, - ): - sample[position:position + len(batch)] = batch - position += len(batch) - else: - vectors = np.fromfile(writer._vector_temp_path, dtype=np.float32).reshape( - writer._vector_count, writer._dimension) - sample = (vectors if count == len(vectors) else - vectors[np.arange(count) * len(vectors) // count]) - return trainer_type.train(writer._training_options(), sample) - - -def worker(args): - with tempfile.TemporaryDirectory(prefix="paimon-stream-train-") as directory: - options = {args.index_type + ".dimension": str(args.dimension), - args.index_type + ".train.sample-ratio": str(args.ratio)} - if args.index_type != "diskann": - options[args.index_type + ".nlist"] = "64" - writer = VindexVectorIndexWriter( - LocalFileIO(), directory, ArrayType(True, AtomicType("FLOAT")), - args.index_type, options, "embedding") - if args.mode != "streaming": - writer._train = lambda np, trainer: one_shot( - writer, np, trainer, args.mode == "sample-matrix") - start = time.perf_counter() - try: - row_id = 0 - with open(args.source, "rb") as stream: - while True: - batch = np.fromfile(stream, dtype=np.float32, count=10000 * args.dimension) - if not batch.size: - break - for vector in batch.reshape(-1, args.dimension): - writer.write(vector, row_id) - row_id += 1 - ingestion_s = time.perf_counter() - start - start_finish = time.perf_counter() - writer.finish() - finish_s = time.perf_counter() - start_finish - peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss - peak_mib = peak / (1024 ** 2 if sys.platform == "darwin" else 1024) - digest = hashlib.sha256() - with open(writer._file_path(), "rb") as stream: - for block in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(block) - finally: - writer.close() - return {"mode": args.mode, "ratio": args.ratio, "ingestion_s": ingestion_s, - "finish_s": finish_s, "total_s": ingestion_s + finish_s, - "peak_rss_mib": peak_mib, "sha256": digest.hexdigest()} - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--rows", type=int, default=200000) - parser.add_argument("--dimension", type=int, default=128) - parser.add_argument("--index-type", default="ivf-flat") - parser.add_argument("--ratios", type=float, nargs="+", default=[1.0, 0.1]) - parser.add_argument("--repeats", type=int, default=3) - parser.add_argument("--output") - parser.add_argument("--source", help=argparse.SUPPRESS) - parser.add_argument("--mode", help=argparse.SUPPRESS) - parser.add_argument("--ratio", type=float, help=argparse.SUPPRESS) - args = parser.parse_args() - if args.mode: - print(json.dumps(worker(args))) - return - if not args.output: - parser.error("--output is required") - records = [] - with tempfile.TemporaryDirectory(prefix="paimon-train-source-") as directory: - source = os.path.join(directory, "vectors") - rng = np.random.default_rng(42) - with open(source, "wb") as stream: - for start in range(0, args.rows, 10000): - rng.standard_normal((min(10000, args.rows - start), args.dimension)).astype( - np.float32).tofile(stream) - for ratio in args.ratios: - expected = None - for _ in range(args.repeats): - for mode in ("baseline", "sample-matrix", "streaming"): - process = subprocess.run( - [sys.executable, "-m", "pypaimon.benchmark.vindex_training_bench", - "--source", source, "--mode", mode, "--ratio", str(ratio), - "--dimension", str(args.dimension), "--index-type", args.index_type], - check=True, capture_output=True, text=True) - record = json.loads(process.stdout) - if expected is None: - expected = record["sha256"] - assert record["sha256"] == expected, "Index bytes changed" - records.append(record) - print(json.dumps(record), flush=True) - with open(args.output, "w") as output: - json.dump({"platform": platform.platform(), "python": platform.python_version(), - "paimon_vindex": version("paimon-vindex"), "parameters": vars(args), - "records": records}, output, indent=2) - - -if __name__ == "__main__": - main()