Implement PIVOT in evaluator (#286)

Author: alancai98

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - Adds some benchmarks for parsing, compiling, planning, & evaluation
+- Implements `PIVOT` operator in evaluator
 
 ### Fixes
 

partiql-eval/src/eval.rs

@@ -383,6 +383,47 @@ impl Evaluable for EvalJoin {
     }
 }
 
+/// Represents an evaluation `Pivot` operator; the `Pivot` enables turning a collection into a
+/// tuple. For `Pivot` operational semantics, see section `6.2` of
+/// [PartiQL Specification — August 1, 2019](https://partiql.org/assets/PartiQL-Specification.pdf).
+#[derive(Debug)]
+pub struct EvalPivot {
+    pub input: Option<Value>,
+    pub key: Box<dyn EvalExpr>,
+    pub value: Box<dyn EvalExpr>,
+}
+
+impl EvalPivot {
+    pub fn new(key: Box<dyn EvalExpr>, value: Box<dyn EvalExpr>) -> Self {
+        EvalPivot {
+            input: None,
+            key,
+            value,
+        }
+    }
+}
+
+impl Evaluable for EvalPivot {
+    fn evaluate(&mut self, ctx: &dyn EvalContext) -> Option<Value> {
+        let input_value = self.input.take().expect("Error in retrieving input value");
+
+        let mut out: Tuple = partiql_tuple![];
+        for binding in input_value.into_iter() {
+            let binding = binding.coerce_to_tuple();
+            let key = self.key.evaluate(&binding, ctx);
+            if let Value::String(s) = key {
+                let value = self.value.evaluate(&binding, ctx);
+                out.insert(&s, value)
+            }
+        }
+        Some(Value::Tuple(Box::new(out)))
+    }
+
+    fn update_input(&mut self, input: Value, _branch_num: u8) {
+        self.input = Some(input);
+    }
+}
+
 /// Represents an evaluation `Unpivot` operator; the `Unpivot` enables ranging over the
 /// attribute-value pairs of a tuple. For `Unpivot` operational semantics, see section `5.2` of
 /// [PartiQL Specification — August 1, 2019](https://partiql.org/assets/PartiQL-Specification.pdf).

partiql-eval/src/plan.rs

@@ -87,6 +87,10 @@ impl EvaluatorPlanner {
             }),
             BindingsOp::Distinct => Box::new(eval::EvalDistinct::new()),
             BindingsOp::Sink => Box::new(eval::EvalSink { input: None }),
+            BindingsOp::Pivot(logical::Pivot { key, value }) => Box::new(eval::EvalPivot::new(
+                self.plan_values(key),
+                self.plan_values(value),
+            )),
             BindingsOp::Unpivot(logical::Unpivot {
                 expr,
                 as_key,

partiql-logical-planner/src/lower.rs

@@ -621,7 +621,13 @@ impl<'ast> Visitor<'ast> for AstToLogical {
 
                 logical::BindingsOp::Project(logical::Project { exprs })
             }
-            ProjectionKind::ProjectPivot(_) => todo!("ProjectionKind::ProjectPivot"),
+            ProjectionKind::ProjectPivot(_) => {
+                assert_eq!(env.len(), 2);
+                let mut iter = env.into_iter();
+                let key = iter.next().unwrap();
+                let value = iter.next().unwrap();
+                logical::BindingsOp::Pivot(logical::Pivot { key, value })
+            }
             ProjectionKind::ProjectValue(_) => {
                 assert_eq!(env.len(), 1);
                 let expr = env.into_iter().next().unwrap();

partiql-logical/src/lib.rs

@@ -162,6 +162,7 @@ impl OpId {
 #[derive(Debug, Clone, Default, Eq, PartialEq)]
 pub enum BindingsOp {
     Scan(Scan),
+    Pivot(Pivot),
     Unpivot(Unpivot),
     Filter(Filter),
     OrderBy,
@@ -187,6 +188,16 @@ pub struct Scan {
     pub at_key: Option<String>,
 }
 
+/// [`Pivot`] represents a PIVOT operator, e.g. `PIVOT sp.price AT sp."symbol` in
+/// `PIVOT sp.price AT sp."symbol" FROM todaysStockPrices sp`. For `Pivot` operational semantics,
+/// see section `6.2` of
+/// [PartiQL Specification — August 1, 2019](https://partiql.org/assets/PartiQL-Specification.pdf).
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct Pivot {
+    pub key: ValueExpr,
+    pub value: ValueExpr,
+}
+
 /// [`Unpivot`] bridges from [`ValueExpr`]s to [`BindingsOp`]s.
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct Unpivot {