DO NOT REVIEW - BENCHMARK WINDOW

datafusion/core/src/physical_planner.rs

@@ -112,13 +112,13 @@ use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule, Sessio
 
 use async_trait::async_trait;
 use datafusion_physical_plan::async_func::{AsyncFuncExec, AsyncMapper};
+use datafusion_physical_plan::windows::PartitionAggregationWindowExec;
 use futures::future::BoxFuture;
 use futures::{StreamExt, TryStreamExt};
 use indexmap::IndexSet;
 use itertools::{Itertools, multiunzip};
 use log::debug;
 use tokio::sync::Mutex;
-
 // Re-export from this module for backwards compatibility.
 pub use datafusion_session::{ExtensionPlanner, PhysicalPlanner};
 
@@ -948,9 +948,17 @@ impl DefaultPhysicalPlanner {
 
                 let uses_bounded_memory =
                     window_expr.iter().all(|e| e.uses_bounded_memory());
+                if PartitionAggregationWindowExec::is_supported(&window_expr) {
+                    Arc::new(PartitionAggregationWindowExec::try_new(
+                        input_exec,
+                        window_expr[0].partition_by().to_vec(),
+                        window_expr,
+                        None,
+                    )?) as _
+                }
                 // If all window expressions can run with bounded memory,
                 // choose the bounded window variant:
-                if uses_bounded_memory {
+                else if uses_bounded_memory {
                     Arc::new(BoundedWindowAggExec::try_new(
                         window_expr,
                         input_exec,

datafusion/physical-expr/src/window/mod.rs

@@ -25,6 +25,7 @@ pub use aggregate::PlainAggregateWindowExpr;
 pub use sliding_aggregate::SlidingAggregateWindowExpr;
 pub use standard::StandardWindowExpr;
 pub use standard_window_function_expr::StandardWindowFunctionExpr;
+pub use window_expr::AggregateWindowExpr;
 pub use window_expr::PartitionBatches;
 pub use window_expr::PartitionKey;
 pub use window_expr::PartitionWindowAggStates;

datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs

@@ -68,7 +68,9 @@ use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext};
 use datafusion_physical_plan::tree_node::PlanContext;
 use datafusion_physical_plan::union::{InterleaveExec, UnionExec, can_interleave};
 use datafusion_physical_plan::windows::WindowAggExec;
-use datafusion_physical_plan::windows::{BoundedWindowAggExec, get_best_fitting_window};
+use datafusion_physical_plan::windows::{
+    BoundedWindowAggExec, PartitionAggregationWindowExec, get_best_fitting_window,
+};
 use datafusion_physical_plan::{
     ChildSatisfactionOptions, Distribution, ExecutionPlan, InputDistributionRequirements,
     Partitioning,
@@ -1196,6 +1198,15 @@ pub fn ensure_distribution(
         )?
     {
         plan = updated_window;
+    } else if let Some(exec) = plan.downcast_ref::<PartitionAggregationWindowExec>()
+        && let Some(updated_window) = get_best_fitting_window(
+            exec.window_expr(),
+            exec.input(),
+            exec.partition_keys(),
+            None,
+        )?
+    {
+        plan = updated_window;
     };
 
     // For joins in partitioned mode, we need exact hash matching between

datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs

@@ -67,7 +67,8 @@ use datafusion_physical_plan::sorts::sort::SortExec;
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion_physical_plan::tree_node::PlanContext;
 use datafusion_physical_plan::windows::{
-    BoundedWindowAggExec, WindowAggExec, get_best_fitting_window,
+    BoundedWindowAggExec, PartitionAggregationWindowExec, WindowAggExec,
+    get_best_fitting_window,
 };
 use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties, InputOrderMode};
 
@@ -584,8 +585,22 @@ fn adjust_window_sort_removal(
             state_observer.clone(),
         )?;
         (window_expr, new_window)
+    } else if let Some(exec) = window_tree
+        .plan
+        .downcast_ref::<PartitionAggregationWindowExec>()
+    {
+        let window_expr = exec.window_expr();
+        let new_window = get_best_fitting_window(
+            window_expr,
+            child_plan,
+            exec.partition_keys(),
+            None,
+        )?;
+        (window_expr, new_window)
     } else {
-        return plan_err!("Expected WindowAggExec or BoundedWindowAggExec");
+        return plan_err!(
+            "Expected WindowAggExec, BoundedWindowAggExec or PartitionAggregationWindowExec"
+        );
     };
 
     window_tree.plan = if let Some(new_window) = new_window {
@@ -604,7 +619,14 @@ fn adjust_window_sort_removal(
         let child_plan = Arc::clone(&child_node.plan);
         window_tree.children.push(child_node);
 
-        if window_expr.iter().all(|e| e.uses_bounded_memory()) {
+        if PartitionAggregationWindowExec::is_supported(window_expr) {
+            Arc::new(PartitionAggregationWindowExec::try_new(
+                child_plan,
+                window_expr[0].partition_by().to_vec(),
+                window_expr.to_vec(),
+                None,
+            )?) as _
+        } else if window_expr.iter().all(|e| e.uses_bounded_memory()) {
             Arc::new(
                 BoundedWindowAggExec::try_new(
                     window_expr.to_vec(),

datafusion/physical-optimizer/src/utils.rs

@@ -26,7 +26,9 @@ use datafusion_physical_plan::sorts::sort::SortExec;
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion_physical_plan::tree_node::PlanContext;
 use datafusion_physical_plan::union::UnionExec;
-use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowAggExec};
+use datafusion_physical_plan::windows::{
+    BoundedWindowAggExec, PartitionAggregationWindowExec, WindowAggExec,
+};
 use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 
 /// This utility function adds a `SortExec` above an operator according to the
@@ -135,7 +137,9 @@ pub fn is_sort(plan: &Arc<dyn ExecutionPlan>) -> bool {
 /// Checks whether the given operator is a window;
 /// i.e. either a [`WindowAggExec`] or a [`BoundedWindowAggExec`].
 pub fn is_window(plan: &Arc<dyn ExecutionPlan>) -> bool {
-    plan.is::<WindowAggExec>() || plan.is::<BoundedWindowAggExec>()
+    plan.is::<WindowAggExec>()
+        || plan.is::<BoundedWindowAggExec>()
+        || plan.is::<PartitionAggregationWindowExec>()
 }
 
 /// Checks whether the given operator is a [`UnionExec`].

datafusion/physical-plan/benches/bounded_window.rs

@@ -71,7 +71,9 @@ use datafusion_functions_window::row_number::row_number_udwf;
 use datafusion_physical_expr::expressions::col;
 use datafusion_physical_expr::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
 use datafusion_physical_plan::test::TestMemoryExec;
-use datafusion_physical_plan::windows::{BoundedWindowAggExec, create_window_expr};
+use datafusion_physical_plan::windows::{
+    BoundedWindowAggExec, PartitionAggregationWindowExec, create_window_expr,
+};
 use datafusion_physical_plan::{ExecutionPlan, InputOrderMode, collect};
 
 const BATCH_SIZE: usize = 8192;
@@ -204,10 +206,25 @@ fn window_exec(
             .expect("window expr")
         })
         .collect::<Vec<_>>();
-    Arc::new(
-        BoundedWindowAggExec::try_new(window_expr, input, mode, true)
-            .expect("bounded window exec"),
-    )
+
+    if mode == InputOrderMode::Sorted
+        && PartitionAggregationWindowExec::is_supported(&window_expr)
+    {
+        Arc::new(
+            PartitionAggregationWindowExec::try_new(
+                input,
+                partitionby_exprs,
+                window_expr.to_vec(),
+                None,
+            )
+            .expect("partition aggregation window exec"),
+        )
+    } else {
+        Arc::new(
+            BoundedWindowAggExec::try_new(window_expr, input, mode, true)
+                .expect("bounded window exec"),
+        )
+    }
 }
 
 fn ts_arg() -> Vec<Arc<dyn PhysicalExpr>> {

datafusion/physical-plan/src/windows/mod.rs

@@ -18,6 +18,8 @@
 //! Physical expressions for window functions
 
 mod bounded_window_agg_exec;
+mod partition_aggregation_window_exec;
+mod partitioning_processor;
 #[cfg(feature = "proto")]
 mod proto;
 mod utils;
@@ -58,6 +60,7 @@ pub use bounded_window_agg_exec::{BoundedWindowAggExec, WindowStateObserver};
 pub use datafusion_physical_expr::window::{
     PlainAggregateWindowExpr, StandardWindowExpr, WindowExpr,
 };
+pub use partition_aggregation_window_exec::PartitionAggregationWindowExec;
 pub use window_agg_exec::WindowAggExec;
 
 /// Build field from window function and add it into schema
@@ -635,7 +638,16 @@ pub fn get_best_fitting_window(
 
     // If all window expressions can run with bounded memory, choose the
     // bounded window variant:
-    if window_expr.iter().all(|e| e.uses_bounded_memory()) {
+    if input_order_mode == InputOrderMode::Sorted
+        && PartitionAggregationWindowExec::is_supported(&window_expr)
+    {
+        Ok(Some(Arc::new(PartitionAggregationWindowExec::try_new(
+            Arc::clone(input),
+            window_expr[0].partition_by().to_vec(),
+            window_expr,
+            None,
+        )?) as _))
+    } else if window_expr.iter().all(|e| e.uses_bounded_memory()) {
         Ok(Some(Arc::new(
             BoundedWindowAggExec::try_new(
                 window_expr,