fix(cubestore): Allowing window functions in cluster send planning

Author: waralexrom

rust/cubestore/cubestore/src/queryplanner/planning.rs

@@ -25,8 +25,7 @@ use datafusion::arrow::datatypes::Field;
 use datafusion::error::DataFusionError;
 use datafusion::physical_plan::empty::EmptyExec;
 use datafusion::physical_plan::{
-    DisplayAs, DisplayFormatType, ExecutionPlan,
-    PlanProperties, SendableRecordBatchStream,
+    DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream,
 };
 use flatbuffers::bitflags::_core::any::Any;
 use flatbuffers::bitflags::_core::fmt::Formatter;
@@ -47,12 +46,12 @@ use crate::queryplanner::providers::InfoSchemaQueryCacheTableProvider;
 use crate::queryplanner::query_executor::{ClusterSendExec, CubeTable, InlineTableProvider};
 use crate::queryplanner::rolling::RollingWindowAggregateSerialized;
 use crate::queryplanner::serialized_plan::PreSerializedPlan;
-use crate::queryplanner::serialized_plan::{
-    IndexSnapshot, InlineSnapshot, PartitionSnapshot,
+use crate::queryplanner::serialized_plan::{IndexSnapshot, InlineSnapshot, PartitionSnapshot};
+use crate::queryplanner::topk::{
+    materialize_topk, ClusterAggregateTopKLowerSerialized, ClusterAggregateTopKUpperSerialized,
 };
 use crate::queryplanner::topk::{plan_topk, DummyTopKLowerExec};
-use crate::queryplanner::topk::{ClusterAggregateTopKUpper, ClusterAggregateTopKLower};
-use crate::queryplanner::topk::{materialize_topk, ClusterAggregateTopKUpperSerialized, ClusterAggregateTopKLowerSerialized};
+use crate::queryplanner::topk::{ClusterAggregateTopKLower, ClusterAggregateTopKUpper};
 use crate::queryplanner::{CubeTableLogical, InfoSchemaTableProvider};
 use crate::table::{cmp_same_types, Row};
 use crate::CubeError;
@@ -64,7 +63,9 @@ use datafusion::execution::{SessionState, TaskContext};
 use datafusion::logical_expr::expr::Alias;
 use datafusion::logical_expr::utils::expr_to_columns;
 use datafusion::logical_expr::{
-    expr, logical_plan, Aggregate, BinaryExpr, Expr, Extension, FetchType, Filter, InvariantLevel, Join, Limit, LogicalPlan, Operator, Projection, SkipType, Sort, SortExpr, SubqueryAlias, TableScan, Union, Unnest, UserDefinedLogicalNode
+    expr, logical_plan, Aggregate, BinaryExpr, Expr, Extension, FetchType, Filter, InvariantLevel,
+    Join, Limit, LogicalPlan, Operator, Projection, SkipType, Sort, SortExpr, SubqueryAlias,
+    TableScan, Union, Unnest, UserDefinedLogicalNode,
 };
 use datafusion::physical_expr::{Distribution, LexRequirement};
 use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner};
@@ -1027,7 +1028,10 @@ fn check_aggregates_expr(table: &IdRow<Table>, aggregates: &Vec<Expr>) -> bool {
 
     for aggr in aggregates.iter() {
         match aggr {
-            Expr::AggregateFunction(expr::AggregateFunction { func, params: expr::AggregateFunctionParams { args, .. } }) => {
+            Expr::AggregateFunction(expr::AggregateFunction {
+                func,
+                params: expr::AggregateFunctionParams { args, .. },
+            }) => {
                 if args.len() != 1 {
                     return false;
                 }
@@ -1529,6 +1533,7 @@ fn pull_up_cluster_send(mut p: LogicalPlan) -> Result<LogicalPlan, DataFusionErr
         LogicalPlan::Extension { .. } => return Ok(p),
         // These nodes collect results from multiple partitions, return unchanged.
         LogicalPlan::Aggregate { .. }
+        | LogicalPlan::Window { .. }
         | LogicalPlan::Sort { .. }
         | LogicalPlan::Limit { .. }
         | LogicalPlan::Repartition { .. } => return Ok(p),
@@ -1710,7 +1715,8 @@ impl ExtensionPlanner for CubeExtensionPlanner {
         } else if let Some(topk_upper) = node.as_any().downcast_ref::<ClusterAggregateTopKUpper>() {
             assert_eq!(inputs.len(), 1);
             assert_eq!(logical_inputs.len(), 1);
-            let msg: &'static str = "ClusterAggregateTopKUpper expects its child to be a ClusterAggregateTopKLower";
+            let msg: &'static str =
+                "ClusterAggregateTopKUpper expects its child to be a ClusterAggregateTopKLower";
             let LogicalPlan::Extension(Extension { node }) = logical_inputs[0] else {
                 return Err(DataFusionError::Internal(msg.to_owned()));
             };
@@ -1719,11 +1725,22 @@ impl ExtensionPlanner for CubeExtensionPlanner {
             };
 
             // The input should be (and must be) a DummyTopKLowerExec node.
-            let Some(DummyTopKLowerExec { schema: _, input: lower_input }) = inputs[0].as_any().downcast_ref::<DummyTopKLowerExec>() else {
+            let Some(DummyTopKLowerExec {
+                schema: _,
+                input: lower_input,
+            }) = inputs[0].as_any().downcast_ref::<DummyTopKLowerExec>()
+            else {
                 return Err(DataFusionError::Internal("ClusterAggregateTopKUpper expects its physical input to be a DummyTopKLowerExec".to_owned()));
             };
 
-            Ok(Some(plan_topk(planner, self, topk_upper, lower_node, lower_input.clone(), state)?))
+            Ok(Some(plan_topk(
+                planner,
+                self,
+                topk_upper,
+                lower_node,
+                lower_input.clone(),
+                state,
+            )?))
         } else if let Some(_) = node.as_any().downcast_ref::<PanicWorkerNode>() {
             assert_eq!(inputs.len(), 0);
             Ok(Some(plan_panic_worker()?))
@@ -1916,7 +1933,9 @@ pub mod tests {
     use crate::queryplanner::pretty_printers::PPOptions;
     use crate::queryplanner::query_executor::ClusterSendExec;
     use crate::queryplanner::serialized_plan::RowRange;
-    use crate::queryplanner::{pretty_printers, sql_to_rel_options, CubeTableLogical, QueryPlannerImpl};
+    use crate::queryplanner::{
+        pretty_printers, sql_to_rel_options, CubeTableLogical, QueryPlannerImpl,
+    };
     use crate::sql::parser::{CubeStoreParser, Statement};
     use crate::table::{Row, TableValue};
     use crate::CubeError;