Skip to content

Commit ea0dd66

Browse files
alancai98alancai98
committed
Adds implementation of path wildcard and path unpivot expressions; also adds IN predicate lowering
1 parent 5629e88 commit ea0dd66

File tree

6 files changed

+136
-41
lines changed

6 files changed

+136
-41
lines changed

‎CHANGELOG.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
1717
- Add `partiql-types` crate that includes data models for PartiQL Types.
1818
- Add `partiql_ast_passes::static_typer` for type annotating the AST.
1919
- Add ability to parse `ORDER BY`, `LIMIT`, `OFFSET` in children of set operators
20+
- Adds evaluation of path wildcard (e.g. `foo[*].a[*].b`) and path unpivot expressions (e.g. `bar.*.c.*.d`)
21+
- Adds AST to logical plan lowering for `IN` expressions
2022

2123
### Fixes
2224

‎partiql-conformance-tests/partiql-tests‎

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

Lines changed: 107 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ use rust_decimal::prelude::FromPrimitive;
1616
use rust_decimal::Decimal;
1717
use std::borrow::{Borrow, Cow};
1818
use std::fmt::Debug;
19+
use std::iter::Peekable;
1920

2021
pub(crate) mod pattern_match;
2122

@@ -107,54 +108,126 @@ pub(crate) enum EvalPathComponent {
107108
KeyExpr(Box<dyn EvalExpr>),
108109
Index(i64),
109110
IndexExpr(Box<dyn EvalExpr>),
111+
PathWildcard,
112+
PathUnpivot,
110113
}
111114

112115
impl EvalExpr for EvalPath {
113116
fn evaluate<'a>(&'a self, bindings: &'a Tuple, ctx: &'a dyn EvalContext) -> Cow<'a, Value> {
114117
#[inline]
115-
fn path_into<'a>(
116-
value: &'a Value,
117-
path: &EvalPathComponent,
118+
fn path<'a, I>(
119+
v: Value,
120+
mut paths: Peekable<I>,
118121
bindings: &'a Tuple,
119122
ctx: &dyn EvalContext,
120-
) -> Option<&'a Value> {
121-
match path {
122-
EvalPathComponent::Key(k) => match value {
123-
Value::Tuple(tuple) => tuple.get(k),
124-
_ => None,
125-
},
126-
EvalPathComponent::Index(idx) => match value {
127-
Value::List(list) if (*idx as usize) < list.len() => list.get(*idx),
128-
_ => None,
129-
},
130-
EvalPathComponent::KeyExpr(ke) => {
131-
let key = ke.evaluate(bindings, ctx);
132-
match (value, key.as_ref()) {
133-
(Value::Tuple(tuple), Value::String(key)) => {
134-
tuple.get(&BindingsName::CaseInsensitive(key.as_ref().clone()))
123+
) -> Value
124+
where
125+
I: Iterator<Item = &'a EvalPathComponent>,
126+
I: Clone,
127+
{
128+
let mut value = v;
129+
while let Some(p) = paths.next() {
130+
match p {
131+
EvalPathComponent::Key(k) => {
132+
value = match value {
133+
Value::Tuple(tuple) => tuple.get(k).unwrap_or_else(|| &Missing).clone(),
134+
_ => Missing,
135135
}
136-
_ => None,
137136
}
138-
}
139-
EvalPathComponent::IndexExpr(ie) => {
140-
if let Value::Integer(idx) = ie.evaluate(bindings, ctx).as_ref() {
141-
match value {
142-
Value::List(list) if (*idx as usize) < list.len() => list.get(*idx),
143-
_ => None,
137+
EvalPathComponent::Index(idx) => {
138+
value = match &value {
139+
Value::List(list) if (*idx as usize) < list.len() => {
140+
list.get(*idx).unwrap_or_else(|| &Missing).clone()
141+
}
142+
_ => Missing,
143+
}
144+
}
145+
EvalPathComponent::KeyExpr(ke) => {
146+
let key = ke.evaluate(bindings, ctx);
147+
value = match (value, key.as_ref()) {
148+
(Value::Tuple(tuple), Value::String(key)) => tuple
149+
.get(&BindingsName::CaseInsensitive(key.as_ref().clone()))
150+
.unwrap_or_else(|| &Missing)
151+
.clone(),
152+
_ => Missing,
144153
}
145-
} else {
146-
None
154+
}
155+
EvalPathComponent::IndexExpr(ie) => {
156+
if let Value::Integer(idx) = ie.evaluate(bindings, ctx).as_ref() {
157+
match &value {
158+
Value::List(list) if (*idx as usize) < list.len() => {
159+
list.get(*idx).unwrap_or_else(|| &Missing).clone()
160+
}
161+
_ => Missing,
162+
}
163+
} else {
164+
Missing
165+
};
166+
}
167+
EvalPathComponent::PathWildcard => {
168+
return match paths.peek().is_some() {
169+
true => {
170+
// iterator is not empty
171+
let other_wildcards_present = paths
172+
.clone()
173+
.any(|_p| matches!(EvalPathComponent::PathWildcard, _p));
174+
if other_wildcards_present {
175+
// other path wildcards so flatten
176+
let values = value
177+
.into_iter()
178+
.flat_map(|v| path(v, paths.clone(), bindings, ctx))
179+
.collect::<Vec<Value>>();
180+
Value::from(Bag::from(values))
181+
} else {
182+
// no other path wildcards
183+
let values = value
184+
.into_iter()
185+
.map(|v| path(v, paths.clone(), bindings, ctx))
186+
.collect::<Vec<Value>>();
187+
Value::from(Bag::from(values))
188+
}
189+
}
190+
false => {
191+
// iterator is empty; path wildcard is last component
192+
Value::from(Bag::from_iter(value.into_iter()))
193+
}
194+
};
195+
}
196+
EvalPathComponent::PathUnpivot => {
197+
return match paths.peek().is_some() {
198+
true => {
199+
// iterator is not empty
200+
let values = value
201+
.coerce_to_tuple()
202+
.into_values()
203+
.flat_map(|v| path(v, paths.clone(), bindings, ctx))
204+
.collect::<Vec<Value>>();
205+
Value::from(Bag::from(values))
206+
}
207+
false =>
208+
// iterator is empty; path unpivot is last component
209+
{
210+
match value {
211+
Value::Tuple(tuple) => {
212+
let values = tuple.into_values().collect::<Vec<Value>>();
213+
Value::from(Bag::from(values))
214+
}
215+
non_tuple => Value::from(Value::coerce_to_bag(non_tuple)),
216+
}
217+
}
218+
};
147219
}
148220
}
149221
}
222+
value
150223
}
151-
let value = self.expr.evaluate(bindings, ctx);
152-
self.components
153-
.iter()
154-
.fold(Some(value.as_ref()), |v, path| {
155-
v.and_then(|v| path_into(v, path, bindings, ctx))
156-
})
157-
.map_or_else(|| Cow::Owned(Value::Missing), |v| Cow::Owned(v.clone()))
224+
let value = self.expr.evaluate(bindings, ctx).into_owned();
225+
Cow::Owned(path(
226+
value,
227+
self.components.iter().peekable(),
228+
bindings,
229+
ctx,
230+
))
158231
}
159232
}
160233

‎partiql-eval/src/plan.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -340,6 +340,8 @@ impl<'c> EvaluatorPlanner<'c> {
340340
PathComponent::IndexExpr(i) => {
341341
eval::expr::EvalPathComponent::IndexExpr(self.plan_values(i))
342342
}
343+
PathComponent::PathWildcard => eval::expr::EvalPathComponent::PathWildcard,
344+
PathComponent::PathUnpivot => eval::expr::EvalPathComponent::PathUnpivot,
343345
})
344346
.collect(),
345347
}),

