lab9: agg extraction

Signed-off-by: Runji Wang <[email protected]>

Author: wangrunji0408

src/agg.rs

@@ -0,0 +1,149 @@
+use egg::Language;
+
+use super::*;
+
+/// The data type of aggragation analysis.
+pub type AggSet = Vec<Expr>;
+
+/// Returns all aggragations in the tree.
+///
+/// Note: if there is an agg over agg, e.g. `sum(count(a))`, only the upper one will be returned.
+pub fn analyze_aggs(egraph: &EGraph, enode: &Expr) -> AggSet {
+    use Expr::*;
+    let x = |i: &Id| egraph[*i].data.aggs.clone();
+    if let Max(_) | Min(_) | Sum(_) | Avg(_) | Count(_) = enode {
+        return vec![enode.clone()];
+    }
+    // merge the set from all children
+    // TODO: ignore plan nodes
+    enode.children().iter().flat_map(x).collect()
+}
+
+#[derive(Debug, PartialEq, Eq)]
+pub enum Error {
+    // #[error("aggregate function calls cannot be nested")]
+    NestedAgg(String),
+    // #[error("WHERE clause cannot contain aggregates")]
+    AggInWhere,
+    // #[error("GROUP BY clause cannot contain aggregates")]
+    AggInGroupBy,
+    // #[error("column {0} must appear in the GROUP BY clause or be used in an aggregate function")]
+    ColumnNotInAgg(String),
+}
+
+/// Converts the SELECT statement into a plan tree.
+///
+/// The nodes of all clauses have been added to the `egraph`.
+/// `from`, `where_`... are the ids of their root node.
+pub fn plan_select(
+    egraph: &mut EGraph,
+    from: Id,
+    where_: Id,
+    having: Id,
+    groupby: Id,
+    orderby: Id,
+    projection: Id,
+) -> Result<Id, Error> {
+    AggExtractor { egraph }.plan_select(from, where_, having, groupby, orderby, projection)
+}
+
+struct AggExtractor<'a> {
+    egraph: &'a mut EGraph,
+}
+
+impl AggExtractor<'_> {
+    fn aggs(&self, id: Id) -> &[Expr] {
+        &self.egraph[id].data.aggs
+    }
+
+    fn node(&self, id: Id) -> &Expr {
+        &self.egraph[id].nodes[0]
+    }
+
+    fn plan_select(
+        &mut self,
+        from: Id,
+        where_: Id,
+        having: Id,
+        groupby: Id,
+        orderby: Id,
+        projection: Id,
+    ) -> Result<Id, Error> {
+        if !self.aggs(where_).is_empty() {
+            return Err(Error::AggInWhere);
+        }
+        if !self.aggs(groupby).is_empty() {
+            return Err(Error::AggInGroupBy);
+        }
+        let mut plan = self.egraph.add(Expr::Filter([where_, from]));
+        let mut to_rewrite = [projection, having, orderby];
+        plan = self.plan_agg(&mut to_rewrite, groupby, plan)?;
+        let [projection, having, orderby] = to_rewrite;
+        plan = self.egraph.add(Expr::Filter([having, plan]));
+        plan = self.egraph.add(Expr::Order([orderby, plan]));
+        plan = self.egraph.add(Expr::Proj([projection, plan]));
+        Ok(plan)
+    }
+
+    /// Extracts all aggregations from `exprs` and generates an [`Agg`](Expr::Agg) plan.
+    /// If no aggregation is found and no `groupby` keys, returns the original `plan`.
+    fn plan_agg(&mut self, exprs: &mut [Id], groupby: Id, plan: Id) -> Result<Id, Error> {
+        let expr_list = self.egraph.add(Expr::List(exprs.to_vec().into()));
+        let aggs = self.aggs(expr_list).to_vec();
+        if aggs.is_empty() && self.node(groupby).as_list().is_empty() {
+            return Ok(plan);
+        }
+        // check nested agg
+        for agg in aggs.iter() {
+            if agg
+                .children()
+                .iter()
+                .any(|child| !self.aggs(*child).is_empty())
+            {
+                return Err(Error::NestedAgg(agg.to_string()));
+            }
+        }
+        let mut list: Vec<_> = aggs.into_iter().map(|agg| self.egraph.add(agg)).collect();
+        // make sure the order of the aggs is deterministic
+        list.sort();
+        list.dedup();
+        let mut schema = list.clone();
+        schema.extend_from_slice(self.node(groupby).as_list());
+        let aggs = self.egraph.add(Expr::List(list.into()));
+        let plan = self.egraph.add(Expr::Agg([aggs, groupby, plan]));
+        // check for not aggregated columns
+        // rewrite the expressions with a wrapper over agg or group keys
+        for id in exprs {
+            *id = self.rewrite_agg_in_expr(*id, &schema)?;
+        }
+        Ok(plan)
+    }
+
+    /// Rewrites the expression `id` with aggs wrapped in a [`Nested`](Expr::Nested) node.
+    /// Returns the new expression.
+    ///
+    /// # Example
+    /// ```text
+    /// id:         (+ (sum a) (+ b 1))
+    /// schema:     (sum a), (+ b 1)
+    /// output:     (+ (`(sum a)) (`(+ b 1)))
+    ///
+    /// so that `id` won't be optimized to:
+    ///             (+ b (+ (sum a) 1))
+    /// which can not be composed by `schema`
+    /// ```
+    fn rewrite_agg_in_expr(&mut self, id: Id, schema: &[Id]) -> Result<Id, Error> {
+        let mut expr = self.node(id).clone();
+        if schema.contains(&id) {
+            // found agg, wrap it with Nested
+            return Ok(self.egraph.add(Expr::Nested(id)));
+        }
+        if let Expr::Column(cid) = &expr {
+            return Err(Error::ColumnNotInAgg(cid.to_string()));
+        }
+        for child in expr.children_mut() {
+            *child = self.rewrite_agg_in_expr(*child, schema)?;
+        }
+        Ok(self.egraph.add(expr))
+    }
+}

