xpark.dataset.connectors.SearchVectorDB

class xpark.dataset.connectors.SearchVectorDB(connection_factory: Callable[[], Any], collection_name: str, *, top_k: int = 10, include_vectors: bool = False, **kwargs: Any)

Operator for vector similarity search with optional metadata filtering.

Supported backends: [‘elasticsearch’, ‘milvus’].

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.

Parameters:

• 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

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())

Methods

__call__(query_vectors)	Execute vector similarity search for a batch of queries.
options(**kwargs)
with_column(query_vectors)	Execute vector similarity search for a batch of queries.

__call__(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.

Parameters:

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.

options(**kwargs: Unpack[ExprUDFOptions]) → Self

with_column(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.

Parameters:

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.