‎partiql-logical-planner/src/lower.rs‎

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -867,6 +867,24 @@ impl<'a, 'ast> Visitor<'ast> for AstToLogical<'a> {
867867
Traverse::Continue
868868
}
869869

870+
fn enter_in(&mut self, _in: &'ast ast::In) -> Traverse {
871+
self.enter_env();
872+
Traverse::Continue
873+
}
874+
fn exit_in(&mut self, _in: &'ast ast::In) -> Traverse {
875+
let mut env = self.exit_env();
876+
eq_or_fault!(self, env.len(), 2, "env.len() != 2");
877+
878+
let rhs = env.pop().unwrap();
879+
let lhs = env.pop().unwrap();
880+
self.push_vexpr(logical::ValueExpr::BinaryExpr(
881+
logical::BinaryOp::In,
882+
Box::new(lhs),
883+
Box::new(rhs),
884+
));
885+
Traverse::Continue
886+
}
887+
870888
fn enter_like(&mut self, _like: &'ast Like) -> Traverse {
871889
self.enter_env();
872890
Traverse::Continue
@@ -1223,12 +1241,8 @@ impl<'a, 'ast> Visitor<'ast> for AstToLogical<'a> {
12231241
}
12241242
}
12251243
}
1226-
PathStep::PathWildCard => {
1227-
not_yet_implemented_fault!(self, "PathStep::PathWildCard".to_string());
1228-
}
1229-
PathStep::PathUnpivot => {
1230-
not_yet_implemented_fault!(self, "PathStep::PathUnpivot".to_string());
1231-
}
1244+
PathStep::PathWildCard => logical::PathComponent::PathWildcard,
1245+
PathStep::PathUnpivot => logical::PathComponent::PathUnpivot,
12321246
};
12331247

12341248
self.push_path_step(step);

‎partiql-logical/src/lib.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -457,6 +457,10 @@ pub enum PathComponent {
457457
Index(i64),
458458
KeyExpr(Box<ValueExpr>),
459459
IndexExpr(Box<ValueExpr>),
460+
/// E.g. `[*]` in `foo[*].a[*].b`
461+
PathWildcard,
462+
/// E.g. `.*` in `bar.*.c.*.d`
463+
PathUnpivot,
460464
}
461465

462466
/// Represents a PartiQL tuple expression, e.g: `{ a.b: a.c * 2, 'count': a.c + 10}`.

0 commit comments

Comments
 (0)