Translate collection literals to Value when lowering (#380)

Author: jpschorr

CHANGELOG.md

@@ -8,8 +8,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 ### Changed
 - *BREAKING:* partiql-eval: `evaluate` on `Evaluable` returns a `Value` rather than an `Option<Value>`
+- *BREAKING:* partiql-ast: changes the modeling of Bag/List/Tuple literals
 ### Added
 - Ability to add and view errors during evaluation with partiql-eval's `EvalContext`
+- AST sub-trees representing literal values are lowered to `Value`s during planning
 ### Fixes
 
 ## [0.4.1] - 2023-05-25

partiql-ast/src/ast.rs

@@ -444,19 +444,16 @@ pub enum Lit {
     #[visit(skip)]
     HexStringLit(String),
     #[visit(skip)]
-    CollectionLit(CollectionLit),
+    StructLit(AstNode<Struct>),
+    #[visit(skip)]
+    BagLit(AstNode<Bag>),
+    #[visit(skip)]
+    ListLit(AstNode<List>),
     /// E.g. `TIME WITH TIME ZONE` in `SELECT TIME WITH TIME ZONE '12:00' FROM ...`
     #[visit(skip)]
     TypedLit(String, Type),
 }
 
-#[derive(Clone, Debug, PartialEq, Eq)]
-#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
-pub enum CollectionLit {
-    ArrayLit(String),
-    BagLit(String),
-}
-
 #[derive(Visit, Clone, Debug, PartialEq, Eq)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 pub struct VarRef {

partiql-logical-planner/src/lower.rs

@@ -6,8 +6,8 @@ use partiql_ast::ast;
 use partiql_ast::ast::{
     Assignment, Bag, Between, BinOp, BinOpKind, Call, CallAgg, CallArg, CallArgNamed,
     CaseSensitivity, CreateIndex, CreateTable, Ddl, DdlOp, Delete, Dml, DmlOp, DropIndex,
-    DropTable, FromClause, FromLet, FromLetKind, GroupByExpr, GroupKey, GroupingStrategy, Insert,
-    InsertValue, Item, Join, JoinKind, JoinSpec, Like, List, Lit, NodeId, NullOrderingSpec,
+    DropTable, Expr, FromClause, FromLet, FromLetKind, GroupByExpr, GroupKey, GroupingStrategy,
+    Insert, InsertValue, Item, Join, JoinKind, JoinSpec, Like, List, Lit, NodeId, NullOrderingSpec,
     OnConflict, OrderByExpr, OrderingSpec, Path, PathStep, ProjectExpr, Projection, ProjectionKind,
     Query, QuerySet, Remove, SearchedCase, Select, Set, SetExpr, SetQuantifier, Sexp, SimpleCase,
     SortSpec, Struct, SymbolPrimitive, UniOp, UniOpKind, VarRef,
@@ -979,48 +979,12 @@ impl<'a, 'ast> Visitor<'ast> for AstToLogical<'a> {
 
     // Values & Value Constructors
 
-    fn enter_lit(&mut self, _lit: &'ast Lit) -> Traverse {
-        let val = match _lit {
-            Lit::Null => Value::Null,
-            Lit::Missing => Value::Missing,
-            Lit::Int8Lit(n) => Value::Integer(*n as i64),
-            Lit::Int16Lit(n) => Value::Integer(*n as i64),
-            Lit::Int32Lit(n) => Value::Integer(*n as i64),
-            Lit::Int64Lit(n) => Value::Integer(*n),
-            Lit::DecimalLit(d) => Value::Decimal(*d),
-            Lit::NumericLit(n) => Value::Decimal(*n),
-            Lit::RealLit(f) => Value::Real(OrderedFloat::from(*f as f64)),
-            Lit::FloatLit(f) => Value::Real(OrderedFloat::from(*f as f64)),
-            Lit::DoubleLit(f) => Value::Real(OrderedFloat::from(*f)),
-            Lit::BoolLit(b) => Value::Boolean(*b),
-            Lit::IonStringLit(s) => match parse_embedded_ion_str(s) {
-                Ok(v) => v,
-                Err(e) => {
-                    // Report error but allow visitor to continue
-                    self.errors.push(e);
-                    Value::Missing
-                }
-            },
-            Lit::CharStringLit(s) => Value::String(Box::new(s.clone())),
-            Lit::NationalCharStringLit(s) => Value::String(Box::new(s.clone())),
-            Lit::BitStringLit(_) => {
-                // Report error but allow visitor to continue
-                not_yet_implemented_err!(self, "Lit::BitStringLit".to_string());
-                Value::Missing
-            }
-            Lit::HexStringLit(_) => {
-                // Report error but allow visitor to continue
-                not_yet_implemented_err!(self, "Lit::HexStringLit".to_string());
-                Value::Missing
-            }
-            Lit::CollectionLit(_) => {
+    fn enter_lit(&mut self, lit: &'ast Lit) -> Traverse {
+        let val = match lit_to_value(lit) {
+            Ok(v) => v,
+            Err(e) => {
                 // Report error but allow visitor to continue
-                not_yet_implemented_err!(self, "Lit::CollectionLit".to_string());
-                Value::Missing
-            }
-            Lit::TypedLit(_, _) => {
-                // Report error but allow visitor to continue
-                not_yet_implemented_err!(self, "Lit::TypedLit".to_string());
+                self.errors.push(e);
                 Value::Missing
             }
         };
@@ -1714,6 +1678,90 @@ impl<'a, 'ast> Visitor<'ast> for AstToLogical<'a> {
     }
 }
 
+fn lit_to_value(lit: &Lit) -> Result<Value, LowerError> {
+    fn expect_lit(v: &Expr) -> Result<Value, LowerError> {
+        match v {
+            Expr::Lit(l) => lit_to_value(&l.node),
+            _ => Err(LowerError::IllegalState(
+                "non literal in literal aggregate".to_string(),
+            )),
+        }
+    }
+
+    fn tuple_pair(pair: &ast::ExprPair) -> Option<Result<(String, Value), LowerError>> {
+        let key = match expect_lit(pair.first.as_ref()) {
+            Ok(Value::String(s)) => s.as_ref().clone(),
+            Ok(_) => {
+                return Some(Err(LowerError::IllegalState(
+                    "non string literal in literal struct key".to_string(),
+                )))
+            }
+            Err(e) => return Some(Err(e)),
+        };
+
+        match expect_lit(pair.second.as_ref()) {
+            Ok(Value::Missing) => None,
+            Ok(val) => Some(Ok((key, val))),
+            Err(e) => Some(Err(e)),
+        }
+    }
+
+    let val = match lit {
+        Lit::Null => Value::Null,
+        Lit::Missing => Value::Missing,
+        Lit::Int8Lit(n) => Value::Integer(*n as i64),
+        Lit::Int16Lit(n) => Value::Integer(*n as i64),
+        Lit::Int32Lit(n) => Value::Integer(*n as i64),
+        Lit::Int64Lit(n) => Value::Integer(*n),
+        Lit::DecimalLit(d) => Value::Decimal(*d),
+        Lit::NumericLit(n) => Value::Decimal(*n),
+        Lit::RealLit(f) => Value::Real(OrderedFloat::from(*f as f64)),
+        Lit::FloatLit(f) => Value::Real(OrderedFloat::from(*f as f64)),
+        Lit::DoubleLit(f) => Value::Real(OrderedFloat::from(*f)),
+        Lit::BoolLit(b) => Value::Boolean(*b),
+        Lit::IonStringLit(s) => parse_embedded_ion_str(s)?,
+        Lit::CharStringLit(s) => Value::String(Box::new(s.clone())),
+        Lit::NationalCharStringLit(s) => Value::String(Box::new(s.clone())),
+        Lit::BitStringLit(_) => {
+            return Err(LowerError::NotYetImplemented(
+                "Lit::BitStringLit".to_string(),
+            ))
+        }
+        Lit::HexStringLit(_) => {
+            return Err(LowerError::NotYetImplemented(
+                "Lit::HexStringLit".to_string(),
+            ))
+        }
+        Lit::BagLit(b) => {
+            let bag: Result<partiql_value::Bag, _> = b
+                .node
+                .values
+                .iter()
+                .map(|l| expect_lit(l.as_ref()))
+                .collect();
+            Value::from(bag?)
+        }
+        Lit::ListLit(l) => {
+            let l: Result<partiql_value::List, _> = l
+                .node
+                .values
+                .iter()
+                .map(|l| expect_lit(l.as_ref()))
+                .collect();
+            Value::from(l?)
+        }
+        Lit::StructLit(s) => {
+            let tuple: Result<partiql_value::Tuple, _> =
+                s.node.fields.iter().filter_map(tuple_pair).collect();
+            Value::from(tuple?)
+        }
+        Lit::TypedLit(_, _) => {
+            return Err(LowerError::NotYetImplemented("Lit::TypedLit".to_string()))
+        }
+    };
+    Ok(val)
+}
+
 fn parse_embedded_ion_str(contents: &str) -> Result<Value, LowerError> {
     fn lit_err(literal: &str, err: impl std::error::Error) -> LowerError {
         LowerError::Literal {

partiql-parser/Cargo.toml

@@ -35,6 +35,8 @@ num-bigint = "~0.4.0"
 bigdecimal = "~0.2.0"
 rust_decimal = { version = "1.25.0", default-features = false, features = ["std"] }
 
+bitflags = "2"
+
 lalrpop-util = "0.20"
 logos = "0.12"
 

partiql-parser/src/parse/parse_util.rs

@@ -1,10 +1,64 @@
 use partiql_ast::ast;
 
+use bitflags::bitflags;
+
+bitflags! {
+    /// Set of AST node attributes to use as synthesized attributes.
+    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+    pub(crate) struct Attrs: u8 {
+        const LIT = 0b00000001;
+
+        const INTERSECTABLE = Self::LIT.bits();
+        const UNIONABLE = 0;
+    }
+}
+
+impl Attrs {
+    /// Combine attributes from two nodes.
+    #[inline]
+    pub fn synthesize(self, other: Self) -> Attrs {
+        ((self & Attrs::INTERSECTABLE) & (other & Attrs::INTERSECTABLE))
+            | ((self & Attrs::UNIONABLE) | (other & Attrs::UNIONABLE))
+    }
+}
+
+/// Wrapper attaching synthesized attributes `Attrs` with an AST node.
+pub(crate) struct Synth<T> {
+    pub(crate) data: T,
+    pub(crate) attrs: Attrs,
+}
+
+impl<T> Synth<T> {
+    #[inline]
+    pub fn new(data: T, attrs: Attrs) -> Self {
+        Synth { data, attrs }
+    }
+
+    #[inline]
+    pub fn empty(data: T) -> Self {
+        Self::new(data, Attrs::empty())
+    }
+}
+
+impl<T> FromIterator<Synth<T>> for Synth<Vec<T>> {
+    #[inline]
+    fn from_iter<I: IntoIterator<Item = Synth<T>>>(iter: I) -> Synth<Vec<T>> {
+        let mut attrs = Attrs::all();
+        let iterator = iter.into_iter().map(|Synth { data, attrs: a }| {
+            attrs = attrs.synthesize(a);
+            data
+        });
+        let data = iterator.collect::<Vec<_>>();
+        Synth { data, attrs }
+    }
+}
+
 pub(crate) enum CallSite {
     Call(ast::Call),
     CallAgg(ast::CallAgg),
 }
 
+#[inline]
 // if this is just a parenthesized expr, lift it out of the query AST, otherwise return input
 //      e.g. `(1+2)` should be a ExprKind::Expr, not wrapped deep in a ExprKind::Query
 pub(crate) fn strip_query(q: Box<ast::Expr>) -> Box<ast::Expr> {