JingsongLi commented on code in PR #8341:
URL: https://github.com/apache/paimon/pull/8341#discussion_r3464878715
##########
paimon-python/pypaimon/globalindex/vindex/vindex_vector_global_index_reader.py:
##########
@@ -70,42 +70,76 @@ def __init__(self, file_io, index_path, io_metas,
options=None):
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"):
+ flat_queries = np.ascontiguousarray(queries).reshape(-1)
+ batch_result = self._reader.search_batch(
+ flat_queries, n, effective_k, nprobe, ef_search,
filter_bytes=filter_bytes)
Review Comment:
This does not match the `paimon-vindex==0.1.0` Python binding.
`VectorIndexReader.search_batch` is defined as `search_batch(self, queries,
top_k, nprobe, ef_search=0, filter_bytes=None)` and it calls
`_float32_matrix(queries, "queries")`, deriving the query count from
`queries.shape[0]`. Passing the flattened array plus the extra `n` argument
shifts all positional parameters, and with `filter_bytes=` set it raises
`TypeError: ... got multiple values for argument filter_bytes`; without a
filter it would also pass the wrong `top_k`. Please pass the 2-D `queries`
array directly and remove the explicit query-count argument, e.g.
`search_batch(queries, effective_k, nprobe, ef_search,
filter_bytes=filter_bytes)`. The test should use a fake reader with the real
0.1.0 signature so this native path is covered.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
