Amazon MemoryDB × ChatGPT Embeddings で 近似値検索してみた

経緯

Amazon MemoryDBでベクトル検索が一般公開で使えるようになったのでChatGPTのEmbedding早速使ってみた。

Amazon MemoryDBについて

詳しくは上記を参照してください。
Amazon ElastiCacheとMemoryDBの違いがわからないのでChatGPTさんに聞いてみた。

ChatGPT: Amazon MemoryDBとAmazon ElastiCacheの違いは？

Amazon MemoryDBとAmazon ElastiCacheは、両方ともインメモリデータベースおよびキャッシュサービスですが、異なるユースケースや設計思想に基づいて提供されています。それぞれの特徴と主な違いについて以下に説明します。

Amazon MemoryDB

特徴

1. フルマネージドサービス: MemoryDBはAWSによって完全に管理されており、クラスタのプロビジョニング、パッチ適用、バックアップなどの運用タスクが自動化されています。
2. 耐久性: メモリベースのデータベースである一方、データはディスクにも保存され、障害からの自動復旧が可能です。データの永続性が高いです。
3. 高い可用性: マルチアベイラビリティゾーンに対応し、自動フェイルオーバーをサポートしているため、高可用性が確保されています。
4. Redis互換: Redis (6.x) に互換性があり、Redisを使った既存のアプリケーションをMemoryDBに移行することが容易です。
5. スループットと低レイテンシ: 高性能なインメモリデータベースであり、低レイテンシかつ高スループットなアクセスを提供します。

Amazon ElastiCache

Amazon ElastiCacheは、MemcachedとRedisの両方をサポートするインメモリキャッシュサービスです。

特徴

1. キャッシング: 主に読み取り負荷を軽減し、アプリケーションのレスポンスを向上させるためのキャッシュとして利用されます。
2. 高パフォーマンス: インメモリストアであるため、非常に高いスループットと低レイテンシが実現できます。
3. 可用性とスケーラビリティ: RedisやMemcachedのクラスタリング機能を活用して、容易にスケールアップ/スケールアウトが可能です。
4. RedisとMemcachedのサポート: ElastiCacheはRedisだけでなく、Memcachedもサポートしているため、具体的なユースケースによってより適切な選択ができます。
5. フルマネージドサービス: MemoryDBと同様にAWSによって完全管理されており、運用の負担が軽減されています。

主な違い

1. データの永続性:

   • MemoryDB: データ永続性が高く、ディスクにデータを保存し、障害時に自動で復旧する機能がある。
   • ElastiCache: 主にキャッシュ用途のため、デフォルトではメモリ内のデータが失われることがありますが、Redisモードではスナップショットや永続化オプションが利用可能です。

2. ユースケース:

   • MemoryDB: 主にインメモリデータベースとして利用され、高い可用性と永続性が求められるアプリケーションに適しています。
   • ElastiCache: 主にキャッシュとして利用され、読み取り負荷を軽減し、データベースのパフォーマンスを向上させるために使用されます。

3. サポートされるソフトウェア:

   • MemoryDB: 現時点ではRedisのみをサポート。
   • ElastiCache: Redisだけでなく、Memcachedもサポート。

4. 可用性と復旧機能:

   • MemoryDB: フルマネージドでマルチAZ対応、高可用性と自動フェイルオーバー機能が統合されています。
   • ElastiCache: Redisモードでのクラスタリングによって高可用性が実現できますが、設定に依存します。

これらの特徴と違いを理解することで、具体的なユースケースに適したサービスを選択することができます。例えば、高いデータ永続性が必要で、かつインメモリの高速データベースが求められる場合はMemoryDBが適しています。一方で、アプリケーションのパフォーマンスを向上させるためのキャッシュが必要な場合はElastiCacheが適していると言えます。

らしいです。はい。
上記以外の違いとしたら金額が違います。

InstanceType	Amazon MemoryDB	Amazon ElastiCache	%diff
*.r7g.large	USD 0.371	USD 0.263	41%
*.r7g.16xlarge	USD 11.811	USD 8.37	41%
*.t4g.small	USD 0.074	USD 0.049	51%
*.t4g.medium	USD 0.147	USD 0.098	50%

