update

.github/workflows/expose_service_ports.sh

@@ -9,6 +9,6 @@ yq eval '.services["pgvecto-rs"].ports += ["5431:5432"]' -i docker/docker-compos
 yq eval '.services["elasticsearch"].ports += ["9200:9200"]' -i docker/docker-compose.yaml
 yq eval '.services.couchbase-server.ports += ["8091-8096:8091-8096"]' -i docker/docker-compose.yaml
 yq eval '.services.couchbase-server.ports += ["11210:11210"]' -i docker/docker-compose.yaml
-yq eval '.services.tidb.ports += ["4000:4000"]' -i docker/docker-compose.yaml
+yq eval '.services.tidb.ports += ["4000:4000"]' -i docker/tidb/docker-compose.yaml
 
 echo "Ports exposed for sandbox, weaviate, tidb, qdrant, chroma, milvus, pgvector, pgvecto-rs, elasticsearch, couchbase"

.github/workflows/vdb-tests.yml

@@ -51,7 +51,15 @@ jobs:
       - name: Expose Service Ports
         run: sh .github/workflows/expose_service_ports.sh
 
-      - name: Set up Vector Stores (TiDB, Weaviate, Qdrant, PGVector, Milvus, PgVecto-RS, Chroma, MyScale, ElasticSearch, Couchbase)
+      - name: Set up Vector Store (TiDB)
+        uses: hoverkraft-tech/[email protected]
+        with:
+          compose-file: docker/tidb/docker-compose.yaml
+          services: |
+            tidb
+            tiflash
+
+      - name: Set up Vector Stores (Weaviate, Qdrant, PGVector, Milvus, PgVecto-RS, Chroma, MyScale, ElasticSearch, Couchbase)
         uses: hoverkraft-tech/[email protected]
         with:
           compose-file: |
@@ -67,7 +75,9 @@ jobs:
             pgvector
             chroma
             elasticsearch
-            tidb
+
+      - name: Check TiDB Ready
+        run: poetry run -C api python api/tests/integration_tests/vdb/tidb_vector/check_tiflash_ready.py
 
       - name: Test Vector Stores
         run: poetry run -C api bash dev/pytest/pytest_vdb.sh

.gitignore

