diff --git a/pyproject.toml b/pyproject.toml index 6706a3d4e..2baeb16e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,8 +39,9 @@ dependencies = [ "environs", "pydantic=0.10.1", + "ujson", ] dynamic = ["version"] diff --git a/vectordb_bench/backend/clients/polardb/polardb.py b/vectordb_bench/backend/clients/polardb/polardb.py index f42b6fca5..d53638b0e 100644 --- a/vectordb_bench/backend/clients/polardb/polardb.py +++ b/vectordb_bench/backend/clients/polardb/polardb.py @@ -109,14 +109,14 @@ def init(self): db_name = self.db_config["database"] hint = "/*+ SET_VAR(imci_enable_fast_vector_search=on) */" - self.insert_sql = f"INSERT INTO {db_name}.{self.table_name} (id, v) VALUES (%s, _binary %s)" # noqa: S608 + self.insert_sql = f"INSERT INTO {db_name}.{self.table_name} (id, v) VALUES (%s, _binary %s)" self.select_sql = ( - f"SELECT {hint} id FROM {db_name}.{self.table_name} " # noqa: S608 + f"SELECT {hint} id FROM {db_name}.{self.table_name} " f"ORDER BY DISTANCE(v, _binary %s, '{metric_type}') " f"LIMIT %s" ) self.select_sql_with_filter = ( - f"SELECT id FROM {db_name}.{self.table_name} " # noqa: S608 + f"SELECT id FROM {db_name}.{self.table_name} " f"WHERE id >= %s " f"ORDER BY DISTANCE(v, _binary %s, '{metric_type}') " f"LIMIT %s" @@ -218,7 +218,7 @@ def _insert_batch(self, embeddings: list[list[float]], metadata: list[int], offs conn, cursor = self._create_connection() try: db_name = self.db_config["database"] - insert_sql = f"INSERT INTO {db_name}.{self.table_name} (id, v) VALUES (%s, _binary %s)" # noqa: S608 + insert_sql = f"INSERT INTO {db_name}.{self.table_name} (id, v) VALUES (%s, _binary %s)" batch_data = [] for i in range(offset, offset + size): batch_data.append((int(metadata[i]), self.vector_to_hex(embeddings[i])))