Add initial GROUP BY implementation (#301)

Author: alancai98

CHANGELOG.md

@@ -16,6 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `serde` feature to `partiql-value` and `partiql-logical` with `Serialize` and `Deserialize` traits.
 - Adds `Display` for `LogicalPlan`
 - Expose `partiql_value::parse_ion` as a public API.
+- Implements `GROUP BY` operator in evaluator
+- Implements `HAVING` operator in evaluator
 
 ### Fixes
 - Fixes Tuple value duplicate equality and hashing

partiql-conformance-tests/partiql-tests

@@ -282,6 +282,77 @@ impl Evaluable for EvalJoin {
     }
 }
 
+/// Represents an evaluation `GROUP BY` operator. For `GROUP BY` operational semantics, see section
+/// `11` of
+/// [PartiQL Specification — August 1, 2019](https://partiql.org/assets/PartiQL-Specification.pdf).
+#[derive(Debug)]
+pub struct EvalGroupBy {
+    pub strategy: EvalGroupingStrategy,
+    pub exprs: HashMap<String, Box<dyn EvalExpr>>,
+    pub group_as_alias: Option<String>,
+    pub input: Option<Value>,
+}
+
+/// Represents the grouping qualifier: ALL or PARTIAL.
+#[derive(Debug)]
+pub enum EvalGroupingStrategy {
+    GroupFull,
+    GroupPartial,
+}
+
+impl EvalGroupBy {
+    #[inline]
+    fn eval_group(&self, bindings: &Tuple, ctx: &dyn EvalContext) -> Tuple {
+        self.exprs
+            .iter()
+            .map(
+                |(alias, expr)| match expr.evaluate(bindings, ctx).into_owned() {
+                    Missing => (alias.as_str(), Value::Null),
+                    val => (alias.as_str(), val),
+                },
+            )
+            .collect::<Tuple>()
+    }
+}
+
+impl Evaluable for EvalGroupBy {
+    fn evaluate(&mut self, ctx: &dyn EvalContext) -> Option<Value> {
+        let group_as_alias = &self.group_as_alias;
+        let input_value = self.input.take().expect("Error in retrieving input value");
+
+        match self.strategy {
+            EvalGroupingStrategy::GroupPartial => todo!(),
+            EvalGroupingStrategy::GroupFull => {
+                let mut groups: HashMap<Tuple, Vec<Value>> = HashMap::new();
+                for v in input_value.into_iter() {
+                    let v_as_tuple = v.coerce_to_tuple();
+                    groups
+                        .entry(self.eval_group(&v_as_tuple, ctx))
+                        .or_insert(vec![])
+                        .push(Value::Tuple(Box::new(v_as_tuple)));
+                }
+
+                let bag = groups
+                    .into_iter()
+                    .map(|(k, v)| match group_as_alias {
+                        None => Value::from(k), // TODO: removing the values here will be insufficient for when aggregations are added since they may have nothing to aggregate over
+                        Some(alias) => {
+                            let mut tuple_with_group = k;
+                            tuple_with_group.insert(alias, Value::Bag(Box::new(Bag::from(v))));
+                            Value::from(tuple_with_group)
+                        }
+                    })
+                    .collect::<Bag>();
+                Some(Value::from(bag))
+            }
+        }
+    }
+
+    fn update_input(&mut self, input: Value, _branch_num: u8) {
+        self.input = Some(input);
+    }
+}
+
 /// 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).
@@ -410,7 +481,7 @@ impl EvalFilter {
             Boolean(bool_val) => *bool_val,
             // Alike SQL, when the expression of the WHERE clause expression evaluates to
             // absent value or a value that is not a Boolean, PartiQL eliminates the corresponding
-            // binding. PartiQL Specification August 1, August 1, 2019 Draft, Section 8. `WHERE clause`
+            // binding. PartiQL Specification August 1, 2019 Draft, Section 8. `WHERE clause`
             _ => false,
         }
     }
@@ -432,6 +503,51 @@ impl Evaluable for EvalFilter {
     }
 }
 
