Implements path wildcard and path unpivot expressions

Author: alancai98

CHANGELOG.md

@@ -17,6 +17,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add `partiql-types` crate that includes data models for PartiQL Types.
 - Add `partiql_ast_passes::static_typer` for type annotating the AST.
 - Add ability to parse `ORDER BY`, `LIMIT`, `OFFSET` in children of set operators
+- Adds evaluation of path wildcard (e.g. `foo[*].a[*].b`) and path unpivot expressions (e.g. `bar.*.c.*.d`)
+- Adds AST to logical plan lowering for `IN` expressions
 
 ### Fixes
 - Fixes parsing of multiple consecutive path wildcards (e.g. `a[*][*][*]`), unpivot (e.g. `a.*.*.*`), and path expressions (e.g. `a[1 + 2][3 + 4][5 + 6]`)—previously these would not parse correctly.

partiql-eval/src/eval/expr/mod.rs

@@ -16,6 +16,7 @@ use rust_decimal::prelude::FromPrimitive;
 use rust_decimal::Decimal;
 use std::borrow::{Borrow, Cow};
 use std::fmt::Debug;
+use std::iter::Peekable;
 
 pub(crate) mod pattern_match;
 
@@ -107,54 +108,126 @@ pub(crate) enum EvalPathComponent {
     KeyExpr(Box<dyn EvalExpr>),
     Index(i64),
     IndexExpr(Box<dyn EvalExpr>),
+    PathWildcard,
+    PathUnpivot,
 }
 
 impl EvalExpr for EvalPath {
     fn evaluate<'a>(&'a self, bindings: &'a Tuple, ctx: &'a dyn EvalContext) -> Cow<'a, Value> {
         #[inline]
-        fn path_into<'a>(
-            value: &'a Value,
-            path: &EvalPathComponent,
+        fn path<'a, I>(
+            v: Value,
+            mut paths: Peekable<I>,
             bindings: &'a Tuple,
             ctx: &dyn EvalContext,
-        ) -> Option<&'a Value> {
-            match path {
-                EvalPathComponent::Key(k) => match value {
-                    Value::Tuple(tuple) => tuple.get(k),
-                    _ => None,
-                },
-                EvalPathComponent::Index(idx) => match value {
-                    Value::List(list) if (*idx as usize) < list.len() => list.get(*idx),
-                    _ => None,
-                },
-                EvalPathComponent::KeyExpr(ke) => {
-                    let key = ke.evaluate(bindings, ctx);
-                    match (value, key.as_ref()) {
-                        (Value::Tuple(tuple), Value::String(key)) => {
-                            tuple.get(&BindingsName::CaseInsensitive(key.as_ref().clone()))
+        ) -> Value
+        where
+            I: Iterator<Item = &'a EvalPathComponent>,
+            I: Clone,
+        {
+            let mut value = v;
+            while let Some(p) = paths.next() {
+                match p {
+                    EvalPathComponent::Key(k) => {
+                        value = match value {
+                            Value::Tuple(tuple) => tuple.get(k).unwrap_or_else(|| &Missing).clone(),
+                            _ => Missing,
                         }
-                        _ => None,
                     }
-                }
-                EvalPathComponent::IndexExpr(ie) => {
-                    if let Value::Integer(idx) = ie.evaluate(bindings, ctx).as_ref() {
-                        match value {
-                            Value::List(list) if (*idx as usize) < list.len() => list.get(*idx),
-                            _ => None,
+                    EvalPathComponent::Index(idx) => {
+                        value = match &value {
+                            Value::List(list) if (*idx as usize) < list.len() => {
+                                list.get(*idx).unwrap_or_else(|| &Missing).clone()
+                            }
+                            _ => Missing,
+                        }
+                    }
+                    EvalPathComponent::KeyExpr(ke) => {
+                        let key = ke.evaluate(bindings, ctx);
+                        value = match (value, key.as_ref()) {
+                            (Value::Tuple(tuple), Value::String(key)) => tuple
+                                .get(&BindingsName::CaseInsensitive(key.as_ref().clone()))
+                                .unwrap_or_else(|| &Missing)
+                                .clone(),
+                            _ => Missing,
+                        }
+                    }
+                    EvalPathComponent::IndexExpr(ie) => {
+                        value = if let Value::Integer(idx) = ie.evaluate(bindings, ctx).as_ref() {
+                            match &value {
+                                Value::List(list) if (*idx as usize) < list.len() => {
+                                    list.get(*idx).unwrap_or_else(|| &Missing).clone()
+                                }
+                                _ => Missing,
+                            }
+                        } else {
+                            Missing
                         }
-                    } else {
-                        None
+                    }
+                    EvalPathComponent::PathWildcard => {
+                        return match paths.peek().is_some() {
+                            true => {
+                                // iterator is not empty
+                                let other_wildcards_present = paths
+                                    .clone()
+                                    .any(|_p| matches!(EvalPathComponent::PathWildcard, _p));
+                                if other_wildcards_present {
+                                    // other path wildcards so flatten
+                                    let values = value
+                                        .into_iter()
+                                        .flat_map(|v| path(v, paths.clone(), bindings, ctx))
+                                        .collect::<Vec<Value>>();
+                                    Value::from(Bag::from(values))
+                                } else {
+                                    // no other path wildcards
+                                    let values = value
+                                        .into_iter()
+                                        .map(|v| path(v, paths.clone(), bindings, ctx))
+                                        .collect::<Vec<Value>>();
+                                    Value::from(Bag::from(values))
+                                }
+                            }
+                            false => {
+                                // iterator is empty; path wildcard is last component
+                                Value::from(Bag::from_iter(value.into_iter()))
+                            }
+                        };
+                    }
+                    EvalPathComponent::PathUnpivot => {
+                        return match paths.peek().is_some() {
+                            true => {
+                                // iterator is not empty
+                                let values = value
+                                    .coerce_to_tuple()
+                                    .into_values()
+                                    .flat_map(|v| path(v, paths.clone(), bindings, ctx))
+                                    .collect::<Vec<Value>>();
+                                Value::from(Bag::from(values))
+                            }
+                            false =>
+                            // iterator is empty; path unpivot is last component
+                            {
+                                match value {
+                                    Value::Tuple(tuple) => {
+                                        let values = tuple.into_values().collect::<Vec<Value>>();
+                                        Value::from(Bag::from(values))
+                                    }
+                                    non_tuple => Value::from(Value::coerce_to_bag(non_tuple)),
+                                }
+                            }
+                        };
                     }
                 }
             }
+            value
         }
-        let value = self.expr.evaluate(bindings, ctx);
-        self.components
-            .iter()
-            .fold(Some(value.as_ref()), |v, path| {
-                v.and_then(|v| path_into(v, path, bindings, ctx))
-            })
-            .map_or_else(|| Cow::Owned(Value::Missing), |v| Cow::Owned(v.clone()))
+        let value = self.expr.evaluate(bindings, ctx).into_owned();
+        Cow::Owned(path(
+            value,
+            self.components.iter().peekable(),
+            bindings,
+            ctx,
+        ))
     }
 }
 

partiql-eval/src/plan.rs

@@ -340,6 +340,8 @@ impl<'c> EvaluatorPlanner<'c> {
                         PathComponent::IndexExpr(i) => {
                             eval::expr::EvalPathComponent::IndexExpr(self.plan_values(i))
                         }
+                        PathComponent::PathWildcard => eval::expr::EvalPathComponent::PathWildcard,
+                        PathComponent::PathUnpivot => eval::expr::EvalPathComponent::PathUnpivot,
                     })
                     .collect(),
             }),

