add window expr

Author: Jiayu Liu

ballista/rust/core/proto/ballista.proto

@@ -39,7 +39,6 @@ message LogicalExprNode {
 
     ScalarValue literal = 3;
 
-
     // binary expressions
     BinaryExprNode binary_expr = 4;
 
@@ -60,6 +59,9 @@ message LogicalExprNode {
     bool wildcard = 15;
     ScalarFunctionNode scalar_function = 16;
     TryCastNode try_cast = 17;
+
+    // window expressions
+    WindowExprNode window_expr = 18;
   }
 }
 
@@ -151,6 +153,25 @@ message AggregateExprNode {
   LogicalExprNode expr = 2;
 }
 
+enum BuiltInWindowFunction {
+  ROW_NUMBER = 0;
+  RANK = 1;
+  DENSE_RANK = 2;
+  LAG = 3;
+  LEAD = 4;
+  FIRST_VALUE = 5;
+  LAST_VALUE = 6;
+}
+
+message WindowExprNode {
+  oneof window_function {
+    AggregateFunction aggr_function = 1;
+    BuiltInWindowFunction built_in_function = 2;
+    // udaf = 3
+  }
+  LogicalExprNode expr = 4;
+}
+
 message BetweenNode {
   LogicalExprNode expr = 1;
   bool negated = 2;
@@ -200,6 +221,7 @@ message LogicalPlanNode {
     EmptyRelationNode empty_relation = 10;
     CreateExternalTableNode create_external_table = 11;
     ExplainNode explain = 12;
+    WindowNode window = 13;
   }
 }
 
@@ -288,6 +310,12 @@ message AggregateNode {
   repeated LogicalExprNode aggr_expr = 3;
 }
 
