🪐

Azure Cosmos DB MongoDB vCoreで事前にデータをフィルタリングした後にベクトル検索をしたい。

2024/08/19に公開

事前にデータを絞ってからベクトル検索したい

ベクトル検索を行うときにAzure Cosmos DB MongoDB vCoreを使用していたのですが、最初にデータをフィルタリングしてからベクトル検索をしたいケースがありました。
全てのデータに対してベクトル検索を行うよりも、あらかじめカテゴリなどを絞ったうえでベクトル検索をしたほうが検索精度が良いからです。

どうやら今まではこの機能がなかったようで、最近プレビュー機能として追加されたようです。
https://devblogs.microsoft.com/cosmosdb/introducing-vcore-based-azure-cosmos-db-for-mongodb-latest-ai-features/

機能を有効にする

この事前フィルタリングベクトル検索機能を有効にするには、まずこちらのページの通りに、サブスクリプション→プレビュー機能に進みます。
https://learn.microsoft.com/en-us/azure/azure-resource-manager/management/preview-features?tabs=azure-portal

「filtering on vector search」を選んで「登録」を行うことで利用可能になります。

ちなみに既に作成済みのリソースには適用されなかったので、機能を有効にした後に新しくCosmosDBを作り直しました。

使い方

以下のようにfilterというキーを追加するだけで使用できました。
例として「Color="赤"」かつ「Fruits="いちご" OR "りんご"」のデータに絞ったうえでベクトル検索をするとこのようになります。

search_query = {
    "$search": { # 検索クエリ
        "cosmosSearch": {
            "vector": query_embedding, #検索ベクトル
            "path": "VectorField", #ベクトルが格納されているフィールド
            "k": 20,  #類似スコア上位20件
+           "filter":{ # ここで事前フィルタリングを行う。
+                 "$and": {
+                       "Color": {
+                             "$in": ["赤"]
+                       },
+                       "Fruits": {
+                             "$in": ["いちご","りんご"]
+                       }
+                 }
+           }
        },
        "returnStoredSource": True
    }
}

今までのクエリにfilterを追加するだけなので簡単でした。

最後に

どうやらAzure Cosmos DB for NoSQLでもベクトル検索ができるようになったようです。
for NoSQLではWHERE句を使用することで同じことが実現できそうですが、私はまだ試せていないです。
https://learn.microsoft.com/ja-jp/azure/cosmos-db/nosql/vector-search

もし試した方がいれば情報いただけると嬉しいです🙇‍♂

Discussion