src/lib.rs

@@ -2,6 +2,7 @@ use std::hash::Hash;
 
 use egg::{define_language, Analysis, DidMerge, Id, Var};
 
+pub mod agg;
 pub mod expr;
 pub mod plan;
 mod value;
@@ -16,9 +17,9 @@ define_language! {
     pub enum Expr {
         // values
         Constant(Value),            // null, true, 1, 'hello'
-        Column(Column),             // t.a, b, c
 
         // utilities
+        "`" = Nested(Id),           // (` expr) a wrapper over expr to prevent optimization
         "list" = List(Box<[Id]>),   // (list ...)
 
         // unary operations
@@ -75,6 +76,8 @@ define_language! {
         "empty" = Empty(Id),                    // (empty child)
                                                     // returns empty chunk
                                                     // with the same schema as `child`
+
+        Column(Column),             // t.a, b, c
     }
 }
 
@@ -101,6 +104,9 @@ pub struct Data {
 
     /// All columns involved in the node.
     pub columns: plan::ColumnSet,
+
+    /// All aggragations in the tree.
+    pub aggs: agg::AggSet,
 }
 
 impl Analysis<Expr> for ExprAnalysis {
@@ -111,6 +117,7 @@ impl Analysis<Expr> for ExprAnalysis {
         Data {
             constant: expr::eval_constant(egraph, enode),
             columns: plan::analyze_columns(egraph, enode),
+            aggs: agg::analyze_aggs(egraph, enode),
         }
     }
 
@@ -125,7 +132,8 @@ impl Analysis<Expr> for ExprAnalysis {
     fn merge(&mut self, to: &mut Self::Data, from: Self::Data) -> DidMerge {
         let merge_const = egg::merge_max(&mut to.constant, from.constant);
         let merge_columns = plan::merge(&mut to.columns, from.columns);
-        merge_const | merge_columns
+        let merge_aggs = egg::merge_max(&mut to.aggs, from.aggs);
+        merge_const | merge_columns | merge_aggs
     }
 
     /// Modify the graph after analyzing a node.

tests/9_agg_extraction.rs

@@ -0,0 +1,161 @@
+use egg::Language;
+use egg_sql_labs::{
+    agg::{plan_select, Error},
+    EGraph, RecExpr,
+};
+
+#[test]
+fn no_agg() {
+    // SELECT a FROM t;
+    test(Case {
+        select: "(list a)",
+        from: "(scan t (list a))",
+        where_: "",
+        having: "",
+        groupby: "",
+        orderby: "",
+        expected: Ok("
+        (proj (list a)
+            (order list
+                (filter true
+                    (filter true
+                        (scan t (list a))
+        ))))"),
+    });
+}
+
+#[test]
+fn agg() {
+    // SELECT sum(a + b) + (a + 1) FROM t
+    // WHERE b > 1
+    // GROUP BY a + 1
+    // HAVING count(a) > 1
+    // ORDER BY max(b)
+    test(Case {
+        select: "(list (+ (sum (+ a b)) (+ a 1)))",
+        from: "(scan t (list a b))",
+        where_: "(> b 1)",
+        having: "(> (count a) 1)",
+        groupby: "(list (+ a 1))",
+        orderby: "(list (asc (max b)))",
+        expected: Ok("
+        (proj (list (+ (` (sum (+ a b))) (` (+ a 1))))
+            (order (list (asc (` (max b))))
+                (filter (> (` (count a)) 1)
+                    (agg (list (sum (+ a b)) (count a) (max b)) (list (+ a 1))
+                        (filter (> b 1)
+                            (scan t (list a b))
+        )))))"),
+    });
+}
+
+#[test]
+fn error_agg_in_where() {
+    // SELECT a FROM t WHERE sum(a) > 1
+    test(Case {
+        select: "(list a)",
+        from: "(scan t (list a b))",
+        where_: "(> (sum a) 1)",
+        having: "",
+        groupby: "",
+        orderby: "",
+        expected: Err(Error::AggInWhere),
+    });
+}
+
+#[test]
+fn error_agg_in_groupby() {
+    // SELECT a FROM t GROUP BY sum(a)
+    test(Case {
+        select: "(list a)",
+        from: "(scan t (list a b))",
+        where_: "",
+        having: "",
+        groupby: "(list (sum a))",
+        orderby: "",
+        expected: Err(Error::AggInGroupBy),
+    });
+}
+
+#[test]
+fn error_nested_agg() {
+    // SELECT count(sum(a)) FROM t
+    test(Case {
+        select: "(list (count (sum a)))",
+        from: "(scan t (list a b))",
+        where_: "",
+        having: "",
+        groupby: "",
+        orderby: "",
+        expected: Err(Error::NestedAgg("count".into())),
+    });
+}
+
+#[test]
+fn error_column_not_in_agg() {
+    // SELECT b FROM t GROUP BY a
+    test(Case {
+        select: "(list b)",
+        from: "(scan t (list a b))",
+        where_: "",
+        having: "",
+        groupby: "(list a)",
+        orderby: "",
+        expected: Err(Error::ColumnNotInAgg("b".into())),
+    });
+}
+
+struct Case {
+    select: &'static str,
+    from: &'static str,
+    where_: &'static str,
+    having: &'static str,
+    groupby: &'static str,
+    orderby: &'static str,
+    expected: Result<&'static str, Error>,
+}
+
+#[track_caller]
+fn test(mut case: Case) {
+    if case.where_.is_empty() {
+        case.where_ = "true";
+    }
+    if case.having.is_empty() {
+        case.having = "true";
+    }
+    if case.groupby.is_empty() {
+        case.groupby = "list";
+    }
+    if case.orderby.is_empty() {
+        case.orderby = "list";
+    }
+    let mut egraph = EGraph::default();
+    let projection = egraph.add_expr(&case.select.parse().unwrap());
+    let from = egraph.add_expr(&case.from.parse().unwrap());
+    let where_ = egraph.add_expr(&case.where_.parse().unwrap());
+    let having = egraph.add_expr(&case.having.parse().unwrap());
+    let groupby = egraph.add_expr(&case.groupby.parse().unwrap());
+    let orderby = egraph.add_expr(&case.orderby.parse().unwrap());
+    match plan_select(
+        &mut egraph,
+        from,
+        where_,
+        having,
+        groupby,
+        orderby,
+        projection,
+    ) {
+        Err(e) => assert_eq!(case.expected, Err(e)),
+        Ok(id) => {
+            let get_node = |id| egraph[id].nodes[0].clone();
+            let actual = get_node(id).build_recexpr(get_node).to_string();
+            let expected = case
+                .expected
+                .expect(&format!("expect error, but got: {actual:?}"))
+                .parse::<RecExpr>()
+                .unwrap()
+                .to_string();
+            assert_eq!(actual, expected);
+        }
+    }
+}