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 812bca965e [python] Support batch vector search for vindex (#8341)
812bca965e is described below
commit 812bca965e67e2eafc14f7dd3098d2c078a3a1d1
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Jun 24 15:17:32 2026 +0800
[python] Support batch vector search for vindex (#8341)
Support batch vector search for the PyPaimon vindex global index reader
so batch vector reads can use the native paimon-vindex batch API instead
of falling back to per-query searches.
---
.../vindex/vindex_vector_global_index_reader.py | 121 +++++++++++++++++----
.../pypaimon/tests/vindex_vector_index_test.py | 118 +++++++++++++++++++-
2 files changed, 215 insertions(+), 24 deletions(-)
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 5b5c93a1b3..4006c7d534 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
@@ -70,42 +70,80 @@ class VindexVectorGlobalIndexReader(GlobalIndexReader):
self._load_lock = threading.Lock()
def visit_vector_search(self, vector_search):
+ results = self._run_search(
+ [vector_search.vector],
+ vector_search.limit,
+ vector_search.include_row_ids,
+ vector_search.options,
+ )
+ return _completed_future(results[0])
+
+ def visit_batch_vector_search(self, batch_vector_search):
+ results = self._run_search(
+ batch_vector_search.vectors,
+ batch_vector_search.limit,
+ batch_vector_search.include_row_ids,
+ batch_vector_search.options,
+ )
+ return _completed_future(results)
+
+ def _run_search(self, vectors, limit, include_row_ids, query_options):
self._ensure_loaded()
- query = np.asarray(vector_search.vector, dtype=np.float32)
- if query.ndim != 1:
- raise ValueError("Query vector must be a one-dimensional float32
array")
- expected_dim = self._metadata.dimension
- if query.shape[0] != expected_dim:
- raise ValueError(
- "Query vector dimension mismatch: expected %d, got %d"
- % (expected_dim, query.shape[0]))
-
- effective_k = self._effective_k(vector_search)
+ queries = self._validate_queries(vectors)
+ n = len(queries)
+ effective_k = self._effective_k(limit, include_row_ids)
if effective_k <= 0:
- return _completed_future(None)
+ return [None] * n
- options = vector_search.options or {}
+ options = query_options or {}
nprobe = _int_parameter(options, NPROBE_PARAMETER, DEFAULT_NPROBE)
ef_search = _int_parameter(options, EF_SEARCH_PARAMETER,
DEFAULT_EF_SEARCH)
- filter_bytes = _filter_bytes(vector_search.include_row_ids)
-
- ids, distances = self._reader.search(
- query, effective_k, nprobe, ef_search, filter_bytes=filter_bytes)
- id_to_scores = _build_scores(ids, distances, self._metadata.metric)
- if not id_to_scores:
- return _completed_future(None)
- return _completed_future(DictBasedScoredIndexResult(id_to_scores))
+ filter_bytes = _filter_bytes(include_row_ids)
+
+ if n == 1:
+ ids, distances = self._reader.search(
+ queries[0], effective_k, nprobe, ef_search,
filter_bytes=filter_bytes)
+ return [_result_from_scores(ids, distances, self._metadata.metric)]
+
+ if hasattr(self._reader, "search_batch"):
+ batch_result = self._reader.search_batch(
+ np.ascontiguousarray(queries),
+ effective_k,
+ nprobe,
+ ef_search,
+ filter_bytes=filter_bytes,
+ )
+ return _batch_results(batch_result, n, effective_k,
self._metadata.metric)
+
+ results = []
+ for query in queries:
+ ids, distances = self._reader.search(
+ query, effective_k, nprobe, ef_search,
filter_bytes=filter_bytes)
+ results.append(_result_from_scores(ids, distances,
self._metadata.metric))
+ return results
+
+ def _validate_queries(self, vectors):
+ queries = []
+ expected_dim = self._metadata.dimension
+ for vector in vectors:
+ query = np.asarray(vector, dtype=np.float32)
+ if query.ndim != 1:
+ raise ValueError("Query vector must be a one-dimensional
float32 array")
+ if query.shape[0] != expected_dim:
+ raise ValueError(
+ "Query vector dimension mismatch: expected %d, got %d"
+ % (expected_dim, query.shape[0]))
+ queries.append(query)
+ return np.asarray(queries, dtype=np.float32)
def vector_metric(self):
self._ensure_loaded()
return self._metadata.metric
- def _effective_k(self, vector_search):
- limit = vector_search.limit
+ def _effective_k(self, limit, include_row_ids):
total_vectors = getattr(self._metadata, "total_vectors", limit)
effective_k = min(limit, int(total_vectors))
- include_row_ids = vector_search.include_row_ids
if include_row_ids is not None:
cardinality = include_row_ids.cardinality()
if cardinality == 0:
@@ -177,6 +215,43 @@ def _build_scores(ids, distances, metric):
return id_to_scores
+def _result_from_scores(ids, distances, metric):
+ id_to_scores = _build_scores(ids, distances, metric)
+ if not id_to_scores:
+ return None
+ return DictBasedScoredIndexResult(id_to_scores)
+
+
+def _batch_results(batch_result, vector_count, effective_k, metric):
+ if hasattr(batch_result, "ids_for_query"):
+ return [
+ _result_from_scores(
+ batch_result.ids_for_query(i),
+ batch_result.distances_for_query(i),
+ metric,
+ )
+ for i in range(vector_count)
+ ]
+
+ ids, distances = batch_result
+ return [
+ _result_from_scores(
+ _values_for_query(ids, i, effective_k),
+ _values_for_query(distances, i, effective_k),
+ metric,
+ )
+ for i in range(vector_count)
+ ]
+
+
+def _values_for_query(values, query_index, effective_k):
+ array = np.asarray(values)
+ if array.ndim >= 2:
+ return array[query_index]
+ start = query_index * effective_k
+ return array[start:start + effective_k]
+
+
def _convert_distance_to_score(distance, metric):
if metric == "l2":
return 1.0 / (1.0 + distance)
diff --git a/paimon-python/pypaimon/tests/vindex_vector_index_test.py
b/paimon-python/pypaimon/tests/vindex_vector_index_test.py
index eae983ca42..ee64512cb2 100644
--- a/paimon-python/pypaimon/tests/vindex_vector_index_test.py
+++ b/paimon-python/pypaimon/tests/vindex_vector_index_test.py
@@ -15,10 +15,74 @@
# specific language governing permissions and limitations
# under the License.
+import io
+import sys
+import types
import unittest
+from pypaimon.globalindex.batch_vector_search import BatchVectorSearch
+from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta
from pypaimon.globalindex.vector_search_result import
DictBasedScoredIndexResult
-from pypaimon.globalindex.vindex.vindex_vector_global_index_reader import
_build_scores
+from pypaimon.globalindex.vindex.vindex_vector_global_index_reader import (
+ VindexVectorGlobalIndexReader,
+ _build_scores,
+)
+from pypaimon.utils.roaring_bitmap import RoaringBitmap64
+
+
+class _FakeMetadata:
+ dimension = 2
+ metric = "l2"
+ total_vectors = 100
+
+
+class _FakeVectorIndexReader:
+ instances = []
+
+ def __init__(self, index_input):
+ self.index_input = index_input
+ self.search_calls = []
+ self.batch_calls = []
+ self.closed = False
+ _FakeVectorIndexReader.instances.append(self)
+
+ def metadata(self):
+ return _FakeMetadata()
+
+ def search(self, query, effective_k, nprobe, ef_search, filter_bytes=None):
+ self.search_calls.append(
+ (list(query), effective_k, nprobe, ef_search, filter_bytes))
+ return list(range(10, 10 + effective_k)), [float(i) for i in
range(effective_k)]
+
+ def search_batch(
+ self, queries, top_k, nprobe, ef_search=0, filter_bytes=None
+ ):
+ self.batch_calls.append(
+ {
+ "queries": queries,
+ "top_k": top_k,
+ "nprobe": nprobe,
+ "ef_search": ef_search,
+ "filter_bytes": filter_bytes,
+ }
+ )
+ ids = []
+ distances = []
+ query_count = queries.shape[0]
+ for query_index in range(query_count):
+ base_id = (query_index + 1) * 10
+ for rank in range(top_k):
+ ids.append(base_id + rank)
+ distances.append(float(rank))
+ return ids, distances
+
+ def close(self):
+ self.closed = True
+
+
+class _BytesFileIO:
+ def new_input_stream(self, path):
+ return io.BytesIO(b"fake-index")
class VindexVectorIndexTest(unittest.TestCase):
@@ -36,6 +100,58 @@ class VindexVectorIndexTest(unittest.TestCase):
top1 = DictBasedScoredIndexResult(id_to_scores).top_k(1)
self.assertEqual([10], top1.results().to_list())
+ def test_batch_search_uses_native_batch_api(self):
+ old_module = sys.modules.get("paimon_vindex")
+ sys.modules["paimon_vindex"] = types.SimpleNamespace(
+ VectorIndexReader=_FakeVectorIndexReader)
+ _FakeVectorIndexReader.instances = []
+
+ try:
+ reader = VindexVectorGlobalIndexReader(
+ file_io=_BytesFileIO(),
+ index_path="/tmp",
+ io_metas=[GlobalIndexIOMeta(file_name="index", file_size=1)],
+ )
+
+ include_ids = RoaringBitmap64()
+ include_ids.add(1)
+ include_ids.add(2)
+
+ results = reader.visit_batch_vector_search(
+ BatchVectorSearch(
+ vectors=[[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]],
+ limit=3,
+ field_name="embedding",
+ options={"ivf.nprobe": "7", "hnsw.ef_search": "11"},
+ ).with_include_row_ids(include_ids)
+ ).result()
+
+ fake_reader = _FakeVectorIndexReader.instances[0]
+ self.assertEqual([], fake_reader.search_calls)
+ self.assertEqual(1, len(fake_reader.batch_calls))
+
+ call = fake_reader.batch_calls[0]
+ self.assertEqual((3, 2), call["queries"].shape)
+ self.assertEqual(2, call["top_k"])
+ self.assertEqual(7, call["nprobe"])
+ self.assertEqual(11, call["ef_search"])
+ self.assertIsNotNone(call["filter_bytes"])
+
+ self.assertEqual(3, len(results))
+ self.assertEqual([10, 11], results[0].results().to_list())
+ self.assertEqual([20, 21], results[1].results().to_list())
+ self.assertEqual([30, 31], results[2].results().to_list())
+ self.assertEqual(1.0, results[0].score_getter()(10))
+ self.assertEqual(0.5, results[0].score_getter()(11))
+
+ reader.close()
+ self.assertTrue(fake_reader.closed)
+ finally:
+ if old_module is None:
+ sys.modules.pop("paimon_vindex", None)
+ else:
+ sys.modules["paimon_vindex"] = old_module
+
if __name__ == '__main__':
unittest.main()
