from __future__ import annotations
import os
import uuid
from collections.abc import Callable
from string import Template
from typing import TYPE_CHECKING, Any
from xpark.dataset.connectors.vectordb.base import (
_CONNECTOR_REGISTRY,
VectorDBClient,
close_connector,
create_connector,
)
from xpark.dataset.connectors.vectordb.types import columns_to_records
from xpark.dataset.constants import IO_WORKER_ENV
from xpark.dataset.datatype import DataType
from xpark.dataset.expressions import BatchColumnClassProtocol, udf
from xpark.dataset.import_utils import lazy_import
if TYPE_CHECKING:
import pyarrow as pa
else:
pa = lazy_import("pyarrow", rename="pa")
[docs]
@udf(return_dtype=DataType.bool())
class InsertVectorDB(BatchColumnClassProtocol):
__doc__ = Template("""Operator for batch-inserting records into a vector database.
Supported backends: $SUPPORTED_BACKENDS.
IDs are **auto-generated** (UUID4) by this operator — callers should **not**
supply an ID column. Internally uses the backend's native ``upsert`` API
(since the auto-generated UUID4 IDs guarantee no conflicts, upsert is
equivalent to insert but is idempotent under retries). If you need to
control record IDs, use :class:`UpsertVectorDB` instead.
The operator is a pure side-effect: it returns a boolean column and does not
produce meaningful output data.
Args:
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:
.. code-block:: python
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")),
)
# metadata => {"category": "tech", "color": "red", "priority": 1}
print(ds.take_all())
""").safe_substitute(SUPPORTED_BACKENDS=sorted(_CONNECTOR_REGISTRY.keys()))
def __init__(
self,
connection_factory: Callable[[], VectorDBClient],
collection_name: str,
metadata_columns: list[str | None] | None = None,
**kwargs: Any,
):
if not os.environ.get(IO_WORKER_ENV):
raise ValueError("InsertVectorDB must run in an IO worker. Use .options(num_workers={'IO': N}).")
self._connector, self._kwargs = create_connector(connection_factory, collection_name, **kwargs)
self._metadata_columns = metadata_columns
def __del__(self) -> None:
close_connector(getattr(self, "_connector", None))
async def __call__(
self,
vectors: pa.ChunkedArray,
*metadata_cols: pa.ChunkedArray,
) -> pa.Array:
"""Execute batch insert.
Args:
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.
"""
num_records = len(vectors)
if not num_records:
return pa.array([], type=pa.bool_())
ids = pa.array([str(uuid.uuid4()) for _ in range(num_records)], type=pa.string())
records = columns_to_records(ids, vectors, *metadata_cols, meta_names=self._metadata_columns)
try:
await self._connector.upsert(records, **self._kwargs)
except Exception as e:
raise RuntimeError(f"Database insert operation failed: {e}") from e
return pa.repeat(False, num_records)
[docs]
@udf(return_dtype=DataType.bool())
class UpsertVectorDB(BatchColumnClassProtocol):
__doc__ = Template("""Operator for batch-upserting records into a vector database.
Supported backends: $SUPPORTED_BACKENDS.
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.
Args:
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:
.. code-block:: python
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())
""").safe_substitute(SUPPORTED_BACKENDS=sorted(_CONNECTOR_REGISTRY.keys()))
def __init__(
self,
connection_factory: Callable[[], VectorDBClient],
collection_name: str,
metadata_columns: list[str | None] | None = None,
**kwargs: Any,
):
if not os.environ.get(IO_WORKER_ENV):
raise ValueError("UpsertVectorDB must run in an IO worker. Use .options(num_workers={'IO': N}).")
self._connector, self._kwargs = create_connector(connection_factory, collection_name, **kwargs)
self._metadata_columns = metadata_columns
def __del__(self) -> None:
close_connector(getattr(self, "_connector", None))
async def __call__(
self,
ids: pa.ChunkedArray,
vectors: pa.ChunkedArray,
*metadata_cols: pa.ChunkedArray,
) -> pa.Array:
"""Execute batch upsert.
Args:
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.
"""
if len(ids) == 0:
return pa.array([], type=pa.bool_())
records = columns_to_records(ids, vectors, *metadata_cols, meta_names=self._metadata_columns)
try:
await self._connector.upsert(records, **self._kwargs)
except Exception as e:
raise RuntimeError(f"Database upsert operation failed: {e}") from e
return pa.repeat(False, len(ids))
[docs]
@udf(return_dtype=DataType.bool())
class DeleteVectorDB(BatchColumnClassProtocol):
__doc__ = Template("""Operator for batch-deleting records from a vector database by ID.
Supported backends: $SUPPORTED_BACKENDS.
Uses the backend's native ``delete`` API. Deleting a non-existent ID is
a no-op (idempotent delete semantics).
The operator is a pure side-effect: it returns a boolean column and does not
produce meaningful output data.
Args:
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.
**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 delete call (e.g. ``timeout``).
Examples:
.. code-block:: python
import pymilvus
from xpark.dataset import from_items
from xpark.dataset.connectors import DeleteVectorDB
from xpark.dataset.expressions import col
ds = from_items([{"id": "1"}, {"id": "2"}, {"id": "3"}])
ds = ds.with_column(
"result",
DeleteVectorDB(
connection_factory=lambda: pymilvus.MilvusClient(uri="http://localhost:19530"),
collection_name="my_col",
)
.options(num_workers={"IO": 1})
.with_column(col("id")),
)
print(ds.take_all())
""").safe_substitute(SUPPORTED_BACKENDS=sorted(_CONNECTOR_REGISTRY.keys()))
def __init__(
self,
connection_factory: Callable[[], VectorDBClient],
collection_name: str,
**kwargs: Any,
):
if not os.environ.get(IO_WORKER_ENV):
raise ValueError("DeleteVectorDB must run in an IO worker. Use .options(num_workers={'IO': N}).")
self._connector, self._kwargs = create_connector(connection_factory, collection_name, **kwargs)
def __del__(self) -> None:
close_connector(getattr(self, "_connector", None))
async def __call__(self, ids: pa.ChunkedArray) -> pa.Array:
"""Execute batch delete.
Args:
ids: A ChunkedArray of IDs to delete.
Returns:
A boolean array (one False per row). The operator is used for its
side-effects only.
"""
if len(ids) == 0:
return pa.array([], type=pa.bool_())
ids_list = ids.to_pylist()
try:
await self._connector.delete(ids_list, **self._kwargs)
except Exception as e:
raise RuntimeError(f"Database delete operation failed: {e}") from e
return pa.repeat(False, len(ids))