Skip to main content

surrealdb_sql/
index.rs

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	/// (Basic) non unique
12	Idx,
13	/// Unique index
14	Uniq,
15	/// HNSW index for distance based metrics
16	Hnsw(HnswParams),
17	/// DiskANN index for distance based metrics
18	DiskAnn(DiskAnnParams),
19	/// Index with Full-Text search capabilities - single writer
20	FullText(FullTextParams),
21	/// Count index
22	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	/// The vector dimension.
52	pub dimension: u16,
53	/// The distance metric used by the DiskANN graph.
54	pub distance: Distance,
55	/// The element type used to encode vectors in the index.
56	pub vector_type: VectorType,
57	/// Target maximum graph degree.
58	pub degree: u16,
59	/// Construction search list size.
60	pub l_build: u16,
61	/// DiskANN pruning alpha.
62	pub alpha: PublicNumber,
63	/// Whether vector-document mappings are keyed by vector hash.
64	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 distance.
71	Chebyshev,
72	/// Cosine distance.
73	Cosine,
74	/// Euclidean distance.
75	#[default]
76	Euclidean,
77	/// Hamming distance.
78	Hamming,
79	/// Jaccard distance.
80	Jaccard,
81	/// Manhattan distance.
82	Manhattan,
83	/// Minkowski distance with the supplied order.
84	Minkowski(PublicNumber),
85	/// Pearson similarity.
86	Pearson,
87	/// Cosine distance for already-normalized vectors.
88	CosineNormalized,
89	/// Inner product similarity, transformed into a distance score.
90	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	/// 64-bit floating point.
114	F64,
115	/// 32-bit floating point.
116	#[default]
117	F32,
118	/// 64-bit signed integer.
119	I64,
120	/// 32-bit signed integer.
121	I32,
122	/// 16-bit signed integer.
123	I16,
124	/// 16-bit floating point.
125	F16,
126	/// 8-bit signed integer.
127	I8,
128	/// 8-bit unsigned integer.
129	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}