xpark.dataset.connectors.UpsertVectorDB

class xpark.dataset.connectors.UpsertVectorDB(connection_factory: Callable[[], Any], collection_name: str, metadata_columns: list[str | None] | None = None, **kwargs: Any)

Operator for batch-upserting records into a vector database.

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

Uses the backend’s native upsert API which inserts new records and updates existing ones with the same ID.

The operator is a pure side-effect: it returns a boolean column and does not produce meaningful output data.

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.

• metadata_columns (list[str | None] | None) – Maps each metadata column passed via .with_column() to a field name, like SQL AS. Each entry can be a str (use as the metadata field name) or None (parse the column value as JSON and merge all its key-value pairs into the metadata dict). Keys are case-sensitive and must be unique. Defaults to None.

• **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 upsert call (e.g. timeout).

Examples

import pymilvus
from xpark.dataset import from_items
from xpark.dataset.connectors import UpsertVectorDB
from xpark.dataset.expressions import col

ds = from_items([
    {"id": "1", "vec": [0.1, 0.2, 0.3], "category": "tech",
     "extra": '{"color": "red", "priority": 1}'},
])
ds = ds.with_column(
    "result",
    UpsertVectorDB(
        connection_factory=lambda: pymilvus.MilvusClient(uri="http://localhost:19530"),
        collection_name="my_col",
        metadata_columns=["category", None],
    )
    .options(num_workers={"IO": 1})
    .with_column(col("id"), col("vec"), col("category"), col("extra")),
)
# metadata => {"category": "tech", "color": "red", "priority": 1}
print(ds.take_all())

Methods

__call__(ids, vectors, *metadata_cols)	Execute batch upsert.
options(**kwargs)
with_column(ids, vectors, *metadata_cols)	Execute batch upsert.

__call__(ids: pa.ChunkedArray, vectors: pa.ChunkedArray, *metadata_cols: pa.ChunkedArray) → pa.Array

Execute batch upsert.

Parameters:

• ids – A ChunkedArray of IDs.

• vectors – A ChunkedArray of list vectors.

• *metadata_cols – Optional metadata columns.

Returns:

A boolean array (one False per row). The operator is used for its side-effects only.

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

with_column(ids: pa.ChunkedArray, vectors: pa.ChunkedArray, *metadata_cols: pa.ChunkedArray) → pa.Array

Execute batch upsert.

Parameters:

• ids – A ChunkedArray of IDs.

• vectors – A ChunkedArray of list vectors.

• *metadata_cols – Optional metadata columns.

Returns:

A boolean array (one False per row). The operator is used for its side-effects only.