mir: window functions initial support

Handle creation and rewrite of Window MIR node.

Addresses: REA-5790

Change-Id: I52c24d9547688b0278b53ae17ed15648c68039e4
Reviewed-on: https://gerrit.readyset.name/c/readyset/+/9647
Reviewed-by: Michael Zink <michael.z@readyset.io>
Tested-by: Buildkite CI

Author: MohamedAbdeen-rs

readyset-dataflow/src/ops/mod.rs

@@ -21,6 +21,7 @@ pub mod project;
 pub mod topk;
 pub mod union;
 pub(crate) mod utils;
+pub mod window;
 
 use crate::ops::grouped::concat::GroupConcat;
 use crate::processing::{

readyset-dataflow/src/ops/window.rs

@@ -0,0 +1,49 @@
+use std::fmt::Display;
+
+use readyset_errors::{unsupported, ReadySetResult};
+use readyset_sql::ast::FunctionExpr;
+use serde::{Deserialize, Serialize};
+
+use crate::ops::grouped::aggregate::Aggregation;
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub enum WindowOperation {
+    RowNumber,
+    Rank,
+    DenseRank,
+    Max,
+    Min,
+    Aggregation(Aggregation),
+}
+
+impl WindowOperation {
+    pub fn from_fn(fn_expr: FunctionExpr) -> ReadySetResult<WindowOperation> {
+        Ok(match fn_expr {
+            FunctionExpr::RowNumber => WindowOperation::RowNumber,
+            FunctionExpr::Max { .. } => WindowOperation::Max,
+            FunctionExpr::Min { .. } => WindowOperation::Min,
+            FunctionExpr::CountStar | FunctionExpr::Count { .. } => {
+                WindowOperation::Aggregation(Aggregation::Count)
+            }
+            FunctionExpr::Sum { .. } => WindowOperation::Aggregation(Aggregation::Sum),
+            FunctionExpr::Avg { .. } => WindowOperation::Aggregation(Aggregation::Avg),
+            FunctionExpr::Rank => WindowOperation::Rank,
+            FunctionExpr::DenseRank => WindowOperation::DenseRank,
+            _ => unsupported!("Function {fn_expr:?} is not supported as a window function"),
+        })
+    }
+}
+
+impl Display for WindowOperation {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            WindowOperation::RowNumber => write!(f, "row_number"),
+            WindowOperation::Rank => write!(f, "rank"),
+            WindowOperation::DenseRank => write!(f, "dense_rank"),
+            WindowOperation::Max => write!(f, "max"),
+            WindowOperation::Min => write!(f, "min"),
+            // TODO: Make Aggregation impl Display
+            WindowOperation::Aggregation(a) => write!(f, "{a:?}"),
+        }
+    }
+}

readyset-mir/src/graph.rs

