Tablestore:Tablestore におけるベクトル検索パフォーマンスの最適化

ベクトル検索のスコアリング式

ベクトル検索は、ユークリッド距離 (euclidean)、コサイン類似度 (cosine)、およびドット積 (dot_product) の距離尺度をサポートしています。各距離尺度には異なるスコアリング式があります。Tablestore はこれらのスコアリング式を使用して、ベクトル間の類似性を評価します。次の表に、スコアリング式を示します。

トラブルシューティングと分析

private static void knnVectorQuery(SyncClient client) {
    SearchQuery searchQuery = new SearchQuery();
    KnnVectorQuery query = new KnnVectorQuery();
    query.setFieldName("Col_Vector");
    query.setTopK(10); // 上位 K 個の最近傍を返します。
    query.setFloat32QueryVector(new float[]{0.1f, 0.2f, 0.3f, 0.4f});
    // 最近傍ベクトルは、Col_Keyword=hangzhou && Col_Long<4 という条件を満たす必要があります。
    query.setFilter(QueryBuilders.bool()
            .must(QueryBuilders.term("Col_Keyword", "hangzhou"))
            .must(QueryBuilders.range("Col_Long").lessThan(4))
    );
    searchQuery.setQuery(query);
    searchQuery.setLimit(10);
    // スコアで並べ替えます。
    searchQuery.setSort(new Sort(Collections.singletonList(new ScoreSort())));
    SearchRequest searchRequest = new SearchRequest("<TABLE_NAME>", "<SEARCH_INDEX_NAME>", searchQuery);
    SearchRequest.ColumnsToGet columnsToGet = new SearchRequest.ColumnsToGet();
    columnsToGet.setColumns(Arrays.asList("Col_Keyword", "Col_Long"));
    searchRequest.setColumnsToGet(columnsToGet);
    // Search インターフェイスにアクセスします。
    SearchResponse resp = client.search(searchRequest);
    for (SearchHit hit : resp.getSearchHits()) {
        // スコアを出力します。
        System.out.println(hit.getScore());
        // データを出力します。
        System.out.println(hit.getRow());
    }
}

付録： ベクトル検索のスコアリング式のデモコード

次の Java コードは、サポートされている距離尺度のスコアリング式を示しています。

import java.util.concurrent.ThreadLocalRandom;

public class CompareVector {

    public static void main(String[] args) {
        // a はクエリベクトルです
        float[] a = randomVector(512);
        // b は、返されることが期待されるインデックス内の行のベクトルです
        float[] b = randomVector(512);
        // 多次元インデックスで設定した類似度測定アルゴリズムを選択し、スコアを出力します
        System.out.println(MetricFunction.COSINE.compare(a, b));
    }

    public static float[] randomVector(int dim) {
        float[] vec = new float[dim];
        for (int i = 0; i < dim; i++) {
            vec[i] = ThreadLocalRandom.current().nextFloat();
            if (ThreadLocalRandom.current().nextBoolean()) {
                vec[i] = -vec[i];
            }
        }
        return l2normalize(vec, true);
    }

    public static float[] l2normalize(float[] v, boolean throwOnZero) {
        double squareSum = 0.0f;
        int dim = v.length;
        for (float x : v) {
            squareSum += x * x;
        }
        if (squareSum == 0) {
            if (throwOnZero) {
                throw new IllegalArgumentException("normalize a zero-length vector");
            } else {
                return v;
            }
        }
        double length = Math.sqrt(squareSum);
        for (int i = 0; i < dim; i++) {
            v[i] /= length;
        }
        return v;
    }

    public enum MetricFunction {
        /**
         * ユークリッド距離。
         */
        EUCLIDEAN {
            @Override
            public float compare(float[] v1, float[] v2) {
                return 1 / (1 + VectorUtil.squareDistance(v1, v2));
            }
        },

        /**
         * ドット積。
         */
        DOT_PRODUCT {
            @Override
            public float compare(float[] v1, float[] v2) {
                return (1 + VectorUtil.dotProduct(v1, v2)) / 2;
            }
        },

        /**
         * コサイン。
         */
        COSINE {
            @Override
            public float compare(float[] v1, float[] v2) {
                return (1 + VectorUtil.cosine(v1, v2)) / 2;
            }
        };

        public abstract float compare(float[] v1, float[] v2);
    }


    static final class VectorUtil {

        private static void checkParam(float[] a, float[] b) {
            if (a.length != b.length) {
                throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
            }
        }

        public static float dotProduct(float[] a, float[] b) {
            checkParam(a, b);
            float res = 0f;
            for (int i = 0; i < a.length; i++) {
                res += b[i] * a[i];
            }
            return res;
        }

        public static float cosine(float[] a, float[] b) {
            checkParam(a, b);
            float sum = 0.0f;
            float norm1 = 0.0f;
            float norm2 = 0.0f;
            for (int i = 0; i < a.length; i++) {
                float elem1 = a[i];
                float elem2 = b[i];
                sum += elem1 * elem2;
                norm1 += elem1 * elem1;
                norm2 += elem2 * elem2;
            }
            return (float) (sum / Math.sqrt((double) norm1 * (double) norm2));
        }

        public static float squareDistance(float[] a, float[] b) {
            checkParam(a, b);
            float sum = 0.0f;
            for (int i = 0; i < a.length; i++) {
                float difference = a[i] - b[i];
                sum += difference * difference;
            }
            return sum;
        }
    }
}