Skip to main content

linkrs_query/parser/ast/
vector.rs

1//! Vector Search AST Definitions
2//!
3//! This module defines the Abstract Syntax Tree (AST) nodes for vector search queries,
4//! including CREATE VECTOR INDEX, SEARCH VECTOR, and related statements.
5
6use linkrs_core::types::expr::contextual::ContextualExpression;
7use linkrs_core::types::span::Span;
8use serde::{Deserialize, Serialize};
9
10use super::stmt::OrderByClause;
11use super::types::{LimitClause, SkipClause};
12
13// ============================================================================
14// Vector Index DDL Statements
15// ============================================================================
16
17/// CREATE VECTOR INDEX statement
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct CreateVectorIndex {
20    pub span: Span,
21    pub index_name: String,
22    pub schema_name: String,
23    pub field_name: String,
24    pub config: VectorIndexConfig,
25    pub if_not_exists: bool,
26}
27
28/// Quantization kind for vector index (mirrors `simvec::QuantizationType` / Qdrant builders)
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
30#[serde(rename_all = "lowercase")]
31pub enum QuantizationKind {
32    Scalar,
33    Binary,
34    Product,
35}
36
37/// Compression ratio for product quantization (mirrors `simvec::CompressionRatio` / Qdrant)
38#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
39#[serde(rename_all = "lowercase")]
40pub enum CompressionRatioKind {
41    X4,
42    X8,
43    X16,
44    X32,
45    X64,
46}
47
48/// Vector index configuration
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct VectorIndexConfig {
51    pub vector_size: usize,
52    pub distance: VectorDistance,
53    pub hnsw_m: Option<usize>,
54    pub hnsw_ef_construct: Option<usize>,
55    /// Quantization kind: None = disabled (exact f32), Some = scalar/binary/product
56    pub quantization: Option<QuantizationKind>,
57    /// Scalar only: quantile in (0,1], default 0.99
58    pub quantile: Option<f32>,
59    /// Product only: compression ratio, default X4
60    pub compression: Option<CompressionRatioKind>,
61    /// Always keep quantized vectors in RAM (scalar/binary/product)
62    pub always_ram: Option<bool>,
63}
64
65/// DROP VECTOR INDEX statement
66#[derive(Debug, Clone, Serialize, Deserialize)]
67pub struct DropVectorIndex {
68    pub span: Span,
69    pub index_name: String,
70    pub if_exists: bool,
71}
72
73// ============================================================================
74// Vector Search DML Statements
75// ============================================================================
76
77/// SEARCH VECTOR statement
78#[derive(Debug, Clone)]
79pub struct SearchVectorStatement {
80    pub span: Span,
81    pub index_name: String,
82    pub query: VectorQueryExpr,
83    pub threshold: Option<f32>,
84    pub where_clause: Option<ContextualExpression>,
85    pub order_clause: Option<OrderByClause>,
86    pub limit: Option<LimitClause>,
87    pub skip: Option<SkipClause>,
88    pub yield_clause: Option<VectorYieldClause>,
89}
90
91/// Vector query expression
92#[derive(Debug, Clone, Serialize, Deserialize)]
93pub struct VectorQueryExpr {
94    pub span: Span,
95    pub query_type: VectorQueryType,
96    pub query_data: String,
97}
98
99/// Vector query type
100#[derive(Debug, Clone, Serialize, Deserialize)]
101pub enum VectorQueryType {
102    /// Direct vector: [0.1, 0.2, ...]
103    Vector,
104    /// Text query (requires embedding service)
105    Text,
106    /// Parameter reference: $param_name
107    Parameter,
108}
109
110/// Vector distance metric
111#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
112#[serde(rename_all = "lowercase")]
113pub enum VectorDistance {
114    Cosine,
115    Euclidean,
116    Dot,
117    Manhattan,
118}
119
120/// YIELD clause for vector search
121#[derive(Debug, Clone)]
122pub struct VectorYieldClause {
123    pub items: Vec<VectorYieldItem>,
124}
125
126/// Yield item for vector search
127#[derive(Debug, Clone)]
128pub struct VectorYieldItem {
129    pub expr: ContextualExpression,
130    pub alias: Option<String>,
131}
132
133// ============================================================================
134// MATCH and LOOKUP extensions for vector search
135// ============================================================================
136
137/// MATCH with vector search
138#[derive(Debug, Clone)]
139pub struct MatchVector {
140    pub span: Span,
141    pub pattern: String,
142    pub vector_condition: VectorMatchCondition,
143    pub yield_clause: Option<VectorYieldClause>,
144}
145
146/// Vector match condition in WHERE clause
147#[derive(Debug, Clone, Serialize, Deserialize)]
148pub struct VectorMatchCondition {
149    pub field: String,
150    pub query: VectorQueryExpr,
151    pub threshold: Option<f32>,
152}
153
154/// LOOKUP with vector search
155#[derive(Debug, Clone)]
156pub struct LookupVector {
157    pub span: Span,
158    pub schema_name: String,
159    pub index_name: String,
160    pub query: VectorQueryExpr,
161    pub yield_clause: Option<VectorYieldClause>,
162    pub limit: Option<usize>,
163}
164
165// ============================================================================
166// Helper Functions
167// ============================================================================
168
169impl CreateVectorIndex {
170    pub fn new(
171        span: Span,
172        index_name: String,
173        schema_name: String,
174        field_name: String,
175        config: VectorIndexConfig,
176    ) -> Self {
177        Self {
178            span,
179            index_name,
180            schema_name,
181            field_name,
182            config,
183            if_not_exists: false,
184        }
185    }
186
187    pub fn with_if_not_exists(mut self, if_not_exists: bool) -> Self {
188        self.if_not_exists = if_not_exists;
189        self
190    }
191}
192
193impl DropVectorIndex {
194    pub fn new(span: Span, index_name: String) -> Self {
195        Self {
196            span,
197            index_name,
198            if_exists: false,
199        }
200    }
201
202    pub fn with_if_exists(mut self, if_exists: bool) -> Self {
203        self.if_exists = if_exists;
204        self
205    }
206}
207
208impl VectorQueryExpr {
209    pub fn vector(span: Span, vector_str: String) -> Self {
210        Self {
211            span,
212            query_type: VectorQueryType::Vector,
213            query_data: vector_str,
214        }
215    }
216
217    pub fn text(span: Span, text: String) -> Self {
218        Self {
219            span,
220            query_type: VectorQueryType::Text,
221            query_data: text,
222        }
223    }
224
225    pub fn parameter(span: Span, param_name: String) -> Self {
226        Self {
227            span,
228            query_type: VectorQueryType::Parameter,
229            query_data: param_name,
230        }
231    }
232}
233
234impl VectorIndexConfig {
235    pub fn new(vector_size: usize, distance: VectorDistance) -> Self {
236        Self {
237            vector_size,
238            distance,
239            hnsw_m: None,
240            hnsw_ef_construct: None,
241            quantization: None,
242            quantile: None,
243            compression: None,
244            always_ram: None,
245        }
246    }
247
248    pub fn with_hnsw(mut self, m: usize, ef_construct: usize) -> Self {
249        self.hnsw_m = Some(m);
250        self.hnsw_ef_construct = Some(ef_construct);
251        self
252    }
253
254    pub fn with_quantization(
255        mut self,
256        quantization: QuantizationKind,
257        quantile: Option<f32>,
258        compression: Option<CompressionRatioKind>,
259        always_ram: Option<bool>,
260    ) -> Self {
261        self.quantization = Some(quantization);
262        self.quantile = quantile;
263        self.compression = compression;
264        self.always_ram = always_ram;
265        self
266    }
267}
268
269impl CompressionRatioKind {
270    pub fn as_str(&self) -> &'static str {
271        match self {
272            CompressionRatioKind::X4 => "x4",
273            CompressionRatioKind::X8 => "x8",
274            CompressionRatioKind::X16 => "x16",
275            CompressionRatioKind::X32 => "x32",
276            CompressionRatioKind::X64 => "x64",
277        }
278    }
279}