TiDB Vector 搶先體驗之用 TiDB 實現以圖搜圖

本文首發自 TiDB 社區專欄：https://tidb.net/blog/0c5672b9
轉載請註明出處！

前言

最早知道 TiDB 要支持向量化的消息應該是在23年10月份左右，到第一次見到 TiDB Vector 的樣子是在今年1月初，當時 dongxu 在朋友圈發了一張圖：

去年我研究了一段時間的向量數據庫，一直對 TiDB 向量特性非常期待，看到這張圖真的就激動萬分，於是第一時間提交了 waitlist 等待體驗 private beta。

苦等幾個月，它終於來了（目前只對 TiDB Serverless 開放）。迫不及待做個小應用嚐嚐鮮。

waitlist申請入口：https://tidb.cloud/ai

體驗入口：https://tidbcloud.com/

創建 TiDB Vector 實例

在收到體驗邀請郵件後，恭喜你可以開始 TiDB Vector 之旅了。

TiDB Serverless 提供了免費試用額度，對於測試用途綽綽有餘，只需要註冊一個 TiDB Cloud 賬號即可。

創建 TiDB Vector 實例和普通的 TiDB 實例並沒有太大區別，在創建集羣頁面可以看到加入瞭如下開關：

不過要注意的是目前 TiDB Vector 只在 AWS 的eu-central-1可用區開放，選到了其他可用區就看不到這個開關。

這裏只需要填一個集羣名稱就可以開始創建，創建成功後的樣子如下所示：

下面開始進入正題。

關於向量的那些事

一些基礎概念

• 向量：向量就是一組浮點數，在編程語言中通常體現爲 float 數組，數組的長度叫做維度（dim），維度越大精度越高，向量的數學表示是多維座標系中的一個點。例如RGB顏色表示法就是一個簡單的向量示例。
• embedding：中文翻譯叫嵌入，感覺不好理解，實質上就是把非結構化數據（文本、語音、圖片、視頻等）通過一系列算法加工變成向量的過程，這裏面的算法叫做模型（model）。
• 向量檢索：計算兩個向量之間的相似度。

向量檢索初體驗

連接到 TiDB Serverless 後，就可以體驗文章開頭圖片中的向量操作。

創建一張帶有向量字段的表，長度是3維。

CREATE TABLE vector_table (
    id int PRIMARY KEY,
    doc TEXT,
    embedding vector < float > (3)
  );

往表中插入向量數據：

INSERT INTO vector_table VALUES (1, 'apple', '[1,1,1]'), (2, 'banana', '[1,1,2]'), (3, 'dog', '[2,2,2]');

根據指定的向量做搜索：

SELECT *, vec_cosine_distance(embedding, '[1,1,3]') as distance FROM vector_table ORDER BY distance LIMIT 3;

+-----------------------+-----------------------+---------------------+
| id      | doc         | embedding             | distance            |
+-----------------------+-----------------------+---------------------+
| 2       | banana      | [1,1,2]               | 0.015268072165338209|
| 3       | dog         | [2,2,2]               | 0.1296117202215108  |
| 1       | apple       | [1,1,1]               | 0.1296117202215108  |
+---------+-------------+-----------------------+---------------------+

這裏的distance就是兩個向量之間的相似度，這個相似度是用vec_cosine_distance函數計算出來的，意味着兩個向量之間的夾角越小相似性越高，夾角大小用餘弦值來衡量。

還有以一種常用的相似度計算方法是比較兩個向量之間的直線距離，稱爲歐式距離。

這也意味着不管兩個向量是否有關聯性，總是能計算出一個相似度，distance越小相似度越高。

向量檢索原理

前面大概也提到了兩種常用的向量檢索方式：餘弦相似度和歐式距離，不妨從從最簡單的二維向量開始推導一下計算過程。

二維向量對應一個平面座標系，一個向量就是座標系中任意一點，要計算兩點之間的直線距離用勾股定理很容易就能得出，兩點夾角的餘弦值也有公式能直接算出來。

拓展到三維座標系，還是套用上一步的數學公式，只是多了一個座標。

以此類推到n維也是一樣的方法。

以上內容來自我去年講的向量數據庫公開課：https://www.bilibili.com/video/BV1YP411t7Do

可以發現維數越多，對算力的要求就越高，計算時間就越長。

第一個 TiDB AI 應用：以圖搜圖

基礎實現

藉助前面介紹的理論知識，一個以圖搜圖的流程應該是這樣子：

下面我用最簡潔直白的代碼演示整個流程，方便大家理解。

