xpark.dataset.connectors.SearchVectorDB#

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

Operator for vector similarity search.

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.

Two operating modes (mutually exclusive):

  • Default mode — pure kNN search using the backend’s built-in field layout (vector + metadata). Accepts the usual connector kwargs (filter, consistency_level, timeout, …). Requires only the query-vector column; any extra bound column raises. For BM25 / hybrid / arbitrary DSL queries, use template mode below.

  • Template mode (body_template provided) — the operator renders a user-supplied request body per row and sends it to the backend as-is via search_with_bodies. body_template accepts either a dict (Python literal, recommended for JSON-based backends like Elasticsearch) or a pre-serialised str (handy for backends whose query language is a string).

Template placeholders (template mode):

  • $vector — whole-value, list[float], the row’s query vector.

  • $top_k — whole-value, int, the operator’s top_k.

  • $collection — whole-value, str.

  • $N / $N.a.b.c — 1-based positional reference to an extra column bound via with_column after the query-vector column. Whole-value form preserves type; interpolated form ("prefix-$1") coerces via str(). Dot-paths drill into dict/struct values and raise on non-dict intermediate values.

  • $ escapes a literal $.

body_template and kwargs are mutually exclusive — bake everything into the template instead. Each invocation issues exactly one request body per row. For multi-leg fan-out + client-side fusion, instantiate one SearchVectorDB per leg (each producing its own column) and pass the per-leg columns into a fusion UDF. For server-side fusion, embed the entire fusion request (e.g. Elasticsearch’s retriever DSL with RRF / linear) inside a single body_template and let the backend return one fused list directly.

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). Also exposed as $top_k inside body_template.

  • 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.

  • body_template – Template mode only. Optional dict | str template rendered per row (opts into template mode). Mutually exclusive with connector kwargs.

  • **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). Default mode only — rejected when body_template is set.

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

Default mode with a metadata filter:

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

Template mode on Elasticsearch — kNN + per-row filter by attribute:

ds = ds.with_column(
    "results",
    SearchVectorDB(
        connection_factory=es_factory,
        collection_name="memories",
        top_k=5,
        body_template={
            "knn": {
                "field": "vector",
                "query_vector": "$vector",
                "k": "$top_k",
                "num_candidates": 50,
                "filter": {"term": {"conv_idx": "$1"}},
            },
            "size": "$top_k",
        },
    )
    .options(num_workers={"IO": 1})
    .with_column(col("query_vec"), col("conv_idx")),
)

Template mode — server-side hybrid retrieval via Elasticsearch’s retriever DSL (RRF fusion of kNN + BM25), one fused list per row:

ds = ds.with_column(
    "results",
    SearchVectorDB(
        connection_factory=es_factory,
        collection_name="memories",
        top_k=10,
        body_template={
            "retriever": {
                "rrf": {
                    "retrievers": [
                        {"standard": {"query": {"match": {"text": "$1"}}}},
                        {"knn": {
                            "field": "vector",
                            "query_vector": "$vector",
                            "k": "$top_k",
                            "num_candidates": 50,
                        }},
                    ],
                    "rank_window_size": 50,
                }
            },
            "size": "$top_k",
        },
    )
    .options(num_workers={"IO": 1})
    .with_column(col("query_vec"), col("question")),
)

Methods

__call__(query_vectors, *bound_columns)

Execute vector similarity search for a batch of queries.

options(**kwargs)

with_column(query_vectors, *bound_columns)

Execute vector similarity search for a batch of queries.
__call__(query_vectors: pa.ChunkedArray, *bound_columns: pa.ChunkedArray) → pa.Array#

Execute vector similarity search for a batch of queries.

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.

  • *bound_columns – Extra columns bound via with_column after the query-vector column. Default mode rejects these; template mode binds them positionally to $1..$N inside body_template (in declaration order — NOT alphabetical).

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, *bound_columns: pa.ChunkedArray) → pa.Array#

Execute vector similarity search for a batch of queries.

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.

  • *bound_columns – Extra columns bound via with_column after the query-vector column. Default mode rejects these; template mode binds them positionally to $1..$N inside body_template (in declaration order — NOT alphabetical).

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.