@@ -163,6 +163,7 @@ docker/volumes/db/data/*
 docker/volumes/redis/data/*
 docker/volumes/weaviate/*
 docker/volumes/qdrant/*
+docker/tidb/volumes/*
 docker/volumes/etcd/*
 docker/volumes/minio/*
 docker/volumes/milvus/*

api/core/rag/datasource/vdb/tidb_vector/tidb_vector.py

@@ -104,14 +104,14 @@ def _create_collection(self, dimension: int):
                         text TEXT NOT NULL,
                         meta JSON NOT NULL,
                         doc_id VARCHAR(64) AS (JSON_UNQUOTE(JSON_EXTRACT(meta, '$.doc_id'))) STORED,
-                        KEY (doc_id),
                         vector VECTOR<FLOAT>({dimension}) NOT NULL COMMENT "hnsw(distance={self._distance_func})",
                         create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
-                        update_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
+                        update_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
+                        KEY (doc_id),
+                        VECTOR INDEX idx_vector ((VEC_COSINE_DISTANCE(vector))) USING HNSW
                     );
                 """)
                 session.execute(create_statement)
-                # tidb vector not support 'CREATE/ADD INDEX' now
                 session.commit()
             redis_client.set(collection_exist_cache_key, 1, ex=3600)
 

api/tests/integration_tests/vdb/tidb_vector/check_tiflash_ready.py

@@ -0,0 +1,57 @@
+import time
+
+import pymysql
+
+
+def check_tiflash_ready() -> bool:
+    try:
+        connection = pymysql.connect(
+            host="localhost",
+            port=4000,
+            user="root",
+            password="",
+        )
+
+        with connection.cursor() as cursor:
+            select_tiflash_query = """
+            SELECT * FROM information_schema.cluster_hardware
+            WHERE TYPE='tiflash'
+            LIMIT 1;
+            """
+            cursor.execute(select_tiflash_query)
+            result = cursor.fetchall()
+            return result is not None and len(result) > 0
+    except Exception as e:
+        print(f"TiFlash is not ready. Exception: {e}")
+        return False
+    finally:
+        if connection:
+            connection.close()
+
+
+def main():
+    max_attempts = 30
+    retry_interval_seconds = 2
+    is_tiflash_ready = False
+    for attempt in range(max_attempts):
+        try:
+            is_tiflash_ready = check_tiflash_ready()
+        except Exception as e:
+            print(f"TiFlash is not ready. Exception: {e}")
+            is_tiflash_ready = False
+
+        if is_tiflash_ready:
+            break
+        else:
+            print(f"Attempt {attempt + 1} failed，retry in {retry_interval_seconds} seconds...")
+            time.sleep(retry_interval_seconds)
+
+    if is_tiflash_ready:
+        print("TiFlash is not ready in TiDB.")
+    else:
+        print(f"TiFlash is not ready in TiDB after {max_attempts} attempting checks.")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()

docker/docker-compose-template.yaml

@@ -207,16 +207,6 @@ services:
       - '${EXPOSE_NGINX_PORT:-80}:${NGINX_PORT:-80}'
       - '${EXPOSE_NGINX_SSL_PORT:-443}:${NGINX_SSL_PORT:-443}'
 
-  # The TiDB vector store.
-  # For production use, please refer to https://github.com/pingcap/tidb-docker-compose
-  tidb:
-    image: pingcap/tidb:v8.4.0
-    profiles:
-      - tidb
-    command:
-      - --store=unistore
-    restart: always
-
   # The Weaviate vector store.
   weaviate:
     image: semitechnologies/weaviate:1.19.0

docker/docker-compose.yaml

@@ -595,16 +595,6 @@ services:
       - '${EXPOSE_NGINX_PORT:-80}:${NGINX_PORT:-80}'
       - '${EXPOSE_NGINX_SSL_PORT:-443}:${NGINX_SSL_PORT:-443}'
 
-  # The TiDB vector store.
-  # For production use, please refer to https://github.com/pingcap/tidb-docker-compose
-  tidb:
-    image: pingcap/tidb:v8.4.0
-    profiles:
-      - tidb
-    command:
-      - --store=unistore
-    restart: always
-
   # The Weaviate vector store.
   weaviate:
     image: semitechnologies/weaviate:1.19.0

docker/tidb/config/pd.toml

@@ -0,0 +1,2 @@
+[replication]
+max-replicas = 1

docker/tidb/config/tiflash-learner.toml

@@ -0,0 +1,10 @@
+log-file = "/logs/tiflash_tikv.log"
+
+[server]
+engine-addr = "tiflash:4030"
+addr = "0.0.0.0:20280"
+advertise-addr = "tiflash:20280"
+status-addr = "tiflash:20292"
+
+[storage]
+data-dir = "/data/flash"

docker/tidb/config/tiflash.toml

@@ -0,0 +1,16 @@
+listen_host = "0.0.0.0"
+path = "/data"
+
+[flash]
+tidb_status_addr = "tidb:10080"
+service_addr = "tiflash:4030"
+
+[flash.proxy]
+config = "/tiflash-learner.toml"
+
+[logger]
+errorlog = "/logs/tiflash_error.log"
+log = "/logs/tiflash.log"
+
+[raft]
+pd_addr = "pd0:2379"

docker/tidb/docker-compose.yaml

@@ -0,0 +1,64 @@
+version: '2.1'
+
+services:
+  pd0:
+    image: pingcap/pd:v8.5.0
+    ports:
+      - "2379"
+    volumes:
+      - ./config/pd.toml:/pd.toml:ro
+      - ./volumes/data:/data
+      - ./volumes/logs:/logs
+    command:
+      - --name=pd0
+      - --client-urls=http://0.0.0.0:2379
+      - --peer-urls=http://0.0.0.0:2380
+      - --advertise-client-urls=http://pd0:2379
+      - --advertise-peer-urls=http://pd0:2380
+      - --initial-cluster=pd0=http://pd0:2380
+      - --data-dir=/data/pd
+      - --config=/pd.toml
+      - --log-file=/logs/pd.log
+    restart: on-failure
+  tikv:
+    image: pingcap/tikv:v8.5.0
+    volumes:
+      - ./volumes/data:/data
+      - ./volumes/logs:/logs
+    command:
+      - --addr=0.0.0.0:20160
+      - --advertise-addr=tikv:20160
+      - --status-addr=tikv:20180
+      - --data-dir=/data/tikv
+      - --pd=pd0:2379
+      - --log-file=/logs/tikv.log
+    depends_on:
+      - "pd0"
+    restart: on-failure
+  tidb:
+    image: pingcap/tidb:v8.5.0
+    ports:
+      - "4000:4000"
+    volumes:
+      - ./volumes/logs:/logs
+    command:
+      - --advertise-address=tidb
+      - --store=tikv
+      - --path=pd0:2379
+      - --log-file=/logs/tidb.log
+    depends_on:
+      - "tikv"
+    restart: on-failure
+  tiflash:
+    image: pingcap/tiflash:v8.5.0
+    volumes:
+      - ./config/tiflash.toml:/tiflash.toml:ro
+      - ./config/tiflash-learner.toml:/tiflash-learner.toml:ro
+      - ./volumes/data:/data
+      - ./volumes/logs:/logs
+    command:
+      - --config=/tiflash.toml
+    depends_on:
+      - "tikv"
+      - "tidb"
+    restart: on-failure