お値段約1.5倍です。

Amazon MemoryDBの構築

では早速構築してみましょう。

サーバー構築

デモで使うので最小構成に。（冗長化もされない）
必ずベクトル検索オプションを有効にすること。

作成してから8分かかりました。

接続してみる

redis-cli -h clustercfg.*****.memorydb.ap-northeast-1.amazonaws.com --tls -p 6379

繋がらない場合はおそらくセキュリティグループの設定を確認してみてください。
最初繋がらなかったです。

clustercfg.*****.memorydb.ap-northeast-1.amazonaws.com:6379>  INFO
# Server
redis_version:7.1.1
used_memory_human:10.45M
maxmemory_human:1.37G

ここまで動けば取り急ぎサーバー構築は完了です。

データベース作成

FT.CREATE : Indexの作成

に従って作成してみます。
サンプル見てもコマンドの命令多すぎて全然意味がわかりません。

FT.CREATE hash_idx1 ON HASH PREFIX 1 hash: SCHEMA vec AS VEC VECTOR HNSW 6 DIM 2 TYPE FLOAT32 DISTANCE_METRIC L2
OK

FT.CREATE json_idx1 ON JSON PREFIX 1 json: SCHEMA $.vec AS VEC VECTOR HNSW 6 DIM 6 TYPE FLOAT32 DISTANCE_METRIC L2
OK

ChatGPTのEmbedding用途以外にあまり思いつきませんし、
AWSはClaude3の開発元のAnthropicに出資してるので忖度してるのかわかりませんが
ChatGPTのEmbeddingの場合どうすればよいのか調べたところ、
下記にたどり着きました。

# 1M件以上の大量データの場合
FT.CREATE embedding ON JSON PREFIX 1 emb: SCHEMA $.v AS VEC VECTOR HNSW 10 TYPE FLOAT32 DIM 1536 DISTANCE_METRIC COSINE M 40 EF_CONSTRUCTION 250

# 1M件以下のデータの場合
FT.CREATE embedding2 ON JSON PREFIX 1 emb: SCHEMA $.v AS VEC VECTOR FLAT 6 TYPE FLOAT32 DIM 1536 DISTANCE_METRIC COSINE

詳しいオプションの説明はChatGPTにしてもらってください。
僕もよーわからんです。
今回投入するデータ量は10K程度なのでFLATオプションで構築してみます。

ちなみにFT.CREATEコマンドをElastiCacheにも投げて見ましたが当然使えませんでした。

*****.ng.0001.apne1.cache.amazonaws.com:6379> FT.CREATE
(error) ERR unknown command `FT.CREATE`, with args beginning with: 

JSON.SET : データ投入

AWSのドキュメント通りJSON.SETでデータ投入できるとのこと。
あとはお手持ちのChatGPTで生成したEmbeddingデータを投入するだけです。

redis-cli -h clustercfg.*****.amazonaws.com --tls -p 6379 JSON.SET emb:11931 . '{"v":[0.01683144,-0.004322319,...,0.005457869,0.0047440086]}'

FT.INFO embedding
17) num_docs
18) (integer) 1

FT.INFOでnum_docsが増えてるのを確認しましょう。
ここの数字が増えてないと正しくIndexされてないです。
[]が間違えてたり次元数が違うと正しく処理されないです。

FT.INFO embedding
11) space_usage
12) (integer) 6,525,849
17) num_docs
18) (integer) 10,422

FT.INFO embedding2
11) space_usage
12) (integer) 6,610,370
17) num_docs
18) (integer) 10,422

とりあえず1万件ぐらい入れたデータで検証してみましょう。

FT.SEARCH : データ検索

ここは盛大にハマりました。
bashでコマンド叩いてたのですが正しく処理できなくて結局pythonで叩いて動きました。
ChatGPTに聞いて作ってるので細かいことは聞かないでください。
AWSドキュメントのサンプルコードの通りにまず実行してみます。

import redis
import struct

