WIP: Cargo fmt topk changes

Author: srh

rust/cubestore/cubestore/src/cluster/mod.rs

@@ -230,7 +230,9 @@ impl WorkerPlanningParams {
     }
     pub fn panic_worker_params() -> WorkerPlanningParams {
         // TODO: Maybe avoid needing to call this function.
-        WorkerPlanningParams { worker_partition_count: 1 }
+        WorkerPlanningParams {
+            worker_partition_count: 1,
+        }
     }
 }
 
@@ -344,7 +346,12 @@ impl WorkerProcessing for WorkerProcessor {
                     let res = services
                         .query_executor
                         .clone()
-                        .execute_worker_plan(plan_node_to_send, worker_planning_params, remote_to_local_names, result)
+                        .execute_worker_plan(
+                            plan_node_to_send,
+                            worker_planning_params,
+                            remote_to_local_names,
+                            result,
+                        )
                         .await;
                     debug!(
                         "Running select in worker completed ({:?})",
@@ -495,7 +502,10 @@ impl Cluster for ClusterImpl {
         worker_planning_params: WorkerPlanningParams,
     ) -> Result<Vec<RecordBatch>, CubeError> {
         let response = self
-            .send_or_process_locally(node_name, NetworkMessage::Select(plan_node, worker_planning_params))
+            .send_or_process_locally(
+                node_name,
+                NetworkMessage::Select(plan_node, worker_planning_params),
+            )
             .await?;
         match response {
             NetworkMessage::SelectResult(r) => {
@@ -512,7 +522,10 @@ impl Cluster for ClusterImpl {
         worker_planning_params: WorkerPlanningParams,
     ) -> Result<String, CubeError> {
         let response = self
-            .send_or_process_locally(node_name, NetworkMessage::ExplainAnalyze(plan, worker_planning_params))
+            .send_or_process_locally(
+                node_name,
+                NetworkMessage::ExplainAnalyze(plan, worker_planning_params),
+            )
             .await?;
         match response {
             NetworkMessage::ExplainAnalyzeResult(r) => r,
@@ -708,7 +721,9 @@ impl Cluster for ClusterImpl {
                 NetworkMessage::SelectResult(res)
             }
             NetworkMessage::ExplainAnalyze(plan, planning_params) => {
-                let res = self.run_local_explain_analyze_worker(plan, planning_params).await;
+                let res = self
+                    .run_local_explain_analyze_worker(plan, planning_params)
+                    .await;
                 NetworkMessage::ExplainAnalyzeResult(res)
             }
             NetworkMessage::WarmupDownload(remote_path, expected_file_size) => {
@@ -1253,7 +1268,9 @@ impl ClusterImpl {
             table_id: None,
             trace_obj: plan_node.trace_obj(),
         };
-        let res = self.run_local_select_worker_impl(plan_node, worker_planning_params).await;
+        let res = self
+            .run_local_select_worker_impl(plan_node, worker_planning_params)
+            .await;
         match res {
             Ok((schema, records, data_loaded_size)) => {
                 self.process_rate_limiter
@@ -1406,7 +1423,12 @@ impl ClusterImpl {
 
         let res = self
             .query_executor
-            .pp_worker_plan(plan_node, worker_planning_params, remote_to_local_names, chunk_id_to_record_batches)
+            .pp_worker_plan(
+                plan_node,
+                worker_planning_params,
+                remote_to_local_names,
+                chunk_id_to_record_batches,
+            )
             .await;
 
         res
@@ -1530,7 +1552,10 @@ impl ClusterImpl {
     async fn start_stream_on_worker(self: Arc<Self>, m: NetworkMessage) -> Box<dyn MessageStream> {
         match m {
             NetworkMessage::SelectStart(p, worker_planning_params) => {
-                let (schema, results) = match self.run_local_select_worker(p, worker_planning_params).await {
+                let (schema, results) = match self
+                    .run_local_select_worker(p, worker_planning_params)
+                    .await
+                {
                     Err(e) => return Box::new(QueryStream::new_error(e)),
                     Ok(x) => x,
                 };

rust/cubestore/cubestore/src/config/mod.rs

@@ -2106,7 +2106,7 @@ impl Config {
                     i.get_service_typed::<dyn CubestoreMetadataCacheFactory>()
                         .await
                         .cache_factory()
-                        .clone()
+                        .clone(),
                 )
             })
             .await;

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

@@ -29,7 +29,10 @@ pub fn push_aggregate_to_workers(
     let p_final_agg: &AggregateExec;
     let p_partial: &Arc<dyn ExecutionPlan>;
     if let Some(a) = p_final.as_any().downcast_ref::<AggregateExec>() {
-        if matches!(a.mode(), AggregateMode::Final | AggregateMode::FinalPartitioned) {
+        if matches!(
+            a.mode(),
+            AggregateMode::Final | AggregateMode::FinalPartitioned
+        ) {
             p_final_agg = a;
             p_partial = a.input();
         } else {
@@ -49,33 +52,45 @@ pub fn push_aggregate_to_workers(
         return Ok(p_final);
     }
 
-    let p_final_input: Arc<dyn ExecutionPlan> = if let Some(cs) = agg.input().as_any().downcast_ref::<ClusterSendExec>() {
-        let clustersend_input = p_partial.clone()
-                    .with_new_children(vec![cs.input_for_optimizations.clone()])?;
+    let p_final_input: Arc<dyn ExecutionPlan> =
+        if let Some(cs) = agg.input().as_any().downcast_ref::<ClusterSendExec>() {
+            let clustersend_input = p_partial
+                .clone()
+                .with_new_children(vec![cs.input_for_optimizations.clone()])?;
 
-        // Note that required_input_ordering is applicable when p_final_agg has a Sorted input mode.
+            // Note that required_input_ordering is applicable when p_final_agg has a Sorted input mode.
 
-        // Router plan, replace partial aggregate with cluster send.
-        Arc::new(
-            cs.with_changed_schema(
-                clustersend_input,
-                p_final_agg.required_input_ordering().into_iter().next().unwrap(),
-            ),
-        )
-    } else if let Some(w) = agg.input().as_any().downcast_ref::<WorkerExec>() {
-        let worker_input = p_partial.clone().with_new_children(vec![w.input.clone()])?;
+            // Router plan, replace partial aggregate with cluster send.
+            Arc::new(
+                cs.with_changed_schema(
+                    clustersend_input,
+                    p_final_agg
+                        .required_input_ordering()
+                        .into_iter()
+                        .next()
+                        .unwrap(),
+                ),
+            )
+        } else if let Some(w) = agg.input().as_any().downcast_ref::<WorkerExec>() {
+            let worker_input = p_partial.clone().with_new_children(vec![w.input.clone()])?;
 
-        // Worker plan, execute partial aggregate inside the worker.
-        Arc::new(WorkerExec::new(
-            worker_input,
-            w.max_batch_rows,
-            w.limit_and_reverse.clone(),
-            p_final_agg.required_input_ordering().into_iter().next().unwrap(),
-            WorkerPlanningParams { worker_partition_count: w.properties().output_partitioning().partition_count() },
-        ))
-    } else {
-        return Ok(p_final);
-    };
+            // Worker plan, execute partial aggregate inside the worker.
+            Arc::new(WorkerExec::new(
+                worker_input,
+                w.max_batch_rows,
+                w.limit_and_reverse.clone(),
+                p_final_agg
+                    .required_input_ordering()
+                    .into_iter()
+                    .next()
+                    .unwrap(),
+                WorkerPlanningParams {
+                    worker_partition_count: w.properties().output_partitioning().partition_count(),
+                },
+            ))
+        } else {
+            return Ok(p_final);
+        };
 
     // We change AggregateMode::FinalPartitioned to AggregateMode::Final, because the ClusterSend
     // node ends up creating an incompatible partitioning for FinalPartitioned.  Some other ideas,
@@ -157,7 +172,10 @@ pub fn ensure_partition_merge_with_acceptable_parent(
     let mut new_children = Vec::new();
 
     for p in parent.children() {
-        new_children.push(ensure_partition_merge_helper(p.clone(), &mut any_new_children)?);
+        new_children.push(ensure_partition_merge_helper(
+            p.clone(),
+            &mut any_new_children,
+        )?);
     }
     if any_new_children {
         parent.with_new_children(new_children)

rust/cubestore/cubestore/src/queryplanner/panic.rs

@@ -161,13 +161,14 @@ pub fn plan_panic_worker() -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
         /* max_batch_rows */ 1,
         /* limit_and_reverse */ None,
         /* required_input_ordering */ None,
-
         // worker_partition_count is generally set to 1 for panic worker messages
         // (SystemCommand::PanicWorker).  What is important is that router and worker nodes have the
         // same plan properties so that DF optimizations run identically -- router node is creating
         // a WorkerExec for some reason. (Also, it's important that DF optimizations run identically
         // when it comes to aggregates pushed down through ClusterSend and the like -- it's actually
         // NOT important for panic worker planning.)
-        WorkerPlanningParams { worker_partition_count: 1 },
+        WorkerPlanningParams {
+            worker_partition_count: 1,
+        },
     )))
 }

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

@@ -51,11 +51,11 @@ use crate::queryplanner::rolling::RollingWindowAggregateSerialized;
 use crate::queryplanner::serialized_plan::{
     IndexSnapshot, InlineSnapshot, PartitionSnapshot, SerializedPlan,
 };
+use crate::queryplanner::topk::plan_topk;
 use crate::queryplanner::topk::ClusterAggregateTopK;
 use crate::queryplanner::{CubeTableLogical, InfoSchemaTableProvider};
 use crate::table::{cmp_same_types, Row};
 use crate::CubeError;
-use crate::queryplanner::topk::plan_topk;
 use datafusion::common;
 use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
 use datafusion::common::DFSchemaRef;
@@ -64,7 +64,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, Filter, Join, Limit, LogicalPlan, Operator, Projection, Sort, SortExpr, SubqueryAlias, TableScan, Union, Unnest, UserDefinedLogicalNode
+    expr, logical_plan, Aggregate, BinaryExpr, Expr, Extension, Filter, Join, Limit, LogicalPlan,
+    Operator, Projection, Sort, SortExpr, SubqueryAlias, TableScan, Union, Unnest,
+    UserDefinedLogicalNode,
 };
 use datafusion::physical_expr::{Distribution, LexRequirement};
 use datafusion::physical_plan::repartition::RepartitionExec;
@@ -1711,7 +1713,10 @@ impl CubeExtensionPlanner {
             Ok(Arc::new(ClusterSendExec::new(
                 c.clone(),
                 if let Some(logical_plan_to_send) = logical_plan_to_send {
-                    Arc::new(self.serialized_plan.replace_logical_plan(logical_plan_to_send.clone())?)
+                    Arc::new(
+                        self.serialized_plan
+                            .replace_logical_plan(logical_plan_to_send.clone())?,
+                    )
                 } else {
                     self.serialized_plan.clone()
                 },
@@ -1752,7 +1757,13 @@ 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));
+        let properties =
+            input
+                .properties()
+                .clone()
+                .with_partitioning(Partitioning::UnknownPartitioning(
+                    worker_planning_params.worker_partition_count,
+                ));
         WorkerExec {
             input,
             max_batch_rows,
@@ -2643,19 +2654,27 @@ pub mod tests {
         }
 
         fn udf_names(&self) -> Vec<String> {
-            self.session_state.scalar_functions().keys().cloned().collect()
+            self.session_state
+                .scalar_functions()
+                .keys()
+                .cloned()
+                .collect()
         }
 
         fn udaf_names(&self) -> Vec<String> {
-            self.session_state.aggregate_functions().keys().cloned().collect()
+            self.session_state
+                .aggregate_functions()
+                .keys()
+                .cloned()
+                .collect()
         }
 
         fn udwf_names(&self) -> Vec<String> {
             self.session_state
-            .window_functions()
-            .keys()
-            .cloned()
-            .collect()
+                .window_functions()
+                .keys()
+                .cloned()
+                .collect()
         }
     }
 

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

@@ -31,8 +31,8 @@ use crate::queryplanner::query_executor::{
 use crate::queryplanner::rolling::RollingWindowAggregate;
 use crate::queryplanner::serialized_plan::{IndexSnapshot, RowRange};
 use crate::queryplanner::tail_limit::TailLimitExec;
-use crate::queryplanner::topk::{AggregateTopKExec, ClusterAggregateTopK};
 use crate::queryplanner::topk::SortColumn;
+use crate::queryplanner::topk::{AggregateTopKExec, ClusterAggregateTopK};
 use crate::queryplanner::trace_data_loaded::TraceDataLoadedExec;
 use crate::queryplanner::{CubeTableLogical, InfoSchemaTableProvider};
 use crate::streaming::topic_table_provider::TopicTableProvider;

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

@@ -1,4 +1,6 @@
-use crate::cluster::{pick_worker_by_ids, pick_worker_by_partitions, Cluster, WorkerPlanningParams};
+use crate::cluster::{
+    pick_worker_by_ids, pick_worker_by_partitions, Cluster, WorkerPlanningParams,
+};
 use crate::config::injection::DIService;
 use crate::config::ConfigObj;
 use crate::metastore::multi_index::MultiPartition;
@@ -75,7 +77,8 @@ use datafusion::physical_plan::sorts::sort::SortExec;
 use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{
-    collect, DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, ExecutionPlanProperties, Partitioning, PhysicalExpr, PlanProperties, SendableRecordBatchStream
+    collect, DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, ExecutionPlanProperties,
+    Partitioning, PhysicalExpr, PlanProperties, SendableRecordBatchStream,
 };
 use datafusion::prelude::{and, SessionConfig, SessionContext};
 use futures_util::{stream, FutureExt, StreamExt, TryStreamExt};
@@ -310,7 +313,8 @@ impl QueryExecutor for QueryExecutorImpl {
         )?;
         let pre_serialized_plan = Arc::new(pre_serialized_plan);
         let ctx = self.router_context(cluster.clone(), pre_serialized_plan.clone())?;
-        let router_plan = ctx.clone()
+        let router_plan = ctx
+            .clone()
             .state()
             .create_physical_plan(pre_serialized_plan.logical_plan())
             .await?;
@@ -337,7 +341,11 @@ impl QueryExecutor for QueryExecutorImpl {
             self.parquet_metadata_cache.cache().clone(),
         )?;
         let pre_serialized_plan = Arc::new(pre_serialized_plan);
-        let ctx = self.worker_context(pre_serialized_plan.clone(), worker_planning_params, data_loaded_size)?;
+        let ctx = self.worker_context(
+            pre_serialized_plan.clone(),
+            worker_planning_params,
+            data_loaded_size,
+        )?;
         let plan_ctx = ctx.clone();
         Ok((
             plan_ctx
@@ -1524,7 +1532,11 @@ impl ClusterSendExec {
         r
     }
 
-    pub fn with_changed_schema(&self, input_for_optimizations: Arc<dyn ExecutionPlan>, new_required_input_ordering: Option<LexRequirement>) -> Self {
+    pub fn with_changed_schema(
+        &self,
+        input_for_optimizations: Arc<dyn ExecutionPlan>,
+        new_required_input_ordering: Option<LexRequirement>,
+    ) -> Self {
         ClusterSendExec {
             properties: Self::compute_properties(
                 input_for_optimizations.properties(),
@@ -1623,7 +1635,11 @@ impl ExecutionPlan for ClusterSendExec {
             // A future that yields a stream
             let fut = async move {
                 cluster
-                    .run_select_stream(&node_name, plan.to_serialized_plan()?, worker_planning_params)
+                    .run_select_stream(
+                        &node_name,
+                        plan.to_serialized_plan()?,
+                        worker_planning_params,
+                    )
                     .await
             };
             // Use TryStreamExt::try_flatten to flatten the stream of streams
@@ -1633,7 +1649,11 @@ impl ExecutionPlan for ClusterSendExec {
         } else {
             let record_batches = async move {
                 cluster
-                    .run_select(&node_name, plan.to_serialized_plan()?, worker_planning_params)
+                    .run_select(
+                        &node_name,
+                        plan.to_serialized_plan()?,
+                        worker_planning_params,
+                    )
                     .await
             };
             let stream = futures::stream::once(record_batches).flat_map(|r| match r {