reformat all

Signed-off-by: min.tian <min.tian.cn@gmail.com>

Author: alwayslove2013

vectordb_bench/backend/clients/__init__.py

@@ -158,7 +158,7 @@ def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901
             from .test.test import Test
 
             return Test
-        
+
         if self == DB.Vespa:
             from .vespa.vespa import Vespa
 
@@ -279,17 +279,16 @@ def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912, C901
             from .test.config import TestConfig
 
             return TestConfig
-        
+
         if self == DB.Vespa:
             from .vespa.config import VespaConfig
 
             return VespaConfig
 
-
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
-    def case_config_cls(  # noqa: PLR0911
+    def case_config_cls(  # noqa: C901, PLR0911, PLR0912
         self,
         index_type: IndexType | None = None,
     ) -> type[DBCaseConfig]:
@@ -377,7 +376,7 @@ def case_config_cls(  # noqa: PLR0911
             from .tidb.config import TiDBIndexConfig
 
             return TiDBIndexConfig
-        
+
         if self == DB.Vespa:
             from .vespa.config import VespaHNSWConfig
 

vectordb_bench/backend/clients/api.py

@@ -162,7 +162,7 @@ def insert_embeddings(
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs,
-    ) -> (int, Exception):
+    ) -> tuple[int, Exception]:
         """Insert the embeddings to the vector database. The default number of embeddings for
         each insert_embeddings is 5000.
 

vectordb_bench/backend/clients/chroma/chroma.py

@@ -65,7 +65,7 @@ def insert_embeddings(
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs: Any,
-    ) -> (int, Exception):
+    ) -> tuple[int, Exception]:
         """Insert embeddings into the database.
 
         Args:
@@ -74,7 +74,7 @@ def insert_embeddings(
             kwargs: other arguments
 
         Returns:
-            (int, Exception): number of embeddings inserted and exception if any
+            tuple[int, Exception]: number of embeddings inserted and exception if any
         """
         ids = [str(i) for i in metadata]
         metadata = [{"id": int(i)} for i in metadata]

vectordb_bench/backend/clients/clickhouse/cli.py

@@ -18,7 +18,7 @@ class ClickhouseTypedDict(TypedDict):
     password: Annotated[str, click.option("--password", type=str, help="DB password")]
     host: Annotated[str, click.option("--host", type=str, help="DB host", required=True)]
     port: Annotated[int, click.option("--port", type=int, default=8123, help="DB Port")]
-    user: Annotated[int, click.option("--user", type=str, default='clickhouse', help="DB user")]
+    user: Annotated[int, click.option("--user", type=str, default="clickhouse", help="DB user")]
     ssl: Annotated[
         bool,
         click.option(

vectordb_bench/backend/clients/clickhouse/clickhouse.py

@@ -1,18 +1,19 @@
 """Wrapper around the Clickhouse vector database over VectorDB"""
 
-import io
 import logging
 from contextlib import contextmanager
 from typing import Any
+
 import clickhouse_connect
-import numpy as np
 
-from ..api import VectorDB, DBCaseConfig
+from ..api import DBCaseConfig, VectorDB
 
 log = logging.getLogger(__name__)
 
+
 class Clickhouse(VectorDB):
     """Use SQLAlchemy instructions"""
+
     def __init__(
         self,
         dim: int,
@@ -32,12 +33,13 @@ def __init__(
         self._vector_field = "embedding"
 
         # construct basic units
-        self.conn =  clickhouse_connect.get_client(
-        host=self.db_config["host"],
-        port=self.db_config["port"],
-        username=self.db_config["user"],
-        password=self.db_config["password"],
-        database=self.db_config["dbname"])
+        self.conn = clickhouse_connect.get_client(
+            host=self.db_config["host"],
+            port=self.db_config["port"],
+            username=self.db_config["user"],
+            password=self.db_config["password"],
+            database=self.db_config["dbname"],
+        )
 
         if drop_old:
             log.info(f"Clickhouse client drop table : {self.table_name}")
@@ -48,20 +50,21 @@ def __init__(
         self.conn = None
 
     @contextmanager
-    def init(self) -> None:
+    def init(self):
         """
         Examples:
             >>> with self.init():
             >>>     self.insert_embeddings()
             >>>     self.search_embedding()
         """
 
-        self.conn =  clickhouse_connect.get_client(
-        host=self.db_config["host"],
-        port=self.db_config["port"],
-        username=self.db_config["user"],
-        password=self.db_config["password"],
-        database=self.db_config["dbname"])
+        self.conn = clickhouse_connect.get_client(
+            host=self.db_config["host"],
+            port=self.db_config["port"],
+            username=self.db_config["user"],
+            password=self.db_config["password"],
+            database=self.db_config["dbname"],
+        )
 
         try:
             yield
@@ -85,9 +88,7 @@ def _create_table(self, dim: int):
             )
 
         except Exception as e:
-            log.warning(
-                f"Failed to create Clickhouse table: {self.table_name} error: {e}"
-            )
+            log.warning(f"Failed to create Clickhouse table: {self.table_name} error: {e}")
             raise e from None
 
     def ready_to_load(self):
@@ -104,16 +105,20 @@ def insert_embeddings(
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs: Any,
-    ) -> (int, Exception):
+    ) -> tuple[int, Exception]:
         assert self.conn is not None, "Connection is not initialized"
 
         try:
             # do not iterate for bulk insert
             items = [metadata, embeddings]
 
-            self.conn.insert(table=self.table_name, data=items,
-                             column_names=['id', 'embedding'], column_type_names=['UInt32', 'Array(Float64)'],
-                             column_oriented=True)
+            self.conn.insert(
+                table=self.table_name,
+                data=items,
+                column_names=["id", "embedding"],
+                column_type_names=["UInt32", "Array(Float64)"],
+                column_oriented=True,
+            )
             return len(metadata), None
         except Exception as e:
             log.warning(f"Failed to insert data into Clickhouse table ({self.table_name}), error: {e}")
@@ -128,22 +133,24 @@ def search_embedding(
     ) -> list[int]:
         assert self.conn is not None, "Connection is not initialized"
 
-        index_param = self.case_config.index_param()
+        index_param = self.case_config.index_param()  # noqa: F841
         search_param = self.case_config.search_param()
 
         if filters:
             gt = filters.get("id")
-            filterSql = (f'SELECT id, {search_param["metric_type"]}(embedding,{query}) AS score '
-                         f'FROM {self.db_config["dbname"]}.{self.table_name}  '
-                         f'WHERE id > {gt} '
-                         f'ORDER BY score LIMIT {k};'
-                         )
-            result = self.conn.query(filterSql).result_rows
+            filter_sql = (
+                f'SELECT id, {search_param["metric_type"]}(embedding,{query}) AS score '  # noqa: S608
+                f'FROM {self.db_config["dbname"]}.{self.table_name}  '
+                f"WHERE id > {gt} "
+                f"ORDER BY score LIMIT {k};"
+            )
+            result = self.conn.query(filter_sql).result_rows
             return [int(row[0]) for row in result]
-        else:
-            selectSql = (f'SELECT id, {search_param["metric_type"]}(embedding,{query}) AS score '
-                         f'FROM {self.db_config["dbname"]}.{self.table_name}  '
-                         f'ORDER BY score LIMIT {k};'
-                         )
-            result = self.conn.query(selectSql).result_rows
+        else:  # noqa: RET505
+            select_sql = (
+                f'SELECT id, {search_param["metric_type"]}(embedding,{query}) AS score '  # noqa: S608
+                f'FROM {self.db_config["dbname"]}.{self.table_name}  '
+                f"ORDER BY score LIMIT {k};"
+            )
+            result = self.conn.query(select_sql).result_rows
             return [int(row[0]) for row in result]

vectordb_bench/backend/clients/elastic_cloud/elastic_cloud.py

@@ -81,7 +81,7 @@ def insert_embeddings(
         embeddings: Iterable[list[float]],
         metadata: list[int],
         **kwargs,
-    ) -> (int, Exception):
+    ) -> tuple[int, Exception]:
         """Insert the embeddings to the elasticsearch."""
         assert self.client is not None, "should self.init() first"
 

vectordb_bench/backend/clients/mariadb/cli.py

@@ -1,4 +1,4 @@
-from typing import Annotated, Optional, Unpack
+from typing import Annotated, Unpack
 
 import click
 from pydantic import SecretStr
@@ -15,68 +15,84 @@
 
 class MariaDBTypedDict(CommonTypedDict):
     user_name: Annotated[
-        str, click.option("--username",
-                          type=str,
-                          help="Username",
-                          required=True,
-                          ),
+        str,
+        click.option(
+            "--username",
+            type=str,
+            help="Username",
+            required=True,
+        ),
     ]
     password: Annotated[
-        str, click.option("--password",
-                     type=str,
-                     help="Password",
-                     required=True,
-                     ),
+        str,
+        click.option(
+            "--password",
+            type=str,
+            help="Password",
+            required=True,
+        ),
     ]
 
     host: Annotated[
-        str, click.option("--host",
-                          type=str,
-                          help="Db host",
-                          default="127.0.0.1",
-                          ),
+        str,
+        click.option(
+            "--host",
+            type=str,
+            help="Db host",
+            default="127.0.0.1",
+        ),
     ]
 
     port: Annotated[
-        int, click.option("--port",
-                          type=int,
-                          default=3306,
-                          help="Db Port",
-                          ),
+        int,
+        click.option(
+            "--port",
+            type=int,
+            default=3306,
+            help="Db Port",
+        ),
     ]
 
     storage_engine: Annotated[
-        int, click.option("--storage-engine",
-                          type=click.Choice(["InnoDB", "MyISAM"]),
-                          help="DB storage engine",
-                          required=True,
-                          ),
+        int,
+        click.option(
+            "--storage-engine",
+            type=click.Choice(["InnoDB", "MyISAM"]),
+            help="DB storage engine",
+            required=True,
+        ),
     ]
 
+
 class MariaDBHNSWTypedDict(MariaDBTypedDict):
-    ...
     m: Annotated[
-        Optional[int], click.option("--m",
-                                    type=int,
-                                    help="M parameter in MHNSW vector indexing",
-                                    required=False,
-                                    ),
+        int | None,
+        click.option(
+            "--m",
+            type=int,
+            help="M parameter in MHNSW vector indexing",
+            required=False,
+        ),
     ]
 
     ef_search: Annotated[
-        Optional[int], click.option("--ef-search",
-                                    type=int,
-                                    help="MariaDB system variable mhnsw_min_limit",
-                                    required=False,
-                                    ),
+        int | None,
+        click.option(
+            "--ef-search",
+            type=int,
+            help="MariaDB system variable mhnsw_min_limit",
+            required=False,
+        ),
     ]
 
     max_cache_size: Annotated[
-        Optional[int], click.option("--max-cache-size",
-                                    type=int,
-                                    help="MariaDB system variable mhnsw_max_cache_size",
-                                    required=False,
-                                    ),
+        int | None,
+        click.option(
+            "--max-cache-size",
+            type=int,
+            help="MariaDB system variable mhnsw_max_cache_size",
+            required=False,
+        ),
     ]
 
 

vectordb_bench/backend/clients/mariadb/config.py

@@ -7,7 +7,7 @@
 
 class MariaDBConfigDict(TypedDict):
     """These keys will be directly used as kwargs in mariadb connection string,
-        so the names must match exactly mariadb API"""
+    so the names must match exactly mariadb API"""
 
     user: str
     password: str
@@ -44,6 +44,7 @@ def parse_metric(self) -> str:
         msg = f"Metric type {self.metric_type} is not supported!"
         raise ValueError(msg)
 
+
 class MariaDBHNSWConfig(MariaDBIndexConfig, DBCaseConfig):
     M: int | None
     ef_search: int | None
@@ -68,7 +69,5 @@ def search_param(self) -> dict:
 
 
 _mariadb_case_config = {
-        IndexType.HNSW: MariaDBHNSWConfig,
+    IndexType.HNSW: MariaDBHNSWConfig,
 }
-
-