veclite_sql/
planner.rs

1use sqlparser::ast::{Expr, Query, SetExpr, Statement, Value};
2
3#[derive(Debug, PartialEq)]
4pub enum LogicalPlan {
5    CreateTable {
6        name: String,
7    },
8    Insert {
9        table: String,
10        id: String,
11        vector: Vec<f32>,
12        metadata: Option<String>,
13    },
14    Search {
15        table: String,
16        query_vector: Vec<f32>,
17        limit: usize,
18        filter: Option<String>,
19    },
20    Compact {
21        table: String,
22    },
23    RetentionPolicy {
24        table: String,
25        interval: String,
26    },
27}
28
29pub struct PlanBuilder;
30
31impl PlanBuilder {
32    pub fn build_from_sql(query: &str) -> Result<LogicalPlan, String> {
33        let q = query.trim();
34        let upper = q.to_uppercase();
35        if upper.starts_with("VACUUM") || upper.starts_with("COMPACT") {
36            let parts: Vec<&str> = q.split_whitespace().collect();
37            if parts.len() >= 2 {
38                let table = parts[1].trim_end_matches(';').to_string();
39                return Ok(LogicalPlan::Compact { table });
40            }
41        }
42
43        let stmts = crate::parser::parse_sql(query)?;
44        if stmts.is_empty() {
45            return Err("No statements found".to_string());
46        }
47        Self::build(stmts[0].clone())
48    }
49
50    pub fn build(stmt: Statement) -> Result<LogicalPlan, String> {
51        match stmt {
52            Statement::CreateTable { name, .. } => {
53                let table_name = name.to_string();
54                Ok(LogicalPlan::CreateTable { name: table_name })
55            }
56            Statement::Insert {
57                table_name, source, ..
58            } => {
59                let table = table_name.to_string();
60                let body = match source.as_ref().map(|s| &*s.body) {
61                    Some(SetExpr::Values(values)) => values,
62                    _ => return Err("Only VALUES inserts are supported".to_string()),
63                };
64
65                if body.rows.len() != 1 {
66                    return Err("Only single row inserts are supported".to_string());
67                }
68
69                let row = &body.rows[0];
70                if row.len() < 2 {
71                    return Err("Insert requires at least id and vector".to_string());
72                }
73
74                let id = Self::extract_string(&row[0])?;
75                let vector = Self::extract_vector(&row[1])?;
76                let metadata = if row.len() > 2 {
77                    Some(Self::extract_string(&row[2])?)
78                } else {
79                    None
80                };
81
82                Ok(LogicalPlan::Insert {
83                    table,
84                    id,
85                    vector,
86                    metadata,
87                })
88            }
89            Statement::Query(query) => Self::build_query(*query),
90            _ => Err(format!("Unsupported statement: {:?}", stmt)),
91        }
92    }
93
94    fn build_query(query: Query) -> Result<LogicalPlan, String> {
95        let select = match *query.body {
96            SetExpr::Select(s) => s,
97            _ => return Err("Only SELECT queries are supported".to_string()),
98        };
99
100        if select.from.is_empty() {
101            return Err("FROM clause is required".to_string());
102        }
103
104        let table = select.from[0].relation.to_string();
105
106        let limit = match query.limit {
107            Some(Expr::Value(Value::Number(n, _))) => n.parse::<usize>().unwrap_or(10),
108            _ => 10,
109        };
110
111        let mut query_vector = vec![];
112
113        if !query.order_by.is_empty() {
114            let order_expr = &query.order_by[0];
115            match &order_expr.expr {
116                Expr::BinaryOp { left: _, op, right } => {
117                    let op_str = op.to_string();
118                    // Some parser versions might format differently, check for basic pgvector ops
119                    if op_str.contains("<->") || op_str.contains("<=>") || op_str.contains("<#>") {
120                        query_vector = Self::extract_vector(right)?;
121                    } else {
122                        return Err(format!("Unsupported order by operator: {}", op_str));
123                    }
124                }
125                _ => return Err("Unsupported ORDER BY expression".to_string()),
126            }
127        } else {
128            return Err("ORDER BY vector distance is required for Search".to_string());
129        }
130
131        Ok(LogicalPlan::Search {
132            table,
133            query_vector,
134            limit,
135            filter: None,
136        })
137    }
138
139    fn extract_string(expr: &Expr) -> Result<String, String> {
140        match expr {
141            Expr::Value(Value::SingleQuotedString(s)) => Ok(s.clone()),
142            Expr::Value(Value::DoubleQuotedString(s)) => Ok(s.clone()),
143            _ => Err("Expected string value".to_string()),
144        }
145    }
146
147    fn extract_vector(expr: &Expr) -> Result<Vec<f32>, String> {
148        match expr {
149            Expr::Value(Value::SingleQuotedString(s))
150            | Expr::Value(Value::DoubleQuotedString(s)) => {
151                let trimmed = s.trim_matches(|c| c == '[' || c == ']');
152                let mut vec = Vec::new();
153                if trimmed.is_empty() {
154                    return Ok(vec);
155                }
156                for part in trimmed.split(',') {
157                    let val: f32 = part
158                        .trim()
159                        .parse()
160                        .map_err(|_| format!("Invalid number in vector: {}", part))?;
161                    vec.push(val);
162                }
163                Ok(vec)
164            }
165            Expr::Array(sqlparser::ast::Array { elem, .. }) => {
166                let mut vec = Vec::new();
167                for e in elem {
168                    match e {
169                        Expr::Value(Value::Number(n, _)) => {
170                            let val: f32 = n
171                                .parse()
172                                .map_err(|_| format!("Invalid number in array: {}", n))?;
173                            vec.push(val);
174                        }
175                        _ => return Err("Expected numbers in array".to_string()),
176                    }
177                }
178                Ok(vec)
179            }
180            _ => Err(format!("Expected vector string or array, got {:?}", expr)),
181        }
182    }
183}
veclite_sql/planner.rs

veclite_sql/
planner.rs
