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

llkv_plan/
conversion.rs