@@ -335,6 +335,10 @@ impl MirGraph {
                 .into_iter()
                 .chain(iter::once(MirColumn::named(&*PAGE_NUMBER_COL)))
                 .collect(),
+            MirNodeInner::Window { output_column, .. } => parent_columns()
+                .into_iter()
+                .chain(iter::once(output_column.clone()))
+                .collect(),
             MirNodeInner::Distinct { group_by } => group_by
                 .iter()
                 .cloned()

readyset-mir/src/node/node_inner.rs

@@ -4,6 +4,7 @@ use common::{DfValue, IndexType};
 use dataflow::ops::grouped::aggregate::Aggregation;
 use dataflow::ops::grouped::extremum::Extremum;
 use dataflow::ops::union;
+use dataflow::ops::window::WindowOperation;
 use dataflow::PostLookupAggregates;
 use derive_more::From;
 use itertools::Itertools;
@@ -98,6 +99,17 @@ pub enum MirNodeInner {
         primary_key: Option<Box<[Column]>>,
         unique_keys: Box<[Box<[Column]>]>,
     },
+    /// Node that evaluates Window Functions over a window.
+    /// PARTITIONS and ORDER BY can be expressions, but should
+    /// be evaluated by a parent node before getting passed here.
+    Window {
+        group_by: Vec<Column>,
+        partition_by: Vec<Column>,
+        order_by: Vec<(Column, OrderType, NullOrder)>,
+        output_column: Column,
+        function: WindowOperation,
+        args: Vec<Column>,
+    },
     /// Node that computes the extreme value (minimum or maximum) of a column grouped by another
     /// set of columns, outputting its result as an additional column.
     ///
@@ -436,6 +448,27 @@ impl MirNodeInner {
                     .join(", ");
                 format!("{op_string} γ[{group_cols}]")
             }
+            MirNodeInner::Window {
+                partition_by,
+                order_by,
+                function,
+                args,
+                ..
+            } => {
+                let op_string = format!("{}({})", function, args.iter().join(", "));
+                let partition_cols = partition_by
+                    .iter()
+                    .map(|c| c.name.as_str())
+                    .collect::<Vec<_>>()
+                    .join(", ");
+                let order_cols = order_by
+                    .iter()
+                    .map(|(c, o, no)| format!("{c} {o}({no})"))
+                    .collect::<Vec<_>>()
+                    .into_iter()
+                    .join(", ");
+                format!("{op_string} Over[{partition_cols}; {order_cols}]")
+            }
             MirNodeInner::Base {
                 column_specs,
                 unique_keys,

readyset-mir/src/rewrite/filters_to_join_keys.rs

@@ -254,6 +254,10 @@ pub(crate) fn convert_filters_to_join_keys(query: &mut MirQuery<'_>) -> ReadySet
                     // TODO: figure out what to do about unions
                     continue 'filter;
                 }
+                MirNodeInner::Window { .. } => {
+                    // TODO: figure out what to do about windows
+                    continue 'filter;
+                }
                 MirNodeInner::Join { .. } => {
                     let join_parents = query.ancestors(ancestor_idx)?;
                     let left_parent = *join_parents

readyset-mir/src/rewrite/predicate_pushup.rs

@@ -16,6 +16,7 @@ fn commutes_with(conditions: &Expr, inner: &MirNodeInner) -> bool {
         | MirNodeInner::Paginate { group_by, .. }
         | MirNodeInner::TopK { group_by, .. }
         | MirNodeInner::Distinct { group_by, .. }
+        | MirNodeInner::Window { group_by, .. }
         | MirNodeInner::Extremum { group_by, .. } => conditions
             .referred_columns()
             .all(|col| group_by.iter().any(|c| c == col)),

readyset-mir/src/rewrite/pull_keys.rs

@@ -39,6 +39,7 @@ fn push_view_key(query: &mut MirQuery<'_>, node_idx: NodeIndex) -> ReadySetResul
         | MirNodeInner::Extremum { group_by, .. }
         | MirNodeInner::Distinct { group_by }
         | MirNodeInner::Paginate { group_by, .. }
+        | MirNodeInner::Window { group_by, .. }
         | MirNodeInner::TopK { group_by, .. } => {
             for ViewKeyColumn { column, op, .. } in &key {
                 invariant_eq!(

readyset-mir/src/visualize.rs

@@ -370,6 +370,27 @@ impl GraphViz for MirNodeInner {
             MirNodeInner::AliasTable { ref table } => {
                 write!(f, "AliasTable [{}]", table.display_unquoted())
             }
+            MirNodeInner::Window {
+                group_by,
+                partition_by,
+                order_by,
+                function,
+                args,
+                ..
+            } => {
+                write!(
+                    f,
+                    "Window [key: {}, p: {}, o: {}, f: {}, args: {}]",
+                    group_by.iter().join(", "),
+                    partition_by.iter().join(", "),
+                    order_by
+                        .iter()
+                        .map(|(c, o, no)| format!("{}: {}({})", c.name.as_str(), o, no))
+                        .join(", "),
+                    function,
+                    args.iter().join(", ")
+                )
+            }
         }
     }
 }

readyset-server/src/controller/mir_to_flow.rs

@@ -136,6 +136,7 @@ pub(super) fn mir_node_to_flow_parts(
                         mig,
                     )?)
                 }
+                MirNodeInner::Window { .. } => unsupported!("lowering window node to flow"),
                 MirNodeInner::Filter { ref conditions } => {
                     invariant_eq!(ancestors.len(), 1);
                     let parent = ancestors[0];

readyset-server/src/controller/sql/mir/mod.rs

@@ -9,6 +9,7 @@ use catalog_tables::is_catalog_table;
 use common::IndexType;
 use dataflow::ops::grouped::aggregate::Aggregation;
 use dataflow::ops::union;
+use dataflow::ops::window::WindowOperation;
 use lazy_static::lazy_static;
 use mir::graph::MirGraph;
 use mir::node::node_inner::MirNodeInner;
@@ -24,7 +25,7 @@ use readyset_errors::{
     unsupported_err, ReadySetError, ReadySetResult,
 };
 use readyset_sql::analysis::visit::{walk_expr, Visitor};
-use readyset_sql::analysis::{self, ReferredColumns};
+use readyset_sql::analysis::{self, is_aggregate, ReferredColumns};
 use readyset_sql::ast::{
     self, BinaryOperator, CaseWhenBranch, ColumnSpecification, CompoundSelectOperator,
     CreateTableBody, Expr, FieldDefinitionExpr, FieldReference, FunctionExpr, GroupByClause,
@@ -47,6 +48,7 @@ use crate::controller::sql::query_graph::{
     to_query_graph, ExprColumn, OutputColumn, Pagination, QueryGraph,
 };
 use crate::controller::sql::query_signature::Signature;
+use crate::sql::query_graph::WindowFunction;
 
 mod grouped;
 mod join;
@@ -2065,6 +2067,134 @@ impl SqlToMirConverter {
         Ok(leaf)
     }
 
+    fn make_window_node(
+        &mut self,
+        query_name: &Relation,
+        name: Relation,
+        mut parent: NodeIndex,
+        funcs: &[WindowFunction],
+        group_by: Vec<Column>,
+    ) -> ReadySetResult<NodeIndex> {
+        if funcs.is_empty() {
+            return Ok(parent);
+        } else if funcs.len() != 1 {
+            // TOOD(mohamed): On paper, it should be possible by simply connecting the
+            // window functions in series (fn1 -> fn2 -> fn3 -> ... -> reader)
+            // However, the performance will be so bad that we rather just not support
+            // it for now. Making the execution "smarter" with how it handles
+            // its inputs and outputs will make this possible, but should be left
+            // as a future improvement
+            unsupported!("Multiple window functions not yet supported");
+        }
+
+        let dialect = readyset_sql::Dialect::MySQL;
+        let WindowFunction {
+            function,
+            partition_by,
+            order_by,
+            alias,
+        } = funcs.first().unwrap().clone();
+
+        let arguments = function.arguments().cloned().collect::<Vec<_>>();
+
+        let function = WindowOperation::from_fn(function)?;
+
+        let order_cols = order_by
+            .iter()
+            .map(|(e, _, _)| e)
+            .cloned()
+            .collect::<Vec<_>>();
+
+        let has_agg = |e: &Expr| matches!(e, Expr::Call(f) if is_aggregate(f));
+        if partition_by.iter().any(has_agg)
+            || arguments.iter().any(has_agg)
+            || order_cols.iter().any(has_agg)
+        {
+            unsupported!("Aggregates in window functions not yet supported");
+        }
+
+        // if the partition, ordering cols, or args require projection,
+        // create a projection node and use that as a parent
+        // TODO: do we need to make this strict? i.e. only project certain
+        // Expr variants?
+        let p = |e: &Expr| !matches!(e, Expr::Column(_));
+        let needs_proj_node =
+            partition_by.iter().any(p) || arguments.iter().any(p) || order_cols.iter().any(p);
+
+        if needs_proj_node {
+            let node_name = format!(
+                "{}_window_project_n{}",
+                name.display_unquoted(),
+                self.mir_graph.node_count()
+            );
+
+            let node = self.make_project_node(
+                query_name,
+                node_name.into(),
+                parent,
+                partition_by
+                    .iter()
+                    .chain(arguments.iter())
+                    .chain(order_cols.iter())
+                    .cloned()
+                    .map(|e| -> ReadySetResult<_> {
+                        Ok(ProjectExpr::Expr {
+                            alias: e
+                                .alias(dialect)
+                                // returns None if e is a placeholder or a variable
+                                .ok_or_else(|| {
+                                    unsupported_err!("Placeholders not allowed in this context")
+                                })?,
+                            expr: e,
+                        })
+                    })
+                    .collect::<ReadySetResult<Vec<_>>>()?,
+            );
+            parent = node;
+        }
+
+        let output_column = Column::named(alias);
+
+        let partition_by = partition_by
+            .iter()
+            .map(|e| e.alias(dialect).unwrap())
+            .map(|e| Column::named(e))
+            .collect();
+
+        let order_by = order_by
+            .into_iter()
+            .map(|(e, order, no)| (Column::named(e.alias(dialect).unwrap()), order, no))
+            .collect();
+
+        let args = arguments
+            .iter()
+            .map(|e| e.alias(dialect).unwrap())
+            .map(|e| Column::named(e))
+            .collect();
+
+        let node_name = format!(
+            "{}_window_n{}",
+            name.display_unquoted(),
+            self.mir_graph.node_count()
+        );
+
+        Ok(self.add_query_node(
+            query_name.clone(),
+            MirNode::new(
+                node_name.into(),
+                MirNodeInner::Window {
+                    group_by,
+                    partition_by,
+                    order_by,
+                    output_column,
+                    function,
+                    args,
+                },
+            ),
+            &[parent],
+        ))
+    }
+
     fn predicates_above_group_by<'a>(
         &mut self,
         query_name: &Relation,
@@ -2370,7 +2500,21 @@ impl SqlToMirConverter {
                 prev_node = subquery_leaf;
             }
 
-            // 8. Add function and grouped nodes
+            // 8. Add window functions or grouped nodes (mutually exclusive for now)
+            if !query_graph.aggregates.is_empty() && !query_graph.window_functions.is_empty() {
+                unsupported!("Mixing window functions and aggregates is not supported yet")
+            };
+
+            let group_by: Vec<_> = view_key.columns.iter().map(|(c, _)| c.clone()).collect();
+
+            prev_node = self.make_window_node(
+                query_name,
+                format!("q_{:x}", query_graph.signature().hash).into(),
+                prev_node,
+                &query_graph.window_functions,
+                group_by,
+            )?;
+
             let mut func_nodes: Vec<NodeIndex> = make_grouped(
                 self,
                 query_name,
@@ -2620,6 +2764,10 @@ impl SqlToMirConverter {
                 let are_repeat_reads_required =
                     query_graph.collapsed_where_in || view_key.index_type == IndexType::BTreeMap;
 
+                if are_repeat_reads_required && !query_graph.window_functions.is_empty() {
+                    unsupported!("Post-lookups are not supported for window functions");
+                }
+
                 let post_lookup_aggregates = if are_repeat_reads_required {
                     // When a query contains WHERE col IN (?, ?, ...), it gets rewritten
                     // (or collapsed) to WHERE col = ? during SQL parsing, with the