[arrow-datafusion] branch master updated: Make window function related struct public for ballista. (#4511)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/master by this push:
     new 73527e361 Make window function related struct public for ballista. (#4511)
73527e361 is described below

commit 73527e36136ee0e557545e2b0f636775d99e79ab
Author: Yang Jiang <ya...@ebay.com>
AuthorDate: Mon Dec 5 22:01:22 2022 +0800

    Make window function related struct public for ballista. (#4511)
    
    Signed-off-by: yangjiang <ya...@ebay.com>
    
    Signed-off-by: yangjiang <ya...@ebay.com>
---
 datafusion/physical-expr/src/expressions/mod.rs  | 5 ++++-
 datafusion/physical-expr/src/window/aggregate.rs | 5 +++++
 datafusion/physical-expr/src/window/built_in.rs  | 5 +++++
 datafusion/physical-expr/src/window/lead_lag.rs  | 7 +++++++
 datafusion/physical-expr/src/window/nth_value.rs | 7 ++++++-
 datafusion/physical-expr/src/window/rank.rs      | 9 ++++++++-
 6 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs
index 26bb3ca1e..8222fb664 100644
--- a/datafusion/physical-expr/src/expressions/mod.rs
+++ b/datafusion/physical-expr/src/expressions/mod.rs
@@ -62,9 +62,12 @@ pub use crate::aggregate::sum_distinct::DistinctSum;
 pub use crate::aggregate::variance::{Variance, VariancePop};
 
 pub use crate::window::cume_dist::cume_dist;
+pub use crate::window::cume_dist::CumeDist;
+pub use crate::window::lead_lag::WindowShift;
 pub use crate::window::lead_lag::{lag, lead};
-pub use crate::window::nth_value::NthValue;
+pub use crate::window::nth_value::{NthValue, NthValueKind};
 pub use crate::window::rank::{dense_rank, percent_rank, rank};
+pub use crate::window::rank::{Rank, RankType};
 pub use crate::window::row_number::RowNumber;
 
 pub use binary::{binary, BinaryExpr};
diff --git a/datafusion/physical-expr/src/window/aggregate.rs b/datafusion/physical-expr/src/window/aggregate.rs
index 17a19b6d7..f835c40a6 100644
--- a/datafusion/physical-expr/src/window/aggregate.rs
+++ b/datafusion/physical-expr/src/window/aggregate.rs
@@ -59,6 +59,11 @@ impl AggregateWindowExpr {
             window_frame,
         }
     }
+
+    /// Get aggregate expr of AggregateWindowExpr
+    pub fn get_aggregate_expr(&self) -> &Arc<dyn AggregateExpr> {
+        &self.aggregate
+    }
 }
 
 /// peer based evaluation based on the fact that batch is pre-sorted given the sort columns
diff --git a/datafusion/physical-expr/src/window/built_in.rs b/datafusion/physical-expr/src/window/built_in.rs
index adf8d5b34..314905bae 100644
--- a/datafusion/physical-expr/src/window/built_in.rs
+++ b/datafusion/physical-expr/src/window/built_in.rs
@@ -55,6 +55,11 @@ impl BuiltInWindowExpr {
             window_frame,
         }
     }
+
+    /// Get BuiltInWindowFunction expr of BuiltInWindowExpr
+    pub fn get_built_in_func_expr(&self) -> &Arc<dyn BuiltInWindowFunctionExpr> {
+        &self.expr
+    }
 }
 
 impl WindowExpr for BuiltInWindowExpr {
diff --git a/datafusion/physical-expr/src/window/lead_lag.rs b/datafusion/physical-expr/src/window/lead_lag.rs
index ef4175d13..c7fc73b9f 100644
--- a/datafusion/physical-expr/src/window/lead_lag.rs
+++ b/datafusion/physical-expr/src/window/lead_lag.rs
@@ -41,6 +41,13 @@ pub struct WindowShift {
     default_value: Option<ScalarValue>,
 }
 
+impl WindowShift {
+    /// Get shift_offset of window shift expression
+    pub fn get_shift_offset(&self) -> i64 {
+        self.shift_offset
+    }
+}
+
 /// lead() window function
 pub fn lead(
     name: String,
diff --git a/datafusion/physical-expr/src/window/nth_value.rs b/datafusion/physical-expr/src/window/nth_value.rs
index e0afb520e..63a2354c9 100644
--- a/datafusion/physical-expr/src/window/nth_value.rs
+++ b/datafusion/physical-expr/src/window/nth_value.rs
@@ -31,7 +31,7 @@ use std::sync::Arc;
 
 /// nth_value kind
 #[derive(Debug, Copy, Clone)]
-enum NthValueKind {
+pub enum NthValueKind {
     First,
     Last,
     Nth(u32),
@@ -94,6 +94,11 @@ impl NthValue {
             }),
         }
     }
+
+    /// Get nth_value kind
+    pub fn get_kind(&self) -> NthValueKind {
+        self.kind
+    }
 }
 
 impl BuiltInWindowFunctionExpr for NthValue {
diff --git a/datafusion/physical-expr/src/window/rank.rs b/datafusion/physical-expr/src/window/rank.rs
index 3447d47b3..8ed0319a1 100644
--- a/datafusion/physical-expr/src/window/rank.rs
+++ b/datafusion/physical-expr/src/window/rank.rs
@@ -37,8 +37,15 @@ pub struct Rank {
     rank_type: RankType,
 }
 
+impl Rank {
+    /// Get rank_type of the rank in window function with order by
+    pub fn get_type(&self) -> RankType {
+        self.rank_type
+    }
+}
+
 #[derive(Debug, Copy, Clone)]
-pub(crate) enum RankType {
+pub enum RankType {
     Basic,
     Dense,
     Percent,