聚类分片

前提条件

已安装Proxima CE包并准备输入表，详情请参见安装Proxima CE包。

基本原理

Proxima CE在检索时有两种划分数据分片的方式：哈希分片与聚类分片。您可以通过设置-sharding_mode参数来选择具体的索引分片模式，值为hash时采用哈希分片，值为cluster时采用聚类分片，当前默认采用哈希分片。

• 哈希分片：在构建索引时，对全量doc集合划分，得到column_num个索引，检索时每条query需要在所有索引分片中查询，最后合并召回结果。

• 聚类分片：核心思路是先对doc进行聚类，将距离接近的doc划分到同一索引分片中，检索时根据query和聚类中心点的距离，选择最近的部分中心点对应的索引分片进行检索。

聚类分片划分索引方式的目的是性能优化，查询时避免查询所有索引分片，只需要检索部分索引分片就能尽可能召回最优结果。聚类分片划分索引时包含如下两个阶段：

1. 索引构建阶段。

   1. 在构建索引时先对doc集合进行kmeans聚类，生成kmeans_cluster_num个中心点。

   2. 将kmeans_cluster_num个中心点，按照空间距离划分成column_num个集合，可以理解为将中心点分配到column_num个索引中。

   3. 在对doc集合进行划分时，将doc划分到距离该doc最近的中心点对应的索引分片中。

2. 索引查询阶段。

   1. query先和所有中心点计算距离。

   2. 根据kmeans_seek_ratio选择一定比例的最近中心点对应的索引分片进行检索。

   3. 对检索的索引分片结果进行结果合并。

适用场景

• 聚类分片方式适用于数据量非常大的情况（十亿数据量级），特别是query数据量极大的场景。

• 适用于构建一次索引，后续多次查询该索引（即一次build，多次seek）的场景。

  说明

  • 聚类索引分片划分方法需要对doc集合进行kmeans聚类，产生时间消耗，并且由于只检索了一部分索引分片，势必会产生一定的召回损失，所以该方式不适用于所有向量检索场景。

  • 聚类分片不支持多类目检索，距离函数不支持除欧式距离、汉明距离以外的其他距离公式。

使用逻辑

1. 指定-sharding_mode为cluster。

2. 在JAR命令的-resources中添加聚类初始中心点表名称。

   说明

   • 此处不是命令行参数，是JAR命令需要的参数。聚类初始中心点名称为用户自定义名称，需要保持唯一，例如foo_init_center_resource。

   • 运行Proxima CE时会创建对应的MaxCompute表存储聚类中心点，由于MaxCompute资源的机制，需要用户手动添加中心点表的具体表名。

3. -kmeans_resource_name参数值需要和-resources中保持一致。因程序无法直接获取-resources的值，所以需要额外的-kmeans_resouce_name命令行参数来传递。

4. 其他参数非必选，可以参考可选参数中名称以kmeans_开头的参数。

导入输入表数据

您可以在DataWorks的SQL节点运行以下命令。

-- 备注：origin_table 来自阿里某业务的 128 维 float 向量数据表
-- 准备 doc 表：
CREATE TABLE cluster_10kw_128f_doc(pk STRING, vector STRING) PARTITIONED BY (pt STRING);
ALTER TABLE cluster_10kw_128f_doc add PARTITION(pt='20221111');
INSERT OVERWRITE TABLE cluster_10kw_128f_doc PARTITION (pt='20221111') SELECT pk, vector FROM origin_table WHERE pt='20221111';

-- 准备 query 表：
CREATE TABLE cluster_10kw_128f_query(pk STRING, vector STRING) PARTITIONED BY (pt STRING);
ALTER TABLE cluster_10kw_128f_query add PARTITION(pt='20221111');
INSERT OVERWRITE TABLE cluster_10kw_128f_query PARTITION (pt='20221111') SELECT pk, vector FROM origin_table WHERE pt='20221111';

使用DataWorks运行

本文以DataWorks运行方式为例，假设已提前创建好了External Volume。

命令如下：

--@resource_reference{"proxima-ce-aliyun-1.0.0.jar"}
jar -resources proxima-ce-aliyun-1.0.0.jar  -- 上传的 proxima-ce jar 包
-classpath proxima-ce-aliyun-1.0.0.jar com.alibaba.proxima2.ce.ProximaCERunner  -- classpath 指定 main 函数入口类
-doc_table cluster_10kw_128f_doc
-doc_table_partition 20221111
-query_table cluster_10kw_128f_query
-query_table_partition 20221111
-output_table cluster_10kw_128f_output
-output_table_partition 20221111
-algo_model hnsw
-data_type float
-pk_type int64
-dimension 128
-column_num 50
-row_num 50
-vector_separator ,
-topk 1,50,100,200 -- 获取 topk 为 1/50/100/200 时各自的召回率
-job_mode train:build:seek:recall
-- -clean_build_volume true -- 保留索引，后续多次运行时可以设置该选项为 true，此时需要设置 job_mode 为 `seek(:recall 可选)` 模式
-external_volume_name udf_proxima_ext
-sharding_mode cluster
-kmeans_resource_name kmeans_center_resource_xxx -- 手动指定 kmeans 资源名称，这里的命名示例为 `kmeans_center_resource_xxx`
-kmeans_cluster_num 1000
-- -kmeans_sample_ratio 0.05 -- 使用默认参数
-- -kmeans_seek_ratio 0.1 -- 使用默认参数
-- -kmeans_iter_num 30 -- 使用默认参数
-- -kmeans_init_center_method "" -- 使用默认参数
-- -kmeans_worker_num 0 -- 使用默认参数
;

运行结果

向量检索  数据类型:4 , 向量维度:128 , 检索方式:HNSW , 计算方法:SquaredEuclidean , 构建模式:train:build:seek:recall
doc表信息 表名: cluster_10kw_128f_doc , 分区:20221111 , doc数量:100000000 , 向量分隔符:,
query表信息 表名: cluster_10kw_128f_query , 分区:20221111 , query数量:100000000 , 向量分隔符:,
输出表信息 表名: cluster_10kw_128f_output , 分区:20221111
行列信息  行数: 50 , 列数:50 , 每列索引doc数量:2000000
是否清除Volume索引:true

各个worker的耗时(单位:秒):
   SegmentationWorker:      3
   TmpTableWorker:      1
   KmeansGraphWorker:       2243
   BuildJobWorker:      4973
   SeekJobWorker:       5922
   TmpResultJoinWorker:     0
   RecallWorker:        986
   CleanUpWorker:       6
总耗时(单位:分钟):235

实际召回率
    Recall@1:   0.999
    Recall@50:  0.9941600000000027
    Recall@100: 0.9902300000000046
    Recall@200: 0.9816199999999914