alopex_sql/planner/
knn_optimizer.rs

1use crate::ast::expr::Literal;
2use crate::executor::evaluator::vector_ops::VectorMetric;
3use crate::planner::logical_plan::LogicalPlan;
4use crate::planner::typed_expr::{Projection, SortExpr, TypedExprKind};
5
6/// ORDER BY + LIMIT から抽出した KNN 最適化パターン。
7#[derive(Debug, Clone, PartialEq)]
8pub struct KnnPattern {
9    pub table: String,
10    pub column: String,
11    pub query_vector: Vec<f32>,
12    pub metric: VectorMetric,
13    pub k: u64,
14    pub sort_direction: SortDirection,
15}
16
17/// ソート方向（ASC / DESC）。
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum SortDirection {
20    Asc,
21    Desc,
22}
23
24/// ORDER BY vector_similarity/vector_distance + LIMIT K の形になっているか検出する。
25///
26/// - LIMIT が存在し、OFFSET が無い（または 0）の場合のみ最適化対象。
27/// - ORDER BY は単一のベクトル関数呼び出しであること。
28/// - メトリクスとソート方向の整合性を満たす場合のみ Some を返す。
29pub fn detect_knn_pattern(plan: &LogicalPlan) -> Option<KnnPattern> {
30    let (sort_plan, k) = extract_limit(plan)?;
31    let (order_expr, input_after_sort) = extract_sort(sort_plan)?;
32    let sort_direction = if order_expr.asc {
33        SortDirection::Asc
34    } else {
35        SortDirection::Desc
36    };
37
38    let (table, _projection, _filter) = extract_scan_context(input_after_sort)?;
39
40    let (func_name, args) = match &order_expr.expr.kind {
41        TypedExprKind::FunctionCall { name, args, .. } => (name.to_ascii_lowercase(), args),
42        _ => return None,
43    };
44
45    if func_name != "vector_similarity" && func_name != "vector_distance" {
46        return None;
47    }
48
49    if args.len() != 3 {
50        return None;
51    }
52
53    let column_name = extract_column_name(&args[0], &table)?;
54    let query_vector = extract_query_vector(&args[1])?;
55    let metric = extract_metric(&args[2])?;
56
57    if !is_valid_knn_direction(metric, sort_direction) {
58        return None;
59    }
60
61    Some(KnnPattern {
62        table,
63        column: column_name,
64        query_vector,
65        metric,
66        k,
67        sort_direction,
68    })
69}
70
71fn extract_limit(plan: &LogicalPlan) -> Option<(&LogicalPlan, u64)> {
72    match plan {
73        // WITH TIES (ties: Some) must never take the index shortcut: the
74        // exact peer set of the boundary row requires the full sort.
75        LogicalPlan::Limit {
76            input,
77            limit: Some(k),
78            offset,
79            ties: None,
80        } if offset.unwrap_or(0) == 0 => Some((input.as_ref(), *k)),
81        _ => None,
82    }
83}
84
85fn extract_sort(plan: &LogicalPlan) -> Option<(&SortExpr, &LogicalPlan)> {
86    if let LogicalPlan::Sort { input, order_by } = plan
87        && order_by.len() == 1
88    {
89        return Some((&order_by[0], input.as_ref()));
90    }
91    None
92}
93
94fn extract_scan_context(
95    plan: &LogicalPlan,
96) -> Option<(
97    String,
98    Projection,
99    Option<crate::planner::typed_expr::TypedExpr>,
100)> {
101    match plan {
102        LogicalPlan::Filter { input, predicate } => {
103            if let LogicalPlan::Scan { table, projection } = input.as_ref() {
104                return Some((table.clone(), projection.clone(), Some(predicate.clone())));
105            }
106            None
107        }
108        LogicalPlan::Scan { table, projection } => Some((table.clone(), projection.clone(), None)),
109        _ => None,
110    }
111}
112
113fn extract_column_name(
114    expr: &crate::planner::typed_expr::TypedExpr,
115    table: &str,
116) -> Option<String> {
117    match &expr.kind {
118        TypedExprKind::ColumnRef {
119            table: tbl, column, ..
120        } if tbl == table => Some(column.clone()),
121        _ => None,
122    }
123}
124
125fn extract_query_vector(expr: &crate::planner::typed_expr::TypedExpr) -> Option<Vec<f32>> {
126    match &expr.kind {
127        TypedExprKind::VectorLiteral(values) if !values.is_empty() => {
128            Some(values.iter().map(|v| *v as f32).collect())
129        }
130        _ => None,
131    }
132}
133
134fn extract_metric(expr: &crate::planner::typed_expr::TypedExpr) -> Option<VectorMetric> {
135    match &expr.kind {
136        TypedExprKind::Literal(Literal::String(s)) => s.parse().ok(),
137        _ => None,
138    }
139}
140
141fn is_valid_knn_direction(metric: VectorMetric, dir: SortDirection) -> bool {
142    matches!(
143        (metric, dir),
144        (VectorMetric::Cosine, SortDirection::Desc)
145            | (VectorMetric::Inner, SortDirection::Desc)
146            | (VectorMetric::L2, SortDirection::Asc)
147    )
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use crate::ast::ddl::VectorMetric as AstVectorMetric;
154    use crate::ast::span::Span;
155    use crate::planner::logical_plan::LogicalPlan;
156    use crate::planner::typed_expr::{Projection, SortExpr, TypedExpr};
157    use crate::planner::types::ResolvedType;
158
159    fn base_vector_type() -> ResolvedType {
160        ResolvedType::Vector {
161            dimension: 2,
162            metric: AstVectorMetric::Cosine,
163        }
164    }
165
166    fn build_plan(order_asc: bool, metric_literal: &str, offset: Option<u64>) -> LogicalPlan {
167        let span = Span::empty();
168        let vector_expr = TypedExpr::function_call(
169            "vector_similarity".to_string(),
170            vec![
171                TypedExpr::column_ref(
172                    "items".to_string(),
173                    "embedding".to_string(),
174                    0,
175                    base_vector_type(),
176                    span,
177                ),
178                TypedExpr::vector_literal(vec![1.0, 0.0], 2, span),
179                TypedExpr::literal(
180                    Literal::String(metric_literal.to_string()),
181                    ResolvedType::Text,
182                    span,
183                ),
184            ],
185            false,
186            false,
187            ResolvedType::Double,
188            span,
189        );
190        let sort = LogicalPlan::Sort {
191            input: Box::new(LogicalPlan::Scan {
192                table: "items".to_string(),
193                projection: Projection::All(vec!["embedding".to_string()]),
194            }),
195            order_by: vec![SortExpr::new(vector_expr, order_asc, false)],
196        };
197
198        LogicalPlan::Limit {
199            input: Box::new(sort),
200            limit: Some(2),
201            offset,
202            ties: None,
203        }
204    }
205
206    #[test]
207    fn detect_knn_pattern_cosine_desc() {
208        let plan = build_plan(false, "cosine", None);
209        let pattern = detect_knn_pattern(&plan).expect("should detect pattern");
210        assert_eq!(pattern.table, "items");
211        assert_eq!(pattern.column, "embedding");
212        assert_eq!(pattern.k, 2);
213        assert_eq!(pattern.metric, VectorMetric::Cosine);
214        assert_eq!(pattern.sort_direction, SortDirection::Desc);
215        assert_eq!(pattern.query_vector, vec![1.0, 0.0]);
216    }
217
218    #[test]
219    fn reject_invalid_direction() {
220        let plan = build_plan(true, "cosine", None);
221        assert!(detect_knn_pattern(&plan).is_none());
222    }
223
224    #[test]
225    fn reject_missing_limit_or_offset() {
226        let plan_no_limit = LogicalPlan::Sort {
227            input: Box::new(LogicalPlan::Scan {
228                table: "items".to_string(),
229                projection: Projection::All(vec!["embedding".to_string()]),
230            }),
231            order_by: vec![],
232        };
233        assert!(detect_knn_pattern(&plan_no_limit).is_none());
234
235        let plan_with_offset = build_plan(false, "cosine", Some(1));
236        assert!(detect_knn_pattern(&plan_with_offset).is_none());
237    }
238
239    #[test]
240    fn reject_unknown_metric() {
241        let plan = build_plan(false, "unknown", None);
242        assert!(detect_knn_pattern(&plan).is_none());
243    }
244
245    #[test]
246    fn reject_with_ties_limit() {
247        // FETCH ... WITH TIES needs the exact peer set of the boundary row,
248        // so the KNN index shortcut must not fire (issue #152, D14).
249        let plan = build_plan(false, "cosine", None);
250        let LogicalPlan::Limit {
251            input,
252            limit,
253            offset,
254            ..
255        } = plan
256        else {
257            panic!("expected Limit plan");
258        };
259        let ties = if let LogicalPlan::Sort { order_by, .. } = input.as_ref() {
260            Some(order_by.clone())
261        } else {
262            panic!("expected Sort input");
263        };
264        let with_ties = LogicalPlan::Limit {
265            input,
266            limit,
267            offset,
268            ties,
269        };
270        assert!(detect_knn_pattern(&with_ties).is_none());
271    }
272}
alopex_sql/planner/knn_optimizer.rs

alopex_sql/planner/
knn_optimizer.rs