首先肯定是先連接到 TiDB 實例，目前官方提供了python SDK包tidb_vector，對SQLAlchemy、Peewee這樣的 ORM 框架也有支持，具體可參考https://github.com/pingcap/tidb-vector-python

這裏簡單起見直接用pymysql手寫 SQL 操作，以下連接參數都可以從 TiDB Cloud 控制檯獲取：

import pymysql

def GetConnection():
    connection = pymysql.connect(
        host = "xxx.xxx.prod.aws.tidbcloud.com",
        port = 4000,
        user = "xxx.root",
        password = "xxx",
        database = "test",
        ssl_verify_cert = True,
        ssl_verify_identity = True,
        ssl_ca = "C:\\Users\\59131\\Downloads\\isrgrootx1.pem"
    )
    return connection

再借助 Towhee 來簡化 embedding 的處理，裏面包含了常用的非結構化數據到向量數據的轉換模型，用流水線（pipeline）的形式清晰構建整個處理過程。

from towhee import ops,pipe,AutoPipes,AutoConfig,DataCollection

image_pipe = AutoPipes.pipeline('text_image_embedding')

這裏使用默認配置構建了一個text_image_embedding流水線，它專門用於對文本和圖片做向量轉換，從引用的源碼中可以看到它使用的模型是clip_vit_base_patch16，默認模態是image。

@AutoConfig.register
class TextImageEmbeddingConfig(BaseModel):
    model: Optional[str] = 'clip_vit_base_patch16'
    modality: Optional[str] = 'image'
    customize_embedding_op: Optional[Any] = None
    normalize_vec: Optional[bool] = True
    device: Optional[int] = -1

clip_vit_base_patch16是一個512維的模型，因此需要在 TiDB 中創建512維的向量字段。

create table if not exists img_list 
(
    id int PRIMARY KEY, 
    path varchar(200) not null, 
    embedding vector<float>(512)
);

我準備了3000張各種各樣的動物圖片用於測試，把它們依次加載到 TiDB 中，完整代碼爲：

def LoadImage(connection):
    cursor = connection.cursor() 
    cursor.execute("create table if not exists img_list (id int PRIMARY KEY, path varchar(200) not null, embedding vector<float>(512));")
    img_dir='D:\\\\test\\\\'
    files = os.listdir(img_dir)
    for i in range(len(files)):
        path=os.path.join(img_dir, files[i])
        embedding = image_pipe(path).get()[0]
        cursor.execute("INSERT INTO img_list VALUE ("+str(i)+",'"+path+"' , '"+np.array2string(embedding, separator=',')+"');")
    connection.commit()

如果用 ORM 框架的話這裏對數據庫和向量加工操作會簡單些，不需要數組到字符串之間的手工轉換。

加載完成後的數據：

下一步定義出根據指定向量在 TiDB 中檢索的函數：

def SearchInTiDB(connection,vector):
    cursor = connection.cursor() 
    begin_time = datetime.datetime.now()
    cursor.execute("select id,path,vec_cosine_distance(embedding, '"+np.array2string(vector, separator=',')+"') as distance from img_list order by distance limit 3;")
    end_time=datetime.datetime.now()
    print("Search time:",(end_time-begin_time).total_seconds())
    df =pd.DataFrame(cursor.fetchall())
    return df[1]

這裏根據餘弦相似度取出結果最相近的3張圖片，返回它們的文件路徑用於預覽顯示。

下一步用相同的 image pipeline 給指定圖片做 embedding 得到向量，把這個向量傳到 TiDB 中去搜索，最後把搜索結果輸出顯示。

def read_images(img_paths):
    imgs = []
    op = ops.image_decode.cv2_rgb()
    for p in img_paths:
        imgs.append(op(p))
    return imgs
    
def ImageSearch(connection,path):    
    emb = image_pipe(path).get()[0]
    res = SearchInTiDB(connection,emb)
    p = (
        pipe.input('path','search_result')
        .map('path', 'img', ops.image_decode.cv2('rgb'))
        .map('search_result','prev',read_images)
        .output('img','prev')
    )
    DataCollection(p(path,res)).show()

看一下最終搜索效果如何。先看一張已經在圖片庫存在的圖（左邊是待搜索的圖，右邊是搜索結果，按相似度由高到低）：

不能說非常相似，只能說是一模一樣，準確度非常高！再看一下不在圖片庫的搜索效果：

圖片庫裏有幾十種動物，能夠準確搜索出需要的是狗，特別是第一張從圖片色彩、畫面角度、動作神態上來說都非常相似。

使用向量索引優化

