chore(cubestore): Upgrade DF: Bugfix from topk: Correct compute_properties in WorkerExec

Author: srh

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

@@ -1757,13 +1757,11 @@ impl WorkerExec {
         required_input_ordering: Option<LexRequirement>,
         worker_planning_params: WorkerPlanningParams,
     ) -> WorkerExec {
-        let properties =
-            input
-                .properties()
-                .clone()
-                .with_partitioning(Partitioning::UnknownPartitioning(
-                    worker_planning_params.worker_partition_count,
-                ));
+        // This, importantly, gives us the same PlanProperties as ClusterSendExec.
+        let properties = ClusterSendExec::compute_properties(
+            input.properties(),
+            worker_planning_params.worker_partition_count,
+        );
         WorkerExec {
             input,
             max_batch_rows,
@@ -1796,12 +1794,16 @@ impl ExecutionPlan for WorkerExec {
     ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
         assert_eq!(children.len(), 1);
         let input = children.into_iter().next().unwrap();
+        let properties: PlanProperties = ClusterSendExec::compute_properties(
+            input.properties(),
+            self.properties.output_partitioning().partition_count(),
+        );
         Ok(Arc::new(WorkerExec {
             input,
             max_batch_rows: self.max_batch_rows,
             limit_and_reverse: self.limit_and_reverse.clone(),
             required_input_ordering: self.required_input_ordering.clone(),
-            properties: self.properties.clone(),
+            properties,
         }))
     }
 
@@ -1831,7 +1833,7 @@ impl ExecutionPlan for WorkerExec {
 
     fn maintains_input_order(&self) -> Vec<bool> {
         // TODO upgrade DF: If the WorkerExec has the number of partitions so it can produce the same output, we could occasionally return true.
-        // vec![self.num_clustersend_partitions <= 1 && self.input_for_optimizations.output_partitioning().partition_count() <= 1]
+        // vec![self.input_for_optimizations.output_partitioning().partition_count() <= 1]
 
         // For now, same as default implementation:
         vec![false]
@@ -1883,7 +1885,7 @@ pub mod tests {
     use datafusion::error::DataFusionError;
     use datafusion::execution::{SessionState, SessionStateBuilder};
     use datafusion::logical_expr::{AggregateUDF, LogicalPlan, ScalarUDF, TableSource, WindowUDF};
-    use datafusion::prelude::{SessionConfig, SessionContext};
+    use datafusion::prelude::SessionConfig;
     use datafusion::sql::TableReference;
     use std::collections::HashMap;
     use std::iter::FromIterator;

rust/cubestore/cubestore/src/queryplanner/query_executor.rs

@@ -1307,12 +1307,19 @@ impl ClusterSendExec {
         })
     }
 
+    /// Also used by WorkerExec (to produce the exact same plan properties so we get the same optimizations).
     pub fn compute_properties(
         input_properties: &PlanProperties,
         partitions_num: usize,
     ) -> PlanProperties {
+        // Coalescing partitions (on the worker side) loses existing orderings:
+        let mut eq_properties = input_properties.eq_properties.clone();
+        if input_properties.output_partitioning().partition_count() > 1 {
+            eq_properties.clear_orderings();
+            eq_properties.clear_per_partition_constants();
+        }
         PlanProperties::new(
-            input_properties.eq_properties.clone(),
+            eq_properties,
             Partitioning::UnknownPartitioning(partitions_num),
             input_properties.execution_mode.clone(),
         )
@@ -1685,6 +1692,7 @@ impl ExecutionPlan for ClusterSendExec {
     }
 
     fn required_input_distribution(&self) -> Vec<Distribution> {
+        // TODO:  If this is in place, and it is obeyed (with EnforceDistribution?), then we don't need to use a CoalescePartitions node in worker exec.
         vec![Distribution::SinglePartition; self.children().len()]
     }
 }