chore(cubestore): Upgrade DF: push down GlobalLimit after adding it to workers

Author: srh

rust/cubestore/cubestore/src/queryplanner/optimizations/distributed_partial_aggregate.rs

@@ -3,8 +3,11 @@ use crate::queryplanner::planning::WorkerExec;
 use crate::queryplanner::query_executor::ClusterSendExec;
 use crate::queryplanner::tail_limit::TailLimitExec;
 use crate::queryplanner::topk::AggregateTopKExec;
+use datafusion::config::ConfigOptions;
 use datafusion::error::DataFusionError;
 use datafusion::physical_expr::LexOrdering;
+use datafusion::physical_optimizer::limit_pushdown::LimitPushdown;
+use datafusion::physical_optimizer::PhysicalOptimizerRule as _;
 use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode};
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::limit::GlobalLimitExec;
@@ -192,10 +195,13 @@ pub fn ensure_partition_merge_with_acceptable_parent(
     }
 }
 
-///Add `GlobalLimitExec` behind worker node if this node has `limit` property set
-///Should be executed after all optimizations which can move `Worker` node or change it input
+/// Add `GlobalLimitExec` behind worker node if this node has `limit` property set and applies DF
+/// `LimitPushdown` optimizer. Should be executed after all optimizations which can move `Worker`
+/// node or change its input.  `config` is ignored -- we pass it to DF's `LimitPushdown` optimizer,
+/// which also ignores it (as of DF 46.0.1).
 pub fn add_limit_to_workers(
     p: Arc<dyn ExecutionPlan>,
+    config: &ConfigOptions,
 ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
     let limit_and_reverse;
     let input;
@@ -217,6 +223,7 @@ pub fn add_limit_to_workers(
         p.with_new_children(vec![limit])
     } else {
         let limit = Arc::new(GlobalLimitExec::new(input.clone(), 0, Some(limit)));
-        p.with_new_children(vec![limit])
+        let limit_optimized = LimitPushdown::new().optimize(limit, config)?;
+        p.with_new_children(vec![limit_optimized])
     }
 }

rust/cubestore/cubestore/src/queryplanner/optimizations/mod.rs

@@ -93,12 +93,13 @@ impl QueryPlanner for CubeQueryPlanner {
         ])
         .create_physical_plan(logical_plan, ctx_state)
         .await?;
-        // TODO: assert there is only a single ClusterSendExec in the plan.  Update: This is no longer true.
-        finalize_physical_plan(
+        let result = finalize_physical_plan(
             p,
             self.memory_handler.clone(),
             self.data_loaded_size.clone(),
-        )
+            ctx_state.config().options(),
+        );
+        result
     }
 }
 
@@ -145,12 +146,14 @@ fn finalize_physical_plan(
     p: Arc<dyn ExecutionPlan>,
     memory_handler: Arc<dyn MemoryHandler>,
     data_loaded_size: Option<Arc<DataLoadedSize>>,
+    config: &ConfigOptions,
 ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
     let p = rewrite_physical_plan(p, &mut |p| add_check_memory_exec(p, memory_handler.clone()))?;
     let p = if let Some(data_loaded_size) = data_loaded_size {
         rewrite_physical_plan(p, &mut |p| add_trace_data_loaded_exec(p, &data_loaded_size))?
     } else {
         p
     };
-    rewrite_physical_plan(p, &mut |p| add_limit_to_workers(p))
+    let p = rewrite_physical_plan(p, &mut |p| add_limit_to_workers(p, config))?;
+    Ok(p)
 }

rust/cubestore/cubestore/src/queryplanner/pretty_printers.rs

@@ -23,6 +23,7 @@ use datafusion_datasource::file_scan_config::FileScanConfig;
 use datafusion_datasource::memory::MemorySourceConfig;
 use datafusion_datasource::source::DataSourceExec;
 use itertools::{repeat_n, Itertools};
+use std::fmt::Write;
 use std::sync::Arc;
 
 use crate::queryplanner::check_memory::CheckMemoryExec;