Implement proto serialization for (Bounded)WindowAggExec. (#7557)

Signed-off-by: Vaibhav <[email protected]>
apache · Sep 22, 2023 · 1c847b4 · 1c847b4
1 parent 70b8620
commit 1c847b4
Show file tree

Hide file tree

Showing 11 changed files with 870 additions and 195 deletions.
diff --git a/datafusion/physical-expr/src/aggregate/regr.rs b/datafusion/physical-expr/src/aggregate/regr.rs
@@ -43,6 +43,12 @@ pub struct Regr {
     expr_x: Arc<dyn PhysicalExpr>,
 }
 
+impl Regr {
+    pub fn get_regr_type(&self) -> RegrType {
+        self.regr_type.clone()
+    }
+}
+
 #[derive(Debug, Clone)]
 #[allow(clippy::upper_case_acronyms)]
 pub enum RegrType {

diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs
@@ -59,7 +59,7 @@ pub use crate::aggregate::grouping::Grouping;
 pub use crate::aggregate::median::Median;
 pub use crate::aggregate::min_max::{Max, Min};
 pub use crate::aggregate::min_max::{MaxAccumulator, MinAccumulator};
-pub use crate::aggregate::regr::Regr;
+pub use crate::aggregate::regr::{Regr, RegrType};
 pub use crate::aggregate::stats::StatsType;
 pub use crate::aggregate::stddev::{Stddev, StddevPop};
 pub use crate::aggregate::sum::Sum;

diff --git a/datafusion/physical-expr/src/window/lead_lag.rs b/datafusion/physical-expr/src/window/lead_lag.rs
@@ -46,6 +46,11 @@ impl WindowShift {
     pub fn get_shift_offset(&self) -> i64 {
         self.shift_offset
     }
+
+    /// Get the default_value for window shift expression.
+    pub fn get_default_value(&self) -> Option<ScalarValue> {
+        self.default_value.clone()
+    }
 }
 
 /// lead() window function

diff --git a/datafusion/physical-expr/src/window/mod.rs b/datafusion/physical-expr/src/window/mod.rs
@@ -31,6 +31,7 @@ pub use aggregate::PlainAggregateWindowExpr;
 pub use built_in::BuiltInWindowExpr;
 pub use built_in_window_function_expr::BuiltInWindowFunctionExpr;
 pub use sliding_aggregate::SlidingAggregateWindowExpr;
+pub use window_expr::NthValueKind;
 pub use window_expr::PartitionBatches;
 pub use window_expr::PartitionKey;
 pub use window_expr::PartitionWindowAggStates;

diff --git a/datafusion/physical-expr/src/window/ntile.rs b/datafusion/physical-expr/src/window/ntile.rs
@@ -41,6 +41,10 @@ impl Ntile {
     pub fn new(name: String, n: u64) -> Self {
         Self { name, n }
     }
+
+    pub fn get_n(&self) -> u64 {
+        self.n
+    }
 }
 
 impl BuiltInWindowFunctionExpr for Ntile {

diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto
@@ -1183,7 +1183,11 @@ message PhysicalWindowExprNode {
     BuiltInWindowFunction built_in_function = 2;
     // udaf = 3
   }
-  PhysicalExprNode expr = 4;
+  repeated PhysicalExprNode args = 4;
+  repeated PhysicalExprNode partition_by = 5;
+  repeated PhysicalSortExprNode order_by = 6;
+  WindowFrame window_frame = 7;
+  string name = 8;
 }
 
 message PhysicalIsNull {
@@ -1388,11 +1392,21 @@ enum AggregateMode {
   SINGLE_PARTITIONED = 4;
 }
 
+message PartiallySortedPartitionSearchMode {
+  repeated uint64 columns = 6;
+}
+
 message WindowAggExecNode {
   PhysicalPlanNode input = 1;
-  repeated PhysicalExprNode window_expr = 2;
-  repeated string window_expr_name = 3;
+  repeated PhysicalWindowExprNode window_expr = 2;
   Schema input_schema = 4;
+  repeated PhysicalExprNode partition_keys = 5;
+  // Set optional to `None` for `BoundedWindowAggExec`.
+  oneof partition_search_mode {
+    EmptyMessage linear = 7;
+    PartiallySortedPartitionSearchMode partially_sorted = 8;
+    EmptyMessage sorted = 9;
+  }
 }
 
 message MaybeFilter {