[python] Read native vector index ranges with bounded concurrency (#9756)

Author: TheR1sing3un

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.

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):