llkv_plan/
conversion.rs

1//! Conversion utilities from SQL AST to Plan types.
2//!
3//! This module provides functions for converting sqlparser AST nodes into
4//! llkv-plan data structures, particularly for literal value conversion
5//! and range SELECT parsing.
6
7use llkv_result::{Error, Result};
8use rustc_hash::FxHashMap;
9use sqlparser::ast::{
10    Expr as SqlExpr, FunctionArg, FunctionArgExpr, GroupByExpr, ObjectName, ObjectNamePart, Select,
11    SelectItem, SelectItemQualifiedWildcardKind, TableAlias, TableFactor, UnaryOperator, Value,
12    ValueWithSpan,
13};
14
15use crate::PlanValue;
16
17/// Convert a SQL expression to a PlanValue literal.
18///
19/// Supports:
20/// - Literal values (numbers, strings, NULL)
21/// - Unary operators (-, +)
22/// - Nested expressions
23/// - Dictionary/struct literals
24///
25/// # Examples
26///
27/// ```ignore
28/// use llkv_plan::conversion::plan_value_from_sql_expr;
29/// use sqlparser::ast::Expr;
30///
31/// let expr = /* parse SQL expression */;
32/// let value = plan_value_from_sql_expr(&expr)?;
33/// ```
34///
35/// # Errors
36///
37/// Returns an error for:
38/// - Non-literal expressions (column references, function calls, etc.)
39/// - Unsupported operators
40/// - Invalid number parsing
41pub fn plan_value_from_sql_expr(expr: &SqlExpr) -> Result<PlanValue> {
42    match expr {
43        SqlExpr::Value(value) => plan_value_from_sql_value(value),
44        SqlExpr::UnaryOp {
45            op: UnaryOperator::Minus,
46            expr,
47        } => match plan_value_from_sql_expr(expr)? {
48            PlanValue::Integer(v) => Ok(PlanValue::Integer(-v)),
49            PlanValue::Float(v) => Ok(PlanValue::Float(-v)),
50            PlanValue::Null | PlanValue::String(_) | PlanValue::Struct(_) => Err(
51                Error::InvalidArgumentError("cannot negate non-numeric literal".into()),
52            ),
53        },
54        SqlExpr::UnaryOp {
55            op: UnaryOperator::Plus,
56            expr,
57        } => plan_value_from_sql_expr(expr),
58        SqlExpr::Nested(inner) => plan_value_from_sql_expr(inner),
59        SqlExpr::Dictionary(fields) => {
60            let mut map = FxHashMap::with_capacity_and_hasher(fields.len(), Default::default());
61            for field in fields {
62                let key = field.key.value.clone();
63                let value = plan_value_from_sql_expr(&field.value)?;
64                map.insert(key, value);
65            }
66            Ok(PlanValue::Struct(map))
67        }
68        other => Err(Error::InvalidArgumentError(format!(
69            "unsupported literal expression: {other:?}"
70        ))),
71    }
72}
73
74/// Convert a SQL value literal to a PlanValue.
75///
76/// Handles:
77/// - NULL
78/// - Numbers (integers and floats)
79/// - Strings
80///
81/// # Examples
82///
83/// ```ignore
84/// use llkv_plan::conversion::plan_value_from_sql_value;
85/// use sqlparser::ast::ValueWithSpan;
86///
87/// let value = /* parse SQL value */;
88/// let plan_value = plan_value_from_sql_value(&value)?;
89/// ```
90///
91/// # Errors
92///
93/// Returns an error for:
94/// - Boolean literals (not yet supported)
95/// - Invalid number formats
96/// - Unsupported value types
97pub fn plan_value_from_sql_value(value: &ValueWithSpan) -> Result<PlanValue> {
98    match &value.value {
99        Value::Null => Ok(PlanValue::Null),
100        Value::Number(text, _) => {
101            if text.contains(['.', 'e', 'E']) {
102                let parsed = text.parse::<f64>().map_err(|err| {
103                    Error::InvalidArgumentError(format!("invalid float literal: {err}"))
104                })?;
105                Ok(PlanValue::Float(parsed))
106            } else {
107                let parsed = text.parse::<i64>().map_err(|err| {
108                    Error::InvalidArgumentError(format!("invalid integer literal: {err}"))
109                })?;
110                Ok(PlanValue::Integer(parsed))
111            }
112        }
113        Value::Boolean(_) => Err(Error::InvalidArgumentError(
114            "BOOLEAN literals are not supported yet".into(),
115        )),
116        other => {
117            if let Some(text) = other.clone().into_string() {
118                Ok(PlanValue::String(text))
119            } else {
120                Err(Error::InvalidArgumentError(format!(
121                    "unsupported literal: {other:?}"
122                )))
123            }
124        }
125    }
126}
127
128// ============================================================================
129// Range SELECT Support
130// ============================================================================
131
132/// Result of parsing a range() SELECT statement.
133#[derive(Clone)]
134pub struct RangeSelectRows {
135    rows: Vec<Vec<PlanValue>>,
136}
137
138impl RangeSelectRows {
139    /// Convert into the underlying rows vector.
140    pub fn into_rows(self) -> Vec<Vec<PlanValue>> {
141        self.rows
142    }
143}
144
145#[derive(Clone)]
146enum RangeProjection {
147    Column,
148    Literal(PlanValue),
149}
150
151#[derive(Clone)]
152struct RangeSpec {
153    start: i64,
154    #[allow(dead_code)] // Used for validation, computed into row_count
155    end: i64,
156    row_count: usize,
157    column_name_lower: String,
158    table_alias_lower: Option<String>,
159}
160
161impl RangeSpec {
162    fn matches_identifier(&self, ident: &str) -> bool {
163        let lower = ident.to_ascii_lowercase();
164        lower == self.column_name_lower || lower == "range"
165    }
166
167    fn matches_table_alias(&self, ident: &str) -> bool {
168        let lower = ident.to_ascii_lowercase();
169        match &self.table_alias_lower {
170            Some(alias) => lower == *alias,
171            None => lower == "range",
172        }
173    }
174
175    fn matches_object_name(&self, name: &ObjectName) -> bool {
176        if name.0.len() != 1 {
177            return false;
178        }
179        match &name.0[0] {
180            ObjectNamePart::Identifier(ident) => self.matches_table_alias(&ident.value),
181            _ => false,
182        }
183    }
184}
185
186/// Extract rows from a range() SELECT statement.
187///
188/// Parses SELECT statements of the form:
189/// - `SELECT * FROM range(10)`
190/// - `SELECT * FROM range(5, 15)`  
191/// - `SELECT range, range * 2 FROM range(100)`
192///
193/// Returns `None` if the SELECT is not a range() query.
194///
195/// # Errors
196///
197/// Returns an error if:
198/// - The range() function has invalid arguments
199/// - Unsupported SELECT clauses are used (WHERE, JOIN, etc.)
200/// - Column references don't match the range source
201pub fn extract_rows_from_range(select: &Select) -> Result<Option<RangeSelectRows>> {
202    let spec = match parse_range_spec(select)? {
203        Some(spec) => spec,
204        None => return Ok(None),
205    };
206
207    if select.selection.is_some() {
208        return Err(Error::InvalidArgumentError(
209            "WHERE clauses are not supported for range() SELECT statements".into(),
210        ));
211    }
212    if select.having.is_some()
213        || !select.named_window.is_empty()
214        || select.qualify.is_some()
215        || select.distinct.is_some()
216        || select.top.is_some()
217        || select.into.is_some()
218        || select.prewhere.is_some()
219        || !select.lateral_views.is_empty()
220        || select.value_table_mode.is_some()
221        || !group_by_is_empty(&select.group_by)
222    {
223        return Err(Error::InvalidArgumentError(
224            "advanced SELECT clauses are not supported for range() SELECT statements".into(),
225        ));
226    }
227
228    let mut projections: Vec<RangeProjection> = Vec::with_capacity(select.projection.len());
229
230    // If projection is empty, treat it as SELECT * (implicit wildcard)
231    if select.projection.is_empty() {
232        projections.push(RangeProjection::Column);
233    } else {
234        for item in &select.projection {
235            let projection = match item {
236                SelectItem::Wildcard(_) => RangeProjection::Column,
237                SelectItem::QualifiedWildcard(kind, _) => match kind {
238                    SelectItemQualifiedWildcardKind::ObjectName(object_name) => {
239                        if spec.matches_object_name(object_name) {
240                            RangeProjection::Column
241                        } else {
242                            return Err(Error::InvalidArgumentError(
243                                "qualified wildcard must reference the range() source".into(),
244                            ));
245                        }
246                    }
247                    SelectItemQualifiedWildcardKind::Expr(_) => {
248                        return Err(Error::InvalidArgumentError(
249                            "expression-qualified wildcards are not supported for range() SELECT statements".into(),
250                        ));
251                    }
252                },
253                SelectItem::UnnamedExpr(expr) => build_range_projection_expr(expr, &spec)?,
254                SelectItem::ExprWithAlias { expr, .. } => build_range_projection_expr(expr, &spec)?,
255            };
256            projections.push(projection);
257        }
258    }
259
260    let mut rows: Vec<Vec<PlanValue>> = Vec::with_capacity(spec.row_count);
261    for idx in 0..spec.row_count {
262        let mut row: Vec<PlanValue> = Vec::with_capacity(projections.len());
263        let value = spec.start + (idx as i64);
264        for projection in &projections {
265            match projection {
266                RangeProjection::Column => row.push(PlanValue::Integer(value)),
267                RangeProjection::Literal(value) => row.push(value.clone()),
268            }
269        }
270        rows.push(row);
271    }
272
273    Ok(Some(RangeSelectRows { rows }))
274}
275
276fn build_range_projection_expr(expr: &SqlExpr, spec: &RangeSpec) -> Result<RangeProjection> {
277    match expr {
278        SqlExpr::Identifier(ident) => {
279            if spec.matches_identifier(&ident.value) {
280                Ok(RangeProjection::Column)
281            } else {
282                Err(Error::InvalidArgumentError(format!(
283                    "unknown column '{}' in range() SELECT",
284                    ident.value
285                )))
286            }
287        }
288        SqlExpr::CompoundIdentifier(parts) => {
289            if parts.len() == 2
290                && spec.matches_table_alias(&parts[0].value)
291                && spec.matches_identifier(&parts[1].value)
292            {
293                Ok(RangeProjection::Column)
294            } else {
295                Err(Error::InvalidArgumentError(
296                    "compound identifiers must reference the range() source".into(),
297                ))
298            }
299        }
300        SqlExpr::Wildcard(_) | SqlExpr::QualifiedWildcard(_, _) => unreachable!(),
301        other => Ok(RangeProjection::Literal(plan_value_from_sql_expr(other)?)),
302    }
303}
304
305fn parse_range_spec(select: &Select) -> Result<Option<RangeSpec>> {
306    if select.from.len() != 1 {
307        return Ok(None);
308    }
309    let item = &select.from[0];
310    if !item.joins.is_empty() {
311        return Err(Error::InvalidArgumentError(
312            "JOIN clauses are not supported for range() SELECT statements".into(),
313        ));
314    }
315
316    match &item.relation {
317        TableFactor::Function {
318            lateral,
319            name,
320            args,
321            alias,
322        } => {
323            if *lateral {
324                return Err(Error::InvalidArgumentError(
325                    "LATERAL range() is not supported".into(),
326                ));
327            }
328            parse_range_spec_from_args(name, args, alias)
329        }
330        TableFactor::Table {
331            name,
332            alias,
333            args: Some(table_args),
334            with_ordinality,
335            ..
336        } => {
337            if *with_ordinality {
338                return Err(Error::InvalidArgumentError(
339                    "WITH ORDINALITY is not supported for range()".into(),
340                ));
341            }
342            if table_args.settings.is_some() {
343                return Err(Error::InvalidArgumentError(
344                    "range() SETTINGS clause is not supported".into(),
345                ));
346            }
347            parse_range_spec_from_args(name, &table_args.args, alias)
348        }
349        _ => Ok(None),
350    }
351}
352
353fn parse_range_spec_from_args(
354    name: &ObjectName,
355    args: &[FunctionArg],
356    alias: &Option<TableAlias>,
357) -> Result<Option<RangeSpec>> {
358    if name.0.len() != 1 {
359        return Ok(None);
360    }
361    let func_name = match &name.0[0] {
362        ObjectNamePart::Identifier(ident) => ident.value.to_ascii_lowercase(),
363        _ => return Ok(None),
364    };
365    if func_name != "range" {
366        return Ok(None);
367    }
368
369    if args.is_empty() || args.len() > 2 {
370        return Err(Error::InvalidArgumentError(
371            "range() requires one or two arguments".into(),
372        ));
373    }
374
375    // Helper to extract integer from argument
376    let extract_int = |arg: &FunctionArg| -> Result<i64> {
377        let arg_expr = match arg {
378            FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => expr,
379            FunctionArg::Unnamed(FunctionArgExpr::QualifiedWildcard(_))
380            | FunctionArg::Unnamed(FunctionArgExpr::Wildcard) => {
381                return Err(Error::InvalidArgumentError(
382                    "range() argument must be an integer literal".into(),
383                ));
384            }
385            FunctionArg::Named { .. } | FunctionArg::ExprNamed { .. } => {
386                return Err(Error::InvalidArgumentError(
387                    "named arguments are not supported for range()".into(),
388                ));
389            }
390        };
391
392        let value = plan_value_from_sql_expr(arg_expr)?;
393        match value {
394            PlanValue::Integer(v) => Ok(v),
395            _ => Err(Error::InvalidArgumentError(
396                "range() argument must be an integer literal".into(),
397            )),
398        }
399    };
400
401    let (start, end, row_count) = if args.len() == 1 {
402        // range(count) - generate [0, count)
403        let count = extract_int(&args[0])?;
404        if count < 0 {
405            return Err(Error::InvalidArgumentError(
406                "range() argument must be non-negative".into(),
407            ));
408        }
409        (0, count, count as usize)
410    } else {
411        // range(start, end) - generate [start, end)
412        let start = extract_int(&args[0])?;
413        let end = extract_int(&args[1])?;
414        if end < start {
415            return Err(Error::InvalidArgumentError(
416                "range() end must be >= start".into(),
417            ));
418        }
419        let row_count = (end - start) as usize;
420        (start, end, row_count)
421    };
422
423    let column_name_lower = alias
424        .as_ref()
425        .and_then(|a| {
426            a.columns
427                .first()
428                .map(|col| col.name.value.to_ascii_lowercase())
429        })
430        .unwrap_or_else(|| "range".to_string());
431    let table_alias_lower = alias.as_ref().map(|a| a.name.value.to_ascii_lowercase());
432
433    Ok(Some(RangeSpec {
434        start,
435        end,
436        row_count,
437        column_name_lower,
438        table_alias_lower,
439    }))
440}
441
442fn group_by_is_empty(expr: &GroupByExpr) -> bool {
443    matches!(
444        expr,
445        GroupByExpr::Expressions(exprs, modifiers)
446            if exprs.is_empty() && modifiers.is_empty()
447    )
448}
449
450#[cfg(test)]
451mod tests {
452    use super::*;
453    use sqlparser::ast::{Expr as SqlExpr, Value, ValueWithSpan};
454
455    fn value_with_span(v: Value) -> ValueWithSpan {
456        ValueWithSpan {
457            value: v,
458            span: sqlparser::tokenizer::Span::empty(),
459        }
460    }
461
462    #[test]
463    fn test_null_value() {
464        let value = value_with_span(Value::Null);
465        assert_eq!(plan_value_from_sql_value(&value).unwrap(), PlanValue::Null);
466    }
467
468    #[test]
469    fn test_integer_value() {
470        let value = value_with_span(Value::Number("42".to_string(), false));
471        assert_eq!(
472            plan_value_from_sql_value(&value).unwrap(),
473            PlanValue::Integer(42)
474        );
475    }
476
477    #[test]
478    fn test_negative_integer() {
479        let value = value_with_span(Value::Number("42".to_string(), false));
480        let expr = SqlExpr::UnaryOp {
481            op: UnaryOperator::Minus,
482            expr: Box::new(SqlExpr::Value(value)),
483        };
484        assert_eq!(
485            plan_value_from_sql_expr(&expr).unwrap(),
486            PlanValue::Integer(-42)
487        );
488    }
489
490    #[test]
491    #[allow(clippy::approx_constant)]
492    fn test_float_value() {
493        let value = value_with_span(Value::Number("3.14".to_string(), false));
494        assert_eq!(
495            plan_value_from_sql_value(&value).unwrap(),
496            PlanValue::Float(3.14)
497        );
498    }
499
500    #[test]
501    fn test_string_value() {
502        let value = value_with_span(Value::SingleQuotedString("hello".to_string()));
503        assert_eq!(
504            plan_value_from_sql_value(&value).unwrap(),
505            PlanValue::String("hello".to_string())
506        );
507    }
508
509    #[test]
510    fn test_nested_expression() {
511        let value = value_with_span(Value::Number("100".to_string(), false));
512        let expr = SqlExpr::Nested(Box::new(SqlExpr::Value(value)));
513        assert_eq!(
514            plan_value_from_sql_expr(&expr).unwrap(),
515            PlanValue::Integer(100)
516        );
517    }
518
519    #[test]
520    fn test_plus_operator() {
521        let value = value_with_span(Value::Number("50".to_string(), false));
522        let expr = SqlExpr::UnaryOp {
523            op: UnaryOperator::Plus,
524            expr: Box::new(SqlExpr::Value(value)),
525        };
526        assert_eq!(
527            plan_value_from_sql_expr(&expr).unwrap(),
528            PlanValue::Integer(50)
529        );
530    }
531
532    #[test]
533    fn test_cannot_negate_string() {
534        let value = value_with_span(Value::SingleQuotedString("test".to_string()));
535        let expr = SqlExpr::UnaryOp {
536            op: UnaryOperator::Minus,
537            expr: Box::new(SqlExpr::Value(value)),
538        };
539        assert!(plan_value_from_sql_expr(&expr).is_err());
540    }
541
542    #[test]
543    fn test_boolean_not_supported() {
544        let value = value_with_span(Value::Boolean(true));
545        assert!(plan_value_from_sql_value(&value).is_err());
546    }
547}