add get_window_frame in window_expr, show frame info in window_agg_exec

Signed-off-by: yangjiang <[email protected]>
apache · Dec 5, 2022 · e87a000 · e87a000
1 parent 34d9bb5
commit e87a000
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 3 deletions.
diff --git a/datafusion/core/src/physical_plan/windows/window_agg_exec.rs b/datafusion/core/src/physical_plan/windows/window_agg_exec.rs
@@ -268,7 +268,14 @@ impl ExecutionPlan for WindowAggExec {
                 let g: Vec<String> = self
                     .window_expr
                     .iter()
-                    .map(|e| format!("{}: {:?}", e.name().to_owned(), e.field()))
+                    .map(|e| {
+                        format!(
+                            "{}: {:?}, frame: {:?}",
+                            e.name().to_owned(),
+                            e.field(),
+                            e.get_window_frame()
+                        )
+                    })
                     .collect();
                 write!(f, "wdw=[{}]", g.join(", "))?;
             }

diff --git a/datafusion/core/tests/sql/window.rs b/datafusion/core/tests/sql/window.rs
@@ -1628,8 +1628,8 @@ async fn test_window_agg_sort() -> Result<()> {
     let expected = {
         vec![
             "ProjectionExec: expr=[c9@3 as c9, SUM(aggregate_test_100.c9) ORDER BY [aggregate_test_100.c9 ASC NULLS LAST]@0 as sum1, SUM(aggregate_test_100.c9) ORDER BY [aggregate_test_100.c9 ASC NULLS LAST, aggregate_test_100.c8 ASC NULLS LAST]@1 as sum2]",
-            "  WindowAggExec: wdw=[SUM(aggregate_test_100.c9): Ok(Field { name: \"SUM(aggregate_test_100.c9)\", data_type: UInt64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} })]",
-            "    WindowAggExec: wdw=[SUM(aggregate_test_100.c9): Ok(Field { name: \"SUM(aggregate_test_100.c9)\", data_type: UInt64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} })]",
+            "  WindowAggExec: wdw=[SUM(aggregate_test_100.c9): Ok(Field { name: \"SUM(aggregate_test_100.c9)\", data_type: UInt64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: None]",
+            "    WindowAggExec: wdw=[SUM(aggregate_test_100.c9): Ok(Field { name: \"SUM(aggregate_test_100.c9)\", data_type: UInt64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: None]",
             "      SortExec: [c9@1 ASC NULLS LAST,c8@0 ASC NULLS LAST]",
         ]
     };

diff --git a/datafusion/physical-expr/src/window/aggregate.rs b/datafusion/physical-expr/src/window/aggregate.rs
@@ -154,4 +154,8 @@ impl WindowExpr for AggregateWindowExpr {
     fn order_by(&self) -> &[PhysicalSortExpr] {
         &self.order_by
     }
+
+    fn get_window_frame(&self) -> Option<&Arc<WindowFrame>> {
+        self.window_frame.as_ref()
+    }
 }
diff --git a/datafusion/physical-expr/src/window/built_in.rs b/datafusion/physical-expr/src/window/built_in.rs
@@ -136,4 +136,8 @@ impl WindowExpr for BuiltInWindowExpr {
         let results = results.iter().map(|i| i.as_ref()).collect::<Vec<_>>();
         concat(&results).map_err(DataFusionError::ArrowError)
     }
+
+    fn get_window_frame(&self) -> Option<&Arc<WindowFrame>> {
+        self.window_frame.as_ref()
+    }
 }
diff --git a/datafusion/physical-expr/src/window/window_expr.rs b/datafusion/physical-expr/src/window/window_expr.rs
@@ -21,6 +21,7 @@ use arrow::compute::kernels::sort::{SortColumn, SortOptions};
 use arrow::record_batch::RecordBatch;
 use arrow::{array::ArrayRef, datatypes::Field};
 use datafusion_common::{DataFusionError, Result};
+use datafusion_expr::WindowFrame;
 use std::any::Any;
 use std::fmt::Debug;
 use std::ops::Range;
@@ -127,4 +128,7 @@ pub trait WindowExpr: Send + Sync + Debug {
             order_by_columns.iter().map(|s| s.values.clone()).collect();
         Ok((values, order_bys))
     }
+
+    // Get window frame of this WindowExpr, None if absent
+    fn get_window_frame(&self) -> Option<&Arc<WindowFrame>>;
 }