ES + FAISS 分布式向量检索引擎的实现原理
本文主要介绍如何基于 ES + FAISS 实现向量检索，并且以 FAISS IndexIVFFlat 索引为例介绍实现方案。对于 IndexIVFFlat 而言，需要先对数据进行聚类，得到若干个聚簇( nlist=1024)。然后在查询时，先计算与查询向量与聚簇之间的距离，选取距离最近的若干个((nprob=128))聚簇进行检索。
离线部分
IndexIVFFlat 聚簇中心点训练
- 索引依赖的数据源存储在 hive 表，先从 hive 下载一部分样本数据，用于聚簇中心点训练。
- 这里基于封装好的 查询hive 表的 client(底层是spark引擎)，将样本数据下载到索引聚簇训练机器的本地磁盘上。
- 然后基于 jni 调用 faiss lib 的 createIndex、train、writeIndex 方法训练中心点聚簇，并将聚簇持久化到磁盘上。
- 将中心点聚簇上传到 S3，用于后续创建 ES 向量索引时，针对自定义的向量字段（FieldMapper），下载解析聚簇文件，为向量字段生成相应的属性（索引类型（indexIVF or ivfpq）、距离计算方式（L2 or IP））。
离线索引构建
目前线上使用的是 faiss indexivfflat 索引，需要训练聚簇中心点，默认nlist=1024个聚簇。
- ES索引依赖的索引数据源存储在hive表，从hive表查出样本数据，用于中心点训练
- 调用faiss lib 的 createIndex、train、writeIndex 方法进行训练，并将训练结果持久化保存到磁盘上
- 将训练好的索引文件，上传到 S3，用于创建 ES 向量索引时，拉取中心点聚簇，存储向量字段所必要的属性。
- 索引构建，通过 spark 作业将包含向量字段的hive表数据，写入 ES，生成 ES 索引。针对向量字段，开启了doc_value，会将float[]保存到doc_value中。
这里通过继承org.elasticsearch.index.mapper.FieldMapper自定义 FieldMapper 创建了一个新的字段类型，专门用于存储向量。在 idnexing doc时， parse 解析向量数组时，先获取当前 doc 最近的中心点聚簇，然后将之保存在向量字段的索引中。
//获取当前 doc 最近的中心点聚簇
int nearestCentroid = centroids.getNearestCentroid(vectors);
//保存当前 doc 所属的最近的中心点聚簇，lucene custom query 在线查询时，用到
if (fieldType().indexOptions() != IndexOptions.NONE || fieldType().stored()) {
BytesRef binaryValue = VectorUtils.intToBytesRef(nearestCentroid);
Field field = new Field(fieldType().name(), binaryValue, fieldType());
fields.add(field);
}
//存储向量float[]数组，在算分时，需要传入 faiss 库函数计算 向量之间的距离
if (index.getVectorType() == VectorType.IVF) {
fields.add(new BinaryDocValuesField(fieldType().name(),
VectorUtils.toBytesRef(handledVector.getOriginVector(), handledVector.getStart(), vectors.length)));
}
在线部分
ES 插件开发
Mapper Plugin
定义一种新的 field type 存储向量数据。
向量数据本质上就是一个128维或者256维的 float 数组，自定义 org.elasticsearch.index.mapper.FieldMapper，在往 ES 索引写入向量数据时，计算该向量与 nlist 个聚簇的距离，得到距离最近的聚簇的位置(中心点位置-int 类型)，存储在该字段中，这样就可用于：倒排检索。并且，以 doc_value 形式存储向量内容，即 float 数组。
这样，基于倒排索引，就能快速定位到需要检索哪个聚簇，然后再读取 doc_value 得到 float 数组，通过 JNI 调用 FAISS 计算待查询向量和该聚簇下的每个 doc 的 float 数组，从而算出各个 doc 与待查询的向量query 的相似度。
Search Plugin
自定义 QueryBuilder 用于向量查询。
继承 org.elasticsearch.index.query.AbstractQueryBuilder 自定义一个专门针对向量字段进行查询的QueryBuilder。
当针对向量字段检索时，该 QueryBuilder 负责创建底层 Lucene 查询。
- threshold
- 聚簇中心点，它是 int 类型，存储在自定义的 Field 类型中。针对聚簇中心点的查询，其实就是一个 Lucene Term 查询。
Lucene custom query
自定义 QueryBuilder 生成向量查询 Query，向量查询Query 定义了 Weight 和 Scorer
Query
继承 org.apache.lucene.search.Query 自定义Lucene Query，用于向量字段的检索。
自定义 org.apache.lucene.queries.function.valuesource.FieldCacheSource 封装 JNI 调用 FAISS 计算两个 float 数组之间的距离（支持 L2 和 IP 两种距离计算）
//multiVector 是当前待查询的向量，获取与它最接近的 topK 个中心点聚簇(topK 就是 nprobe 参数)
this.centroidTerms = centroids.getTopKCentroid(multiVector);
//centroidTerms 代表 topK 中心点聚簇，重写query，只查询中心点聚簇下的向量。中心点聚簇之间使用 should
BooleanQuery.Builder bq = new BooleanQuery.Builder();
for (BytesRef term : centroidTerms) {
if (termsEnum != null && termsEnum.seekExact(term)) {
bq.add(new TermQuery(new Term(field, term)), BooleanClause.Occur.SHOULD);
}
}
从这里可以看出，查多个中心点聚簇，其实质是改写成了 多个 term query 之间的 should 查询。
Weight
通过 rewrite 方法将向量字段的查询改写成 Lucene Bool term query。每个聚簇中心点对应一个 Term 查询，如果要查询多个聚簇中心点，则是多个 Term 查询语句，使用 Bool 查询包装，使用 SHOULD 定义多个聚簇查询关系。
Score
FAISS 计算出来的距离，作为 score 分数
算分时，有个 threshold，低于阈值的向量会被认为低相关的，不被返回
if (metricType == MetricType.IP) {
return score() >= threshold;
} else {
return score() <= threshold;
}
