Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions paimon-python/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@

import os
import threading
from concurrent.futures import ThreadPoolExecutor, wait
from urllib.parse import urlparse

import numpy as np

Expand All @@ -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):
Expand Down Expand Up @@ -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()
Expand All @@ -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):
Expand All @@ -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):
Expand Down
Loading
Loading