+/// Represents an evaluation `Having` operator; for an input bag of binding tuples the `Having`
+/// operator filters out the binding tuples that does not meet the condition expressed as `expr`,
+/// e.g. `a = 10` in `HAVING a = 10` expression.
+#[derive(Debug)]
+pub struct EvalHaving {
+    pub expr: Box<dyn EvalExpr>,
+    pub input: Option<Value>,
+}
+
+impl EvalHaving {
+    pub fn new(expr: Box<dyn EvalExpr>) -> Self {
+        EvalHaving { expr, input: None }
+    }
+
+    #[inline]
+    fn eval_having(&self, bindings: &Tuple, ctx: &dyn EvalContext) -> bool {
+        let result = self.expr.evaluate(bindings, ctx);
+        match result.as_ref() {
+            Boolean(bool_val) => *bool_val,
+            // Alike SQL, when the expression of the HAVING clause expression evaluates to
+            // absent value or a value that is not a Boolean, PartiQL eliminates the corresponding
+            // binding. PartiQL Specification August 1, 2019 Draft, Section 11.1.
+            // > HAVING behaves identical to a WHERE, once groups are already formulated earlier
+            // See Section 8 on WHERE semantics
+            _ => false,
+        }
+    }
+}
+
+impl Evaluable for EvalHaving {
+    fn evaluate(&mut self, ctx: &dyn EvalContext) -> Option<Value> {
+        let input_value = self.input.take().expect("Error in retrieving input value");
+
+        let filtered = input_value
+            .into_iter()
+            .map(Value::coerce_to_tuple)
+            .filter_map(|v| self.eval_having(&v, ctx).then_some(v));
+        Some(Value::from(filtered.collect::<Bag>()))
+    }
+
+    fn update_input(&mut self, input: Value, _branch_num: u8) {
+        self.input = Some(input);
+    }
+}
+
 /// Represents an evaluation `LIMIT` and/or `OFFSET` operator.
 #[derive(Debug)]
 pub struct EvalLimitOffset {

partiql-eval/src/eval/evaluable.rs

@@ -5,12 +5,12 @@ use std::collections::HashMap;
 use partiql_logical as logical;
 
 use partiql_logical::{
-    BinaryOp, BindingsOp, CallName, IsTypeExpr, JoinKind, LogicalPlan, OpId, PathComponent,
-    Pattern, PatternMatchExpr, SearchedCase, Type, UnaryOp, ValueExpr,
+    BinaryOp, BindingsOp, CallName, GroupingStrategy, IsTypeExpr, JoinKind, LogicalPlan, OpId,
+    PathComponent, Pattern, PatternMatchExpr, SearchedCase, Type, UnaryOp, ValueExpr,
 };
 
 use crate::eval;
-use crate::eval::evaluable::{EvalJoinKind, EvalSubQueryExpr, Evaluable};
+use crate::eval::evaluable::{EvalGroupingStrategy, EvalJoinKind, EvalSubQueryExpr, Evaluable};
 use crate::eval::expr::pattern_match::like_to_re_pattern;
 use crate::eval::expr::{
     EvalBagExpr, EvalBetweenExpr, EvalBinOp, EvalBinOpExpr, EvalDynamicLookup, EvalExpr, EvalFnAbs,
@@ -91,6 +91,10 @@ impl EvaluatorPlanner {
                 expr: self.plan_values(expr),
                 input: None,
             }),
+            BindingsOp::Having(logical::Having { expr }) => Box::new(eval::evaluable::EvalHaving {
+                expr: self.plan_values(expr),
+                input: None,
+            }),
             BindingsOp::Distinct => Box::new(eval::evaluable::EvalDistinct::new()),
             BindingsOp::Sink => Box::new(eval::evaluable::EvalSink { input: None }),
             BindingsOp::Pivot(logical::Pivot { key, value }) => Box::new(
@@ -129,6 +133,27 @@ impl EvaluatorPlanner {
                     on,
                 ))
             }
+            BindingsOp::GroupBy(logical::GroupBy {
+                strategy,
+                exprs,
+                group_as_alias,
+            }) => {
+                let strategy = match strategy {
+                    GroupingStrategy::GroupFull => EvalGroupingStrategy::GroupFull,
+                    GroupingStrategy::GroupPartial => EvalGroupingStrategy::GroupPartial,
+                };
+                let exprs: HashMap<_, _> = exprs
+                    .iter()
+                    .map(|(k, v)| (k.clone(), self.plan_values(v)))
+                    .collect();
+                let group_as_alias = group_as_alias.as_ref().map(|alias| alias.to_string());
+                Box::new(eval::evaluable::EvalGroupBy {
+                    strategy,
+                    exprs,
+                    group_as_alias,
+                    input: None,
+                })
+            }
             BindingsOp::ExprQuery(logical::ExprQuery { expr }) => {
                 let expr = self.plan_values(expr);
                 Box::new(eval::evaluable::EvalExprQuery::new(expr))
@@ -143,7 +168,6 @@ impl EvaluatorPlanner {
             }
 
             BindingsOp::SetOp => todo!("SetOp"),
-            BindingsOp::GroupBy => todo!("GroupBy"),
         }
     }
 

partiql-eval/src/plan.rs

@@ -6,10 +6,11 @@ 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, Insert, InsertValue, Item, Join,
-    JoinKind, JoinSpec, Like, List, Lit, NodeId, OnConflict, OrderByExpr, Path, PathStep,
-    ProjectExpr, Projection, ProjectionKind, Query, QuerySet, Remove, SearchedCase, Select, Set,
-    SetExpr, SetQuantifier, Sexp, SimpleCase, Struct, SymbolPrimitive, UniOp, UniOpKind, VarRef,
+    DropTable, FromClause, FromLet, FromLetKind, GroupByExpr, GroupKey, GroupingStrategy, Insert,
+    InsertValue, Item, Join, JoinKind, JoinSpec, Like, List, Lit, NodeId, OnConflict, OrderByExpr,
+    Path, PathStep, ProjectExpr, Projection, ProjectionKind, Query, QuerySet, Remove, SearchedCase,
+    Select, Set, SetExpr, SetQuantifier, Sexp, SimpleCase, Struct, SymbolPrimitive, UniOp,
+    UniOpKind, VarRef,
 };
 use partiql_ast::visit::{Visit, Visitor};
 use partiql_logical as logical;
