vectordb-bench 0.0.21__py3-none-any.whl → 0.0.23__py3-none-any.whl

vectordb_bench/backend/clients/milvus/cli.py

@@ -194,6 +194,56 @@ def MilvusGPUIVFFlat(**parameters: Unpack[MilvusGPUIVFTypedDict]):
         **parameters,
     )
 
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusGPUBruteForceTypedDict)
+def MilvusGPUBruteForce(**parameters: Unpack[MilvusGPUBruteForceTypedDict]):
+    from .config import GPUBruteForceConfig, MilvusConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]),
+        ),
+        db_case_config=GPUBruteForceConfig(
+            metric_type=parameters["metric_type"],
+            limit=parameters["limit"],  # top-k for search
+        ),
+        **parameters,
+    )
+
+class MilvusGPUBruteForceTypedDict(CommonTypedDict, MilvusTypedDict):
+    metric_type: Annotated[
+        str,
+        click.option("--metric-type", type=str, required=True, help="Metric type for brute force search"),
+    ]
+    limit: Annotated[
+        int,
+        click.option("--limit", type=int, required=True, help="Top-k limit for search"),
+    ]
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusGPUBruteForceTypedDict)
+def MilvusGPUBruteForce(**parameters: Unpack[MilvusGPUBruteForceTypedDict]):
+    from .config import GPUBruteForceConfig, MilvusConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]),
+        ),
+        db_case_config=GPUBruteForceConfig(
+            metric_type=parameters["metric_type"],
+            limit=parameters["limit"],  # top-k for search
+        ),
+        **parameters,
+    )
+
 
 class MilvusGPUIVFPQTypedDict(
     CommonTypedDict,

vectordb_bench/backend/clients/milvus/config.py

@@ -40,6 +40,7 @@ class MilvusIndexConfig(BaseModel):
             IndexType.GPU_CAGRA,
             IndexType.GPU_IVF_FLAT,
             IndexType.GPU_IVF_PQ,
+            IndexType.GPU_BRUTE_FORCE,
         ]
 
     def parse_metric(self) -> str:
@@ -184,6 +185,37 @@ class GPUIVFFlatConfig(MilvusIndexConfig, DBCaseConfig):
         }
 
 
+class GPUBruteForceConfig(MilvusIndexConfig, DBCaseConfig):
+    limit: int = 10  # Default top-k for search
+    metric_type: str  # Metric type (e.g., 'L2', 'IP', etc.)
+    index: IndexType = IndexType.GPU_BRUTE_FORCE  # Index type set to GPU_BRUTE_FORCE
+
+    def index_param(self) -> dict:
+        """
+        Returns the parameters for creating the GPU_BRUTE_FORCE index.
+        No additional parameters required for index building.
+        """
+        return {
+            "metric_type": self.parse_metric(),  # Metric type for distance calculation (L2, IP, etc.)
+            "index_type": self.index.value,  # GPU_BRUTE_FORCE index type
+            "params": {},  # No additional parameters for GPU_BRUTE_FORCE
+        }
+
+    def search_param(self) -> dict:
+        """
+        Returns the parameters for performing a search on the GPU_BRUTE_FORCE index.
+        Only metric_type and top-k (limit) are needed for search.
+        """
+        return {
+            "metric_type": self.parse_metric(),  # Metric type for search
+            "params": {
+                "nprobe": 1,  # For GPU_BRUTE_FORCE, set nprobe to 1 (brute force search)
+                "limit": self.limit,  # Top-k for search
+            },
+        }
+
+
+
 class GPUIVFPQConfig(MilvusIndexConfig, DBCaseConfig):
     nlist: int = 1024
     m: int = 0
@@ -261,4 +293,5 @@ _milvus_case_config = {
     IndexType.GPU_IVF_FLAT: GPUIVFFlatConfig,
     IndexType.GPU_IVF_PQ: GPUIVFPQConfig,
     IndexType.GPU_CAGRA: GPUCAGRAConfig,
+    IndexType.GPU_BRUTE_FORCE: GPUBruteForceConfig,
 }