partiql-logical-planner/src/lower.rs

@@ -867,6 +867,24 @@ impl<'a, 'ast> Visitor<'ast> for AstToLogical<'a> {
         Traverse::Continue
     }
 
+    fn enter_in(&mut self, _in: &'ast ast::In) -> Traverse {
+        self.enter_env();
+        Traverse::Continue
+    }
+    fn exit_in(&mut self, _in: &'ast ast::In) -> Traverse {
+        let mut env = self.exit_env();
+        eq_or_fault!(self, env.len(), 2, "env.len() != 2");
+
+        let rhs = env.pop().unwrap();
+        let lhs = env.pop().unwrap();
+        self.push_vexpr(logical::ValueExpr::BinaryExpr(
+            logical::BinaryOp::In,
+            Box::new(lhs),
+            Box::new(rhs),
+        ));
+        Traverse::Continue
+    }
+
     fn enter_like(&mut self, _like: &'ast Like) -> Traverse {
         self.enter_env();
         Traverse::Continue
@@ -1223,12 +1241,8 @@ impl<'a, 'ast> Visitor<'ast> for AstToLogical<'a> {
                     }
                 }
             }
-            PathStep::PathWildCard => {
-                not_yet_implemented_fault!(self, "PathStep::PathWildCard".to_string());
-            }
-            PathStep::PathUnpivot => {
-                not_yet_implemented_fault!(self, "PathStep::PathUnpivot".to_string());
-            }
+            PathStep::PathWildCard => logical::PathComponent::PathWildcard,
+            PathStep::PathUnpivot => logical::PathComponent::PathUnpivot,
         };
 
         self.push_path_step(step);

partiql-logical/src/lib.rs

@@ -457,6 +457,10 @@ pub enum PathComponent {
     Index(i64),
     KeyExpr(Box<ValueExpr>),
     IndexExpr(Box<ValueExpr>),
+    /// E.g. `[*]` in `foo[*].a[*].b`
+    PathWildcard,
+    /// E.g. `.*` in `bar.*.c.*.d`
+    PathUnpivot,
 }
 
 /// Represents a PartiQL tuple expression, e.g: `{ a.b: a.c * 2, 'count': a.c + 10}`.