沒錯，向量也能加索引，但這個索引和傳統的 B+ Tree 索引有些區別。前面提到向量相似度計算是一個非常消耗 CPU 的過程，如果每次計算都採用全量暴力搜索的方式那麼無疑效率非常低。上一節演示的案例就是用指定的向量與表裏的3000個向量逐一計算，最簡單粗暴的辦法。

向量索引犧牲了一定的準確度來提升性能，通常採用 ANN（近似最近鄰搜索） 算法，HNSW 是最知名的算法之一。TiDB Vector 目前對它已經有了支持：

create table if not exists img_list_hnsw 
(
    id int PRIMARY KEY, 
    path varchar(200) not null, 
    embedding vector<float>(512) COMMENT "hnsw(distance=cosine)"
);

重新把3000張圖片加載到新的img_list_hnsw表做搜索測試。

以下分別是不帶索引和帶索引的查詢耗時，第二次明顯要快很多，如果數據量越大這個差距會越明顯，只是目前還無法通過執行計劃或其他方式區分出索引使用情況。

E:\GitLocal\AITester>python tidb_vec.py
Search time: 0.320241
+------------------------------------+------------------------------------------------------------------------------------------------------+
| img                                | prev                                                                                                 |
+====================================+======================================================================================================+
| Image shape=(900, 900, 3) mode=RGB | [Image shape=(84, 84, 3) mode=RGB,Image shape=(84, 84, 3) mode=RGB,Image shape=(84, 84, 3) mode=RGB] |
+------------------------------------+------------------------------------------------------------------------------------------------------+

E:\GitLocal\AITester>python tidb_vec.py
Search time: 0.239746
+------------------------------------+------------------------------------------------------------------------------------------------------+
| img                                | prev                                                                                                 |
+====================================+======================================================================================================+
| Image shape=(900, 900, 3) mode=RGB | [Image shape=(84, 84, 3) mode=RGB,Image shape=(84, 84, 3) mode=RGB,Image shape=(84, 84, 3) mode=RGB] |
+------------------------------------+------------------------------------------------------------------------------------------------------+

實際在本次測試中發現，使用 HNSW 索引對搜索結果準確度沒有任何影響。

自然語言實現圖片搜索

本來到這裏測試目的已經達到了，突發奇想想試一下用自然語言也來實現圖片搜索。於是對代碼稍加改造：

def TextSearch(connection,text):
    text_conf = AutoConfig.load_config('text_image_embedding')
    text_conf.modality = 'text'

    text_pipe = AutoPipes.pipeline('text_image_embedding', text_conf)
    embedding = text_pipe(text).get()[0]
    
    res=SearchInTiDB(connection,embedding)
    p = (
        pipe.input('text','search_result')
        .map('search_result','prev',read_images)
        .output('text','prev')
    )
    DataCollection(p(text,res)).show()

還是用的clip_vit_base_patch16模型，只是使用模態改成了文本。通過對文本做 embedding 後得到向量數據送到 TiDB 中進行搜索，流程和前面基本一樣。

看一下最終效果：

可以發現英文的搜索效果要很多，這個主要是因爲模型對於中文理解能力比較差，英文語義下 TiDB 的向量搜索準確度依然非常高。

基於 TiDB Vector，前後不到100行代碼就實現了以圖搜圖和自然語言搜圖。

未來展望

反正第一時間體驗完的感受就是：太香了，強烈推薦給大家！

在以往，想在關係型數據庫中對非結構化數據實現搜索是一件不敢想象的事，哪怕是號稱無所不能的 PostgreSQL 在向量插件的加持下也沒有獲得太多關注，這其中有場景、性能、生態等各方面的因素制約。而如今在 AI 大浪潮中，應用場景變得多樣化，生態鏈變得更豐富，TiDB Vector 的誕生恰逢其時。

但是不可忽視的是，傳統數據庫集成向量化的能力已經是大勢所趨，哪怕是 Redis 這樣的產品也擁有了向量能力。前有專門的向量數據庫阻擊，後有各種傳統數據庫追趕，這注定是一個慘烈的賽道，希望 TiDB 能深度打磨產品，突圍成功。

期待的功能：更多的索引類型、GPU加速等。

當然了，最大的願望必須是 TiDB On-Premises 中能儘快看到 Vector 的身影。

給 TiDB 點贊！

作者介紹：hey-hoho，來自神州數碼鈦合金戰隊，是一支致力於爲企業提供分佈式數據庫TiDB整體解決方案的專業技術團隊。團隊成員擁有豐富的數據庫從業背景，全部擁有TiDB高級資格證書，並活躍於TiDB開源社區，是官方認證合作伙伴。目前已爲10+客戶提供了專業的TiDB交付服務，涵蓋金融、證券、物流、電力、政府、零售等重點行業。