xpark.dataset.connectors.InsertVectorDB#

class xpark.dataset.connectors.InsertVectorDB(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-inserting records into a vector database.

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

IDs are auto-generated (UUID4) unless an id_template is provided. Internally uses the backend’s native upsert API (since UUID4 IDs guarantee no conflicts, upsert is equivalent to insert but is idempotent under retries). If you need caller-supplied IDs, use UpsertVectorDB instead.

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

Two operating modes (mutually exclusive):

  • Default mode (no source_template) — uses the backend’s built-in _source layout (vector / metadata / …). metadata_columns controls SQL-like column-to-field naming.

  • Template mode (source_template provided) — renders a user-supplied dict or str template per row to produce the raw _source document. metadata_columns must be None in this mode — metadata fields are embedded directly in the template via positional $N.

Template placeholders (for source_template and id_template):

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

  • $id — whole-value, the per-row document id (either UUID4 or the id_template result); available inside source_template so a body can embed the document id.

  • $collection — whole-value, str, the target collection name.

  • $N / $N.a.b.c — 1-based positional reference to an extra column bound via with_column after the 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 $.

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() 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 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. E.g. "$1-$2" builds the id from the first two bound columns. Requires source_template to also be set. Defaults to UUID4.

  • **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 InsertVectorDB
from xpark.dataset.expressions import col

ds = from_items([
    {"vec": [0.1, 0.2, 0.3], "category": "tech",
     "extra": '{"color": "red", "priority": 1}'},
])
ds = ds.with_column(
    "result",
    InsertVectorDB(
        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("vec"), col("category"), col("extra")),
)
# Resulting metadata dict => {"category": "tech", "color": "red", "priority": 1}
print(ds.take_all())

Template mode on Elasticsearch — custom _source with typed fields and a deterministic id built from two columns:

from xpark.dataset.connectors import InsertVectorDB
from xpark.dataset.expressions import col

ds = ds.with_column(
    "result",
    InsertVectorDB(
        connection_factory=es_factory,
        collection_name="memories",
        source_template={
            "vector": "$vector",
            "doc_id": "$id",
            "conv_idx": "$1",
            "memory_text": "$2",
            "user": "$3.name",
            "tag": "conv-$1",
        },
        id_template="$1-$2",
    )
    .options(num_workers={"IO": 1})
    .with_column(col("vec"), col("conv_idx"), col("memory_text"), col("user")),
)

Methods

__call__(vectors, *metadata_cols)

Execute batch insert.

options(**kwargs)

with_column(vectors, *metadata_cols)

Execute batch insert.
__call__(vectors: pa.ChunkedArray, *metadata_cols: pa.ChunkedArray) → pa.Array#

Execute batch insert.

Parameters:

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

Returns:

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

options(**kwargs: Unpack[ExprUDFOptions]) → Self#
with_column(vectors: pa.ChunkedArray, *metadata_cols: pa.ChunkedArray) → pa.Array#

Execute batch insert.

Parameters:

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

Returns:

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