@@ -1087,17 +1088,71 @@ impl<'ast> Visitor<'ast> for AstToLogical {
     }
 
     fn exit_having_clause(&mut self, _having_clause: &'ast ast::HavingClause) {
-        let _env = self.exit_env();
-        todo!("having clause");
+        let mut env = self.exit_env();
+        assert_eq!(env.len(), 1);
+
+        let having = BindingsOp::Having(logical::Having {
+            expr: env.pop().unwrap(),
+        });
+        let id = self.plan.add_operator(having);
+
+        self.current_clauses_mut().having_clause.replace(id);
     }
 
     fn enter_group_by_expr(&mut self, _group_by_expr: &'ast GroupByExpr) {
+        self.enter_benv();
         self.enter_env();
     }
 
     fn exit_group_by_expr(&mut self, _group_by_expr: &'ast GroupByExpr) {
-        let _env = self.exit_env();
-        todo!("group by clause");
+        let benv = self.exit_benv();
+        assert_eq!(benv.len(), 0); // TODO sub-query
+        let env = self.exit_env();
+
+        let group_as_alias = _group_by_expr
+            .group_as_alias
+            .as_ref()
+            .map(|SymbolPrimitive { value, case: _ }| value.clone());
+
+        let strategy = match _group_by_expr.strategy {
+            GroupingStrategy::GroupFull => logical::GroupingStrategy::GroupFull,
+            GroupingStrategy::GroupPartial => logical::GroupingStrategy::GroupPartial,
+        };
+        let mut exprs = HashMap::with_capacity(env.len() / 2);
+        let mut iter = env.into_iter();
+        while let Some(value) = iter.next() {
+            let alias = iter.next().unwrap();
+            let alias = match alias {
+                ValueExpr::Lit(lit) => match *lit {
+                    Value::String(s) => (*s).clone(),
+                    _ => panic!("unexpected literal"),
+                },
+                _ => panic!("unexpected alias type"),
+            };
+            exprs.insert(alias, value);
+        }
+        let group_by: BindingsOp = BindingsOp::GroupBy(logical::GroupBy {
+            strategy,
+            exprs,
+            group_as_alias,
+        });
+
+        let id = self.plan.add_operator(group_by);
+        self.current_clauses_mut().group_by_clause.replace(id);
+    }
+
+    fn exit_group_key(&mut self, _group_key: &'ast GroupKey) {
+        let as_key: &name_resolver::Symbol = self
+            .key_registry
+            .aliases
+            .get(self.current_node())
+            .expect("alias");
+        // TODO intern strings
+        let as_key = match as_key {
+            name_resolver::Symbol::Known(sym) => sym.value.clone(),
+            name_resolver::Symbol::Unknown(id) => format!("_{id}"),
+        };
+        self.push_value(as_key.into());
     }
 
     fn enter_order_by_expr(&mut self, _order_by_expr: &'ast OrderByExpr) {

partiql-logical-planner/src/lower.rs

@@ -1,6 +1,7 @@
 use fnv::FnvBuildHasher;
 use indexmap::{IndexMap, IndexSet};
 use partiql_ast::ast;
+use partiql_ast::ast::{GroupByExpr, GroupKey};
 use partiql_ast::visit::{Visit, Visitor};
 use std::sync::atomic::{AtomicU32, Ordering};
 
@@ -327,6 +328,36 @@ impl<'ast> Visitor<'ast> for NameResolver {
             .produce_required
             .insert(as_alias);
     }
+
+    fn exit_group_key(&mut self, group_key: &'ast GroupKey) {
+        let id = *self.current_node();
+        // get the "as" alias for each `GROUP BY` expr
+        // 1. if explicitly given
+        // 2. else try to infer if a simple variable reference or path
+        // 3. else it is currently 'Unknown'
+        let as_alias = if let Some(sym) = &group_key.as_alias {
+            Symbol::Known(sym.clone())
+        } else if let Some(sym) = infer_alias(&group_key.expr) {
+            Symbol::Known(sym)
+        } else {
+            Symbol::Unknown(self.id_gen.next_id())
+        };
+        self.aliases.insert(id, as_alias.clone());
+        self.keyref_stack
+            .last_mut()
+            .unwrap()
+            .produce_required
+            .insert(as_alias);
+    }
+
+    fn exit_group_by_expr(&mut self, group_by_expr: &'ast GroupByExpr) {
+        // add the `GROUP AS` alias
+        if let Some(sym) = &group_by_expr.group_as_alias {
+            let id = *self.current_node();
+            let as_alias = Symbol::Known(sym.clone());
+            self.aliases.insert(id, as_alias);
+        }
+    }
 }
 
 /// Attempt to infer an alias for a simple variable reference expression.