chore(cubestore): Upgrade DF: Pretty-print sort column improvements for MergeSort and TopK Aggregates

Author: srh

rust/cubestore/cubestore-sql-tests/src/tests.rs

@@ -3403,13 +3403,13 @@ async fn planning_inplace_aggregate2(service: Box<dyn SqlClient>) {
     assert_eq!(
         pp_phys_plan_ext(p.router.as_ref(), &verbose),
         "Projection, [url, sum(Data.hits)@1:hits]\
-           \n  AggregateTopK, limit: 10, sortBy: [2 desc null last]\
+           \n  AggregateTopK, limit: 10, sortBy: [2 desc nulls last]\
            \n    ClusterSend, partitions: [[1, 2]], sort_order: [1]"
     );
     assert_eq!(
         pp_phys_plan_ext(p.worker.as_ref(), &verbose),
         "Projection, [url, sum(Data.hits)@1:hits]\
-           \n  AggregateTopK, limit: 10, sortBy: [2 desc null last]\
+           \n  AggregateTopK, limit: 10, sortBy: [2 desc nulls last]\
            \n    Worker, sort_order: [1]\
            \n      Sort, by: [sum(Data.hits)@1 desc nulls last], sort_order: [1]\
            \n        LinearSingleAggregate\

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

@@ -2229,7 +2229,7 @@ pub mod tests {
         assert_eq!(
             pretty_printers::pp_plan_ext(&plan, &with_sort_by),
             "Projection, [amount, customer]\
-           \n  ClusterAggregateTopK, limit: 10, sortBy: [2 desc null last]\
+           \n  ClusterAggregateTopK, limit: 10, sortBy: [2 desc nulls last]\
            \n    Scan s.Orders, source: CubeTable(index: by_customer:3:[]:sort_on[order_customer]), fields: [order_customer, order_amount]"
         );
 
@@ -2243,7 +2243,7 @@ pub mod tests {
         assert_eq!(
             pretty_printers::pp_plan_ext(&plan, &with_sort_by),
             "Projection, [customer, amount]\
-           \n  ClusterAggregateTopK, limit: 10, sortBy: [2 null last]\
+           \n  ClusterAggregateTopK, limit: 10, sortBy: [2 nulls last]\
            \n    Scan s.Orders, source: CubeTable(index: by_customer:3:[]:sort_on[order_customer]), fields: [order_customer, order_amount]"
         );
 
@@ -2261,7 +2261,7 @@ pub mod tests {
         assert_eq!(
             pretty_printers::pp_plan_ext(&plan, &verbose),
             "Projection, [customer, amount, min_amount, max_amount]\
-           \n  ClusterAggregateTopK, limit: 10, aggs: [sum(s.Orders.order_amount), min(s.Orders.order_amount), max(s.Orders.order_amount)], sortBy: [3 desc null last, 2 null last]\
+           \n  ClusterAggregateTopK, limit: 10, aggs: [sum(s.Orders.order_amount), min(s.Orders.order_amount), max(s.Orders.order_amount)], sortBy: [3 desc nulls last, 2 nulls last]\
            \n    Scan s.Orders, source: CubeTable(index: by_customer:3:[]:sort_on[order_customer]), fields: [order_customer, order_amount]"
         );
 

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

@@ -11,7 +11,7 @@ use datafusion::logical_expr::{
     Aggregate, EmptyRelation, Explain, Extension, FetchType, Filter, Join, Limit, LogicalPlan,
     Projection, Repartition, SkipType, Sort, TableScan, Union, Window,
 };
-use datafusion::physical_expr::{AcrossPartitions, ConstExpr};
+use datafusion::physical_expr::{AcrossPartitions, ConstExpr, LexOrdering};
 use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode};
 use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec;
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
@@ -479,14 +479,34 @@ pub fn pp_sort_columns(first_agg: usize, cs: &[SortColumn]) -> String {
                     r += " desc";
                 }
                 if !c.nulls_first {
-                    r += " null last";
+                    r += " nulls last";
                 }
                 r
             })
             .join(", ")
     )
 }
 
+fn pp_append_sort_by(out: &mut String, ordering: &LexOrdering) {
+    let _ = write!(
+        out,
+        ", by: [{}]",
+        ordering
+            .iter()
+            .map(|e| {
+                let mut r = format!("{}", e.expr);
+                if e.options.descending {
+                    r += " desc";
+                }
+                if !e.options.nulls_first {
+                    r += " nulls last";
+                }
+                r
+            })
+            .join(", "),
+    );
+}
+
 fn pp_phys_plan_indented(p: &dyn ExecutionPlan, indent: usize, o: &PPOptions, out: &mut String) {
     if p.as_any().is::<CheckMemoryExec>() && !o.show_check_memory_nodes {
         //We don't show CheckMemoryExec in plan by default
@@ -588,23 +608,9 @@ fn pp_phys_plan_indented(p: &dyn ExecutionPlan, indent: usize, o: &PPOptions, ou
             }
         } else if let Some(s) = a.downcast_ref::<SortExec>() {
             *out += "Sort";
+
             if o.show_sort_by {
-                *out += &format!(
-                    ", by: [{}]",
-                    s.expr()
-                        .iter()
-                        .map(|e| {
-                            let mut r = format!("{}", e.expr);
-                            if e.options.descending {
-                                r += " desc";
-                            }
-                            if !e.options.nulls_first {
-                                r += " nulls last";
-                            }
-                            r
-                        })
-                        .join(", ")
-                );
+                pp_append_sort_by(out, s.expr());
             }
             if let Some(fetch) = s.fetch() {
                 *out += &format!(", fetch: {}", fetch);
@@ -656,8 +662,9 @@ fn pp_phys_plan_indented(p: &dyn ExecutionPlan, indent: usize, o: &PPOptions, ou
             *out += "CoalescePartitions";
         } else if let Some(s) = a.downcast_ref::<SortPreservingMergeExec>() {
             *out += "MergeSort";
-            // } else if let Some(_) = a.downcast_ref::<MergeReSortExec>() {
-            //     *out += "MergeResort";
+            if o.show_sort_by {
+                pp_append_sort_by(out, s.expr());
+            }
             if let Some(fetch) = s.fetch() {
                 *out += &format!(", fetch: {}", fetch);
             }

rust/cubestore/cubestore/src/queryplanner/topk/execute.rs

@@ -1294,7 +1294,7 @@ mod tests {
         .unwrap();
         assert_eq!(r, vec![vec![0, 0]]);
 
-        // Ascending, null first.
+        // Ascending, nulls first.
         proto.change_order(vec![SortColumn {
             agg_index: 0,
             asc: true,
@@ -1315,7 +1315,7 @@ mod tests {
         .unwrap();
         assert_eq!(r, vec![vec![Some(2), None]]);
 
-        // Ascending, null last.
+        // Ascending, nulls last.
         proto.change_order(vec![SortColumn {
             agg_index: 0,
             asc: true,