vectordb_bench/backend/clients/mongodb/config.py

@@ -0,0 +1,53 @@
+from enum import Enum
+
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
+
+
+class QuantizationType(Enum):
+    NONE = "none"
+    BINARY = "binary"
+    SCALAR = "scalar"
+
+
+class MongoDBConfig(DBConfig, BaseModel):
+    connection_string: SecretStr = "mongodb+srv://<user>:<password>@<cluster_name>.heatl.mongodb.net"
+    database: str = "vdb_bench"
+
+    def to_dict(self) -> dict:
+        return {
+            "connection_string": self.connection_string.get_secret_value(),
+            "database": self.database,
+        }
+
+
+class MongoDBIndexConfig(BaseModel, DBCaseConfig):
+    index: IndexType = IndexType.HNSW  # MongoDB uses HNSW for vector search
+    metric_type: MetricType = MetricType.COSINE
+    num_candidates_ratio: int = 10  # Default numCandidates ratio for vector search
+    quantization: QuantizationType = QuantizationType.NONE  # Quantization type if applicable
+
+    def parse_metric(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "euclidean"
+        if self.metric_type == MetricType.IP:
+            return "dotProduct"
+        return "cosine"  # Default to cosine similarity
+
+    def index_param(self) -> dict:
+        return {
+            "type": "vectorSearch",
+            "fields": [
+                {
+                    "type": "vector",
+                    "similarity": self.parse_metric(),
+                    "numDimensions": None,  # Will be set in MongoDB class
+                    "path": "vector",  # Vector field name
+                    "quantization": self.quantization.value,
+                }
+            ],
+        }
+
+    def search_param(self) -> dict:
+        return {"num_candidates_ratio": self.num_candidates_ratio}

vectordb_bench/backend/clients/mongodb/mongodb.py

@@ -0,0 +1,200 @@
+import logging
+import time
+from contextlib import contextmanager
+
+from pymongo import MongoClient
+from pymongo.operations import SearchIndexModel
+
+from ..api import VectorDB
+from .config import MongoDBIndexConfig
+
+log = logging.getLogger(__name__)
+
+
+class MongoDBError(Exception):
+    """Custom exception class for MongoDB client errors."""
+
+
+class MongoDB(VectorDB):
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: MongoDBIndexConfig,
+        collection_name: str = "vdb_bench_collection",
+        id_field: str = "id",
+        vector_field: str = "vector",
+        drop_old: bool = False,
+        **kwargs,
+    ):
+        self.dim = dim
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.collection_name = collection_name
+        self.id_field = id_field
+        self.vector_field = vector_field
+        self.drop_old = drop_old
+
+        # Update index dimensions
+        index_params = self.case_config.index_param()
+        log.info(f"index params: {index_params}")
+        index_params["fields"][0]["numDimensions"] = dim
+        self.index_params = index_params
+
+        # Initialize  - they'll also be set in init()
+        uri = self.db_config["connection_string"]
+        self.client = MongoClient(uri)
+        self.db = self.client[self.db_config["database"]]
+        self.collection = self.db[self.collection_name]
+        if self.drop_old and self.collection_name in self.db.list_collection_names():
+            log.info(f"MongoDB client dropping old collection: {self.collection_name}")
+            self.db.drop_collection(self.collection_name)
+        self.client = None
+        self.db = None
+        self.collection = None
+
+    @contextmanager
+    def init(self):
+        """Initialize MongoDB client and cleanup when done"""
+        try:
+            uri = self.db_config["connection_string"]
+            self.client = MongoClient(uri)
+            self.db = self.client[self.db_config["database"]]
+            self.collection = self.db[self.collection_name]
+
+            yield
+        finally:
+            if self.client is not None:
+                self.client.close()
+                self.client = None
+                self.db = None
+                self.collection = None
+
+    def _create_index(self) -> None:
+        """Create vector search index"""
+        index_name = "vector_index"
+        index_params = self.index_params
+        log.info(f"index params {index_params}")
+        # drop index if already exists
+        if self.collection.list_indexes():
+            all_indexes = self.collection.list_search_indexes()
+            if any(idx.get("name") == index_name for idx in all_indexes):
+                log.info(f"Drop index: {index_name}")
+                try:
+                    self.collection.drop_search_index(index_name)
+                    while True:
+                        indices = list(self.collection.list_search_indexes())
+                        indices = [idx for idx in indices if idx["name"] == index_name]
+                        log.debug(f"index status {indices}")
+                        if len(indices) == 0:
+                            break
+                        log.info(f"index deleting {indices}")
+                except Exception:
+                    log.exception(f"Error dropping index {index_name}")
+        try:
+            # Create vector search index
+            search_index = SearchIndexModel(definition=index_params, name=index_name, type="vectorSearch")
+
+            self.collection.create_search_index(search_index)
+            log.info(f"Created vector search index: {index_name}")
+            self._wait_for_index_ready(index_name)
+
+            # Create regular index on id field for faster lookups
+            self.collection.create_index(self.id_field)
+            log.info(f"Created index on {self.id_field} field")
+
+        except Exception:
+            log.exception(f"Error creating index {index_name}")
+            raise
+
+    def _wait_for_index_ready(self, index_name: str, check_interval: int = 5) -> None:
+        """Wait for index to be ready"""
+        while True:
+            indices = list(self.collection.list_search_indexes())
+            log.debug(f"index status {indices}")
+            if indices and any(idx.get("name") == index_name and idx.get("queryable") for idx in indices):
+                break
+            for idx in indices:
+                if idx.get("name") == index_name and idx.get("status") == "FAILED":
+                    error_msg = f"Index {index_name} failed to build"
+                    raise MongoDBError(error_msg)
+
+            time.sleep(check_interval)
+        log.info(f"Index {index_name} is ready")
+
+    def need_normalize_cosine(self) -> bool:
+        return False
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs,
+    ) -> (int, Exception | None):
+        """Insert embeddings into MongoDB"""
+
+        # Prepare documents in bulk
+        documents = [
+            {
+                self.id_field: id_,
+                self.vector_field: embedding,
+            }
+            for id_, embedding in zip(metadata, embeddings, strict=False)
+        ]
+
+        # Use ordered=False for better insert performance
+        try:
+            self.collection.insert_many(documents, ordered=False)
+        except Exception as e:
+            return 0, e
+        return len(documents), None
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+        **kwargs,
+    ) -> list[int]:
+        """Search for similar vectors"""
+        search_params = self.case_config.search_param()
+
+        vector_search = {"queryVector": query, "index": "vector_index", "path": self.vector_field, "limit": k}
+
+        # Add exact search parameter if specified
+        if search_params["exact"]:
+            vector_search["exact"] = True
+        else:
+            # Set numCandidates based on k value and data size
+            # For 50K dataset, use higher multiplier for better recall
+            num_candidates = min(10000, k * search_params["num_candidates_ratio"])
+            vector_search["numCandidates"] = num_candidates
+
+        # Add filter if specified
+        if filters:
+            log.info(f"Applying filter: {filters}")
+            vector_search["filter"] = {
+                "id": {"gte": filters["id"]},
+            }
+        pipeline = [
+            {"$vectorSearch": vector_search},
+            {
+                "$project": {
+                    "_id": 0,
+                    self.id_field: 1,
+                    "score": {"$meta": "vectorSearchScore"},  # Include similarity score
+                }
+            },
+        ]
+
+        results = list(self.collection.aggregate(pipeline))
+        return [doc[self.id_field] for doc in results]
+
+    def optimize(self, data_size: int | None = None) -> None:
+        """MongoDB vector search indexes are self-optimizing"""
+        log.info("optimize for search")
+        self._create_index()
+        self._wait_for_index_ready("vector_index")
+
+    def ready_to_load(self) -> None:
+        """MongoDB is always ready to load"""

