本文主要介紹如何基於 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;
}