+message WindowNode {
+  LogicalPlanNode input = 1;
+  repeated LogicalExprNode partition_by_expr = 2;
+  repeated LogicalExprNode order_by_expr = 3;
+}
+
 enum JoinType {
   INNER = 0;
   LEFT = 1;

ballista/rust/core/src/serde/logical_plan/from_proto.rs

@@ -33,6 +33,7 @@ use datafusion::logical_plan::{
 };
 use datafusion::physical_plan::aggregates::AggregateFunction;
 use datafusion::physical_plan::csv::CsvReadOptions;
+use datafusion::physical_plan::windows::BuiltInWindowFunction;
 use datafusion::scalar::ScalarValue;
 use protobuf::logical_plan_node::LogicalPlanType;
 use protobuf::{logical_expr_node::ExprType, scalar_type};
@@ -75,6 +76,23 @@ impl TryInto<LogicalPlan> for &protobuf::LogicalPlanNode {
                     .build()
                     .map_err(|e| e.into())
             }
+            LogicalPlanType::Window(window) => {
+                let input: LogicalPlan = convert_box_required!(window.input)?;
+                let partition_by_expr = window
+                    .partition_by_expr
+                    .iter()
+                    .map(|expr| expr.try_into())
+                    .collect::<Result<Vec<_>, _>>()?;
+                let order_by_expr = window
+                    .order_by_expr
+                    .iter()
+                    .map(|expr| expr.try_into())
+                    .collect::<Result<Vec<_>, _>>()?;
+                LogicalPlanBuilder::from(&input)
+                    .window(partition_by_expr, order_by_expr)?
+                    .build()
+                    .map_err(|e| e.into())
+            }
             LogicalPlanType::Aggregate(aggregate) => {
                 let input: LogicalPlan = convert_box_required!(aggregate.input)?;
                 let group_expr = aggregate
@@ -872,6 +890,8 @@ impl TryInto<Expr> for &protobuf::LogicalExprNode {
 
     fn try_into(self) -> Result<Expr, Self::Error> {
         use protobuf::logical_expr_node::ExprType;
+        use protobuf::window_expr_node;
+        use protobuf::WindowExprNode;
 
         let expr_type = self
             .expr_type
@@ -889,6 +909,48 @@ impl TryInto<Expr> for &protobuf::LogicalExprNode {
                 let scalar_value: datafusion::scalar::ScalarValue = literal.try_into()?;
                 Ok(Expr::Literal(scalar_value))
             }
+            ExprType::WindowExpr(expr) => {
+                let window_function = expr
+                    .window_function
+                    .as_ref()
+                    .ok_or_else(|| proto_error("Received empty window function"))?;
+                match window_function {
+                    window_expr_node::WindowFunction::AggrFunction(i) => {
+                        let aggr_function = protobuf::AggregateFunction::from_i32(*i)
+                            .ok_or_else(|| {
+                                proto_error(format!(
+                                    "Received an unknown aggregate window function: {}",
+                                    i
+                                ))
+                            })?;
+
+                        Ok(Expr::WindowFunction {
+                            fun: windows::WindowFunction::AggregateFunction(
+                                AggregateFunction::from(aggr_function),
+                            ),
+                            args: vec![parse_required_expr(&expr.expr)?],
+                        })
+                    }
+                    window_expr_node::WindowFunction::BuiltInFunction(i) => {
+                        let built_in_function =
+                            protobuf::BuiltInWindowFunction::from_i32(*i).ok_or_else(
+                                || {
+                                    proto_error(format!(
+                                        "Received an unknown aggregate window function: {}",
+                                        i
+                                    ))
+                                },
+                            )?;
+
+                        Ok(Expr::WindowFunction {
+                            fun: windows::WindowFunction::BuiltInWindowFunction(
+                                BuiltInWindowFunction::from(built_in_function),
+                            ),
+                            args: vec![parse_required_expr(&expr.expr)?],
+                        })
+                    }
+                }
+            }
             ExprType::AggregateExpr(expr) => {
                 let aggr_function =
                     protobuf::AggregateFunction::from_i32(expr.aggr_function)
@@ -898,13 +960,7 @@ impl TryInto<Expr> for &protobuf::LogicalExprNode {
                                 expr.aggr_function
                             ))
                         })?;
-                let fun = match aggr_function {
-                    protobuf::AggregateFunction::Min => AggregateFunction::Min,
-                    protobuf::AggregateFunction::Max => AggregateFunction::Max,
-                    protobuf::AggregateFunction::Sum => AggregateFunction::Sum,
-                    protobuf::AggregateFunction::Avg => AggregateFunction::Avg,
-                    protobuf::AggregateFunction::Count => AggregateFunction::Count,
-                };
+                let fun = AggregateFunction::from(aggr_function);
 
                 Ok(Expr::AggregateFunction {
                     fun,
@@ -1152,6 +1208,7 @@ impl TryInto<arrow::datatypes::Field> for &protobuf::Field {
 }
 
 use datafusion::physical_plan::datetime_expressions::{date_trunc, to_timestamp};
+use datafusion::physical_plan::{aggregates, windows};
 use datafusion::prelude::{
     array, length, lower, ltrim, md5, rtrim, sha224, sha256, sha384, sha512, trim, upper,
 };
@@ -1202,3 +1259,37 @@ fn parse_optional_expr(
         None => Ok(None),
     }
 }
+
+impl From<protobuf::AggregateFunction> for AggregateFunction {
+    fn from(aggr_function: protobuf::AggregateFunction) -> Self {
+        match aggr_function {
+            protobuf::AggregateFunction::Min => AggregateFunction::Min,
+            protobuf::AggregateFunction::Max => AggregateFunction::Max,
+            protobuf::AggregateFunction::Sum => AggregateFunction::Sum,
+            protobuf::AggregateFunction::Avg => AggregateFunction::Avg,
+            protobuf::AggregateFunction::Count => AggregateFunction::Count,
+        }
+    }
+}
+
+impl From<protobuf::BuiltInWindowFunction> for BuiltInWindowFunction {
+    fn from(built_in_function: protobuf::BuiltInWindowFunction) -> Self {
+        match built_in_function {
+            protobuf::BuiltInWindowFunction::RowNumber => {
+                BuiltInWindowFunction::RowNumber
+            }
+            protobuf::BuiltInWindowFunction::Rank => BuiltInWindowFunction::Rank,
+            protobuf::BuiltInWindowFunction::DenseRank => {
+                BuiltInWindowFunction::DenseRank
+            }
+            protobuf::BuiltInWindowFunction::Lag => BuiltInWindowFunction::Lag,
+            protobuf::BuiltInWindowFunction::Lead => BuiltInWindowFunction::Lead,
+            protobuf::BuiltInWindowFunction::FirstValue => {
+                BuiltInWindowFunction::FirstValue
+            }
+            protobuf::BuiltInWindowFunction::LastValue => {
+                BuiltInWindowFunction::LastValue
+            }
+        }
+    }
+}

ballista/rust/core/src/serde/logical_plan/to_proto.rs

@@ -31,6 +31,7 @@ use arrow::datatypes::{DataType, Schema};
 use datafusion::datasource::CsvFile;
 use datafusion::logical_plan::{Expr, JoinType, LogicalPlan};
 use datafusion::physical_plan::aggregates::AggregateFunction;
+use datafusion::physical_plan::windows::{BuiltInWindowFunction, WindowFunction};
 use datafusion::{datasource::parquet::ParquetTable, logical_plan::exprlist_to_fields};
 use protobuf::{
     arrow_type, logical_expr_node::ExprType, scalar_type, DateUnit, Field,
@@ -772,6 +773,29 @@ impl TryInto<protobuf::LogicalPlanNode> for &LogicalPlan {
                     ))),
                 })
             }
