diff --git a/paimon-python/README.md b/paimon-python/README.md index f864ff265bf0..ed96878ce459 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -286,3 +286,19 @@ 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. 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)])