From bd4036283f654a5be6600de5908e15287f714833 Mon Sep 17 00:00:00 2001 From: chaoyang Date: Sat, 12 Sep 2026 16:09:15 +0800 Subject: [PATCH 1/2] [python] Read native vector index ranges with bounded concurrency --- paimon-python/README.md | 29 +++ .../pypaimon/benchmark/vindex_io_bench.py | 209 ++++++++++++++++ .../vindex_vector_global_index_reader.py | 110 +++++++-- .../pypaimon/tests/vindex_input_test.py | 233 ++++++++++++++++++ 4 files changed, 564 insertions(+), 17 deletions(-) create mode 100644 paimon-python/pypaimon/benchmark/vindex_io_bench.py create mode 100644 paimon-python/pypaimon/tests/vindex_input_test.py diff --git a/paimon-python/README.md b/paimon-python/README.md index f864ff265bf0..706467382e76 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -286,3 +286,32 @@ 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. + + +# Vector index range reads + +Native vector indexes (`ivf-flat`, `ivf-pq`, `ivf-sq`, `ivf-rq`, and `diskann`) +read multiple file ranges concurrently when the input stream supports +thread-safe positional reads. Set the table option `vindex.read.parallelism` +to a positive integer to control the maximum number of concurrent reads per +index reader, including reads from concurrent native query callbacks. + +The default is **4** for remote index paths and **1** for local paths (including +`file://`). Setting it to **1** disables range-level concurrency. Streams that +only support `seek` and `read` remain serialized. Workers are created lazily +and released when the index reader closes; separate readers have separate +budgets. This option controls index I/O, not shard search or native compute +threads. + +A reproducible serial/concurrent comparison is available with `pypaimon[vindex]` +installed: + +```shell +python -m pypaimon.benchmark.vindex_io_bench --output /tmp/vindex-io.json +``` + +The benchmark compares the original serial adapter with parallelism 1/2/4/8, +checks byte-for-byte range results and identical native search row IDs/scores, +and reports P50/P95 latency, read count, bytes read, and peak concurrent reads. +It uses local files with optional injected per-read latency, not a live object +store. Native query timings include reader open, initialization, and close. diff --git a/paimon-python/pypaimon/benchmark/vindex_io_bench.py b/paimon-python/pypaimon/benchmark/vindex_io_bench.py new file mode 100644 index 000000000000..b6da9f37b6d7 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/vindex_io_bench.py @@ -0,0 +1,209 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Compare serial and concurrent vindex I/O with local or injected read latency. + +Run with pypaimon[vindex] installed: + python -m pypaimon.benchmark.vindex_io_bench --output /tmp/vindex-io.json + +Latency is simulated per positional read; this is not an S3 benchmark. +The native search measurement includes reader open, initialization, search and close. +Index construction and correctness assertions are outside the timed region. +""" + +import argparse +from importlib.metadata import version +import json +import os +import platform +import tempfile +import threading +import time +from unittest import mock + +import numpy as np + +from pypaimon.common.file_io import pread +from pypaimon.globalindex.batch_vector_search import BatchVectorSearch +from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta +from pypaimon.globalindex.vindex import vindex_vector_global_index_reader as adapter + + +class MeasuredStream: + def __init__(self, stream, delay): + self.stream = stream + self.delay = delay + self.lock = threading.Lock() + self.calls = self.bytes = self.active = self.peak = 0 + + def read_at(self, length, offset): + with self.lock: + self.calls += 1 + self.active += 1 + self.peak = max(self.peak, self.active) + try: + if self.delay: + time.sleep(self.delay) + data = pread(self.stream, length, offset) + with self.lock: + self.bytes += len(data) + return data + finally: + with self.lock: + self.active -= 1 + + def close(self): + self.stream.close() + + +class SerialInput: + """Original position-read path, used as the baseline.""" + def __init__(self, stream, parallelism=1): + self.stream = stream + + def pread_many(self, ranges): + return [pread(self.stream, length, offset) for offset, length in ranges] + + def close(self): + pass + + +class MeasuredFileIO: + def __init__(self, delay): + self.delay = delay + self.streams = [] + + def new_input_stream(self, path): + stream = MeasuredStream(open(path, "rb"), self.delay) + self.streams.append(stream) + return stream + + +def summary(times, streams): + return { + "p50_ms": float(np.percentile(times, 50) * 1000), + "p95_ms": float(np.percentile(times, 95) * 1000), + "reads_per_iteration": sum(s.calls for s in streams) / len(times), + "bytes_per_iteration": sum(s.bytes for s in streams) / len(times), + "peak_concurrent_reads": max(s.peak for s in streams), + } + + +def signature(results): + return [sorted((row_id, result.score_getter()(row_id)) + for row_id in result.results().to_list()) for result in results] + + +def run(args, directory): + from paimon_vindex import VectorIndexTrainer, VectorIndexWriter + + rng = np.random.default_rng(42) + vectors = rng.standard_normal((args.rows, args.dimension)).astype(np.float32) + queries = rng.standard_normal((args.batch_size, args.dimension)).astype(np.float32) + path = os.path.join(directory, "index") + options = {"index.type": args.index_type, "metric": "l2"} + if args.index_type.startswith("ivf_"): + options["nlist"] = "64" + with VectorIndexTrainer.train(options, vectors) as training: + with VectorIndexWriter(training) as writer: + writer.add_vectors(np.arange(args.rows, dtype=np.int64), vectors) + with open(path, "wb") as output: + writer.write(output) + with open(path, "rb") as stream: + payload = stream.read() + ranges = [(i * 4096, 4096) for i in range(args.range_count)] + if len(payload) < args.range_count * 4096: + raise ValueError("Index too small for requested microbenchmark ranges") + expected_chunks = [payload[o:o + n] for o, n in ranges] + query = BatchVectorSearch(vectors=queries.tolist(), limit=10, field_name="embedding", + options=({"diskann.l_search": "100"} if args.index_type == "diskann" + else {"ivf.nprobe": "16"})) + records = [] + input_class = adapter.PaimonVindexInput + for delay_ms in args.latency_ms: + # DiskANN can choose a different read plan from the header-read latency. + expected = None + for parallelism in [0] + args.parallelism: + cls = SerialInput if parallelism == 0 else input_class + label = "baseline" if parallelism == 0 else str(parallelism) + delay = delay_ms / 1000 + stream = MeasuredStream(open(path, "rb"), delay) + input_ = cls(stream, max(1, parallelism)) + try: + # Warm the reusable executor; native timings below include cold startup. + assert input_.pread_many(ranges) == expected_chunks + stream.calls = stream.bytes = stream.peak = 0 + times = [] + for _ in range(args.iterations): + start = time.perf_counter() + chunks = input_.pread_many(ranges) + times.append(time.perf_counter() - start) + assert chunks == expected_chunks + micro = summary(times, [stream]) + finally: + input_.close() + stream.close() + file_io = MeasuredFileIO(delay) + times = [] + with mock.patch.object(adapter, "PaimonVindexInput", cls): + for _ in range(args.iterations): + reader = adapter.VindexVectorGlobalIndexReader( + file_io, directory, + [GlobalIndexIOMeta(file_name="index", file_size=len(payload))], + options={"vindex.read.parallelism": str(max(1, parallelism))}) + start = time.perf_counter() + try: + results = reader.visit_batch_vector_search(query).result() + finally: + reader.close() + times.append(time.perf_counter() - start) + actual = signature(results) + if expected is None: + expected = actual + assert actual == expected, "Native row IDs or scores changed" + record = {"latency_ms": delay_ms, "parallelism": label, + "ranges": micro, "native_search": summary(times, file_io.streams)} + records.append(record) + print(json.dumps(record), flush=True) + return records + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--index-type", default="ivf_flat", + choices=["ivf_flat", "ivf_pq", "ivf_sq", "ivf_rq", "diskann"]) + parser.add_argument("--iterations", type=int, default=20) + parser.add_argument("--latency-ms", type=float, nargs="+", default=[0, 2, 10]) + parser.add_argument("--parallelism", type=int, nargs="+", default=[1, 2, 4, 8]) + parser.add_argument("--range-count", type=int, default=32) + parser.add_argument("--rows", type=int, default=16384) + parser.add_argument("--dimension", type=int, default=64) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--output", required=True) + args = parser.parse_args() + with tempfile.TemporaryDirectory(prefix="paimon-vindex-io-") as directory: + records = run(args, directory) + report = {"platform": platform.platform(), "python": platform.python_version(), + "paimon_vindex": version("paimon-vindex"), "cpu_count": os.cpu_count(), + "rayon_num_threads": os.environ.get("RAYON_NUM_THREADS"), + "parameters": vars(args), "records": records} + with open(args.output, "w") as output: + json.dump(report, output, indent=2) + + +if __name__ == "__main__": + main() diff --git a/paimon-python/pypaimon/globalindex/vindex/vindex_vector_global_index_reader.py b/paimon-python/pypaimon/globalindex/vindex/vindex_vector_global_index_reader.py index d2576e2f2e05..3a7efb7d830e 100644 --- a/paimon-python/pypaimon/globalindex/vindex/vindex_vector_global_index_reader.py +++ b/paimon-python/pypaimon/globalindex/vindex/vindex_vector_global_index_reader.py @@ -19,6 +19,8 @@ import os import threading +from concurrent.futures import ThreadPoolExecutor, wait +from urllib.parse import urlparse import numpy as np @@ -30,26 +32,69 @@ NPROBE_PARAMETER = "ivf.nprobe" L_SEARCH_PARAMETER = "diskann.l_search" +READ_PARALLELISM_PARAMETER = "vindex.read.parallelism" class PaimonVindexInput: """Input adapter required by paimon_vindex.VectorIndexReader.""" - def __init__(self, stream): + def __init__(self, stream, parallelism=1): + if isinstance(parallelism, bool) or not isinstance(parallelism, int) or parallelism < 1: + raise ValueError("Vector index read parallelism must be a positive integer") self._stream = stream + self._parallelism = parallelism + self._read_slots = threading.BoundedSemaphore(parallelism) self._supports_pread = supports_pread(stream) self._lock = threading.Lock() + self._executor = None + self._closed = False def pread_many(self, ranges): - if self._supports_pread: - return [pread(self._stream, length, offset) for offset, length in ranges] - - chunks = [] + ranges = list(ranges) with self._lock: + if self._closed: + raise ValueError("Vector index input is closed") + if not self._supports_pread: + chunks = [] + for offset, length in ranges: + self._stream.seek(offset) + chunks.append(self._stream.read(length)) + return chunks + if self._parallelism > 1 and len(ranges) > 1: + if self._executor is None: + self._executor = ThreadPoolExecutor( + max_workers=self._parallelism, + thread_name_prefix="paimon-vindex-io") + executor = self._executor + + if self._parallelism == 1: + with self._read_slots: + return [pread(self._stream, length, offset) for offset, length in ranges] + if executor is None or len(ranges) <= 1: + return [self._pread(offset, length) for offset, length in ranges] + + futures = [] + try: for offset, length in ranges: - self._stream.seek(offset) - chunks.append(self._stream.read(length)) - return chunks + futures.append(executor.submit(self._pread, offset, length)) + return [future.result() for future in futures] + finally: + # A failed range must not leave reads using a stream the caller may close. + wait(futures) + + def _pread(self, offset, length): + # Native query workers may also issue single-range callbacks concurrently. + with self._read_slots: + return pread(self._stream, length, offset) + + def close(self): + """Release workers; the owner remains responsible for closing the stream.""" + with self._lock: + self._closed = True + executor = self._executor + self._executor = None + if executor is not None: + executor.shutdown(wait=True) class VindexVectorGlobalIndexReader(GlobalIndexReader): @@ -165,10 +210,12 @@ def _ensure_loaded(self): file_path = (self._io_meta.external_path if self._io_meta.external_path else os.path.join(self._index_path, self._io_meta.file_name)) + parallelism = _read_parallelism(self._options, file_path) stream = self._file_io.new_input_stream(file_path) reader = None + index_input = None try: - index_input = PaimonVindexInput(stream) + index_input = PaimonVindexInput(stream, parallelism) reader = VectorIndexReader(index_input) self._metadata = reader.metadata() reader.optimize_for_search() @@ -177,9 +224,15 @@ def _ensure_loaded(self): self._search_params_type = SearchParams self._stream = stream except Exception: - if reader is not None: - reader.close() - stream.close() + try: + if reader is not None: + reader.close() + finally: + try: + if index_input is not None: + index_input.close() + finally: + stream.close() raise def __enter__(self): @@ -190,12 +243,35 @@ def __exit__(self, exc_type, exc_val, exc_tb): return False def close(self): - if self._reader is not None: - self._reader.close() + try: + if self._reader is not None: + self._reader.close() + finally: self._reader = None - if self._stream is not None: - self._stream.close() - self._stream = None + try: + if self._index_input is not None: + self._index_input.close() + finally: + self._index_input = None + if self._stream is not None: + self._stream.close() + self._stream = None + + +def _read_parallelism(options, file_path): + value = options.get(READ_PARALLELISM_PARAMETER) + if value is None: + # Avoid thread scheduling overhead for local files, including Windows paths. + scheme = urlparse(file_path).scheme + return 1 if scheme in ("", "file") or len(scheme) == 1 else 4 + try: + parallelism = int(str(value)) + except (ValueError, TypeError): + parallelism = 0 + if parallelism < 1: + raise ValueError("'%s' must be a positive integer, got: %s" + % (READ_PARALLELISM_PARAMETER, value)) + return parallelism def _filter_bytes(include_row_ids): diff --git a/paimon-python/pypaimon/tests/vindex_input_test.py b/paimon-python/pypaimon/tests/vindex_input_test.py new file mode 100644 index 000000000000..892231462d4d --- /dev/null +++ b/paimon-python/pypaimon/tests/vindex_input_test.py @@ -0,0 +1,233 @@ +# 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 io +import tempfile +import threading +import unittest +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError +from unittest import mock + +from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta +from pypaimon.globalindex.vindex.vindex_vector_global_index_reader import ( + PaimonVindexInput, + VindexVectorGlobalIndexReader, + _read_parallelism, +) + + +class VindexInputTest(unittest.TestCase): + + def test_positional_reads_preserve_order_and_cursor(self): + data = bytes(range(256)) * 100 + ranges = [(100, 20), (0, 12), (100, 20), (250, 30), (len(data) - 2, 10), (0, 0)] + with tempfile.TemporaryFile() as stream: + stream.write(data) + stream.flush() + stream.seek(7) + input_ = PaimonVindexInput(stream, parallelism=4) + try: + self.assertEqual([data[o:o + n] for o, n in ranges], input_.pread_many(ranges)) + self.assertEqual(7, stream.tell()) + self.assertEqual([], input_.pread_many([])) + finally: + input_.close() + self.assertFalse(stream.closed) + input_.close() + with self.assertRaisesRegex(ValueError, "closed"): + input_.pread_many([(0, 1)]) + + def test_concurrent_callbacks_share_worker_limit(self): + entered = threading.Event() + release = threading.Event() + lock = threading.Lock() + active = 0 + peak = 0 + + class Stream: + def read_at(self, length, offset): + nonlocal active, peak + with lock: + active += 1 + peak = max(peak, active) + if active == 2: + entered.set() + try: + if not release.wait(5): + raise TimeoutError("Readers were not released") + return bytes([offset]) * length + finally: + with lock: + active -= 1 + + input_ = PaimonVindexInput(Stream(), parallelism=2) + with ThreadPoolExecutor(2) as callers: + try: + first = callers.submit(input_.pread_many, [(3, 2), (1, 4), (2, 1)]) + second = callers.submit(input_.pread_many, [(4, 1)]) + self.assertTrue(entered.wait(5), "Position reads did not overlap") + release.set() + self.assertEqual([b"\x03" * 2, b"\x01" * 4, b"\x02"], first.result(5)) + self.assertEqual([b"\x04"], second.result(5)) + self.assertEqual(2, peak) + finally: + release.set() + input_.close() + + def test_failure_waits_for_other_reads_before_returning(self): + started = threading.Event() + release = threading.Event() + finished = threading.Event() + error = OSError("range read failed") + + class Stream: + def read_at(self, length, offset): + if offset == 0: + if not started.wait(5): + raise TimeoutError("Second range did not start") + raise error + started.set() + if not release.wait(5): + raise TimeoutError("Second range was not released") + finished.set() + return b"x" + + input_ = PaimonVindexInput(Stream(), parallelism=2) + with ThreadPoolExecutor(1) as caller: + try: + result = caller.submit(input_.pread_many, [(0, 1), (1, 1)]) + self.assertTrue(started.wait(5)) + with self.assertRaises(FutureTimeoutError): + result.result(timeout=0.05) + release.set() + with self.assertRaises(OSError) as raised: + result.result(5) + self.assertIs(error, raised.exception) + self.assertTrue(finished.is_set()) + finally: + release.set() + input_.close() + + def test_seek_read_fallback_is_serial_across_callbacks(self): + class Stream(io.BytesIO): + def __init__(self): + super().__init__(b"abcdefgh") + self.guard = threading.Lock() + + def seek(self, offset): + if not self.guard.acquire(blocking=False): + raise AssertionError("Concurrent seek/read") + return super().seek(offset) + + def read(self, length): + try: + return super().read(length) + finally: + self.guard.release() + + input_ = PaimonVindexInput(Stream(), parallelism=4) + try: + with ThreadPoolExecutor(4) as pool: + results = list(pool.map(input_.pread_many, [[(3, 2), (0, 3)]] * 20)) + self.assertEqual([[b"de", b"abc"]] * 20, results) + self.assertIsNone(input_._executor) + finally: + input_.close() + + def test_serial_and_single_range_reads_do_not_start_workers(self): + stream = mock.Mock(spec=["read_at"]) + stream.read_at.return_value = b"x" + for parallelism, ranges in ((1, [(0, 1), (1, 1)]), (4, [(0, 1)]), (4, [])): + input_ = PaimonVindexInput(stream, parallelism) + try: + self.assertEqual([b"x"] * len(ranges), input_.pread_many(ranges)) + self.assertIsNone(input_._executor) + finally: + input_.close() + + def test_parallelism_defaults_and_validation(self): + for path in ("/tmp/index", "file:///tmp/index", "C:/index"): + self.assertEqual(1, _read_parallelism({}, path)) + for path in ("s3://bucket/index", "hdfs://host/index", "oss://bucket/index"): + self.assertEqual(4, _read_parallelism({}, path)) + for value in (1, "2", 8): + self.assertEqual(int(value), _read_parallelism({"vindex.read.parallelism": value}, "x")) + for value in (0, -1, "invalid", "1.5", 1.5, True): + with self.assertRaisesRegex(ValueError, "positive integer"): + _read_parallelism({"vindex.read.parallelism": value}, "x") + + def test_reader_releases_workers_and_stream_on_open_failure(self): + self._check_reader_cleanup("initialize") + + def test_reader_releases_workers_and_stream_on_close(self): + self._check_reader_cleanup("success") + + def test_reader_releases_workers_when_native_constructor_fails(self): + self._check_reader_cleanup("constructor") + + def test_reader_releases_workers_when_native_close_fails(self): + self._check_reader_cleanup("close") + + def _check_reader_cleanup(self, phase): + stream = mock.Mock(spec=["read_at", "close"]) + stream.read_at.return_value = b"x" + io_ = mock.Mock() + io_.new_input_stream.return_value = stream + native = mock.Mock() + inputs = [] + workers = [] + + def open_reader(input_): + inputs.append(input_) + self.assertEqual([b"x", b"x"], input_.pread_many([(0, 1), (1, 1)])) + workers.extend(input_._executor._threads) + if phase == "constructor": + raise error + return native + + error = OSError("native reader failed") + if phase == "initialize": + native.optimize_for_search.side_effect = error + if phase == "close": + native.close.side_effect = error + module = mock.Mock() + module.VectorIndexReader.side_effect = open_reader + with mock.patch.dict("sys.modules", {"paimon_vindex": module}): + reader = VindexVectorGlobalIndexReader( + io_, "s3://bucket", [GlobalIndexIOMeta(file_name="index", file_size=2)]) + if phase in ("constructor", "initialize"): + with self.assertRaises(OSError) as raised: + reader._ensure_loaded() + self.assertIs(error, raised.exception) + else: + reader._ensure_loaded() + if phase == "close": + with self.assertRaises(OSError) as raised: + reader.close() + self.assertIs(error, raised.exception) + else: + reader.close() + reader.close() + if phase == "constructor": + native.close.assert_not_called() + else: + native.close.assert_called_once_with() + stream.close.assert_called_once_with() + self.assertTrue(workers) + self.assertTrue(all(not worker.is_alive() for worker in workers)) + with self.assertRaisesRegex(ValueError, "closed"): + inputs[0].pread_many([(0, 1)]) From 87d8ef7d7dea7266891f3e97b4a4ca151c8b28bd 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 | 13 -- .../pypaimon/benchmark/vindex_io_bench.py | 209 ------------------ 2 files changed, 222 deletions(-) delete mode 100644 paimon-python/pypaimon/benchmark/vindex_io_bench.py diff --git a/paimon-python/README.md b/paimon-python/README.md index 706467382e76..ed96878ce459 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -302,16 +302,3 @@ only support `seek` and `read` remain serialized. Workers are created lazily and released when the index reader closes; separate readers have separate budgets. This option controls index I/O, not shard search or native compute threads. - -A reproducible serial/concurrent comparison is available with `pypaimon[vindex]` -installed: - -```shell -python -m pypaimon.benchmark.vindex_io_bench --output /tmp/vindex-io.json -``` - -The benchmark compares the original serial adapter with parallelism 1/2/4/8, -checks byte-for-byte range results and identical native search row IDs/scores, -and reports P50/P95 latency, read count, bytes read, and peak concurrent reads. -It uses local files with optional injected per-read latency, not a live object -store. Native query timings include reader open, initialization, and close. diff --git a/paimon-python/pypaimon/benchmark/vindex_io_bench.py b/paimon-python/pypaimon/benchmark/vindex_io_bench.py deleted file mode 100644 index b6da9f37b6d7..000000000000 --- a/paimon-python/pypaimon/benchmark/vindex_io_bench.py +++ /dev/null @@ -1,209 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -"""Compare serial and concurrent vindex I/O with local or injected read latency. - -Run with pypaimon[vindex] installed: - python -m pypaimon.benchmark.vindex_io_bench --output /tmp/vindex-io.json - -Latency is simulated per positional read; this is not an S3 benchmark. -The native search measurement includes reader open, initialization, search and close. -Index construction and correctness assertions are outside the timed region. -""" - -import argparse -from importlib.metadata import version -import json -import os -import platform -import tempfile -import threading -import time -from unittest import mock - -import numpy as np - -from pypaimon.common.file_io import pread -from pypaimon.globalindex.batch_vector_search import BatchVectorSearch -from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta -from pypaimon.globalindex.vindex import vindex_vector_global_index_reader as adapter - - -class MeasuredStream: - def __init__(self, stream, delay): - self.stream = stream - self.delay = delay - self.lock = threading.Lock() - self.calls = self.bytes = self.active = self.peak = 0 - - def read_at(self, length, offset): - with self.lock: - self.calls += 1 - self.active += 1 - self.peak = max(self.peak, self.active) - try: - if self.delay: - time.sleep(self.delay) - data = pread(self.stream, length, offset) - with self.lock: - self.bytes += len(data) - return data - finally: - with self.lock: - self.active -= 1 - - def close(self): - self.stream.close() - - -class SerialInput: - """Original position-read path, used as the baseline.""" - def __init__(self, stream, parallelism=1): - self.stream = stream - - def pread_many(self, ranges): - return [pread(self.stream, length, offset) for offset, length in ranges] - - def close(self): - pass - - -class MeasuredFileIO: - def __init__(self, delay): - self.delay = delay - self.streams = [] - - def new_input_stream(self, path): - stream = MeasuredStream(open(path, "rb"), self.delay) - self.streams.append(stream) - return stream - - -def summary(times, streams): - return { - "p50_ms": float(np.percentile(times, 50) * 1000), - "p95_ms": float(np.percentile(times, 95) * 1000), - "reads_per_iteration": sum(s.calls for s in streams) / len(times), - "bytes_per_iteration": sum(s.bytes for s in streams) / len(times), - "peak_concurrent_reads": max(s.peak for s in streams), - } - - -def signature(results): - return [sorted((row_id, result.score_getter()(row_id)) - for row_id in result.results().to_list()) for result in results] - - -def run(args, directory): - from paimon_vindex import VectorIndexTrainer, VectorIndexWriter - - rng = np.random.default_rng(42) - vectors = rng.standard_normal((args.rows, args.dimension)).astype(np.float32) - queries = rng.standard_normal((args.batch_size, args.dimension)).astype(np.float32) - path = os.path.join(directory, "index") - options = {"index.type": args.index_type, "metric": "l2"} - if args.index_type.startswith("ivf_"): - options["nlist"] = "64" - with VectorIndexTrainer.train(options, vectors) as training: - with VectorIndexWriter(training) as writer: - writer.add_vectors(np.arange(args.rows, dtype=np.int64), vectors) - with open(path, "wb") as output: - writer.write(output) - with open(path, "rb") as stream: - payload = stream.read() - ranges = [(i * 4096, 4096) for i in range(args.range_count)] - if len(payload) < args.range_count * 4096: - raise ValueError("Index too small for requested microbenchmark ranges") - expected_chunks = [payload[o:o + n] for o, n in ranges] - query = BatchVectorSearch(vectors=queries.tolist(), limit=10, field_name="embedding", - options=({"diskann.l_search": "100"} if args.index_type == "diskann" - else {"ivf.nprobe": "16"})) - records = [] - input_class = adapter.PaimonVindexInput - for delay_ms in args.latency_ms: - # DiskANN can choose a different read plan from the header-read latency. - expected = None - for parallelism in [0] + args.parallelism: - cls = SerialInput if parallelism == 0 else input_class - label = "baseline" if parallelism == 0 else str(parallelism) - delay = delay_ms / 1000 - stream = MeasuredStream(open(path, "rb"), delay) - input_ = cls(stream, max(1, parallelism)) - try: - # Warm the reusable executor; native timings below include cold startup. - assert input_.pread_many(ranges) == expected_chunks - stream.calls = stream.bytes = stream.peak = 0 - times = [] - for _ in range(args.iterations): - start = time.perf_counter() - chunks = input_.pread_many(ranges) - times.append(time.perf_counter() - start) - assert chunks == expected_chunks - micro = summary(times, [stream]) - finally: - input_.close() - stream.close() - file_io = MeasuredFileIO(delay) - times = [] - with mock.patch.object(adapter, "PaimonVindexInput", cls): - for _ in range(args.iterations): - reader = adapter.VindexVectorGlobalIndexReader( - file_io, directory, - [GlobalIndexIOMeta(file_name="index", file_size=len(payload))], - options={"vindex.read.parallelism": str(max(1, parallelism))}) - start = time.perf_counter() - try: - results = reader.visit_batch_vector_search(query).result() - finally: - reader.close() - times.append(time.perf_counter() - start) - actual = signature(results) - if expected is None: - expected = actual - assert actual == expected, "Native row IDs or scores changed" - record = {"latency_ms": delay_ms, "parallelism": label, - "ranges": micro, "native_search": summary(times, file_io.streams)} - records.append(record) - print(json.dumps(record), flush=True) - return records - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--index-type", default="ivf_flat", - choices=["ivf_flat", "ivf_pq", "ivf_sq", "ivf_rq", "diskann"]) - parser.add_argument("--iterations", type=int, default=20) - parser.add_argument("--latency-ms", type=float, nargs="+", default=[0, 2, 10]) - parser.add_argument("--parallelism", type=int, nargs="+", default=[1, 2, 4, 8]) - parser.add_argument("--range-count", type=int, default=32) - parser.add_argument("--rows", type=int, default=16384) - parser.add_argument("--dimension", type=int, default=64) - parser.add_argument("--batch-size", type=int, default=1) - parser.add_argument("--output", required=True) - args = parser.parse_args() - with tempfile.TemporaryDirectory(prefix="paimon-vindex-io-") as directory: - records = run(args, directory) - report = {"platform": platform.platform(), "python": platform.python_version(), - "paimon_vindex": version("paimon-vindex"), "cpu_count": os.cpu_count(), - "rayon_num_threads": os.environ.get("RAYON_NUM_THREADS"), - "parameters": vars(args), "records": records} - with open(args.output, "w") as output: - json.dump(report, output, indent=2) - - -if __name__ == "__main__": - main()