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

vectordb_bench/backend/clients/tidb/cli.py

@@ -0,0 +1,98 @@
+from typing import Annotated, Unpack
+
+import click
+from pydantic import SecretStr
+
+from vectordb_bench.backend.clients import DB
+
+from ....cli.cli import CommonTypedDict, cli, click_parameter_decorators_from_typed_dict, run
+
+
+class TiDBTypedDict(CommonTypedDict):
+    user_name: Annotated[
+        str,
+        click.option(
+            "--username",
+            type=str,
+            help="Username",
+            default="root",
+            show_default=True,
+            required=True,
+        ),
+    ]
+    password: Annotated[
+        str,
+        click.option(
+            "--password",
+            type=str,
+            default="",
+            show_default=True,
+            help="Password",
+        ),
+    ]
+    host: Annotated[
+        str,
+        click.option(
+            "--host",
+            type=str,
+            default="127.0.0.1",
+            show_default=True,
+            required=True,
+            help="Db host",
+        ),
+    ]
+    port: Annotated[
+        int,
+        click.option(
+            "--port",
+            type=int,
+            default=4000,
+            show_default=True,
+            required=True,
+            help="Db Port",
+        ),
+    ]
+    db_name: Annotated[
+        str,
+        click.option(
+            "--db-name",
+            type=str,
+            default="test",
+            show_default=True,
+            required=True,
+            help="Db name",
+        ),
+    ]
+    ssl: Annotated[
+        bool,
+        click.option(
+            "--ssl/--no-ssl",
+            default=False,
+            show_default=True,
+            is_flag=True,
+            help="Enable or disable SSL, for TiDB Serverless SSL must be enabled",
+        ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(TiDBTypedDict)
+def TiDB(
+    **parameters: Unpack[TiDBTypedDict],
+):
+    from .config import TiDBConfig, TiDBIndexConfig
+
+    run(
+        db=DB.TiDB,
+        db_config=TiDBConfig(
+            db_label=parameters["db_label"],
+            user_name=parameters["username"],
+            password=SecretStr(parameters["password"]),
+            host=parameters["host"],
+            port=parameters["port"],
+            db_name=parameters["db_name"],
+            ssl=parameters["ssl"],
+        ),
+        db_case_config=TiDBIndexConfig(),
+        **parameters,
+    )

vectordb_bench/backend/clients/tidb/config.py

@@ -0,0 +1,49 @@
+from pydantic import SecretStr, BaseModel, validator
+from ..api import DBConfig, DBCaseConfig, MetricType
+
+
+class TiDBConfig(DBConfig):
+    user_name: str = "root"
+    password: SecretStr
+    host: str = "127.0.0.1"
+    port: int = 4000
+    db_name: str = "test"
+    ssl: bool = False
+
+    @validator("*")
+    def not_empty_field(cls, v: any, field: any):
+        return v
+
+    def to_dict(self) -> dict:
+        pwd_str = self.password.get_secret_value()
+        return {
+            "host": self.host,
+            "port": self.port,
+            "user": self.user_name,
+            "password": pwd_str,
+            "database": self.db_name,
+            "ssl_verify_cert": self.ssl,
+            "ssl_verify_identity": self.ssl,
+        }
+
+
+class TiDBIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+
+    def get_metric_fn(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "vec_l2_distance"
+        elif self.metric_type == MetricType.COSINE:
+            return "vec_cosine_distance"
+        else:
+            raise ValueError(f"Unsupported metric type: {self.metric_type}")
+
+    def index_param(self) -> dict:
+        return {
+            "metric_fn": self.get_metric_fn(),
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_fn": self.get_metric_fn(),
+        }

vectordb_bench/backend/clients/tidb/tidb.py

@@ -0,0 +1,234 @@
+import concurrent.futures
+import io
+import logging
+import time
+from contextlib import contextmanager
+from typing import Any, Optional, Tuple
+
+import pymysql
+
+from ..api import VectorDB
+from .config import TiDBIndexConfig
+
+log = logging.getLogger(__name__)
+
+
+class TiDB(VectorDB):
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: TiDBIndexConfig,
+        collection_name: str = "vector_bench_test",
+        drop_old: bool = False,
+        **kwargs,
+    ):
+        self.name = "TiDB"
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.table_name = collection_name
+        self.dim = dim
+        self.conn = None  # To be inited by init()
+        self.cursor = None  # To be inited by init()
+
+        self.search_fn = db_case_config.search_param()["metric_fn"]
+
+        if drop_old:
+            self._drop_table()
+            self._create_table()
+
+    @contextmanager
+    def init(self):
+        with self._get_connection() as (conn, cursor):
+            self.conn = conn
+            self.cursor = cursor
+            try:
+                yield
+            finally:
+                self.conn = None
+                self.cursor = None
+
+    @contextmanager
+    def _get_connection(self):
+        with pymysql.connect(**self.db_config) as conn:
+            conn.autocommit = False
+            with conn.cursor() as cursor:
+                yield conn, cursor
+
+    def _drop_table(self):
+        try:
+            with self._get_connection() as (conn, cursor):
+                cursor.execute(f"DROP TABLE IF EXISTS {self.table_name}")
+                conn.commit()
+        except Exception as e:
+            log.warning("Failed to drop table: %s error: %s", self.table_name, e)
+            raise e
+
+    def _create_table(self):
+        try:
+            index_param = self.case_config.index_param()
+            with self._get_connection() as (conn, cursor):
+                cursor.execute(
+                    f"""
+                    CREATE TABLE {self.table_name} (
+                        id BIGINT PRIMARY KEY,
+                        embedding VECTOR({self.dim}) NOT NULL,
+                        VECTOR INDEX (({index_param["metric_fn"]}(embedding)))
+                    );
+                    """
+                )
+                conn.commit()
+        except Exception as e:
+            log.warning("Failed to create table: %s error: %s", self.table_name, e)
+            raise e
+
+    def ready_to_load(self) -> bool:
+        pass
+
+    def optimize(self, data_size: int | None = None) -> None:
+        while True:
+            progress = self._optimize_check_tiflash_replica_progress()
+            if progress != 1:
+                log.info("Data replication not ready, progress: %d", progress)
+                time.sleep(2)
+            else:
+                break
+
+        log.info("Waiting TiFlash to catch up...")
+        self._optimize_wait_tiflash_catch_up()
+
+        log.info("Start compacting TiFlash replica...")
+        self._optimize_compact_tiflash()
+
+        log.info("Waiting index build to finish...")
+        log_reduce_seq = 0
+        while True:
+            pending_rows = self._optimize_get_tiflash_index_pending_rows()
+            if pending_rows > 0:
+                if log_reduce_seq % 15 == 0:
+                    log.info("Index not fully built, pending rows: %d", pending_rows)
+                log_reduce_seq += 1
+                time.sleep(2)
+            else:
+                break
+
+        log.info("Index build finished successfully.")
+
+    def _optimize_check_tiflash_replica_progress(self):
+        try:
+            database = self.db_config["database"]
+            with self._get_connection() as (_, cursor):
+                cursor.execute(
+                    f"""
+                    SELECT PROGRESS FROM information_schema.tiflash_replica
+                    WHERE TABLE_SCHEMA = "{database}" AND TABLE_NAME = "{self.table_name}"
+                    """
+                )
+                result = cursor.fetchone()
+                return result[0]
+        except Exception as e:
+            log.warning("Failed to check TiFlash replica progress: %s", e)
+            raise e
+
+    def _optimize_wait_tiflash_catch_up(self):
+        try:
+            with self._get_connection() as (conn, cursor):
+                cursor.execute('SET @@TIDB_ISOLATION_READ_ENGINES="tidb,tiflash"')
+                conn.commit()
+                cursor.execute(f"SELECT COUNT(*) FROM {self.table_name}")
+                result = cursor.fetchone()
+                return result[0]
+        except Exception as e:
+            log.warning("Failed to wait TiFlash to catch up: %s", e)
+            raise e
+
+    def _optimize_compact_tiflash(self):
+        try:
+            with self._get_connection() as (conn, cursor):
+                cursor.execute(f"ALTER TABLE {self.table_name} COMPACT")
+                conn.commit()
+        except Exception as e:
+            log.warning("Failed to compact table: %s", e)
+            raise e
+
+    def _optimize_get_tiflash_index_pending_rows(self):
+        try:
+            database = self.db_config["database"]
+            with self._get_connection() as (_, cursor):
+                cursor.execute(
+                    f"""
+                    SELECT SUM(ROWS_STABLE_NOT_INDEXED)
+                    FROM information_schema.tiflash_indexes
+                    WHERE TIDB_DATABASE = "{database}" AND TIDB_TABLE = "{self.table_name}"
+                    """
+                )
+                result = cursor.fetchone()
+                return result[0]
+        except Exception as e:
+            log.warning("Failed to read TiFlash index pending rows: %s", e)
+            raise e
+
+    def _insert_embeddings_serial(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        offset: int,
+        size: int,
+    ) -> Exception:
+        try:
+            with self._get_connection() as (conn, cursor):
+                buf = io.StringIO()
+                buf.write(f"INSERT INTO {self.table_name} (id, embedding) VALUES ")
+                for i in range(offset, offset + size):
+                    if i > offset:
+                        buf.write(",")
+                    buf.write(f'({metadata[i]}, "{str(embeddings[i])}")')
+                cursor.execute(buf.getvalue())
+                conn.commit()
+        except Exception as e:
+            log.warning("Failed to insert data into table: %s", e)
+            raise e
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> Tuple[int, Optional[Exception]]:
+        workers = 10
+        # Avoid exceeding MAX_ALLOWED_PACKET (default=64MB)
+        max_batch_size = 64 * 1024 * 1024 // 24 // self.dim
+        batch_size = len(embeddings) // workers
+        if batch_size > max_batch_size:
+            batch_size = max_batch_size
+        with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
+            futures = []
+            for i in range(0, len(embeddings), batch_size):
+                offset = i
+                size = min(batch_size, len(embeddings) - i)
+                future = executor.submit(self._insert_embeddings_serial, embeddings, metadata, offset, size)
+                futures.append(future)
+            done, pending = concurrent.futures.wait(futures, return_when=concurrent.futures.FIRST_EXCEPTION)
+            executor.shutdown(wait=False)
+            for future in done:
+                future.result()
+            for future in pending:
+                future.cancel()
+        return len(metadata), None
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+        timeout: int | None = None,
+        **kwargs: Any,
+    ) -> list[int]:
+        self.cursor.execute(
+            f"""
+            SELECT id FROM {self.table_name}
+            ORDER BY {self.search_fn}(embedding, "{str(query)}") LIMIT {k};
+            """
+        )
+        result = self.cursor.fetchall()
+        return [int(i[0]) for i in result]

vectordb_bench/cli/vectordbbench.py

@@ -1,5 +1,6 @@
 from ..backend.clients.alloydb.cli import AlloyDBScaNN
 from ..backend.clients.aws_opensearch.cli import AWSOpenSearch
+from ..backend.clients.mariadb.cli import MariaDBHNSW
 from ..backend.clients.memorydb.cli import MemoryDB
 from ..backend.clients.milvus.cli import MilvusAutoIndex
 from ..backend.clients.pgdiskann.cli import PgDiskAnn
@@ -10,6 +11,7 @@ from ..backend.clients.redis.cli import Redis
 from ..backend.clients.test.cli import Test
 from ..backend.clients.weaviate_cloud.cli import Weaviate
 from ..backend.clients.zilliz_cloud.cli import ZillizAutoIndex
+from ..backend.clients.tidb.cli import TiDB
 from .cli import cli
 
 cli.add_command(PgVectorHNSW)
@@ -25,6 +27,8 @@ cli.add_command(AWSOpenSearch)
 cli.add_command(PgVectorScaleDiskAnn)
 cli.add_command(PgDiskAnn)
 cli.add_command(AlloyDBScaNN)
+cli.add_command(MariaDBHNSW)
+cli.add_command(TiDB)
 
 
 if __name__ == "__main__":

vectordb_bench/frontend/components/custom/displaypPrams.py

@@ -3,7 +3,7 @@ def displayParams(st):
         """
 - `Folder Path` - The path to the folder containing all the files. Please ensure that all files in the folder are in the `Parquet` format.
   - Vectors data files: The file must be named `train.parquet` and should have two columns: `id` as an incrementing `int` and `emb` as an array of `float32`.
-  - Query test vectors: The file must be named `test.parquet` and should have two columns: `id` as an incrementing `int` and `emb` as an array of `float32`.
+  - Query test vectors: The file must be named `test.parquet` and should have two columns: `id` as an incrementing `int` and `emb` as an array of `float32`. 
   - Ground truth file: The file must be named `neighbors.parquet` and should have two columns: `id` corresponding to query vectors and `neighbors_id` as an array of `int`.
 
 - `Train File Count` - If the vector file is too large, you can consider splitting it into multiple files. The naming format for the split files should be `train-[index]-of-[file_count].parquet`. For example, `train-01-of-10.parquet` represents the second file (0-indexed) among 10 split files.
@@ -11,3 +11,14 @@ def displayParams(st):
 - `Use Shuffled Data` - If you check this option, the vector data files need to be modified. VectorDBBench will load the data labeled with `shuffle`. For example, use `shuffle_train.parquet` instead of `train.parquet` and `shuffle_train-04-of-10.parquet` instead of `train-04-of-10.parquet`. The `id` column in the shuffled data can be in any order.
 """
     )