vectordb_bench/backend/clients/pgvector/cli.py

@@ -82,7 +82,17 @@ class PgVectorTypedDict(CommonTypedDict):
         click.option(
             "--quantization-type",
             type=click.Choice(["none", "bit", "halfvec"]),
-            help="quantization type for vectors",
+            help="quantization type for vectors (in index)",
+            required=False,
+        ),
+    ]
+    table_quantization_type: Annotated[
+        str | None,
+        click.option(
+            "--table-quantization-type",
+            type=click.Choice(["none", "bit", "halfvec"]),
+            help="quantization type for vectors (in table). "
+            "If equal to bit, the parameter quantization_type will be set to bit too.",
             required=False,
         ),
     ]
@@ -146,6 +156,7 @@ def PgVectorIVFFlat(
             lists=parameters["lists"],
             probes=parameters["probes"],
             quantization_type=parameters["quantization_type"],
+            table_quantization_type=parameters["table_quantization_type"],
             reranking=parameters["reranking"],
             reranking_metric=parameters["reranking_metric"],
             quantized_fetch_limit=parameters["quantized_fetch_limit"],
@@ -182,6 +193,7 @@ def PgVectorHNSW(
             maintenance_work_mem=parameters["maintenance_work_mem"],
             max_parallel_workers=parameters["max_parallel_workers"],
             quantization_type=parameters["quantization_type"],
+            table_quantization_type=parameters["table_quantization_type"],
             reranking=parameters["reranking"],
             reranking_metric=parameters["reranking_metric"],
             quantized_fetch_limit=parameters["quantized_fetch_limit"],

vectordb_bench/backend/clients/pgvector/config.py

@@ -80,7 +80,12 @@ class PgVectorIndexConfig(BaseModel, DBCaseConfig):
 
         if d.get(self.quantization_type) is None:
             return d.get("_fallback").get(self.metric_type)
-        return d.get(self.quantization_type).get(self.metric_type)
+        metric = d.get(self.quantization_type).get(self.metric_type)
+        # If using binary quantization for the index, use a bit metric
+        # no matter what metric was selected for vector or halfvec data
+        if self.quantization_type == "bit" and metric is None:
+            return "bit_hamming_ops"
+        return metric
 
     def parse_metric_fun_op(self) -> LiteralString:
         if self.quantization_type == "bit":
@@ -168,14 +173,19 @@ class PgVectorIVFFlatConfig(PgVectorIndexConfig):
     maintenance_work_mem: str | None = None
     max_parallel_workers: int | None = None
     quantization_type: str | None = None
+    table_quantization_type: str | None
     reranking: bool | None = None
     quantized_fetch_limit: int | None = None
     reranking_metric: str | None = None
 
     def index_param(self) -> PgVectorIndexParam:
         index_parameters = {"lists": self.lists}
-        if self.quantization_type == "none":
-            self.quantization_type = None
+        if self.quantization_type == "none" or self.quantization_type is None:
+            self.quantization_type = "vector"
+        if self.table_quantization_type == "none" or self.table_quantization_type is None:
+            self.table_quantization_type = "vector"
+        if self.table_quantization_type == "bit":
+            self.quantization_type = "bit"
         return {
             "metric": self.parse_metric(),
             "index_type": self.index.value,
@@ -183,6 +193,7 @@ class PgVectorIVFFlatConfig(PgVectorIndexConfig):
             "maintenance_work_mem": self.maintenance_work_mem,
             "max_parallel_workers": self.max_parallel_workers,
             "quantization_type": self.quantization_type,
+            "table_quantization_type": self.table_quantization_type,
         }
 
     def search_param(self) -> PgVectorSearchParam:
@@ -212,14 +223,19 @@ class PgVectorHNSWConfig(PgVectorIndexConfig):
     maintenance_work_mem: str | None = None
     max_parallel_workers: int | None = None
     quantization_type: str | None = None
+    table_quantization_type: str | None
     reranking: bool | None = None
     quantized_fetch_limit: int | None = None
     reranking_metric: str | None = None
 
     def index_param(self) -> PgVectorIndexParam:
         index_parameters = {"m": self.m, "ef_construction": self.ef_construction}
-        if self.quantization_type == "none":
-            self.quantization_type = None
+        if self.quantization_type == "none" or self.quantization_type is None:
+            self.quantization_type = "vector"
+        if self.table_quantization_type == "none" or self.table_quantization_type is None:
+            self.table_quantization_type = "vector"
+        if self.table_quantization_type == "bit":
+            self.quantization_type = "bit"
         return {
             "metric": self.parse_metric(),
             "index_type": self.index.value,
@@ -227,6 +243,7 @@ class PgVectorHNSWConfig(PgVectorIndexConfig):
             "maintenance_work_mem": self.maintenance_work_mem,
             "max_parallel_workers": self.max_parallel_workers,
             "quantization_type": self.quantization_type,
+            "table_quantization_type": self.table_quantization_type,
         }
 
     def search_param(self) -> PgVectorSearchParam:

vectordb_bench/backend/clients/pgvector/pgvector.py

@@ -94,7 +94,7 @@ class PgVector(VectorDB):
         reranking = self.case_config.search_param()["reranking"]
         column_name = (
             sql.SQL("binary_quantize({0})").format(sql.Identifier("embedding"))
-            if index_param["quantization_type"] == "bit"
+            if index_param["quantization_type"] == "bit" and index_param["table_quantization_type"] != "bit"
             else sql.SQL("embedding")
         )
         search_vector = (
@@ -104,7 +104,8 @@ class PgVector(VectorDB):
         )
 
         # The following sections assume that the quantization_type value matches the quantization function name
-        if index_param["quantization_type"] is not None:
+        if index_param["quantization_type"] != index_param["table_quantization_type"]:
+            # Reranking makes sense only if table quantization is not "bit"
             if index_param["quantization_type"] == "bit" and reranking:
                 # Embeddings needs to be passed to binary_quantize function if quantization_type is bit
                 search_query = sql.Composed(
@@ -113,7 +114,7 @@ class PgVector(VectorDB):
                             """
                             SELECT i.id
                             FROM (
-                                SELECT id, embedding {reranking_metric_fun_op} %s::vector AS distance
+                                SELECT id, embedding {reranking_metric_fun_op} %s::{table_quantization_type} AS distance
                                 FROM public.{table_name} {where_clause}
                                 ORDER BY {column_name}::{quantization_type}({dim})
                             """,
@@ -123,6 +124,8 @@ class PgVector(VectorDB):
                             reranking_metric_fun_op=sql.SQL(
                                 self.case_config.search_param()["reranking_metric_fun_op"],
                             ),
+                            search_vector=search_vector,
+                            table_quantization_type=sql.SQL(index_param["table_quantization_type"]),
                             quantization_type=sql.SQL(index_param["quantization_type"]),
                             dim=sql.Literal(self.dim),
                             where_clause=sql.SQL("WHERE id >= %s") if filtered else sql.SQL(""),
@@ -130,7 +133,7 @@ class PgVector(VectorDB):
                         sql.SQL(self.case_config.search_param()["metric_fun_op"]),
                         sql.SQL(
                             """
-                                {search_vector}
+                                {search_vector}::{quantization_type}({dim})
                                 LIMIT {quantized_fetch_limit}
                             ) i
                             ORDER BY i.distance
@@ -138,6 +141,8 @@ class PgVector(VectorDB):
                             """,
                         ).format(
                             search_vector=search_vector,
+                            quantization_type=sql.SQL(index_param["quantization_type"]),
+                            dim=sql.Literal(self.dim),
                             quantized_fetch_limit=sql.Literal(
                                 self.case_config.search_param()["quantized_fetch_limit"],
                             ),
@@ -160,10 +165,12 @@ class PgVector(VectorDB):
                             where_clause=sql.SQL("WHERE id >= %s") if filtered else sql.SQL(""),
                         ),
                         sql.SQL(self.case_config.search_param()["metric_fun_op"]),
-                        sql.SQL(" {search_vector} LIMIT %s::int").format(
+                        sql.SQL(" {search_vector}::{quantization_type}({dim}) LIMIT %s::int").format(
                             search_vector=search_vector,
+                            quantization_type=sql.SQL(index_param["quantization_type"]),
+                            dim=sql.Literal(self.dim),
                         ),
-                    ],
+                    ]
                 )
         else:
             search_query = sql.Composed(
@@ -175,8 +182,12 @@ class PgVector(VectorDB):
                         where_clause=sql.SQL("WHERE id >= %s") if filtered else sql.SQL(""),
                     ),
                     sql.SQL(self.case_config.search_param()["metric_fun_op"]),
-                    sql.SQL(" %s::vector LIMIT %s::int"),
-                ],
+                    sql.SQL(" {search_vector}::{quantization_type}({dim}) LIMIT %s::int").format(
+                        search_vector=search_vector,
+                        quantization_type=sql.SQL(index_param["quantization_type"]),
+                        dim=sql.Literal(self.dim),
+                    ),
+                ]
             )
 
         return search_query
