🐙

Databricks ベクトル検索のfilter機能

2024/01/31に公開

Databricks

以前、こちらの記事にDatabricksのベクトル検索とFaiss/ChromaDBとのコードの書き方の違いを書きました。

Databricksのベクトル検索について簡単におさらいをすると、Python SDKを使うと以下のようにsimilarity_search関数を使ってインデックスを検索できます。

results = vs_index.similarity_search(
  query_vector=[0.4345,0.9861,0.2419,0.9271,0.7619,・・・],
  columns=["id", "contents"],
  num_results=5)

この際、この関数の引数として、以下のようにFiltersというパラメーターを指定できます。Filtersの役割はその名の通り、インデックスデータのフィルタリングです。従って、インデックス検索の対象データを絞り込むことが可能です。

例えば、以下のようなインデックスがあるとします。

id	category	version	contents	embedding
1	food	1	たんぱく質を含む食べ物というと、動物性のものでは肉、魚、牛乳、チーズ、卵、植物性のものでは大豆を使った豆腐、納豆などが思い浮かびますが、ご飯、パン、麺などにもたんぱく質は含まれています。さらに、量は少ないですが、野菜、果物、海草などにも含まれます。	[0.4422,0.6267,0.1659,0.3990,0.9824,・・・
2	food	2	たんぱく質を含む食べ物というと、動物性のものでは肉、魚、牛乳、チーズ、卵、植物性のものでは大豆を使った豆腐、納豆などが思い浮かびますが、ご飯、パン、麺などの主食にもたんぱく質は含まれています。これらの食品は、毎日の食事で欠かせないものであり、たんぱく質の良い供給源になり得ます。例えば、全粒粉を使ったパンや、雑穀を含むご飯は、加工されていないため、より多くのたんぱく質を含んでいます。また、パスタや麺類も、選ぶ種類によってはたんぱく質の良い供給源となります。	[0.8512,0.6964,0.5832,0.3062,0.0685,・・・
3	sport	1	サッカーのワールドカップは、1930年に始まりましたが、第二次世界大戦のために1942年と1946年の大会は開催されませんでした。このため、サッカーのワールドカップは4年に1度開催される伝統が途絶えた唯一の期間となります。戦後、1950年にブラジルで大会が再開され、以降は世界中で最も注目されるスポーツイベントの一つとして成長し続けています。	[0.9078,0.1906,0.0220,0.7822,0.2821,・・・

このインデックスに対して、下記のようにCategoryは「food」である、かつ、Versionは２以上という条件でフィルタリングしたデータに対して、インデックス検索を行うことができます。

results = vs_index.similarity_search(
  query_vector=list(query_vector.astype(float)),
  columns=["source", "response"],
  num_results=5,
  filters={"category": "food", "version >=": 2})

他にも、下記のような指定が可能です。

### 等しい ###
filters={"category": "food"}

### 以上/以下、より大きい/より小さい ###
filters={"version >=": 2}
filters={"version <=": 2}
filters={"version >": 2}
filters={"version <": 2}

### IN ###
filters={"category": ("food", "sports")} # categoryがfoodまたはsports

### OR ###
filters={"category OR version": ["food", 2]} # "categoryがfood"、または、"versionが2"

### AND ###
filters={"category": "food", "version >=": 2} # "categoryがfood"、かつ、"versionが2以上"

### NOT ###
filters={"category NOT": "food"} # categoryがfoodでないもの

BFN!

Discussion