+            LogicalPlan::Window {
+                input,
+                partition_by_expr,
+                order_by_expr,
+                ..
+            } => {
+                let input: protobuf::LogicalPlanNode = input.as_ref().try_into()?;
+                Ok(protobuf::LogicalPlanNode {
+                    logical_plan_type: Some(LogicalPlanType::Window(Box::new(
+                        protobuf::WindowNode {
+                            input: Some(Box::new(input)),
+                            partition_by_expr: partition_by_expr
+                                .iter()
+                                .map(|expr| expr.try_into())
+                                .collect::<Result<Vec<_>, BallistaError>>()?,
+                            order_by_expr: order_by_expr
+                                .iter()
+                                .map(|expr| expr.try_into())
+                                .collect::<Result<Vec<_>, BallistaError>>()?,
+                        },
+                    ))),
+                })
+            }
             LogicalPlan::Aggregate {
                 input,
                 group_expr,
@@ -997,6 +1021,30 @@ impl TryInto<protobuf::LogicalExprNode> for &Expr {
                     expr_type: Some(ExprType::BinaryExpr(binary_expr)),
                 })
             }
+            Expr::WindowFunction {
+                ref fun, ref args, ..
+            } => {
+                let window_function = match fun {
+                    WindowFunction::AggregateFunction(fun) => {
+                        protobuf::window_expr_node::WindowFunction::AggrFunction(
+                            protobuf::AggregateFunction::from(fun).into(),
+                        )
+                    }
+                    WindowFunction::BuiltInWindowFunction(fun) => {
+                        protobuf::window_expr_node::WindowFunction::BuiltInFunction(
+                            protobuf::BuiltInWindowFunction::from(fun).into(),
+                        )
+                    }
+                };
+                let arg = &args[0];
+                let window_expr = Box::new(protobuf::WindowExprNode {
+                    expr: Some(Box::new(arg.try_into()?)),
+                    window_function: Some(window_function),
+                });
+                Ok(protobuf::LogicalExprNode {
+                    expr_type: Some(ExprType::WindowExpr(window_expr)),
+                })
+            }
             Expr::AggregateFunction {
                 ref fun, ref args, ..
             } => {
@@ -1178,6 +1226,32 @@ impl Into<protobuf::Schema> for &Schema {
     }
 }
 
+impl From<&AggregateFunction> for protobuf::AggregateFunction {
+    fn from(value: &AggregateFunction) -> Self {
+        match value {
+            AggregateFunction::Min => Self::Min,
+            AggregateFunction::Max => Self::Max,
+            AggregateFunction::Sum => Self::Sum,
+            AggregateFunction::Avg => Self::Avg,
+            AggregateFunction::Count => Self::Count,
+        }
+    }
+}
+
+impl From<&BuiltInWindowFunction> for protobuf::BuiltInWindowFunction {
+    fn from(value: &BuiltInWindowFunction) -> Self {
+        match value {
+            BuiltInWindowFunction::FirstValue => Self::FirstValue,
+            BuiltInWindowFunction::LastValue => Self::LastValue,
+            BuiltInWindowFunction::RowNumber => Self::RowNumber,
+            BuiltInWindowFunction::Rank => Self::Rank,
+            BuiltInWindowFunction::Lag => Self::Lag,
+            BuiltInWindowFunction::Lead => Self::Lead,
+            BuiltInWindowFunction::DenseRank => Self::DenseRank,
+        }
+    }
+}
+
 impl TryFrom<&arrow::datatypes::DataType> for protobuf::ScalarType {
     type Error = BallistaError;
     fn try_from(value: &arrow::datatypes::DataType) -> Result<Self, Self::Error> {

datafusion/src/logical_plan/builder.rs

@@ -289,6 +289,31 @@ impl LogicalPlanBuilder {
         }))
     }
 
+    /// Apply a window: partition by the `partition_by_expr` expressions
+    /// and calculating the window spec.
+    pub fn window(
+        &self,
+        partition_by_expr: impl IntoIterator<Item = Expr>,
+        order_by_expr: impl IntoIterator<Item = Expr>,
+    ) -> Result<Self> {
+        let partition_by_expr = partition_by_expr.into_iter().collect::<Vec<Expr>>();
+        let order_by_expr = order_by_expr.into_iter().collect::<Vec<Expr>>();
+
+        let all_expr = partition_by_expr.iter().chain(order_by_expr.iter());
+
+        validate_unique_names("Windows", all_expr.clone(), self.plan.schema())?;
+
+        let window_schema =
+            DFSchema::new(exprlist_to_fields(all_expr, self.plan.schema())?)?;
+
+        Ok(Self::from(&LogicalPlan::Window {
+            input: Arc::new(self.plan.clone()),
+            partition_by_expr,
+            order_by_expr,
+            schema: DFSchemaRef::new(window_schema),
+        }))
+    }
+
     /// Apply an aggregate: grouping on the `group_expr` expressions
     /// and calculating `aggr_expr` aggregates for each distinct
     /// value of the `group_expr`;