Skip to main content

surql_parser/upstream/sql/
index.rs

1use crate::compat::types::PublicNumber;
2use crate::upstream::fmt::EscapeKwFreeIdent;
3use crate::upstream::sql::Cond;
4use crate::upstream::sql::scoring::Scoring;
5use surrealdb_types::{SqlFormat, ToSql, write_sql};
6#[derive(Clone, Debug, Eq, PartialEq)]
7#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
8pub enum Index {
9	/// (Basic) non unique
10	Idx,
11	/// Unique index
12	Uniq,
13	/// HNSW index for distance based metrics
14	Hnsw(HnswParams),
15	/// Index with Full-Text search capabilities - single writer
16	FullText(FullTextParams),
17	/// Count index
18	Count(Option<Cond>),
19}
20impl From<Index> for crate::compat::catalog::Index {
21	fn from(v: Index) -> Self {
22		match v {
23			Index::Idx => Self::Idx,
24			Index::Uniq => Self::Uniq,
25			Index::Hnsw(p) => Self::Hnsw(p.into()),
26			Index::FullText(p) => Self::FullText(p.into()),
27			Index::Count(c) => Self::Count(c.map(Into::into)),
28		}
29	}
30}
31impl From<crate::compat::catalog::Index> for Index {
32	fn from(v: crate::compat::catalog::Index) -> Self {
33		match v {
34			crate::compat::catalog::Index::Idx => Self::Idx,
35			crate::compat::catalog::Index::Uniq => Self::Uniq,
36			crate::compat::catalog::Index::Hnsw(p) => Self::Hnsw(p.into()),
37			crate::compat::catalog::Index::FullText(p) => Self::FullText(p.into()),
38			crate::compat::catalog::Index::Count(c) => Self::Count(c.map(Into::into)),
39		}
40	}
41}
42#[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Hash)]
43#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
44pub struct FullTextParams {
45	pub az: String,
46	pub hl: bool,
47	pub sc: Scoring,
48}
49impl From<FullTextParams> for crate::compat::catalog::FullTextParams {
50	fn from(v: FullTextParams) -> Self {
51		crate::compat::catalog::FullTextParams {
52			analyzer: v.az.clone(),
53			highlight: v.hl,
54			scoring: v.sc.into(),
55		}
56	}
57}
58impl From<crate::compat::catalog::FullTextParams> for FullTextParams {
59	fn from(v: crate::compat::catalog::FullTextParams) -> Self {
60		Self {
61			az: v.analyzer,
62			hl: v.highlight,
63			sc: v.scoring.into(),
64		}
65	}
66}
67#[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Hash)]
68#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
69pub struct HnswParams {
70	pub dimension: u16,
71	pub distance: Distance,
72	pub vector_type: VectorType,
73	pub m: u8,
74	pub m0: u8,
75	pub ef_construction: u16,
76	pub extend_candidates: bool,
77	pub keep_pruned_connections: bool,
78	pub ml: PublicNumber,
79	pub use_hashed_vector: bool,
80}
81impl From<HnswParams> for crate::compat::catalog::HnswParams {
82	fn from(v: HnswParams) -> Self {
83		crate::compat::catalog::HnswParams {
84			dimension: v.dimension,
85			distance: v.distance.into(),
86			vector_type: v.vector_type.into(),
87			m: v.m,
88			m0: v.m0,
89			ef_construction: v.ef_construction,
90			ml: v.ml.into(),
91			extend_candidates: v.extend_candidates,
92			keep_pruned_connections: v.keep_pruned_connections,
93			use_hashed_vector: v.use_hashed_vector,
94		}
95	}
96}
97impl From<crate::compat::catalog::HnswParams> for HnswParams {
98	fn from(v: crate::compat::catalog::HnswParams) -> Self {
99		Self {
100			dimension: v.dimension,
101			distance: v.distance.into(),
102			vector_type: v.vector_type.into(),
103			m: v.m,
104			m0: v.m0,
105			ef_construction: v.ef_construction,
106			ml: v.ml.into(),
107			extend_candidates: v.extend_candidates,
108			keep_pruned_connections: v.keep_pruned_connections,
109			use_hashed_vector: v.use_hashed_vector,
110		}
111	}
112}
113#[derive(Clone, Default, Debug, Eq, PartialEq, PartialOrd, Hash)]
114#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
115pub enum Distance {
116	Chebyshev,
117	Cosine,
118	#[default]
119	Euclidean,
120	Hamming,
121	Jaccard,
122	Manhattan,
123	Minkowski(PublicNumber),
124	Pearson,
125}
126impl ToSql for Distance {
127	fn fmt_sql(&self, f: &mut String, fmt: SqlFormat) {
128		match self {
129			Self::Chebyshev => f.push_str("CHEBYSHEV"),
130			Self::Cosine => f.push_str("COSINE"),
131			Self::Euclidean => f.push_str("EUCLIDEAN"),
132			Self::Hamming => f.push_str("HAMMING"),
133			Self::Jaccard => f.push_str("JACCARD"),
134			Self::Manhattan => f.push_str("MANHATTAN"),
135			Self::Minkowski(order) => write_sql!(f, fmt, "MINKOWSKI {}", order),
136			Self::Pearson => f.push_str("PEARSON"),
137		}
138	}
139}
140impl From<Distance> for crate::compat::catalog::Distance {
141	fn from(v: Distance) -> Self {
142		match v {
143			Distance::Chebyshev => crate::compat::catalog::Distance::Chebyshev,
144			Distance::Cosine => crate::compat::catalog::Distance::Cosine,
145			Distance::Euclidean => crate::compat::catalog::Distance::Euclidean,
146			Distance::Hamming => crate::compat::catalog::Distance::Hamming,
147			Distance::Jaccard => crate::compat::catalog::Distance::Jaccard,
148			Distance::Manhattan => crate::compat::catalog::Distance::Manhattan,
149			Distance::Minkowski(n) => crate::compat::catalog::Distance::Minkowski(n.into()),
150			Distance::Pearson => crate::compat::catalog::Distance::Pearson,
151		}
152	}
153}
154impl From<crate::compat::catalog::Distance> for Distance {
155	fn from(v: crate::compat::catalog::Distance) -> Self {
156		match v {
157			crate::compat::catalog::Distance::Chebyshev => Self::Chebyshev,
158			crate::compat::catalog::Distance::Cosine => Self::Cosine,
159			crate::compat::catalog::Distance::Euclidean => Self::Euclidean,
160			crate::compat::catalog::Distance::Hamming => Self::Hamming,
161			crate::compat::catalog::Distance::Jaccard => Self::Jaccard,
162			crate::compat::catalog::Distance::Manhattan => Self::Manhattan,
163			crate::compat::catalog::Distance::Minkowski(n) => Self::Minkowski(n.into()),
164			crate::compat::catalog::Distance::Pearson => Self::Pearson,
165		}
166	}
167}
168#[derive(Clone, Copy, Default, Debug, Eq, PartialEq, PartialOrd, Hash)]
169#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
170pub enum VectorType {
171	F64,
172	#[default]
173	F32,
174	I64,
175	I32,
176	I16,
177}
178impl ToSql for VectorType {
179	fn fmt_sql(&self, f: &mut String, _fmt: SqlFormat) {
180		match self {
181			Self::F64 => f.push_str("F64"),
182			Self::F32 => f.push_str("F32"),
183			Self::I64 => f.push_str("I64"),
184			Self::I32 => f.push_str("I32"),
185			Self::I16 => f.push_str("I16"),
186		}
187	}
188}
189impl ToSql for Index {
190	fn fmt_sql(&self, f: &mut String, fmt: SqlFormat) {
191		match self {
192			Self::Idx => {}
193			Self::Uniq => f.push_str("UNIQUE"),
194			Self::Count(c) => {
195				f.push_str("COUNT");
196				if let Some(v) = c {
197					write_sql!(f, fmt, " {}", v)
198				}
199			}
200			Self::FullText(p) => {
201				write_sql!(
202					f,
203					fmt,
204					"FULLTEXT ANALYZER {} {}",
205					EscapeKwFreeIdent(&p.az),
206					p.sc
207				);
208				if p.hl {
209					f.push_str(" HIGHLIGHTS")
210				}
211			}
212			Self::Hnsw(p) => {
213				write_sql!(
214					f,
215					fmt,
216					"HNSW DIMENSION {} DIST {} TYPE {} EFC {} M {} M0 {} LM {}",
217					p.dimension,
218					p.distance,
219					p.vector_type,
220					p.ef_construction,
221					p.m,
222					p.m0,
223					p.ml
224				);
225				if p.extend_candidates {
226					f.push_str(" EXTEND_CANDIDATES")
227				}
228				if p.keep_pruned_connections {
229					f.push_str(" KEEP_PRUNED_CONNECTIONS")
230				}
231				if p.use_hashed_vector {
232					f.push_str(" HASHED_VECTOR")
233				}
234			}
235		}
236	}
237}
238impl From<VectorType> for crate::compat::catalog::VectorType {
239	fn from(v: VectorType) -> Self {
240		match v {
241			VectorType::F64 => Self::F64,
242			VectorType::F32 => Self::F32,
243			VectorType::I64 => Self::I64,
244			VectorType::I32 => Self::I32,
245			VectorType::I16 => Self::I16,
246		}
247	}
248}
249impl From<crate::compat::catalog::VectorType> for VectorType {
250	fn from(v: crate::compat::catalog::VectorType) -> Self {
251		match v {
252			crate::compat::catalog::VectorType::F64 => Self::F64,
253			crate::compat::catalog::VectorType::F32 => Self::F32,
254			crate::compat::catalog::VectorType::I64 => Self::I64,
255			crate::compat::catalog::VectorType::I32 => Self::I32,
256			crate::compat::catalog::VectorType::I16 => Self::I16,
257		}
258	}
259}