// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// 移植自 milvus-io/milvus client/index/{hnsw,ivf,flat,auto}.go(Apache-2.0)。
///|
/// `index_type` 的字面量。与上游 `index.IndexType` 逐条对齐。
pub(all) enum IndexType {
Flat
BinFlat
IvfFlat
BinIvfFlat
IvfPQ
IvfSQ8
IvfRabitQ
HNSW
HNSWSQ
HNSWPQ
HNSWPRQ
IvfHNSW
AutoIndex
DiskANN
AISAQ
SCANN
MinHashLSH
SparseInverted
SparseWAND
GpuIvfFlat
GpuIvfPQ
GpuCagra
GpuBruteForce
Trie
Sorted
Inverted
Bitmap
Ngram
FmIndex
RTree
} derive(Eq, Debug)
///|
/// 上游字面量。
pub fn IndexType::to_string(self : IndexType) -> String {
match self {
Flat => "FLAT"
BinFlat => "BIN_FLAT"
IvfFlat => "IVF_FLAT"
BinIvfFlat => "BIN_IVF_FLAT"
IvfPQ => "IVF_PQ"
IvfSQ8 => "IVF_SQ8"
IvfRabitQ => "IVF_RABITQ"
HNSW => "HNSW"
HNSWSQ => "HNSW_SQ"
HNSWPQ => "HNSW_PQ"
HNSWPRQ => "HNSW_PRQ"
IvfHNSW => "IVF_HNSW"
AutoIndex => "AUTOINDEX"
DiskANN => "DISKANN"
AISAQ => "AISAQ"
SCANN => "SCANN"
MinHashLSH => "MINHASH_LSH"
SparseInverted => "SPARSE_INVERTED_INDEX"
SparseWAND => "SPARSE_WAND"
GpuIvfFlat => "GPU_IVF_FLAT"
GpuIvfPQ => "GPU_IVF_PQ"
GpuCagra => "GPU_CAGRA"
GpuBruteForce => "GPU_BRUTE_FORCE"
Trie => "Trie"
Sorted => "STL_SORT"
Inverted => "INVERTED"
Bitmap => "BITMAP"
Ngram => "NGRAM"
FmIndex => "FMINDEX"
RTree => "RTREE"
}
}
///|
/// 反解上游字面量。
pub fn IndexType::from_string(s : String) -> IndexType? {
match s {
"FLAT" => Some(Flat)
"BIN_FLAT" => Some(BinFlat)
"IVF_FLAT" => Some(IvfFlat)
"BIN_IVF_FLAT" => Some(BinIvfFlat)
"IVF_PQ" => Some(IvfPQ)
"IVF_SQ8" => Some(IvfSQ8)
"IVF_RABITQ" => Some(IvfRabitQ)
"HNSW" => Some(HNSW)
"HNSW_SQ" => Some(HNSWSQ)
"HNSW_PQ" => Some(HNSWPQ)
"HNSW_PRQ" => Some(HNSWPRQ)
"IVF_HNSW" => Some(IvfHNSW)
"AUTOINDEX" => Some(AutoIndex)
"DISKANN" => Some(DiskANN)
"AISAQ" => Some(AISAQ)
"SCANN" => Some(SCANN)
"MINHASH_LSH" => Some(MinHashLSH)
"SPARSE_INVERTED_INDEX" => Some(SparseInverted)
"SPARSE_WAND" => Some(SparseWAND)
"GPU_IVF_FLAT" => Some(GpuIvfFlat)
"GPU_IVF_PQ" => Some(GpuIvfPQ)
"GPU_CAGRA" => Some(GpuCagra)
"GPU_BRUTE_FORCE" => Some(GpuBruteForce)
"Trie" => Some(Trie)
"STL_SORT" => Some(Sorted)
"INVERTED" => Some(Inverted)
"BITMAP" => Some(Bitmap)
"NGRAM" => Some(Ngram)
"FMINDEX" => Some(FmIndex)
"RTREE" => Some(RTree)
_ => None
}
}
///|
/// 一个索引定义:可选名字、类型,以及构建参数。
pub(all) struct Index {
/// 索引名。空串表示让服务端用默认名(上游把 index_name 留空即为默认)。
name : String
metric_type : MetricType
index_type : IndexType
params : IndexParams
}
///|
/// 组装出完整参数:`metric_type` / `index_type` 固定在最前,其余按插入序。
/// 上游每个具体构建器都把这两个键写进 `Params()`,位置在 Go 的 map 里无序;
/// 这里定死顺序,便于测试快照稳定。
fn finish(
name : String,
metric_type : MetricType,
index_type : IndexType,
params : IndexParams,
) -> Index {
let full = IndexParams::new()
full.set(metric_type_key, metric_type.to_string())
full.set(index_type_key, index_type.to_string())
for entry in params.to_pairs() {
// 这两个键由构建器掌管,调用方经 WithExtraIndexParams 传入时必须被忽略,
// 否则 Params() 的 index_type 会与 Index::index_type 打架。
if entry.0 == index_type_key || entry.0 == metric_type_key {
continue
}
full.set(entry.0, entry.1)
}
{ name, metric_type, index_type, params: full, }
}
///|
/// 索引的构建参数(含 `index_type` / `metric_type`)。
pub fn Index::params(self : Index) -> IndexParams {
self.params
}
///|
/// 索引名。
pub fn Index::name(self : Index) -> String {
self.name
}
///|
/// 索引类型。
pub fn Index::index_type(self : Index) -> IndexType {
self.index_type
}
///|
/// 度量类型。
pub fn Index::metric_type(self : Index) -> MetricType {
self.metric_type
}
///|
/// 合并额外构建参数,键名不做白名单,原样透传(服务端校验)。
/// `index_type` / `metric_type` 是保留键,静默忽略。
pub fn Index::with_extra_params(self : Index, extra : IndexParams) -> Index {
let merged = IndexParams::from_pairs(self.params.to_pairs())
for entry in extra.to_pairs() {
merged.set(entry.0, entry.1)
}
finish(self.name, self.metric_type, self.index_type, merged)
}
///|
/// 设置索引名。
pub fn Index::with_name(self : Index, name : String) -> Index {
{ ..self, name, }
}
///|
/// 通用索引:调用方自己给全参数,`index_type` 从入参里取。
/// 上游 `NewGenericIndex` 同名同义,留作「构建器还没覆盖的参数」的出口。
pub fn new_generic_index(
name : String,
params : IndexParams,
) -> Index raise IndexParamError {
let raw = match params.get(index_type_key) {
Some(s) => s
None => raise InvalidArgument(reason="generic index requires index_type")
}
let index_type = match IndexType::from_string(raw) {
Some(t) => t
None => raise InvalidArgument(reason="unknown index_type: \{raw}")
}
let metric_type = match params.get(metric_type_key) {
Some(s) =>
match MetricType::from_string(s) {
Some(m) => m
None => raise InvalidArgument(reason="unknown metric_type: \{s}")
}
None => raise InvalidArgument(reason="generic index requires metric_type")
}
finish(name, metric_type, index_type, params)
}
///|
/// HNSW:图索引,`M` 是每层邻居数,`efConstruction` 是建图时的候选队列长度。
///
/// 上游默认值 `M=16`、`efConstruction=200`(Milvus 文档与 Go SDK 示例一致);
/// 越界检查与 knowhere 的 `CheckAndAdjust` 对齐:`M ≥ 2`,`efConstruction ≥ 1`。
pub fn new_hnsw_index(
metric_type : MetricType,
m? : Int = 16,
ef_construction? : Int = 200,
) -> Index raise IndexParamError {
if m < 2 {
raise OutOfRange(key="M", value=m, expected=">= 2")
}
if ef_construction < 1 {
raise OutOfRange(
key="efConstruction",
value=ef_construction,
expected=">= 1",
)
}
let params = IndexParams::new()
params.set_int("M", m)
params.set_int("efConstruction", ef_construction)
finish("", metric_type, HNSW, params)
}
///|
/// IVF_FLAT:倒排 + 原始向量,`nlist` 是聚类中心数。
/// 上游 `NewIvfFlatIndex` 不校验,服务端要求 `nlist ∈ [1, 65536]`。
pub fn new_ivf_flat_index(
metric_type : MetricType,
nlist? : Int = 128,
) -> Index raise IndexParamError {
if nlist < 1 || nlist > 65536 {
raise OutOfRange(key="nlist", value=nlist, expected="[1, 65536]")
}
let params = IndexParams::new()
params.set_int("nlist", nlist)
finish("", metric_type, IvfFlat, params)
}
///|
/// AUTOINDEX:把参数完全交给服务端调优。
/// 上游 `NewAutoIndex` 只写 `index_type` 与 `metric_type`。
pub fn new_auto_index(metric_type : MetricType) -> Index {
finish("", metric_type, AutoIndex, IndexParams::new())
}
///|
/// FLAT:暴力检索,只有 `metric_type` 与 `index_type`。
pub fn new_flat_index(metric_type : MetricType) -> Index {
finish("", metric_type, Flat, IndexParams::new())
}
///|
/// 给浮点向量的索引补一道度量校验。
///
/// Milvus 用 `HAMMING` 配 `FloatVector` 只会在服务端被拒(`IllegalMetricType`),
/// 客户端提前判可以省一次往返。**不改变服务端语义**:这里放过的组合,服务端仍可能拒绝。
pub fn check_float_vector_metric(index : Index) -> Unit raise IndexParamError {
if !index.metric_type.valid_for_float_vector() {
raise InvalidCombination(
index_type=index.index_type.to_string(),
reason="metric_type \{index.metric_type.to_string()} is not valid for float vectors",
)
}
}