from __future__ import annotations
import logging
import os
from collections.abc import Callable
from string import Template
from typing import TYPE_CHECKING, Any
from xpark.dataset.connectors.vectordb.base import (
_CONNECTOR_REGISTRY,
VectorDBClient,
close_connector,
create_connector,
)
from xpark.dataset.connectors.vectordb.types import (
VectorDBError,
search_results_to_arrow,
vectors_to_numpy,
)
from xpark.dataset.constants import IO_WORKER_ENV
from xpark.dataset.datatype import DataType
from xpark.dataset.expressions import BatchColumnClassProtocol, udf
from xpark.dataset.import_utils import lazy_import
if TYPE_CHECKING:
import numpy as np
import pyarrow as pa
else:
np = lazy_import("numpy", rename="np")
pa = lazy_import("pyarrow", rename="pa")
logger = logging.getLogger("ray")
[docs]
@udf(
return_dtype=DataType.from_arrow(
pa.list_(
pa.struct(
[
pa.field("id", pa.string()),
pa.field("score", pa.float64()),
pa.field("metadata", pa.string()),
pa.field("vector", pa.list_(pa.float64())),
]
)
)
)
)
class SearchVectorDB(BatchColumnClassProtocol):
__doc__ = Template("""Operator for vector similarity search with optional metadata filtering.
Supported backends: $SUPPORTED_BACKENDS.
Accepts a column of query vectors and performs vector similarity search
against the configured vector database backend. Each result includes an
``id``, ``score``, a JSON-encoded ``metadata`` string, and optionally the
stored ``vector`` when ``include_vectors=True`` is passed.
Args:
connection_factory: A callable that returns a vector database client
(e.g. ``lambda: pymilvus.MilvusClient(uri="...")``). The backend
is auto-detected from the returned client type.
collection_name: Name of the target collection.
top_k: Maximum number of results per query (default ``10``).
include_vectors: If ``True``, each result will contain the stored
``vector`` field as a list of floats. If ``False`` (default),
the ``vector`` field is ``None`` in every result.
**connector_kwargs: Extra keyword arguments. The connector-init
reserved keys ``max_retries``, ``retry_delay``, ``retry_backoff``
are consumed by the connector itself; all remaining kwargs are
forwarded to the underlying search call (e.g. ``timeout``,
``filter``, ``consistency_level``).
Returns:
A ``pa.Array`` of type
``list<struct<id: string, score: float64, metadata: string, vector: list<float64>>>``,
one list of result structs per query. The ``id`` field type matches
the IDs stored in the collection (``string`` or ``int64``).
``vector`` is ``None`` when ``include_vectors=False``.
Examples:
.. code-block:: python
import pymilvus
from xpark.dataset import from_items
from xpark.dataset.connectors import SearchVectorDB
from xpark.dataset.expressions import col
queries = [{"vector": [0.1, 0.2, 0.3]}]
ds = from_items(queries)
ds = ds.with_column(
"results",
SearchVectorDB(
connection_factory=lambda: pymilvus.MilvusClient(uri="http://localhost:19530"),
collection_name="my_col",
top_k=10,
include_vectors=True,
filter='category == "tech"',
)
.options(num_workers={"IO": 1})
.with_column(col("vector")),
)
# include_vectors=False (default):
# {"id": "1", "score": 0.99, "metadata": '{"category": "tech"}', "vector": None}
# include_vectors=True:
# {"id": "1", "score": 0.99, "metadata": '{"category": "tech"}', "vector": [0.1, 0.2, 0.3]}
print(ds.take_all())
""").safe_substitute(SUPPORTED_BACKENDS=sorted(_CONNECTOR_REGISTRY.keys()))
def __init__(
self,
connection_factory: Callable[[], VectorDBClient],
collection_name: str,
*,
top_k: int = 10,
include_vectors: bool = False,
**kwargs: Any,
):
if not os.environ.get(IO_WORKER_ENV):
raise ValueError("SearchVectorDB must run in an IO worker. Use .options(num_workers={'IO': N}).")
self._connector, self._search_kwargs = create_connector(connection_factory, collection_name, **kwargs)
self._top_k = top_k
self._include_vector = include_vectors
def __del__(self) -> None:
close_connector(getattr(self, "_connector", None))
async def __call__(
self,
query_vectors: pa.ChunkedArray,
) -> pa.Array:
"""Execute vector similarity search for a batch of queries.
All queries in the batch share the metadata filter specified at
construction time. Valid (non-null) query vectors are collected and
sent in a single batch call to the backend's search API.
Args:
query_vectors: A ChunkedArray where each element is a list of floats
(the query vector). Null vectors are not accepted and will
raise a ``VectorDBError``.
Returns:
A ``pa.Array`` of type
``list<struct<id, score: float64, metadata: string, vector: list<float64>>>``
– one list of result structs per query. The ``id`` field type
matches the IDs stored in the collection (``string`` or ``int64``).
``vector`` is ``None`` when ``include_vectors=False``.
"""
if query_vectors.null_count > 0:
raise VectorDBError(
"query_vectors must not contain null entries",
operation="search",
)
if query_vectors.length() == 0:
return search_results_to_arrow([])
vectors_np = vectors_to_numpy(
query_vectors,
column_name="query_vectors",
)
if not isinstance(vectors_np, np.ndarray):
raise VectorDBError(
"query_vectors must be fixed-size vectors for search",
operation="search",
)
try:
all_results = await self._connector.search(
vectors=vectors_np,
top_k=self._top_k,
include_vectors=self._include_vector,
**self._search_kwargs,
)
except Exception as e:
raise VectorDBError(
f"Database search failed: {str(e)}",
operation="search",
)
return search_results_to_arrow(all_results)