🐙

TiDB for AI : pytidb AI に特化したTiDB用 Python SDK その3 全文検索

に公開

今まで2回にわたりpytidbという AI に特化したTiDB用 Python SDK を触ってきました。
https://zenn.dev/kameoncloud/articles/6494cd5c51ad96
https://zenn.dev/kameoncloud/articles/a6cacc25f85862

今日は全文検索インデックスを見ていきます。
https://zenn.dev/kameoncloud/articles/7892342f622ffe
過去にこのブログで一度TiDBの全文検索インデックスを試した記事を書きました。pytidbでもこの機能をサポートしています。
https://pingcap.github.io/ai/guides/fulltext-search/

現在以下のリージョンでプレビューサポートされています。
・フランクフルト
・シンガポール

さっそくやってみる

1.クラスターの作成とフルテキストインデックス付テーブルの作成

ではTiDB Serverlessクラスターのフランクフルトシンガポールのどちらかに作成します。
過去作成したconnect.pyを以下に置き換えて実行します。

connect.py
from pytidb import TiDBClient
from pytidb.embeddings import EmbeddingFunction
from pytidb.schema import TableModel, Field, VectorField
from pytidb.schema import FullTextField

# データベース接続
client = TiDBClient.connect(
    host="gateway01.eu-central-1.prod.aws.tidbcloud.com",
    port=4000,
    username="3CLxN311CtrMaPB.root",
    password="VyCy0xxxxx",
    database="test",
)

# 接続確認
result = client.execute("SELECT 1 as test")
print("接続成功")
print(f"結果: {result}")

class Item(TableModel):
    __tablename__ = "items"
    id: int = Field(primary_key=True)
    title: str = FullTextField(fts_parser="MULTILINGUAL")

table = client.create_table(schema=Item, mode="overwrite")

まず新しくFullTextFieldライブラリをインポートしています。接続が完了した後テーブルを作成しています。

class Item(TableModel):
    __tablename__ = "items"
    id: int = Field(primary_key=True)
    title: str = FullTextField(fts_parser="MULTILINGUAL")

table = client.create_table(schema=Item, mode="overwrite")

これはSQLでいえば以下に相当します。

CREATE TABLE items(
    id INT PRIMARY KEY,
    title TEXT,
    FULLTEXT INDEX (title) WITH PARSER MULTILINGUAL
);

1点違いがあるのはmode="overwrite"です。これは既存のテーブルが存在している場合、テーブルを削除してテーブルを再作成します。すでにテーブル存在している場合オペレーションをスキップさせる場合はmode="exists_ok"を指定します。

2.データの投入

次にconnect.pyに以下の部分を最後に追加します。

connect.py

table.bulk_insert([
    Item(id=1, title="Bluetooth Earphones, HiFi sound, 48h battery, Fast charge, Low latency"),
    Item(id=2, title="Bluetooth 5.3 Headphones, Noise Cancelling, Immersive sound, Comfortable"),
    Item(id=3, title="IPX7 Waterproof Earbuds, Sport ready, Touch control, High-quality music"),
    Item(id=4, title="Sports Earbuds, Secure fit, Sweatproof, Long battery, Workout support"),
    Item(id=5, title="Wired Headphones, Studio-grade, HD sound, Comfortable, Pro music experience"),
    Item(id=6, title="Bluetoothイヤホン HiFi音質 48hバッテリー 急速充電 低遅延"),
    Item(id=7, title="Bluetooth5.3ヘッドホン ノイズキャンセリング 没入サウンド 快適装着"),
    Item(id=8, title="IPX7防水イヤホン スポーツ対応 タッチ操作 高音質音楽"),
    Item(id=9, title="スポーツイヤホン 安定装着 防汗 長持ちバッテリー ワークアウト対応"),
    Item(id=10, title="有線ヘッドホン スタジオ級 HDサウンド 快適装着 プロ音楽体験"),
    Item(id=11, title="无线蓝牙耳机 HiFi音质 48小时超长续航 快速充电 低延迟"),
    Item(id=12, title="蓝牙5.3降噪头戴式耳机 杜比全景声 沉浸音效 舒适佩戴 畅享静谧音乐时光"),
    Item(id=13, title="IPX7防水真无线耳机 运动无忧 智能触控 随时畅听高品质音乐"),
    Item(id=14, title="运动专用耳机 稳固佩戴 防汗设计 超长续航 低延迟音频 高清通话"),
    Item(id=15, title="录音室级有线耳机 高清音质 舒适佩戴 可拆卸线材 多设备兼容 降噪麦克风"),
])

SQL Editorでselect * from items;を実行すると以下の通りデータがは挿入されていることがわかります。

3.全文検索の実行

ではさらに以下をconnect.pyに追記します。

results = table.search("Bluetooth Headphones", search_type="fulltext").limit(3).to_list()
print(json.dumps(results, indent=2, ensure_ascii=False))
python connect.py
接続成功
結果: rowcount=1 success=True message=None
[
  {
    "id": 2,
    "title": "Bluetooth 5.3 Headphones, Noise Cancelling, Immersive sound, Comfortable",
    "_match_score": 3.7390857,
    "_score": 3.7390857
  },
  {
    "id": 5,
    "title": "Wired Headphones, Studio-grade, HD sound, Comfortable, Pro music experience",
    "_match_score": 1.9798478,
    "_score": 1.9798478
  },
  {
    "id": 1,
    "title": "Bluetooth Earphones, HiFi sound, 48h battery, Fast charge, Low latency",
    "_match_score": 1.620981,
    "_score": 1.620981
  }
]

Bluetooth Headphonesに対するマッチ度合が高いものが表示されます。このマッチ度合はBM25というアルゴリズムで評価されてています。
BM25は広く使われ、高い 頑健性(robustness) と 効率性(efficiency) を持つテキスト検索アルゴリズムです。

Discussion