# Redisに接続設定
client = redis.StrictRedis(
    host="clustercfg.*****.memorydb.ap-northeast-1.amazonaws.com",
    port=6379,
    ssl=True
)

# Sampleのを試してみる
embedding = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
query_vec = struct.pack(f'{len(embedding)}f', *embedding)
response = client.execute_command(
  'FT.SEARCH', 'json_idx1', '*=>[KNN 100 @VEC $query_vec]', 'PARAMS', 2, 'query_vec', query_vec, 'DIALECT', 2
)
print(response)
[3, b'json:1', [b'$', b'[{"vec":[10.0, 20.0, 30.0, 40.0, 50.0, 60.0]}]', b'__VEC_score', b'95.3939208984375'], b'json:0', [b'$', b'[{"vec":[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]}]', b'__VEC_score', b'9.5393924713134766'], b'json:2', [b'$', b'[{"vec":[1.1, 1.2, 1.3, 1.4, 1.5, 1.6]}]', b'__VEC_score', b'3.3331665992736816']]

動いた！

ChatGPTのEmbeddingも試してみます。

embedding = [0.024686571, -0.024063775, ..., 0.02222189, 0.02222189]
len(embedding)
# 念の為次元があってるか念の為確認する
# 1536

query_vec = struct.pack(f'{len(embedding)}f', *embedding)
response = client.execute_command(
  'FT.SEARCH', 'embedding', '*=>[KNN 100 @VEC $query_vec]', 'PARAMS', 2, 'query_vec', query_vec, 'DIALECT', 2
)

# 1,000件出したい場合
query_vec = struct.pack(f'{len(embedding)}f', *embedding)
response = client.execute_command(
  'FT.SEARCH', 'embedding', '*=>[KNN 1000 @VEC $query_vec]', 'PARAMS', 2, 'query_vec', query_vec, 'DIALECT', 4, 'LIMIT', 0, 1000
)

# 結果からキーだけを抽出
def extract_keys(response):
    # レスポンスの形式を確認
    if isinstance(response, list) and len(response) > 1:
        # レスポンスはリストで、最初の要素は件数、残りは検索結果
        _, *results = response
        # 検索結果のキーを抽出
        keys = [results[i] for i in range(0, len(results), 2)]
        return keys
    return []

# 抽出したキーを表示
keys = extract_keys(response)
print(keys)

# キーと__VEC_scoreを抽出
def extract_keys_and_vec_scores(response):
    keys_and_scores = []
    if isinstance(response, list) and len(response) > 1:
        # レスポンスはリストで、最初の要素は件数、続く要素は検索結果
        _, *results = response
        for i in range(0, len(results), 2):
            key = results[i]
            score_data = results[i + 1]
            # __VEC_score がキーとなる場合、スコアを辞書から取り出す
            if isinstance(score_data, list):
                # スコアがリストの場合
                vec_score = None
                for j in range(0, len(score_data), 2):
                    if score_data[j] == b'__VEC_score':
                        vec_score = float(score_data[j + 1])
                        break
                keys_and_scores.append((key, vec_score))
            else:
                keys_and_scores.append((key, None))
    return keys_and_scores

# 抽出したキーと __VEC_score を表示
keys_and_scores = extract_keys_and_vec_scores(response)
for key, vec_score in keys_and_scores:
    print(f"Key: {key}, __VEC_score: {vec_score}")

Key: b'emb:7831627', __VEC_score: 0.18631483614444733
Key: b'emb:7831406', __VEC_score: 0.1861632615327835
Key: b'emb:7922159', __VEC_score: 0.18594330549240112
Key: b'emb:769', __VEC_score: 0.18583698570728302
Key: b'emb:7876977', __VEC_score: 0.18565867841243744
Key: b'emb:7952667', __VEC_score: 0.18547238409519196
Key: b'emb:8017', __VEC_score: 0.18525390326976776
Key: b'emb:7875609', __VEC_score: 0.18519586324691772
Key: b'emb:7767184', __VEC_score: 0.18516741693019867
Key: b'emb:25470', __VEC_score: 0.18494324386119843

