xpark.dataset.connectors.UpsertVectorDB#

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

Operator for batch-upserting records into a vector database.

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

Uses the backend’s native upsert API: inserts new records and updates existing ones with the same id. Unlike InsertVectorDB, callers must supply an id column (with_column(col('id'), col('vec'), ...)).

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

Two operating modes match InsertVectorDB:

  • Default mode (no source_template) — uses the backend’s built-in _source layout. The first bound column is the id column; remaining bound columns are metadata, named via metadata_columns.

  • Template mode (source_template provided) — renders the _source document per row. The first bound column is still the id column (used when id_template is not set); if id_template is set, it overrides the bound id column. Extra bound columns are bound positionally to $1..$N.

Template placeholders are identical to InsertVectorDB: $vector, $id, $collection, positional $N / $N.path, and $ for a literal $.

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) – Default mode only. Maps each metadata column passed via .with_column() (after the id and vector columns) to a field name. Each entry can be a str (use as the metadata field name) or None (parse the column value as a JSON string and merge all its key-value pairs into the metadata dict). Keys are case-sensitive and must be unique. Defaults to None. Rejected under template mode.

  • source_template – Template mode only. Optional dict | str template rendered to the _source document per row (opts into template mode).

  • id_template – Template mode only. Optional str template for the document id. When set, it overrides the bound id column. Requires source_template to also be set.

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

Examples

Default mode with per-column metadata + a JSON-expansion column:

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")),
)
# Resulting metadata dict => {"category": "tech", "color": "red", "priority": 1}
print(ds.take_all())

Template mode: the bound id column is the source of $id (no id_template needed):

ds = ds.with_column(
    "result",
    UpsertVectorDB(
        connection_factory=es_factory,
        collection_name="memories",
        source_template={
            "vector": "$vector",
            "doc_id": "$id",
            "memory_text": "$1",
            "entities": "$2",
        },
    )
    .options(num_workers={"IO": 1})
    .with_column(col("id"), col("vec"), col("memory_text"), col("entities")),
)

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. Under template mode, used as $id unless id_template is set (in which case it is overridden).

  • vectors – A ChunkedArray of list vectors.

  • *metadata_cols – Extra columns bound via with_column. In default mode they are interpreted as metadata per metadata_columns. In template mode they are bound positionally to $1..$N inside source_template / id_template (in declaration order — NOT alphabetical).

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. Under template mode, used as $id unless id_template is set (in which case it is overridden).

  • vectors – A ChunkedArray of list vectors.

  • *metadata_cols – Extra columns bound via with_column. In default mode they are interpreted as metadata per metadata_columns. In template mode they are bound positionally to $1..$N inside source_template / id_template (in declaration order — NOT alphabetical).

Returns:

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