This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 38b3c1637d [python] Read native vector index ranges with bounded
concurrency (#9756)
38b3c1637d is described below
commit 38b3c1637dd95450e07752b9fb9dc6df74a0418a
Author: chaoyang <[email protected]>
AuthorDate: Sun Sep 13 16:25:31 2026 +0800
[python] Read native vector index ranges with bounded concurrency (#9756)
---
paimon-python/README.md | 16 ++
.../vindex/vindex_vector_global_index_reader.py | 110 ++++++++--
paimon-python/pypaimon/tests/vindex_input_test.py | 233 +++++++++++++++++++++
3 files changed, 342 insertions(+), 17 deletions(-)
diff --git a/paimon-python/README.md b/paimon-python/README.md
index f864ff265b..ed96878ce4 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 d2576e2f2e..3a7efb7d83 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 @@ VINDEX_IDENTIFIERS = ("ivf-flat", "ivf-pq", "ivf-sq",
"ivf-rq", "diskann")
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 @@ class VindexVectorGlobalIndexReader(GlobalIndexReader):
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 @@ class VindexVectorGlobalIndexReader(GlobalIndexReader):
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 @@ class VindexVectorGlobalIndexReader(GlobalIndexReader):
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 0000000000..892231462d
--- /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)])