+    st.caption(
+        """We recommend limiting the number of test query vectors, like 1,000.""",
+        help="""
+When conducting concurrent query tests, Vdbbench creates a large number of processes. 
+To minimize additional communication overhead during testing, 
+we prepare a complete set of test queries for each process, allowing them to run independently.\n
+However, this means that as the number of concurrent processes increases, 
+the number of copied query vectors also increases significantly, 
+which can place substantial pressure on memory resources.
+""",
+    )

vectordb_bench/frontend/components/run_test/submitTask.py

@@ -1,6 +1,8 @@
 from datetime import datetime
+from vectordb_bench import config
 from vectordb_bench.frontend.config import styles
 from vectordb_bench.interface import benchmark_runner
+from vectordb_bench.models import TaskConfig
 
 
 def submitTask(st, tasks, isAllValid):
@@ -47,16 +49,31 @@ def advancedSettings(st):
     k = container[0].number_input("k", min_value=1, value=100, label_visibility="collapsed")
     container[1].caption("K value for number of nearest neighbors to search")
 
-    return index_already_exists, use_aliyun, k
+    container = st.columns([1, 2])
+    defaultconcurrentInput = ",".join(map(str, config.NUM_CONCURRENCY))
+    concurrentInput = container[0].text_input(
+        "Concurrent Input", value=defaultconcurrentInput, label_visibility="collapsed"
+    )
+    container[1].caption("num of concurrencies for search tests to get max-qps")
+    return index_already_exists, use_aliyun, k, concurrentInput
 
 
-def controlPanel(st, tasks, taskLabel, isAllValid):
-    index_already_exists, use_aliyun, k = advancedSettings(st)
+def controlPanel(st, tasks: list[TaskConfig], taskLabel, isAllValid):
+    index_already_exists, use_aliyun, k, concurrentInput = advancedSettings(st)
 
     def runHandler():
         benchmark_runner.set_drop_old(not index_already_exists)
