| # 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 ctypes |
| import operator |
| import threading |
| from dataclasses import dataclass |
| from enum import IntEnum |
| from typing import Mapping, Optional |
| |
| import numpy as np |
| |
| from . import _ffi |
| from ._ffi import lib |
| |
| _SIZE_T_MAX = ctypes.c_size_t(-1).value |
| _UINT64_MAX = ctypes.c_uint64(-1).value |
| |
| |
| def _size_t(value, name: str, *, allow_zero: bool) -> int: |
| try: |
| value = operator.index(value) |
| except TypeError as exc: |
| raise ValueError(f"{name} must be an integer") from exc |
| lower_bound = 0 if allow_zero else 1 |
| if not lower_bound <= value <= _SIZE_T_MAX: |
| raise ValueError( |
| f"{name} must be in [{lower_bound}, {_SIZE_T_MAX}]" |
| ) |
| return value |
| |
| |
| def _uint64(value, name: str) -> int: |
| try: |
| value = operator.index(value) |
| except TypeError as exc: |
| raise ValueError(f"{name} must be an integer") from exc |
| if not 0 <= value <= _UINT64_MAX: |
| raise ValueError(f"{name} must be in [0, {_UINT64_MAX}]") |
| return value |
| |
| |
| class _NativeHandleLock: |
| """Serialize a native handle while rejecting reentry from its callbacks.""" |
| |
| def __init__(self): |
| self._lock = threading.Lock() |
| self._state_lock = threading.Lock() |
| self._owner = None |
| self._callback_depths = {} |
| |
| def __enter__(self): |
| current = threading.get_ident() |
| with self._state_lock: |
| if ( |
| self._owner == current |
| or self._callback_depths.get(current, 0) > 0 |
| ): |
| raise RuntimeError( |
| "reentrant native-handle operation is not allowed" |
| ) |
| self._lock.acquire() |
| with self._state_lock: |
| self._owner = current |
| return self |
| |
| def __exit__(self, exc_type, exc_val, exc_tb): |
| with self._state_lock: |
| self._owner = None |
| self._lock.release() |
| return False |
| |
| def _enter_callback(self): |
| current = threading.get_ident() |
| with self._state_lock: |
| self._callback_depths[current] = ( |
| self._callback_depths.get(current, 0) + 1 |
| ) |
| |
| def _exit_callback(self): |
| current = threading.get_ident() |
| with self._state_lock: |
| depth = self._callback_depths[current] |
| if depth == 1: |
| del self._callback_depths[current] |
| else: |
| self._callback_depths[current] = depth - 1 |
| |
| INDEX_TYPES = { |
| 0: "ivf_flat", |
| 1: "ivf_pq", |
| 4: "ivf_rq", |
| 5: "diskann", |
| 6: "ivf_sq", |
| } |
| |
| METRICS = { |
| 0: "l2", |
| 1: "inner_product", |
| 2: "cosine", |
| } |
| |
| |
| class SearchWidth(IntEnum): |
| AUTO = 0 |
| IVF_NPROBE = 1 |
| DISKANN_L_SEARCH = 2 |
| |
| |
| @dataclass(frozen=True) |
| class VectorIndexMetadata: |
| index_type: str |
| dimension: int |
| nlist: int |
| metric: str |
| total_vectors: int |
| pq_m: Optional[int] = None |
| pq_bits: Optional[int] = None |
| rq_bits: Optional[int] = None |
| diskann_max_degree: Optional[int] = None |
| diskann_build_search_list_size: Optional[int] = None |
| diskann_alpha: Optional[float] = None |
| |
| |
| @dataclass(frozen=True) |
| class VectorIndexReadPlan: |
| random_read_latency_nanos: int |
| window_bytes: int |
| max_ranges_per_read: int |
| graph_beam_width: int |
| filtered_graph_beam_width: int |
| adjacency_preload_bytes: int |
| adjacency_cache_bytes: int |
| raw_vector_cache_bytes: int |
| memory_budget_bytes: int |
| |
| |
| @dataclass(frozen=True) |
| class SearchParams: |
| top_k: int |
| search_width: SearchWidth = SearchWidth.AUTO |
| width: int = 0 |
| |
| def __post_init__(self): |
| try: |
| search_width = SearchWidth(self.search_width) |
| except (TypeError, ValueError) as exc: |
| raise ValueError("search_width is invalid") from exc |
| top_k = _size_t(self.top_k, "top_k", allow_zero=False) |
| if search_width == SearchWidth.AUTO: |
| width = _size_t(self.width, "automatic search width", allow_zero=True) |
| if width: |
| raise ValueError("automatic search width must be zero") |
| else: |
| name = ( |
| "nprobe" |
| if search_width == SearchWidth.IVF_NPROBE |
| else "l_search" |
| ) |
| width = _size_t(self.width, name, allow_zero=False) |
| object.__setattr__(self, "top_k", top_k) |
| object.__setattr__(self, "search_width", search_width) |
| object.__setattr__(self, "width", width) |
| |
| @classmethod |
| def automatic(cls, top_k: int): |
| return cls(top_k=top_k) |
| |
| @classmethod |
| def ivf(cls, top_k: int, nprobe: int): |
| return cls( |
| top_k=top_k, |
| search_width=SearchWidth.IVF_NPROBE, |
| width=nprobe, |
| ) |
| |
| @classmethod |
| def diskann(cls, top_k: int, l_search: int): |
| return cls( |
| top_k=top_k, |
| search_width=SearchWidth.DISKANN_L_SEARCH, |
| width=l_search, |
| ) |
| |
| def to_ffi(self): |
| return _ffi.PaimonVindexSearchParams( |
| self.top_k, |
| int(self.search_width), |
| self.width, |
| ) |
| |
| |
| def _check_error(message="operation failed"): |
| err = lib.paimon_vindex_last_error() |
| if err: |
| raise RuntimeError(err.decode("utf-8", errors="replace")) |
| raise RuntimeError(message) |
| |
| |
| def _metadata_from_ffi(raw): |
| return VectorIndexMetadata( |
| index_type=INDEX_TYPES.get(raw.index_type, f"unknown_{raw.index_type}"), |
| dimension=raw.dimension, |
| nlist=raw.nlist, |
| metric=METRICS.get(raw.metric, f"unknown_{raw.metric}"), |
| total_vectors=raw.total_vectors, |
| pq_m=raw.pq_m or None, |
| pq_bits=raw.pq_bits or None, |
| rq_bits=raw.rq_bits or None, |
| diskann_max_degree=raw.diskann_max_degree or None, |
| diskann_build_search_list_size=raw.diskann_build_search_list_size or None, |
| diskann_alpha=raw.diskann_alpha or None, |
| ) |
| |
| |
| def _float32_matrix(value, name): |
| array = np.asarray(value, dtype=np.float32) |
| if array.ndim != 2: |
| raise ValueError(f"{name} must be a two-dimensional float32 array") |
| return np.ascontiguousarray(array) |
| |
| |
| def _float32_vector(value, name): |
| array = np.asarray(value, dtype=np.float32) |
| if array.ndim != 1: |
| raise ValueError(f"{name} must be a one-dimensional float32 array") |
| return np.ascontiguousarray(array) |
| |
| |
| def _int64_vector(value, name): |
| array = np.asarray(value, dtype=np.int64) |
| if array.ndim != 1: |
| raise ValueError(f"{name} must be a one-dimensional int64 array") |
| return np.ascontiguousarray(array) |
| |
| |
| def _bytes_buffer(value, name): |
| if isinstance(value, memoryview): |
| value = value.tobytes() |
| if not isinstance(value, (bytes, bytearray)): |
| raise ValueError(f"{name} must be bytes") |
| data = bytes(value) |
| if not data: |
| return None, 0, data |
| buf = (ctypes.c_uint8 * len(data)).from_buffer_copy(data) |
| return buf, len(data), data |
| |
| |
| def _option_arrays(options: Mapping[str, str]): |
| option_items = list(options.items()) |
| key_bytes = [] |
| value_bytes = [] |
| for key, value in option_items: |
| if not isinstance(key, str) or not isinstance(value, str): |
| raise ValueError("options must be a mapping of str to str") |
| key_bytes.append(key.encode("utf-8")) |
| value_bytes.append(value.encode("utf-8")) |
| keys = (ctypes.c_char_p * len(key_bytes))(*key_bytes) |
| values = (ctypes.c_char_p * len(value_bytes))(*value_bytes) |
| return option_items, key_bytes, value_bytes, keys, values |
| |
| |
| def _make_read_ranges_callback(input, native_handle_lock=None): |
| @_ffi.READ_RANGES_FN |
| def read_ranges_callback(ctx, requests, request_count): |
| if native_handle_lock is not None: |
| native_handle_lock._enter_callback() |
| try: |
| ranges = [ |
| (requests[i].offset, requests[i].len) |
| for i in range(request_count) |
| ] |
| chunks = input.pread_many(ranges) |
| if len(chunks) != request_count: |
| return -1 |
| for i, chunk in enumerate(chunks): |
| data = bytes(chunk) |
| if len(data) != requests[i].len: |
| return -1 |
| ctypes.memmove(requests[i].buf, data, len(data)) |
| return 0 |
| except Exception: |
| return -1 |
| finally: |
| if native_handle_lock is not None: |
| native_handle_lock._exit_callback() |
| |
| return read_ranges_callback |
| |
| |
| class VectorIndexTraining: |
| def __init__(self, handle): |
| self._native_handle_lock = _NativeHandleLock() |
| self._closed = False |
| self._handle = handle |
| |
| def _require_open(self): |
| if self._closed or not self._handle: |
| raise RuntimeError("VectorIndexTraining is closed") |
| |
| def _take_handle(self): |
| with self._native_handle_lock: |
| self._require_open() |
| handle = self._handle |
| self._handle = None |
| self._closed = True |
| return handle |
| |
| def close(self): |
| with self._native_handle_lock: |
| if self._handle: |
| lib.paimon_vindex_training_free(self._handle) |
| self._handle = None |
| self._closed = True |
| |
| def __enter__(self): |
| with self._native_handle_lock: |
| self._require_open() |
| return self |
| |
| def __exit__(self, exc_type, exc_val, exc_tb): |
| self.close() |
| return False |
| |
| def __del__(self): |
| try: |
| self.close() |
| except Exception: |
| pass |
| |
| |
| class VectorIndexTrainer: |
| def __init__(self, options: Mapping[str, str]): |
| self._native_handle_lock = _NativeHandleLock() |
| self._closed = False |
| ( |
| option_items, |
| self._key_bytes, |
| self._value_bytes, |
| self._keys, |
| self._values, |
| ) = _option_arrays(options) |
| self._handle = lib.paimon_vindex_trainer_open( |
| self._keys, |
| self._values, |
| len(option_items), |
| ) |
| if not self._handle: |
| _check_error("failed to open trainer") |
| self._dimension = self._read_dimension() |
| |
| @classmethod |
| def create(cls, options: Mapping[str, str]): |
| return cls(options) |
| |
| @classmethod |
| def train(cls, options: Mapping[str, str], data): |
| data = _float32_matrix(data, "data") |
| resolved_options = dict(options) |
| if resolved_options.get("dimension") in (None, "auto"): |
| resolved_options["dimension"] = str(data.shape[1]) |
| if ( |
| resolved_options.get("index.type", "").startswith("ivf_") |
| and resolved_options.get("nlist") in (None, "auto") |
| and "expected-vector-count" not in resolved_options |
| ): |
| resolved_options["expected-vector-count"] = str(data.shape[0]) |
| with cls(resolved_options) as trainer: |
| return trainer.add_training_vectors(data).finish_training() |
| |
| def _require_open(self): |
| if self._closed or not self._handle: |
| raise RuntimeError("VectorIndexTrainer is closed") |
| |
| def _read_dimension(self): |
| with self._native_handle_lock: |
| self._require_open() |
| out = ctypes.c_size_t(0) |
| rc = lib.paimon_vindex_trainer_dimension( |
| self._handle, ctypes.byref(out) |
| ) |
| if rc != 0: |
| _check_error("trainer dimension failed") |
| return out.value |
| |
| @property |
| def dimension(self): |
| with self._native_handle_lock: |
| self._require_open() |
| return self._dimension |
| |
| def add_training_vectors(self, data): |
| data = _float32_matrix(data, "data") |
| with self._native_handle_lock: |
| self._require_open() |
| if data.shape[1] != self._dimension: |
| raise RuntimeError( |
| f"training data length {data.size} does not match vector count " |
| f"* dimension {data.shape[0] * self._dimension}" |
| ) |
| rc = lib.paimon_vindex_trainer_add_training_vectors( |
| self._handle, |
| data.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), |
| data.shape[0], |
| ) |
| if rc != 0: |
| _check_error("add training vectors failed") |
| return self |
| |
| def finish_training(self): |
| with self._native_handle_lock: |
| self._require_open() |
| handle = self._handle |
| training = lib.paimon_vindex_trainer_finish(handle) |
| lib.paimon_vindex_trainer_free(handle) |
| self._handle = None |
| self._closed = True |
| if not training: |
| _check_error("finish training failed") |
| return VectorIndexTraining(training) |
| |
| def close(self): |
| with self._native_handle_lock: |
| if self._handle: |
| lib.paimon_vindex_trainer_free(self._handle) |
| self._handle = None |
| self._closed = True |
| |
| def __enter__(self): |
| with self._native_handle_lock: |
| self._require_open() |
| return self |
| |
| def __exit__(self, exc_type, exc_val, exc_tb): |
| self.close() |
| return False |
| |
| def __del__(self): |
| try: |
| self.close() |
| except Exception: |
| pass |
| |
| |
| class VectorIndexWriter: |
| def __init__(self, training: VectorIndexTraining): |
| if not isinstance(training, VectorIndexTraining): |
| raise TypeError("training must be a VectorIndexTraining") |
| self._native_handle_lock = _NativeHandleLock() |
| self._closed = False |
| training_handle = training._take_handle() |
| self._handle = lib.paimon_vindex_writer_open(training_handle) |
| lib.paimon_vindex_training_free(training_handle) |
| if not self._handle: |
| _check_error("failed to open writer") |
| self._dimension = self._read_dimension() |
| |
| def _require_open(self): |
| if self._closed or not self._handle: |
| raise RuntimeError("VectorIndexWriter is closed") |
| |
| def _read_dimension(self): |
| with self._native_handle_lock: |
| self._require_open() |
| out = ctypes.c_size_t(0) |
| rc = lib.paimon_vindex_writer_dimension( |
| self._handle, ctypes.byref(out) |
| ) |
| if rc != 0: |
| _check_error("writer dimension failed") |
| return out.value |
| |
| @property |
| def dimension(self): |
| with self._native_handle_lock: |
| self._require_open() |
| return self._dimension |
| |
| def add_vectors(self, ids, data): |
| data = _float32_matrix(data, "data") |
| ids = _int64_vector(ids, "ids") |
| with self._native_handle_lock: |
| self._require_open() |
| if data.shape[1] != self._dimension: |
| raise RuntimeError( |
| f"vector data length {data.size} does not match vector count " |
| f"* dimension {data.shape[0] * self._dimension}" |
| ) |
| if ids.shape[0] != data.shape[0]: |
| raise RuntimeError( |
| f"ids length {ids.shape[0]} does not match vector count " |
| f"{data.shape[0]}" |
| ) |
| rc = lib.paimon_vindex_writer_add_vectors( |
| self._handle, |
| ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), |
| data.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), |
| data.shape[0], |
| ) |
| if rc != 0: |
| _check_error("add_vectors failed") |
| |
| def write(self, file): |
| pos = 0 |
| |
| @_ffi.WRITE_FN |
| def write_callback(ctx, data, length): |
| nonlocal pos |
| try: |
| payload = ctypes.string_at(data, length) |
| written = file.write(payload) |
| if written is not None and written != length: |
| return -1 |
| pos += length |
| return 0 |
| except Exception: |
| return -1 |
| |
| @_ffi.FLUSH_FN |
| def flush_callback(ctx): |
| try: |
| flush = getattr(file, "flush", None) |
| if flush is not None: |
| flush() |
| return 0 |
| except Exception: |
| return -1 |
| |
| @_ffi.GET_POS_FN |
| def pos_callback(ctx): |
| return pos |
| |
| output = _ffi.PaimonVindexOutputFile() |
| output.ctx = None |
| output.write_fn = write_callback |
| output.flush_fn = flush_callback |
| output.get_pos_fn = pos_callback |
| |
| with self._native_handle_lock: |
| self._require_open() |
| rc = lib.paimon_vindex_writer_write_index(self._handle, output) |
| if rc != 0: |
| _check_error("write index failed") |
| |
| def close(self): |
| with self._native_handle_lock: |
| if self._handle: |
| lib.paimon_vindex_writer_free(self._handle) |
| self._handle = None |
| self._closed = True |
| |
| def __enter__(self): |
| with self._native_handle_lock: |
| self._require_open() |
| return self |
| |
| def __exit__(self, exc_type, exc_val, exc_tb): |
| self.close() |
| return False |
| |
| def __del__(self): |
| try: |
| self.close() |
| except Exception: |
| pass |
| |
| |
| class VectorIndexReader: |
| def __init__( |
| self, |
| input, |
| memory_budget_bytes: int = 4 * 1024 * 1024 * 1024, |
| ): |
| self._native_handle_lock = _NativeHandleLock() |
| self._input = input |
| self._closed = False |
| |
| memory_budget_bytes = _size_t( |
| memory_budget_bytes, "memory_budget_bytes", allow_zero=True |
| ) |
| |
| self._read_ranges_callback = _make_read_ranges_callback( |
| self._input, self._native_handle_lock |
| ) |
| input_file = _ffi.PaimonVindexInputFile() |
| input_file.ctx = None |
| input_file.read_ranges_fn = self._read_ranges_callback |
| capability_names = ( |
| "estimated_random_read_latency_nanos", |
| "preferred_window_bytes", |
| "max_ranges_per_read", |
| ) |
| capabilities = {} |
| for name in capability_names: |
| value = getattr(self._input, name, 0) |
| if name == "estimated_random_read_latency_nanos": |
| capabilities[name] = _uint64(value, f"input.{name}") |
| else: |
| capabilities[name] = _size_t( |
| value, f"input.{name}", allow_zero=True |
| ) |
| input_file.estimated_random_read_latency_nanos = capabilities[ |
| "estimated_random_read_latency_nanos" |
| ] |
| input_file.preferred_window_bytes = capabilities["preferred_window_bytes"] |
| input_file.max_ranges_per_read = capabilities["max_ranges_per_read"] |
| options = _ffi.PaimonVindexReaderOptions(memory_budget_bytes) |
| self._handle = lib.paimon_vindex_reader_open_with_options(input_file, options) |
| if not self._handle: |
| _check_error("failed to open reader") |
| self._metadata = self.metadata() |
| |
| def _require_open(self): |
| if self._closed or not self._handle: |
| raise RuntimeError("VectorIndexReader is closed") |
| |
| @property |
| def index_type(self): |
| return self.metadata().index_type |
| |
| @property |
| def dimension(self): |
| return self.metadata().dimension |
| |
| @property |
| def nlist(self): |
| return self.metadata().nlist |
| |
| @property |
| def total_vectors(self): |
| return self.metadata().total_vectors |
| |
| def metadata(self): |
| with self._native_handle_lock: |
| self._require_open() |
| raw = _ffi.PaimonVindexMetadata() |
| rc = lib.paimon_vindex_reader_metadata( |
| self._handle, ctypes.byref(raw) |
| ) |
| if rc != 0: |
| _check_error("metadata failed") |
| return _metadata_from_ffi(raw) |
| |
| def optimize_for_search(self): |
| with self._native_handle_lock: |
| self._require_open() |
| rc = lib.paimon_vindex_reader_optimize_for_search(self._handle) |
| if rc != 0: |
| _check_error("optimize_for_search failed") |
| |
| def warmup_queries(self, queries, l_search: int = 0): |
| queries = _float32_matrix(queries, "queries") |
| if queries.shape[1] != self._metadata.dimension: |
| raise RuntimeError( |
| f"queries length {queries.size} does not match nq * dimension " |
| f"{queries.shape[0] * self._metadata.dimension}" |
| ) |
| l_search = _size_t(l_search, "l_search", allow_zero=True) |
| with self._native_handle_lock: |
| self._require_open() |
| rc = lib.paimon_vindex_reader_warmup_queries( |
| self._handle, |
| queries.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), |
| queries.shape[0], |
| l_search, |
| ) |
| if rc != 0: |
| _check_error("warmup_queries failed") |
| |
| def calibrate_search_width(self, queries, top_k: int = 10): |
| queries = _float32_matrix(queries, "queries") |
| if queries.shape[1] != self._metadata.dimension: |
| raise RuntimeError( |
| f"queries length {queries.size} does not match nq * dimension " |
| f"{queries.shape[0] * self._metadata.dimension}" |
| ) |
| top_k = _size_t(top_k, "top_k", allow_zero=False) |
| with self._native_handle_lock: |
| self._require_open() |
| out = ctypes.c_size_t() |
| rc = lib.paimon_vindex_reader_calibrate_search_width( |
| self._handle, |
| queries.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), |
| queries.shape[0], |
| top_k, |
| ctypes.byref(out), |
| ) |
| if rc != 0: |
| _check_error("calibrate_search_width failed") |
| return out.value |
| |
| def read_plan(self): |
| with self._native_handle_lock: |
| self._require_open() |
| raw = _ffi.PaimonVindexReadPlan() |
| rc = lib.paimon_vindex_reader_read_plan( |
| self._handle, ctypes.byref(raw) |
| ) |
| if rc != 0: |
| _check_error("read_plan failed") |
| return VectorIndexReadPlan( |
| random_read_latency_nanos=raw.random_read_latency_nanos, |
| window_bytes=raw.window_bytes, |
| max_ranges_per_read=raw.max_ranges_per_read, |
| graph_beam_width=raw.graph_beam_width, |
| filtered_graph_beam_width=raw.filtered_graph_beam_width, |
| adjacency_preload_bytes=raw.adjacency_preload_bytes, |
| adjacency_cache_bytes=raw.adjacency_cache_bytes, |
| raw_vector_cache_bytes=raw.raw_vector_cache_bytes, |
| memory_budget_bytes=raw.memory_budget_bytes, |
| ) |
| |
| def _filter_args(self, filter_bytes): |
| if filter_bytes is None: |
| return None, 0, None |
| return _bytes_buffer(filter_bytes, "filter_bytes") |
| |
| def search(self, query, params: SearchParams, filter_bytes=None): |
| query = _float32_vector(query, "query") |
| if query.shape[0] != self._metadata.dimension: |
| raise RuntimeError( |
| f"query length {query.shape[0]} does not match index dimension " |
| f"{self._metadata.dimension}" |
| ) |
| ffi_params = params.to_ffi() |
| ids = np.empty(params.top_k, dtype=np.int64) |
| distances = np.empty(params.top_k, dtype=np.float32) |
| |
| with self._native_handle_lock: |
| self._require_open() |
| if filter_bytes is None: |
| rc = lib.paimon_vindex_reader_search( |
| self._handle, |
| query.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), |
| ffi_params, |
| ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), |
| distances.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), |
| params.top_k, |
| ) |
| else: |
| filter_buf, filter_len, _ = self._filter_args(filter_bytes) |
| rc = lib.paimon_vindex_reader_search_with_roaring_filter( |
| self._handle, |
| query.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), |
| ffi_params, |
| filter_buf, |
| filter_len, |
| ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), |
| distances.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), |
| params.top_k, |
| ) |
| if rc != 0: |
| _check_error("search failed") |
| return ids, distances |
| |
| def search_batch(self, queries, params: SearchParams, filter_bytes=None): |
| queries = _float32_matrix(queries, "queries") |
| if queries.shape[1] != self._metadata.dimension: |
| raise RuntimeError( |
| f"queries length {queries.size} does not match nq * dimension " |
| f"{queries.shape[0] * self._metadata.dimension}" |
| ) |
| ffi_params = params.to_ffi() |
| result_len = queries.shape[0] * params.top_k |
| ids = np.empty((queries.shape[0], params.top_k), dtype=np.int64) |
| distances = np.empty((queries.shape[0], params.top_k), dtype=np.float32) |
| |
| with self._native_handle_lock: |
| self._require_open() |
| if filter_bytes is None: |
| rc = lib.paimon_vindex_reader_search_batch( |
| self._handle, |
| queries.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), |
| queries.shape[0], |
| ffi_params, |
| ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), |
| distances.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), |
| result_len, |
| ) |
| else: |
| filter_buf, filter_len, _ = self._filter_args(filter_bytes) |
| rc = lib.paimon_vindex_reader_search_batch_with_roaring_filter( |
| self._handle, |
| queries.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), |
| queries.shape[0], |
| ffi_params, |
| filter_buf, |
| filter_len, |
| ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), |
| distances.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), |
| result_len, |
| ) |
| if rc != 0: |
| _check_error("batch search failed") |
| return ids, distances |
| |
| def close(self): |
| with self._native_handle_lock: |
| if self._handle: |
| lib.paimon_vindex_reader_free(self._handle) |
| self._handle = None |
| self._closed = True |
| |
| def __enter__(self): |
| with self._native_handle_lock: |
| self._require_open() |
| return self |
| |
| def __exit__(self, exc_type, exc_val, exc_tb): |
| self.close() |
| return False |
| |
| def __del__(self): |
| try: |
| self.close() |
| except Exception: |
| pass |
| |
| |
| __all__ = [ |
| "SearchParams", |
| "VectorIndexMetadata", |
| "VectorIndexReadPlan", |
| "VectorIndexReader", |
| "VectorIndexTrainer", |
| "VectorIndexTraining", |
| "VectorIndexWriter", |
| ] |
