From b5d256cdae0a37b1a4cd0500fff43eccfc9ea829 Mon Sep 17 00:00:00 2001 From: Alan Cai Date: Wed, 5 Jul 2023 15:30:35 -0700 Subject: [PATCH] Implements path wildcard and path unpivot expressions --- CHANGELOG.md | 2 + partiql-eval/src/eval/expr/mod.rs | 141 ++++++++++++++++++++------- partiql-eval/src/plan.rs | 2 + partiql-logical-planner/src/lower.rs | 26 +++-- partiql-logical/src/lib.rs | 4 + 5 files changed, 135 insertions(+), 40 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 834ae525..fb61d242 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/partiql-eval/src/eval/expr/mod.rs b/partiql-eval/src/eval/expr/mod.rs index ac3f82f8..55c05f88 100644 --- a/partiql-eval/src/eval/expr/mod.rs +++ b/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), Index(i64), IndexExpr(Box), + 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, 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, + 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::>(); + Value::from(Bag::from(values)) + } else { + // no other path wildcards + let values = value + .into_iter() + .map(|v| path(v, paths.clone(), bindings, ctx)) + .collect::>(); + 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::>(); + 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::>(); + 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, + )) } } diff --git a/partiql-eval/src/plan.rs b/partiql-eval/src/plan.rs index b04eda22..64d8a44d 100644 --- a/partiql-eval/src/plan.rs +++ b/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(), }), diff --git a/partiql-logical-planner/src/lower.rs b/partiql-logical-planner/src/lower.rs index 1094a8cd..ef675b38 100644 --- a/partiql-logical-planner/src/lower.rs +++ b/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); diff --git a/partiql-logical/src/lib.rs b/partiql-logical/src/lib.rs index c6d4a2ea..6e7865ba 100644 --- a/partiql-logical/src/lib.rs +++ b/partiql-logical/src/lib.rs @@ -457,6 +457,10 @@ pub enum PathComponent { Index(i64), KeyExpr(Box), IndexExpr(Box), + /// 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}`.