chore(cubestore): Upgrade DF: pretty_printer adjustments: show_partitions, show_schema

Author: srh

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

@@ -4,7 +4,6 @@ use crate::queryplanner::query_executor::ClusterSendExec;
 use crate::queryplanner::tail_limit::TailLimitExec;
 use crate::queryplanner::topk::AggregateTopKExec;
 use datafusion::error::DataFusionError;
-use datafusion::physical_optimizer::topk_aggregation::TopKAggregation;
 use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode};
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::limit::GlobalLimitExec;

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

@@ -1,7 +1,9 @@
 //! Presentation of query plans for use in tests.
 
 use bigdecimal::ToPrimitive;
+use datafusion::arrow::datatypes::Schema;
 use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
+use datafusion::common::DFSchema;
 use datafusion::datasource::physical_plan::ParquetExec;
 use datafusion::datasource::{DefaultTableSource, TableProvider};
 use datafusion::error::DataFusionError;
@@ -34,7 +36,7 @@ use crate::queryplanner::tail_limit::TailLimitExec;
 use crate::queryplanner::topk::SortColumn;
 use crate::queryplanner::topk::{AggregateTopKExec, ClusterAggregateTopKUpper, ClusterAggregateTopKLower};
 use crate::queryplanner::trace_data_loaded::TraceDataLoadedExec;
-use crate::queryplanner::{CubeTableLogical, InfoSchemaTableProvider};
+use crate::queryplanner::{CubeTableLogical, InfoSchemaTableProvider, QueryPlan};
 use crate::streaming::topic_table_provider::TopicTableProvider;
 use datafusion::physical_plan::empty::EmptyExec;
 use datafusion::physical_plan::expressions::Column;
@@ -51,29 +53,24 @@ pub struct PPOptions {
     pub show_filters: bool,
     pub show_sort_by: bool,
     pub show_aggregations: bool,
-    // TODO: Maybe prettify output, name this show_schema.
-    pub debug_schema: bool,
+    pub show_schema: bool,
     // Applies only to physical plan.
     pub show_output_hints: bool,
     pub show_check_memory_nodes: bool,
+    pub show_partitions: bool,
 }
 
 impl PPOptions {
-    pub fn not_everything() -> PPOptions {
+    #[allow(unused)]
+    pub fn everything() -> PPOptions {
         PPOptions {
             show_filters: true,
             show_sort_by: true,
             show_aggregations: true,
-            debug_schema: false,
+            show_schema: true,
             show_output_hints: true,
             show_check_memory_nodes: true,
-        }
-    }
-
-    pub fn truly_everything() -> PPOptions {
-        PPOptions {
-            debug_schema: true,
-            ..PPOptions::not_everything()
+            show_partitions: true,
         }
     }
 
@@ -93,7 +90,18 @@ pub fn pp_phys_plan_ext(p: &dyn ExecutionPlan, o: &PPOptions) -> String {
 }
 
 pub fn pp_plan(p: &LogicalPlan) -> String {
-    pp_plan_ext(p, &PPOptions::default())
+    pp_plan_ext(p, &PPOptions::none())
+}
+
+pub fn pp_query_plan_ext(qp: &QueryPlan, o: &PPOptions) -> String {
+    pp_plan_ext(match qp {
+        QueryPlan::Meta(p) => p,
+        QueryPlan::Select(pre_serialized_plan, _) => pre_serialized_plan.logical_plan()
+    }, o)
+}
+
+pub fn pp_query_plan(p: &QueryPlan) -> String {
+    pp_query_plan_ext(p, &PPOptions::none())
 }
 
 pub fn pp_plan_ext(p: &LogicalPlan, opts: &PPOptions) -> String {
@@ -178,7 +186,7 @@ pub fn pp_plan_ext(p: &LogicalPlan, opts: &PPOptions) -> String {
                     }
                 }
                 LogicalPlan::Union(Union { schema, .. }) => {
-                    self.output += &format!("Union, schema: {}", schema)
+                    self.output += &format!("Union, schema: {}", pp_df_schema(schema.as_ref()))
                 }
                 LogicalPlan::Join(Join { on, .. }) => {
                     self.output += &format!(
@@ -372,8 +380,8 @@ pub fn pp_plan_ext(p: &LogicalPlan, opts: &PPOptions) -> String {
                 }
             }
 
-            if self.opts.debug_schema {
-                self.output += &format!(", debug_schema: {:?}", plan.schema());
+            if self.opts.show_schema {
+                self.output += &format!(", schema: {}", pp_df_schema(plan.schema().as_ref()));
             }
 
             if !saw_expected_topk_lower {
@@ -475,6 +483,8 @@ fn pp_phys_plan_indented(p: &dyn ExecutionPlan, indent: usize, o: &PPOptions, ou
         }
         out.extend(repeat_n(' ', indent));
 
+        let mut skip_show_partitions = false;
+
         let a = p.as_any();
         if let Some(t) = a.downcast_ref::<CubeTableExec>() {
             *out += &format!("Scan, index: {}", pp_index(&t.index_snapshot));
@@ -588,6 +598,7 @@ fn pp_phys_plan_indented(p: &dyn ExecutionPlan, indent: usize, o: &PPOptions, ou
                     })
                     .join(", ")
             );
+            skip_show_partitions = true;
         } else if let Some(topk) = a.downcast_ref::<AggregateTopKExec>() {
             *out += &format!("AggregateTopK, limit: {:?}", topk.limit);
             if o.show_aggregations {
@@ -661,14 +672,6 @@ fn pp_phys_plan_indented(p: &dyn ExecutionPlan, indent: usize, o: &PPOptions, ou
             *out += &to_string.split(" ").next().unwrap_or(&to_string);
         }
 
-        // TODO upgrade DF - remove
-        // *out += &format!(", schema: {}", p.schema());
-        // *out += &format!(
-        //     ", partitions: {}, output_ordering: {:?}",
-        //     p.properties().partitioning.partition_count(),
-        //     p.output_ordering()
-        // );
-
         if o.show_output_hints {
             let properties: &PlanProperties = p.properties();
 
@@ -728,8 +731,12 @@ fn pp_phys_plan_indented(p: &dyn ExecutionPlan, indent: usize, o: &PPOptions, ou
             }
         }
 
-        if o.debug_schema {
-            *out += &format!(", debug_schema: {:?}", p.schema());
+        if o.show_schema {
+            *out += &format!(", schema: {}", pp_schema(p.schema().as_ref()));
+        }
+
+        if o.show_partitions && !skip_show_partitions {
+            *out += &format!(", partitions: {}", p.properties().output_partitioning().partition_count());
         }
     }
 }
@@ -752,3 +759,15 @@ fn pp_row_range(r: &RowRange) -> String {
 fn pp_exprs(v: &Vec<Expr>) -> String {
     "[".to_owned() + &v.iter().map(|e: &Expr| format!("{}", e)).join(", ") + "]"
 }
+
+fn pp_df_schema(schema: &DFSchema) -> String {
+    // Like pp_schema but with qualifiers.
+    format!("{}", schema)
+}
+
+fn pp_schema(schema: &Schema) -> String {
+    // Mimicking DFSchema's Display
+    format!("fields:[{}], metadata:{:?}",
+        schema.fields.iter().map(|f| f.name()).join(", "),
+        schema.metadata)
+}

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

@@ -1692,7 +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.
+        // TODO:  Ensure this is obeyed... or allow worker partitions to be sent separately.
         vec![Distribution::SinglePartition; self.children().len()]
     }
 }