adding window spec

Author: Jiayu Liu

ballista/rust/core/proto/ballista.proto

@@ -314,6 +314,42 @@ message WindowNode {
   LogicalPlanNode input = 1;
   repeated LogicalExprNode partition_by_expr = 2;
   repeated LogicalExprNode order_by_expr = 3;
+  // "optional" keyword is stable in protoc 3.15 but prost is still on 3.14 (see https://github.com/danburkert/prost/issues/430)
+  // this syntax is ugly but is binary compatible with the "optional" keyword (see https://stackoverflow.com/questions/42622015/how-to-define-an-optional-field-in-protobuf-3)
+  oneof window_frame {
+    WindowFrame frame = 4;
+  }
+}
+
+enum WindowFrameUnits {
+  ROWS = 0;
+  RANGE = 1;
+  GROUPS = 2;
+}
+
+message WindowFrame {
+  WindowFrameUnits window_frame_units = 1;
+  WindowFrameBound start_bound = 2;
+  // "optional" keyword is stable in protoc 3.15 but prost is still on 3.14 (see https://github.com/danburkert/prost/issues/430)
+  // this syntax is ugly but is binary compatible with the "optional" keyword (see https://stackoverflow.com/questions/42622015/how-to-define-an-optional-field-in-protobuf-3)
+  oneof end_bound {
+    WindowFrameBound bound = 3;
+  }
+}
+
+enum WindowFrameBoundType {
+  CURRENT_ROW = 0;
+  PRECEDING = 1;
+  FOLLOWING = 2;
+}
+
+message WindowFrameBound {
+  WindowFrameBoundType window_frame_bound_type = 1;
+  // "optional" keyword is stable in protoc 3.15 but prost is still on 3.14 (see https://github.com/danburkert/prost/issues/430)
+  // this syntax is ugly but is binary compatible with the "optional" keyword (see https://stackoverflow.com/questions/42622015/how-to-define-an-optional-field-in-protobuf-3)
+  oneof bound_value {
+    uint64 value = 2;
+  }
 }
 
 enum JoinType {

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

@@ -17,15 +17,15 @@
 
 //! Serde code to convert from protocol buffers to Rust data structures.
 
+use crate::error::BallistaError;
+use crate::serde::{proto_error, protobuf};
+use crate::{convert_box_required, convert_required};
+use sqlparser::ast::{WindowFrame, WindowFrameBound, WindowFrameUnits};
 use std::{
     convert::{From, TryInto},
     unimplemented,
 };
 
-use crate::error::BallistaError;
-use crate::serde::{proto_error, protobuf};
-use crate::{convert_box_required, convert_required};
-
 use arrow::datatypes::{DataType, Field, Schema};
 use datafusion::logical_plan::{
     abs, acos, asin, atan, ceil, cos, exp, floor, ln, log10, log2, round, signum, sin,
@@ -88,8 +88,10 @@ impl TryInto<LogicalPlan> for &protobuf::LogicalPlanNode {
                     .iter()
                     .map(|expr| expr.try_into())
                     .collect::<Result<Vec<_>, _>>()?;
+                // FIXME parse window frame
+                let window_frame = None;
                 LogicalPlanBuilder::from(&input)
-                    .window(partition_by_expr, order_by_expr)?
+                    .window(partition_by_expr, order_by_expr, window_frame)?
                     .build()
                     .map_err(|e| e.into())
             }
@@ -1260,6 +1262,72 @@ fn parse_optional_expr(
     }
 }
 
+impl From<protobuf::WindowFrameUnits> for WindowFrameUnits {
+    fn from(units: protobuf::WindowFrameUnits) -> Self {
+        match units {
+            protobuf::WindowFrameUnits::Rows => WindowFrameUnits::Rows,
+            protobuf::WindowFrameUnits::Range => WindowFrameUnits::Range,
+            protobuf::WindowFrameUnits::Groups => WindowFrameUnits::Groups,
+        }
+    }
+}
+
+impl TryFrom<protobuf::WindowFrameBound> for WindowFrameBound {
+    type Error = BallistaError;
+
+    fn try_from(bound: protobuf::WindowFrameBound) -> Result<Self, Self::Error> {
+        let bound_type = protobuf::WindowFrameBoundType::from_i32(bound.window_frame_bound_type).ok_or_else(|| {
+            proto_error(format!(
+                "Received a WindowFrameBound message with unknown WindowFrameBoundType {}",
+                bound.window_frame_bound_type
+            ))
+        })?.into();
+        match bound_type {
+            protobuf::WindowFrameBoundType::CurrentRow => {
+                Ok(WindowFrameBound::CurrentRow)
+            }
+            protobuf::WindowFrameBoundType::Preceding => {
+                // FIXME implement bound value parsing
+                Ok(WindowFrameBound::Preceding(Some(1)))
+            }
+            protobuf::WindowFrameBoundType::Following => {
+                // FIXME implement bound value parsing
+                Ok(WindowFrameBound::Following(Some(1)))
+            }
+        }
+    }
+}
+
+impl TryFrom<protobuf::WindowFrame> for WindowFrame {
+    type Error = BallistaError;
+
+    fn try_from(window: protobuf::WindowFrame) -> Result<Self, Self::Error> {
+        let units = protobuf::WindowFrameUnits::from_i32(window.window_frame_units)
+            .ok_or_else(|| {
+                proto_error(format!(
+                    "Received a WindowFrame message with unknown WindowFrameUnits {}",
+                    window.window_frame_units
+                ))
+            })?
+            .into();
+        let start_bound = window
+            .start_bound
+            .ok_or_else(|| {
+                proto_error(
+                    "Received a WindowFrame message with no start_bound".to_owned(),
+                )
+            })?
+            .try_into()?;
+        // FIXME parse end bound
+        let end_bound = None;
+        Ok(WindowFrame {
+            units,
+            start_bound,
+            end_bound,
+        })
+    }
+}
+
 impl From<protobuf::AggregateFunction> for AggregateFunction {
     fn from(aggr_function: protobuf::AggregateFunction) -> Self {
         match aggr_function {

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

@@ -26,7 +26,6 @@ use std::{
 
 use crate::datasource::DfTableAdapter;
 use crate::serde::{protobuf, BallistaError};
-
 use arrow::datatypes::{DataType, Schema};
 use datafusion::datasource::CsvFile;
 use datafusion::logical_plan::{Expr, JoinType, LogicalPlan};
@@ -37,6 +36,7 @@ use protobuf::{
     arrow_type, logical_expr_node::ExprType, scalar_type, DateUnit, Field,
     PrimitiveScalarType, ScalarListValue, ScalarType,
 };
+use sqlparser::ast::{WindowFrame, WindowFrameBound, WindowFrameUnits};
 
 use super::super::proto_error;
 use datafusion::physical_plan::functions::BuiltinScalarFunction;
@@ -777,9 +777,12 @@ impl TryInto<protobuf::LogicalPlanNode> for &LogicalPlan {
                 input,
                 partition_by_expr,
                 order_by_expr,
+                window_frame,
                 ..
             } => {
                 let input: protobuf::LogicalPlanNode = input.as_ref().try_into()?;
+                // FIXME: implement
+                let window_frame = None;
                 Ok(protobuf::LogicalPlanNode {
                     logical_plan_type: Some(LogicalPlanType::Window(Box::new(
                         protobuf::WindowNode {
@@ -792,6 +795,7 @@ impl TryInto<protobuf::LogicalPlanNode> for &LogicalPlan {
                                 .iter()
                                 .map(|expr| expr.try_into())
                                 .collect::<Result<Vec<_>, BallistaError>>()?,
+                            window_frame,
                         },
                     ))),
                 })
@@ -1252,6 +1256,32 @@ impl From<&BuiltInWindowFunction> for protobuf::BuiltInWindowFunction {
     }
 }
 
+impl From<WindowFrameUnits> for protobuf::WindowFrameUnits {
+    fn from(units: WindowFrameUnits) -> Self {
+        match units {
+            WindowFrameUnits::Rows => protobuf::WindowFrameUnits::Rows,
+            WindowFrameUnits::Range => protobuf::WindowFrameUnits::Range,
+            WindowFrameUnits::Groups => protobuf::WindowFrameUnits::Groups,
+        }
+    }
+}
+
+impl TryFrom<WindowFrameBound> for protobuf::WindowFrameBound {
+    type Error = BallistaError;
+
+    fn try_from(bound: WindowFrameBound) -> Result<Self, Self::Error> {
+        unimplemented!("not implemented")
+    }
+}
+
+impl TryFrom<WindowFrame> for protobuf::WindowFrame {
+    type Error = BallistaError;
+
+    fn try_from(window: WindowFrame) -> Result<Self, Self::Error> {
+        unimplemented!("not implemented")
+    }
+}
+
 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

@@ -24,18 +24,18 @@ use arrow::{
     record_batch::RecordBatch,
 };
 
+use super::dfschema::ToDFSchema;
+use super::{
+    col, exprlist_to_fields, Expr, JoinType, LogicalPlan, PlanType, StringifiedPlan,
+};
 use crate::datasource::TableProvider;
 use crate::error::{DataFusionError, Result};
+use crate::logical_plan::{DFField, DFSchema, DFSchemaRef, Partitioning};
 use crate::{
     datasource::{empty::EmptyTable, parquet::ParquetTable, CsvFile, MemTable},
     prelude::CsvReadOptions,
 };
-
-use super::dfschema::ToDFSchema;
-use super::{
-    col, exprlist_to_fields, Expr, JoinType, LogicalPlan, PlanType, StringifiedPlan,
-};
-use crate::logical_plan::{DFField, DFSchema, DFSchemaRef, Partitioning};
+use sqlparser::ast::WindowFrame;
 use std::collections::HashSet;
 
 /// Builder for logical plans
@@ -295,21 +295,19 @@ impl LogicalPlanBuilder {
         &self,
         partition_by_expr: impl IntoIterator<Item = Expr>,
         order_by_expr: impl IntoIterator<Item = Expr>,
+        window_frame: Option<WindowFrame>,
     ) -> 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,
+            window_frame,
             schema: DFSchemaRef::new(window_schema),
         }))
     }

datafusion/src/logical_plan/expr.rs

@@ -197,7 +197,6 @@ pub enum Expr {
         fun: windows::WindowFunction,
         /// List of expressions to feed to the functions as arguments
         args: Vec<Expr>,
-        // TODO windowexec
     },
     /// aggregate function
     AggregateUDF {

datafusion/src/logical_plan/plan.rs

@@ -23,10 +23,10 @@ use std::{
     sync::Arc,
 };
 
-use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
-
 use crate::datasource::TableProvider;
 use crate::sql::parser::FileType;
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use sqlparser::ast::WindowFrame;
 
 use super::expr::Expr;
 use super::extension::UserDefinedLogicalNode;
@@ -91,6 +91,8 @@ pub enum LogicalPlan {
         partition_by_expr: Vec<Expr>,
         /// Order by expressions
         order_by_expr: Vec<Expr>,
+        /// Window Frame
+        window_frame: Option<WindowFrame>,
         /// The schema description of the window output
         schema: DFSchemaRef,
     },

datafusion/src/optimizer/projection_push_down.rs

@@ -198,6 +198,7 @@ fn optimize_plan(
             input,
             partition_by_expr,
             order_by_expr,
+            window_frame,
             ..
         } => {
             let new_schema = DFSchema::new(
@@ -211,6 +212,7 @@ fn optimize_plan(
             Ok(LogicalPlan::Window {
                 partition_by_expr: partition_by_expr.clone(),
                 order_by_expr: order_by_expr.clone(),
+                window_frame: window_frame.clone(),
                 input: Arc::new(optimize_plan(
                     optimizer,
                     &input,

datafusion/src/optimizer/utils.rs

@@ -192,11 +192,13 @@ pub fn from_plan(
         LogicalPlan::Window {
             partition_by_expr,
             order_by_expr,
+            window_frame,
             schema,
             ..
         } => Ok(LogicalPlan::Window {
             partition_by_expr: expr[0..partition_by_expr.len()].to_vec(),
             order_by_expr: expr[order_by_expr.len()..].to_vec(),
+            window_frame: window_frame.clone(),
             input: Arc::new(inputs[0].clone()),
             schema: schema.clone(),
         }),