AstraDB/Cassandraで実現するベクトル検索〜最小限のコードによる理解

# はじめに

## 本記事の目的

ベクトル検索についてご存知でしょうか？

この記事では「ベクトル検索」については、すでにご存知であることを前提として、それを実行するための、新しい手法を紹介することを目的として、技術者にとって、実質的なイメージを持つことのできる情報を、簡潔に、伝えたいと思います。

CassandraとAstraDBについての基本的な理解も前提とします。

この記事を読むことによって：
 - ベクトル検索の実行について、他のテクノロジーでの経験のある人は、新しい実行手段として、イメージを持つことができるでしょう。
 - CassandraとPythonライブラリについての知識があれば、何がおこなわれているか、さらに具体的に理解できるでしょう。
 - AstraDBについての知識、あるいは関心があれば、実際にAstraDBを使って、ベクトル検索を実行することも難しくないでしょう。

# 紹介（2023年6月時点）

### Apache Caddandre: CEP-30: Approximate Nearest Neighbor(ANN) Vector Search via Storage-Attached Indexes

https://cwiki.apache.org/confluence/display/CASSANDRA/CEP-30%3A+Approximate+Nearest+Neighbor%28ANN%29+Vector+Search+via+Storage-Attached+Indexes

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/176567/bd2f2b4d-7d32-ed2c-a4c9-052540911652.png)

ステータスは「Adopted」になりました。

### AstraDB

上記の機能が、プレビューとして利用できるようになっています。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/176567/dc6730dc-e278-57d4-e5cd-4649272bfbbf.png)

上記のボタンからデータベースを新規作成するとPreviewモードのデータベースが作成され、下記のようにVector Searchのドキュメントのリンクがオーバービューに表示れます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/176567/9701f1ac-ec81-c9aa-943b-1ef226214ceb.png)

https://docs.datastax.com/en/astra-serverless/docs/vector-search/overview.html

https://www.datastax.com/blog/introducing-vector-search-empowering-cassandra-astra-db-developers-to-build-generative-ai-applications

# 実践

## 基本的な流れ

1. 値のベクトル化
1. データセットの登録とインデックス化
1. ベクトル検索実行

## 値のベクトル化

ここでは、SentenceTransformersのMiniLMを利用する。

https://www.sbert.net/

```
!pip sentence-transformers
```
```python
from sentence_transformers import SentenceTransformer
transformer = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
```
```python
embedding = transformer.encode(text).tolist()
```
## データセットの登録とインデックス化

### テーブル定義

ベクトルデータを格納するためのテーブル定義。

```python
session.execute(f"""CREATE TABLE IF NOT EXISTS vector_preview.products_table
(product_id int,
 chunk_id int,
 
 product_name text,
 description text,
 price text,
 
 description_embedding vector<float, 384>,
 
 PRIMARY KEY (product_id, chunk_id))""")
```

### インデックス定義

インデックスを作成。

```python
session.execute(f"""CREATE CUSTOM INDEX IF NOT EXISTS minilm_desc ON vector_preview.products_table (description_embedding) USING 'org.apache.cassandra.index.sai.StorageAttachedIndex'""")
```
### データ登録

```python
query = SimpleStatement(
                f"""
                INSERT INTO vector_preview.products_table
                (product_id, chunk_id, product_name, description, price, description_embedding)
                VALUES (%s, %s, %s, %s, %s, %s)
                """
            )
    display(row)

    session.execute(query, (product_id, chunk_id, product_name, description, pricevalue, embedding))
```

## ベクトル検索実行

```python
query = SimpleStatement(
    f"""
    SELECT *
    FROM demo.products_table
    ORDER BY description_embedding ANN OF {embedding} LIMIT 5;
    """
    )
```

```python
results = session.execute(query)
top_5_products = results._current_rows

for row in top_5_products:
  print(f"""{row.product_id}, {row.product_name}, {row.description}, {row.price}\n""")
```

## 最後に

本記事の執筆は、多分に自分自身の理解の整理のためではありましたが、説明こそ、不親切と思われる方も多かったかもしれませんが、はじめに整理したように、前提を置けば、かえって簡潔に内容が伝わったことを期待します。


## 追記（2023年７月19日）

### AstraDB ベクトル検索機能 GA

2023年7月18日に、上記ではプレビューのステータスとして紹介したAstraDBのベクトル検索機能がGAになり、プロダクションで用いることができるようになりました。

### サンプル/チュートリアル

上記では、そのエッセンスのみを紹介したApache Cassandraのベクトル検索の拡張機能ですが、下記のGitHubリポジトリで、実際に動作を確認することのできるサンプル/チュートリアルを公開しました(解説・コメントは英語の他、日本語でも用意しています）。

https://github.com/YoshiyukiKono/semantic-text-search