1use common::fmt::EscapeKwFreeIdent;
2use surrealdb_strand::Strand;
3use surrealdb_types::{Number as PublicNumber, SqlFormat, ToSql, write_sql};
4
5use crate::Cond;
6use crate::scoring::Scoring;
7
8#[derive(Clone, Debug, Eq, PartialEq)]
9#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
10pub enum Index {
11 Idx,
13 Uniq,
15 Hnsw(HnswParams),
17 DiskAnn(DiskAnnParams),
19 FullText(FullTextParams),
21 Count(Option<Cond>),
23}
24
25#[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Hash)]
26#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
27pub struct FullTextParams {
28 pub az: Strand,
29 pub hl: bool,
30 pub sc: Scoring,
31}
32
33#[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Hash)]
34#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
35pub struct HnswParams {
36 pub dimension: u16,
37 pub distance: Distance,
38 pub vector_type: VectorType,
39 pub m: u8,
40 pub m0: u8,
41 pub ef_construction: u16,
42 pub extend_candidates: bool,
43 pub keep_pruned_connections: bool,
44 pub ml: PublicNumber,
45 pub use_hashed_vector: bool,
46}
47
48#[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Hash)]
49#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
50pub struct DiskAnnParams {
51 pub dimension: u16,
53 pub distance: Distance,
55 pub vector_type: VectorType,
57 pub degree: u16,
59 pub l_build: u16,
61 pub alpha: PublicNumber,
63 pub use_hashed_vector: bool,
65}
66
67#[derive(Clone, Default, Debug, Eq, PartialEq, PartialOrd, Hash)]
68#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
69pub enum Distance {
70 Chebyshev,
72 Cosine,
74 #[default]
76 Euclidean,
77 Hamming,
79 Jaccard,
81 Manhattan,
83 Minkowski(PublicNumber),
85 Pearson,
87 CosineNormalized,
89 InnerProduct,
91}
92
93impl ToSql for Distance {
94 fn fmt_sql(&self, f: &mut String, fmt: SqlFormat) {
95 match self {
96 Self::Chebyshev => f.push_str("CHEBYSHEV"),
97 Self::Cosine => f.push_str("COSINE"),
98 Self::CosineNormalized => f.push_str("COSINE_NORMALIZED"),
99 Self::Euclidean => f.push_str("EUCLIDEAN"),
100 Self::Hamming => f.push_str("HAMMING"),
101 Self::InnerProduct => f.push_str("INNER_PRODUCT"),
102 Self::Jaccard => f.push_str("JACCARD"),
103 Self::Manhattan => f.push_str("MANHATTAN"),
104 Self::Minkowski(order) => write_sql!(f, fmt, "MINKOWSKI {}", order),
105 Self::Pearson => f.push_str("PEARSON"),
106 }
107 }
108}
109
110#[derive(Clone, Copy, Default, Debug, Eq, PartialEq, PartialOrd, Hash)]
111#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
112pub enum VectorType {
113 F64,
115 #[default]
117 F32,
118 I64,
120 I32,
122 I16,
124 F16,
126 I8,
128 U8,
130}
131
132impl ToSql for VectorType {
133 fn fmt_sql(&self, f: &mut String, _fmt: SqlFormat) {
134 match self {
135 Self::F64 => f.push_str("F64"),
136 Self::F16 => f.push_str("F16"),
137 Self::F32 => f.push_str("F32"),
138 Self::I64 => f.push_str("I64"),
139 Self::I32 => f.push_str("I32"),
140 Self::I16 => f.push_str("I16"),
141 Self::I8 => f.push_str("I8"),
142 Self::U8 => f.push_str("U8"),
143 }
144 }
145}
146
147impl ToSql for Index {
148 fn fmt_sql(&self, f: &mut String, fmt: SqlFormat) {
149 match self {
150 Self::Idx => {}
151 Self::Uniq => f.push_str("UNIQUE"),
152 Self::Count(c) => {
153 f.push_str("COUNT");
154 if let Some(v) = c {
155 write_sql!(f, fmt, " {}", v)
156 }
157 }
158 Self::FullText(p) => {
159 write_sql!(f, fmt, "FULLTEXT ANALYZER {} {}", EscapeKwFreeIdent(&p.az), p.sc);
160 if p.hl {
161 f.push_str(" HIGHLIGHTS")
162 }
163 }
164 Self::Hnsw(p) => {
165 write_sql!(
166 f,
167 fmt,
168 "HNSW DIMENSION {} DIST {} TYPE {} EFC {} M {} M0 {} LM {}",
169 p.dimension,
170 p.distance,
171 p.vector_type,
172 p.ef_construction,
173 p.m,
174 p.m0,
175 p.ml
176 );
177 if p.extend_candidates {
178 f.push_str(" EXTEND_CANDIDATES")
179 }
180 if p.keep_pruned_connections {
181 f.push_str(" KEEP_PRUNED_CONNECTIONS")
182 }
183 if p.use_hashed_vector {
184 f.push_str(" HASHED_VECTOR")
185 }
186 }
187 Self::DiskAnn(p) => {
188 write_sql!(
189 f,
190 fmt,
191 "DISKANN DIMENSION {} DIST {} TYPE {} DEGREE {} L_BUILD {} ALPHA {}",
192 p.dimension,
193 p.distance,
194 p.vector_type,
195 p.degree,
196 p.l_build,
197 p.alpha
198 );
199 if p.use_hashed_vector {
200 f.push_str(" HASHED_VECTOR")
201 }
202 }
203 }
204 }
205}