+
+        try:
+            concurrentInput_list = [int(item.strip()) for item in concurrentInput.split(",")]
+        except ValueError:
+            st.write("please input correct number")
+            return None
+
         for task in tasks:
             task.case_config.k = k
+            task.case_config.concurrency_search_config.num_concurrency = concurrentInput_list
+
         benchmark_runner.set_download_address(use_aliyun)
         benchmark_runner.run(tasks, taskLabel)
 

vectordb_bench/frontend/config/dbCaseConfigs.py

@@ -173,6 +173,7 @@ CaseConfigParamInput_IndexType = CaseConfigInput(
             IndexType.GPU_IVF_FLAT.value,
             IndexType.GPU_IVF_PQ.value,
             IndexType.GPU_CAGRA.value,
+            IndexType.GPU_BRUTE_FORCE.value,
         ],
     },
 )
@@ -562,6 +563,7 @@ CaseConfigParamInput_Nlist = CaseConfigInput(
         IndexType.IVFSQ8.value,
         IndexType.GPU_IVF_FLAT.value,
         IndexType.GPU_IVF_PQ.value,
+        IndexType.GPU_BRUTE_FORCE.value,
     ],
 )
 
@@ -579,6 +581,7 @@ CaseConfigParamInput_Nprobe = CaseConfigInput(
         IndexType.IVFSQ8.value,
         IndexType.GPU_IVF_FLAT.value,
         IndexType.GPU_IVF_PQ.value,
+        IndexType.GPU_BRUTE_FORCE.value,
     ],
 )
 