@@ -323,7 +334,7 @@ class PgVector(VectorDB):
                 )
         with_clause = sql.SQL("WITH ({});").format(sql.SQL(", ").join(options)) if any(options) else sql.Composed(())
 
-        if index_param["quantization_type"] is not None:
+        if index_param["quantization_type"] != index_param["table_quantization_type"]:
             index_create_sql = sql.SQL(
                 """
                 CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name}
@@ -365,14 +376,23 @@ class PgVector(VectorDB):
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
 
+        index_param = self.case_config.index_param()
+
         try:
             log.info(f"{self.name} client create table : {self.table_name}")
 
             # create table
             self.cursor.execute(
                 sql.SQL(
-                    "CREATE TABLE IF NOT EXISTS public.{table_name} (id BIGINT PRIMARY KEY, embedding vector({dim}));",
-                ).format(table_name=sql.Identifier(self.table_name), dim=dim),
+                    """
+                    CREATE TABLE IF NOT EXISTS public.{table_name}
+                    (id BIGINT PRIMARY KEY, embedding {table_quantization_type}({dim}));
+                    """
+                ).format(
+                    table_name=sql.Identifier(self.table_name),
+                    table_quantization_type=sql.SQL(index_param["table_quantization_type"]),
+                    dim=dim,
+                )
             )
             self.cursor.execute(
                 sql.SQL(
@@ -393,18 +413,41 @@ class PgVector(VectorDB):
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
 
+        index_param = self.case_config.index_param()
+
         try:
             metadata_arr = np.array(metadata)
             embeddings_arr = np.array(embeddings)
 
-            with self.cursor.copy(
-                sql.SQL("COPY public.{table_name} FROM STDIN (FORMAT BINARY)").format(
-                    table_name=sql.Identifier(self.table_name),
-                ),
-            ) as copy:
-                copy.set_types(["bigint", "vector"])
-                for i, row in enumerate(metadata_arr):
-                    copy.write_row((row, embeddings_arr[i]))
+            if index_param["table_quantization_type"] == "bit":
+                with self.cursor.copy(
+                    sql.SQL("COPY public.{table_name} FROM STDIN (FORMAT TEXT)").format(
+                        table_name=sql.Identifier(self.table_name)
+                    )
+                ) as copy:
+                    # Same logic as pgvector binary_quantize
+                    for i, row in enumerate(metadata_arr):
+                        embeddings_bit = ""
+                        for embedding in embeddings_arr[i]:
+                            if embedding > 0:
+                                embeddings_bit += "1"
+                            else:
+                                embeddings_bit += "0"
+                        copy.write_row((str(row), embeddings_bit))
+            else:
+                with self.cursor.copy(
+                    sql.SQL("COPY public.{table_name} FROM STDIN (FORMAT BINARY)").format(
+                        table_name=sql.Identifier(self.table_name)
+                    )
+                ) as copy:
+                    if index_param["table_quantization_type"] == "halfvec":
+                        copy.set_types(["bigint", "halfvec"])
+                        for i, row in enumerate(metadata_arr):
+                            copy.write_row((row, np.float16(embeddings_arr[i])))
+                    else:
+                        copy.set_types(["bigint", "vector"])
+                        for i, row in enumerate(metadata_arr):
+                            copy.write_row((row, embeddings_arr[i]))
             self.conn.commit()
 
             if kwargs.get("last_batch"):