Dandandan commented on code in PR #242: URL: https://github.com/apache/arrow-ballista/pull/242#discussion_r981308484
########## ballista/rust/scheduler/src/state/execution_graph_dot.rs: ########## @@ -0,0 +1,537 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Utilities for producing dot diagrams from execution graphs + +use crate::state::execution_graph::{ExecutionGraph, ExecutionStage}; +use ballista_core::execution_plans::{ + ShuffleReaderExec, ShuffleWriterExec, UnresolvedShuffleExec, +}; +use datafusion::datasource::listing::PartitionedFile; +use datafusion::physical_plan::aggregates::AggregateExec; +use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::cross_join::CrossJoinExec; +use datafusion::physical_plan::file_format::{ + AvroExec, CsvExec, FileScanConfig, NdJsonExec, ParquetExec, +}; +use datafusion::physical_plan::filter::FilterExec; +use datafusion::physical_plan::hash_join::HashJoinExec; +use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; +use datafusion::physical_plan::memory::MemoryExec; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::union::UnionExec; +use datafusion::physical_plan::{ExecutionPlan, Partitioning, PhysicalExpr}; +use log::debug; +use object_store::path::Path; +use std::collections::HashMap; +use std::fmt::{self, Write}; +use std::sync::Arc; + +/// Utility for producing dot diagrams from execution graphs +pub struct ExecutionGraphDot { + graph: Arc<ExecutionGraph>, +} + +impl ExecutionGraphDot { + /// Create a DOT graph from the provided ExecutionGraph + pub fn generate(graph: Arc<ExecutionGraph>) -> Result<String, fmt::Error> { + let mut dot = Self { graph }; + dot._generate() + } + + fn _generate(&mut self) -> Result<String, fmt::Error> { + // sort the stages by key for deterministic output for tests + let stages = self.graph.stages(); + let mut stage_ids: Vec<usize> = stages.keys().cloned().collect(); + stage_ids.sort(); + + let mut dot = String::new(); + + writeln!(&mut dot, "digraph G {{")?; + + let mut cluster = 0; + let mut stage_meta = vec![]; + + #[allow(clippy::explicit_counter_loop)] + for id in &stage_ids { + let stage = stages.get(id).unwrap(); // safe unwrap + let stage_name = format!("stage_{}", id); + writeln!(&mut dot, "\tsubgraph cluster{} {{", cluster)?; + match stage { + ExecutionStage::UnResolved(stage) => { + writeln!(&mut dot, "\t\tlabel = \"Stage {} [UnResolved]\";", id)?; + stage_meta.push(write_stage_plan( + &mut dot, + &stage_name, + &stage.plan, + 0, + )?); + } + ExecutionStage::Resolved(stage) => { + writeln!(&mut dot, "\t\tlabel = \"Stage {} [Resolved]\";", id)?; + stage_meta.push(write_stage_plan( + &mut dot, + &stage_name, + &stage.plan, + 0, + )?); + } + ExecutionStage::Running(stage) => { + writeln!(&mut dot, "\t\tlabel = \"Stage {} [Running]\";", id)?; + stage_meta.push(write_stage_plan( + &mut dot, + &stage_name, + &stage.plan, + 0, + )?); + } + ExecutionStage::Completed(stage) => { + writeln!(&mut dot, "\t\tlabel = \"Stage {} [Completed]\";", id)?; + stage_meta.push(write_stage_plan( + &mut dot, + &stage_name, + &stage.plan, + 0, + )?); + } + ExecutionStage::Failed(stage) => { + writeln!(&mut dot, "\t\tlabel = \"Stage {} [FAILED]\";", id)?; + stage_meta.push(write_stage_plan( + &mut dot, + &stage_name, + &stage.plan, + 0, + )?); + } + } + cluster += 1; + writeln!(&mut dot, "\t}}")?; // end of subgraph + } + + // write links between stages + for meta in &stage_meta { + let mut links = vec![]; + for (reader_node, parent_stage_id) in &meta.readers { + // shuffle write node is always node zero + let parent_shuffle_write_node = format!("stage_{}_0", parent_stage_id); + links.push(format!("{} -> {}", parent_shuffle_write_node, reader_node,)); + } + // keep the order deterministic + links.sort(); + for link in links { + writeln!(&mut dot, "\t{}", link)?; + } + } + + writeln!(&mut dot, "}}")?; // end of digraph + + Ok(dot) + } +} + +/// Write the query tree for a single stage and build metadata needed to later draw +/// the links between the stages +fn write_stage_plan( + f: &mut String, + prefix: &str, + plan: &Arc<dyn ExecutionPlan>, + i: usize, +) -> Result<StagePlanState, fmt::Error> { + let mut state = StagePlanState { + readers: HashMap::new(), + }; + write_plan_recursive(f, prefix, plan, i, &mut state)?; + Ok(state) +} + +fn write_plan_recursive( Review Comment: We could use/move something similar to DataFusion (minus the stage part, currently). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