@@ -703,6 +706,7 @@ CaseConfigParamInput_cache_dataset_on_device = CaseConfigInput(
         IndexType.GPU_CAGRA.value,
         IndexType.GPU_IVF_PQ.value,
         IndexType.GPU_IVF_FLAT.value,
+        IndexType.GPU_BRUTE_FORCE.value,
     ],
 )
 
@@ -720,6 +724,7 @@ CaseConfigParamInput_refine_ratio = CaseConfigInput(
         IndexType.GPU_CAGRA.value,
         IndexType.GPU_IVF_PQ.value,
         IndexType.GPU_IVF_FLAT.value,
+        IndexType.GPU_BRUTE_FORCE.value,
     ],
 )
 
@@ -818,6 +823,19 @@ CaseConfigParamInput_QuantizationRatio_PgVectoRS = CaseConfigInput(
     ],
 )
 
+CaseConfigParamInput_TableQuantizationType_PgVector = CaseConfigInput(
+    label=CaseConfigParamType.tableQuantizationType,
+    inputType=InputType.Option,
+    inputConfig={
+        "options": ["none", "bit", "halfvec"],
+    },
+    isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
+    in [
+        IndexType.HNSW.value,
+        IndexType.IVFFlat.value,
+    ],
+)
+
 CaseConfigParamInput_max_parallel_workers_PgVectorRS = CaseConfigInput(
     label=CaseConfigParamType.max_parallel_workers,
     displayLabel="Max parallel workers",
@@ -1040,6 +1058,84 @@ CaseConfigParamInput_NumCandidates_AliES = CaseConfigInput(
     },
 )
 