いい感じに取れました！

性能評価

手動計算の場合約8秒かかった計算が0.4秒で終わりました。
Python内だけだ0.1秒以下でした。

Redisコマンドの処理時間: 0.008824348449707031秒
Redisコマンドの処理時間: 0.009100198745727539秒
Redisコマンドの処理時間: 0.008658885955810547秒
Redisコマンドの処理時間: 0.009626150131225586秒
Redisコマンドの処理時間: 0.008827686309814453秒
平均処理時間: 0.00900705337524414秒

手動計算処理時間: 4.6729121208191
手動計算処理時間: 5.1756479740143
手動計算処理時間: 4.8089139461517
手動計算処理時間: 5.0742518901825
手動計算処理時間: 4.9225928783417
平均処理時間: 4.93086336129986秒

ざっと500倍早くなった！

最後に

ベクトルの近似値計算は件数が多くなると計算コストが増えちゃってレスポンス悪くなるのでこのタイミングでMemoryDBに切り替えたいなと思います。
Elasticsearch / OpenSearch でもベクトル計算できるようなので調べてみようと思います。

BitStarでは絶賛OpenAIをはじめとした生成AIを使ったデータ分析/プロダクト開発に力を入れています。
エンジニアを募集中です。詳しくは下記を参照ください。

Appendix

KEYの一括削除

KEYを一括削除する際に一般的にxargsで削除するっぽいのですがエラーになって削除できません。

redis-cli -h clustercfg.*****.memorydb.ap-northeast-1.amazonaws.com --tls -p 6379 KEYS emb:* | xargs redis-cli -h clustercfg.*****.memorydb.ap-northeast-1.amazonaws.com --tls -p 6379 DEL
(error) CROSSSLOT Keys in request don't hash to the same slot

ChatGPTに聞いてみた。

Redis ClusterやAmazon MemoryDB for Redisで遭遇する「CROSSSLOT Keys in request don't hash to the same slot」エラーは、あるコマンドが複数のキーを扱おうとしたとき、そのキーが異なるスロットに属している場合に発生します。Redis Clusterでは、各キーは特定のスロットにマップされ、そのスロットは特定のノードに割り当てられます。複数のキーが異なるスロットに属する場合、そのキーに対する操作を一度に行うことはできません。

どうやらclusterでサーバーが作成されているとキーの一括削除ができないようです。
回避案としてChatGPTにpythonのコード作ってもらったサンプルです。

# 削除したいキーのパターン
pattern = 'emb:*'

def delete_keys_one_by_one(pattern):
    cursor = 0
    while True:
        cursor, keys = client.scan(cursor=cursor, match=pattern)
        for key in keys:
            client.delete(key)
            print(f"Deleted key: {key}")
        if cursor == 0:
            break

# キー削除を実行
delete_keys_one_by_one(pattern)

PHPのコード

ベクトルの近似計算コード

function similarity($u, $v) {
    $dotProduct = 0;
    $uLength = 0;
    $vLength = 0;
    for ($i = 0; $i < count($u); $i++) {
        $dotProduct += $u[$i] * $v[$i];
        $uLength += $u[$i] * $u[$i];
        $vLength += $v[$i] * $v[$i];
    }
    $uLength = sqrt($uLength);
    $vLength = sqrt($vLength);
    return $dotProduct / ($uLength * $vLength);
}

Redisコマンド

$redis = new Redis([
    'host' => 'clustercfg.*****.memorydb.ap-northeast-1.amazonaws.com',
    'port' => 6379,
    'connectTimeout' => 2.5,
]);
$emb = [0.024686571, -0.024063775, ..., 0.02222189, 0.02222189];
$query_vec = pack('f*', ...$emb);
$response = $redis->rawCommand(
    'FT.SEARCH',
    'embedding', // インデックス名
    '*=>[KNN 100 @VEC $query_vec]', // 検索クエリ
    'PARAMS', 2, 'query_vec', $query_vec, // パラメータ
    'DIALECT', 2, // ダイアレクトバージョン
    'LIMIT', 0, 100 // 取得件数
);