milvus/index/
mod.rs

1use strum_macros::{Display, EnumString};
2
3use crate::proto::{
4    common::{IndexState, KeyValuePair},
5    milvus::IndexDescription,
6};
7use std::{collections::HashMap, str::FromStr};
8
9#[derive(Debug, Clone, Copy, EnumString, Display)]
10pub enum IndexType {
11    #[strum(serialize = "FLAT")]
12    Flat,
13    #[strum(serialize = "BIN_FLAT")]
14    BinFlat,
15    #[strum(serialize = "IVF_FLAT")]
16    IvfFlat,
17    #[strum(serialize = "BIN_IVF_FLAT")]
18    BinIvfFlat,
19    #[strum(serialize = "IVF_PQ")]
20    IvfPQ,
21    #[strum(serialize = "IVF_SQ8")]
22    IvfSQ8,
23    #[strum(serialize = "IVF_SQ8_HYBRID")]
24    IvfSQ8H,
25    #[strum(serialize = "NSG")]
26    NSG,
27    #[strum(serialize = "HNSW")]
28    HNSW,
29    #[strum(serialize = "RHNSW_FLAT")]
30    RHNSWFlat,
31    #[strum(serialize = "RHNSW_PQ")]
32    RHNSWPQ,
33    #[strum(serialize = "RHNSW_SQ")]
34    RHNSWSQ,
35    #[strum(serialize = "IVF_HNSW")]
36    IvfHNSW,
37    #[strum(serialize = "ANNOY")]
38    ANNOY,
39    #[strum(serialize = "NGT_PANNG")]
40    NGTPANNG,
41    #[strum(serialize = "NGT_ONNG")]
42    NGTONNG,
43}
44
45#[derive(Debug, Clone, Copy, EnumString, Display)]
46pub enum MetricType {
47    L2,
48    IP,
49    HAMMING,
50    JACCARD,
51    TANIMOTO,
52    SUBSTRUCTURE,
53    SUPERSTRUCTURE,
54}
55
56#[derive(Debug, Clone)]
57pub struct IndexParams {
58    name: String,
59    index_type: IndexType,
60    metric_type: MetricType,
61    params: HashMap<String, String>,
62}
63
64impl IndexParams {
65    pub fn new(
66        name: String,
67        index_type: IndexType,
68        metric_type: MetricType,
69        params: HashMap<String, String>,
70    ) -> Self {
71        Self {
72            name,
73            index_type,
74            metric_type,
75            params,
76        }
77    }
78
79    pub fn name(&self) -> &String {
80        &self.name
81    }
82
83    pub fn index_type(&self) -> IndexType {
84        self.index_type
85    }
86
87    pub fn metric_type(&self) -> MetricType {
88        self.metric_type
89    }
90
91    pub fn params(&self) -> &HashMap<String, String> {
92        &self.params
93    }
94
95    pub fn extra_params(&self) -> HashMap<String, String> {
96        HashMap::from([
97            ("index_type".to_owned(), self.index_type().to_string()),
98            ("metric_type".to_owned(), self.metric_type().to_string()),
99            (
100                "params".to_owned(),
101                serde_json::to_string(&self.params()).unwrap(),
102            ),
103        ])
104    }
105
106    pub fn extra_kv_params(&self) -> Vec<KeyValuePair> {
107        self.extra_params()
108            .into_iter()
109            .map(|(k, v)| KeyValuePair { key: k, value: v })
110            .collect()
111    }
112}
113
114#[derive(Debug, Clone)]
115pub struct IndexInfo {
116    field_name: String,
117    id: i64,
118    params: IndexParams,
119    state: IndexState,
120}
121
122impl IndexInfo {
123    pub fn field_name(&self) -> &str {
124        &self.field_name
125    }
126
127    pub fn id(&self) -> i64 {
128        self.id
129    }
130
131    pub fn params(&self) -> &IndexParams {
132        &self.params
133    }
134
135    pub fn state(&self) -> IndexState {
136        self.state
137    }
138}
139
140impl From<IndexDescription> for IndexInfo {
141    fn from(description: IndexDescription) -> Self {
142        let mut params: HashMap<String, String> = HashMap::from_iter(
143            description
144                .params
145                .iter()
146                .map(|kv| (kv.key.clone(), kv.value.clone())),
147        );
148
149        let index_type = IndexType::from_str(&params.remove("index_type").unwrap()).unwrap();
150        let metric_type = MetricType::from_str(&params.remove("metric_type").unwrap()).unwrap();
151        let params = serde_json::from_str(params.get("params").unwrap()).unwrap();
152
153        let params = IndexParams::new(
154            description.index_name.clone(),
155            index_type,
156            metric_type,
157            params,
158        );
159        Self {
160            field_name: description.field_name.clone(),
161            id: description.index_id,
162            params: params,
163            state: description.state(),
164        }
165    }
166}
milvus/index/mod.rs

milvus/index/
mod.rs