+CaseConfigParamInput_IndexType_MariaDB = CaseConfigInput(
+    label=CaseConfigParamType.IndexType,
+    inputHelp="Select Index Type",
+    inputType=InputType.Option,
+    inputConfig={
+        "options": [
+            IndexType.HNSW.value,
+        ],
+    },
+)
+
+CaseConfigParamInput_StorageEngine_MariaDB = CaseConfigInput(
+    label=CaseConfigParamType.storage_engine,
+    inputHelp="Select Storage Engine",
+    inputType=InputType.Option,
+    inputConfig={
+        "options": ["InnoDB", "MyISAM"],
+    },
+)
+
+CaseConfigParamInput_M_MariaDB = CaseConfigInput(
+    label=CaseConfigParamType.M,
+    inputHelp="M parameter in MHNSW vector indexing",
+    inputType=InputType.Number,
+    inputConfig={
+        "min": 3,
+        "max": 200,
+        "value": 6,
+    },
+    isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
+    == IndexType.HNSW.value,
+)
+
+CaseConfigParamInput_EFSearch_MariaDB = CaseConfigInput(
+    label=CaseConfigParamType.ef_search,
+    inputHelp="mhnsw_ef_search",
+    inputType=InputType.Number,
+    inputConfig={
+        "min": 1,
+        "max": 10000,
+        "value": 20,
+    },
+    isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
+    == IndexType.HNSW.value,
+)
+
+CaseConfigParamInput_CacheSize_MariaDB = CaseConfigInput(
+    label=CaseConfigParamType.max_cache_size,
+    inputHelp="mhnsw_max_cache_size",
+    inputType=InputType.Number,
+    inputConfig={
+        "min": 1048576,
+        "max": (1 << 53) - 1,
+        "value": 16 * 1024 ** 3,
+    },
+    isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
+    == IndexType.HNSW.value,
+)
+
+CaseConfigParamInput_MongoDBQuantizationType = CaseConfigInput(
+    label=CaseConfigParamType.mongodb_quantization_type,
+    inputType=InputType.Option,
+    inputConfig={
+        "options": ["none", "scalar", "binary"],
+    },
+)
+
+
+CaseConfigParamInput_MongoDBNumCandidatesRatio = CaseConfigInput(
+    label=CaseConfigParamType.mongodb_num_candidates_ratio,
+    inputType=InputType.Number,
+    inputConfig={
+        "min": 10,
+        "max": 20,
+        "value": 10,
+    },
+)
+
 
 MilvusLoadConfig = [
     CaseConfigParamInput_IndexType,
@@ -1113,6 +1209,7 @@ PgVectorLoadingConfig = [
     CaseConfigParamInput_m,
     CaseConfigParamInput_EFConstruction_PgVector,
     CaseConfigParamInput_QuantizationType_PgVector,
+    CaseConfigParamInput_TableQuantizationType_PgVector,
     CaseConfigParamInput_maintenance_work_mem_PgVector,
     CaseConfigParamInput_max_parallel_workers_PgVector,
 ]
@@ -1124,6 +1221,7 @@ PgVectorPerformanceConfig = [
     CaseConfigParamInput_Lists_PgVector,
     CaseConfigParamInput_Probes_PgVector,
     CaseConfigParamInput_QuantizationType_PgVector,
+    CaseConfigParamInput_TableQuantizationType_PgVector,
     CaseConfigParamInput_maintenance_work_mem_PgVector,
     CaseConfigParamInput_max_parallel_workers_PgVector,
     CaseConfigParamInput_reranking_PgVector,
@@ -1224,6 +1322,28 @@ AliyunElasticsearchPerformanceConfig = [
     CaseConfigParamInput_NumCandidates_AliES,
 ]
 
+MongoDBLoadingConfig = [
+    CaseConfigParamInput_MongoDBQuantizationType,
+]
+MongoDBPerformanceConfig = [
+    CaseConfigParamInput_MongoDBQuantizationType,
+    CaseConfigParamInput_MongoDBNumCandidatesRatio,
+]
+
+MariaDBLoadingConfig = [
+    CaseConfigParamInput_IndexType_MariaDB,
+    CaseConfigParamInput_StorageEngine_MariaDB,
+    CaseConfigParamInput_M_MariaDB,
+    CaseConfigParamInput_CacheSize_MariaDB,
+]
+MariaDBPerformanceConfig = [
+    CaseConfigParamInput_IndexType_MariaDB,
+    CaseConfigParamInput_StorageEngine_MariaDB,
+    CaseConfigParamInput_M_MariaDB,
+    CaseConfigParamInput_CacheSize_MariaDB,
+    CaseConfigParamInput_EFSearch_MariaDB,
+]
+
 CASE_CONFIG_MAP = {
     DB.Milvus: {
         CaseLabel.Load: MilvusLoadConfig,
@@ -1272,4 +1392,12 @@ CASE_CONFIG_MAP = {
         CaseLabel.Load: AliyunOpensearchLoadingConfig,
         CaseLabel.Performance: AliyunOpenSearchPerformanceConfig,
     },
+    DB.MongoDB: {
+        CaseLabel.Load: MongoDBLoadingConfig,
+        CaseLabel.Performance: MongoDBPerformanceConfig,
+    },
+    DB.MariaDB: {
+        CaseLabel.Load: MariaDBLoadingConfig,
+        CaseLabel.Performance: MariaDBPerformanceConfig,
+    },
 }

vectordb_bench/frontend/config/styles.py

@@ -47,6 +47,7 @@ DB_TO_ICON = {
     DB.Redis: "https://assets.zilliz.com/Redis_Cloud_74b8bfef39.png",
     DB.Chroma: "https://assets.zilliz.com/chroma_ceb3f06ed7.png",
     DB.AWSOpenSearch: "https://assets.zilliz.com/opensearch_1eee37584e.jpeg",
+    DB.TiDB: "https://img2.pingcap.com/forms/3/d/3d7fd5f9767323d6f037795704211ac44b4923d6.png",
 }
 
 # RedisCloud color: #0D6EFD
@@ -61,4 +62,5 @@ COLOR_MAP = {
     DB.PgVector.value: "#4C779A",
     DB.Redis.value: "#0D6EFD",
     DB.AWSOpenSearch.value: "#0DCAF0",
+    DB.TiDB.value: "#0D6EFD",
 }