diff --git a/Cargo.lock b/Cargo.lock index dbe6096c6..79cf460d9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -287,7 +287,7 @@ dependencies = [ "arrow-schema", "arrow-select", "atoi", - "base64", + "base64 0.22.1", "chrono", "comfy-table", "half", @@ -335,7 +335,7 @@ dependencies = [ "arrow-cast", "arrow-ipc", "arrow-schema", - "base64", + "base64 0.22.1", "bytes", "futures", "prost", @@ -1064,6 +1064,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b25655df2c3cdd83c5e5b293b88acd880332b2ddadd7c30ac43144fdc0033da9" + [[package]] name = "base64-simd" version = "0.8.0" @@ -2193,6 +2199,7 @@ dependencies = [ "arrow-ipc", "arrow-select", "async-trait", + "base64 0.23.0", "bincode", "bytes", "chrono", @@ -2282,6 +2289,7 @@ version = "0.1.0" dependencies = [ "arrow", "async-trait", + "base64 0.23.0", "color-eyre", "crossterm", "datafusion", @@ -2368,7 +2376,7 @@ checksum = "14872c47bfc3d21e53ec82f57074e6987a15941c1e2f43cde4ac6ae2746634e3" dependencies = [ "arrow", "arrow-buffer", - "base64", + "base64 0.22.1", "blake2", "blake3", "chrono", @@ -3599,7 +3607,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-util", @@ -4436,7 +4444,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "622acbc9100d3c10e2ee15804b0caa40e55c933d5aa53814cd520805b7958a49" dependencies = [ "async-trait", - "base64", + "base64 0.22.1", "bytes", "chrono", "form_urlencoded", @@ -4617,7 +4625,7 @@ dependencies = [ "arrow-ipc", "arrow-schema", "arrow-select", - "base64", + "base64 0.22.1", "brotli", "bytes", "chrono", @@ -5347,7 +5355,7 @@ version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-core", @@ -6072,7 +6080,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4676b37242ccbd1aabf56edb093a4827dc49086c0ffd764a5705899e0f35f8f7" dependencies = [ "anyhow", - "base64", + "base64 0.22.1", "bitflags 2.11.1", "fancy-regex", "filedescriptor", @@ -6376,7 +6384,7 @@ checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" dependencies = [ "async-trait", "axum 0.8.9", - "base64", + "base64 0.22.1", "bytes", "h2 0.4.14", "http 1.4.0", diff --git a/Cargo.toml b/Cargo.toml index 79366ed33..901600717 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -59,6 +59,7 @@ insta = { version = "1.46.0", features = ["filters"], optional = true } parquet = { version = "58", optional = true } arrow = { version = "58", optional = true, features = ["test_utils"] } hyper-util = { version = "0.1.16", optional = true } +base64 = "0.23.0" [features] default = ["grpc"] diff --git a/console/Cargo.toml b/console/Cargo.toml index 5267842e2..c81ed99d5 100644 --- a/console/Cargo.toml +++ b/console/Cargo.toml @@ -15,6 +15,7 @@ tonic = "0.14.2" datafusion-distributed = { path = "..", features = ["integration", "system-metrics"] } url = "2.5.7" tokio-stream = "0.1.18" +base64 = "0.23.0" [dev-dependencies] arrow = "58" diff --git a/console/src/main.rs b/console/src/main.rs index 731b504ba..dbf0cce97 100644 --- a/console/src/main.rs +++ b/console/src/main.rs @@ -5,6 +5,8 @@ mod ui; mod worker; use app::App; +use base64::engine::general_purpose::STANDARD; +use base64::engine::Engine; use crossterm::event::{self, Event}; use ratatui::DefaultTerminal; use std::time::{Duration, Instant}; @@ -24,6 +26,10 @@ struct Args { /// Polling interval in milliseconds #[structopt(long = "poll-interval", default_value = "100")] poll_interval: u64, + + /// Decode and print a base64-encoded plan string produced by explain_analyze. + #[structopt(long = "encoded-plan")] + encoded_plan: Option, } #[tokio::main] @@ -32,6 +38,14 @@ async fn main() -> color_eyre::Result<()> { let args = Args::from_args(); + if let Some(encoded) = args.encoded_plan { + let decoded = Engine::decode(&STANDARD, &encoded) + .map(|b| String::from_utf8_lossy(&b).into_owned()) + .unwrap_or(encoded); + println!("{decoded}"); + return Ok(()); + } + let seed_url = Url::parse(&format!("http://localhost:{}", args.port)).expect("valid URL"); let poll_interval = Duration::from_millis(args.poll_interval); diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index cb126a884..ba90050c6 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -1,3 +1,5 @@ +use base64::engine::general_purpose::STANDARD; +use base64::engine::Engine; use crate::common::require_one_child; use crate::coordinator::metrics_store::MetricsStore; use crate::coordinator::prepare_dynamic_plan::prepare_dynamic_plan; @@ -134,6 +136,12 @@ impl DistributedExec { .clone() .ok_or_else(|| internal_datafusion_err!("No head stage found. Was execute() called?")) } + /// Decodes a base64-encoded plan string produced by [`explain_analyze`](crate::explain_analyze). + pub fn extract_encoded_plan(&self, encoded_plan: &str) -> String { + Engine::decode(&STANDARD, encoded_plan) + .map(|bytes| String::from_utf8_lossy(&bytes).into_owned()) + .unwrap_or_else(|_| encoded_plan.to_string()) + } } impl DisplayAs for DistributedExec { diff --git a/src/stage.rs b/src/stage.rs index 7903a48a9..6765b34d4 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -1,4 +1,7 @@ +use base64::engine::general_purpose::STANDARD; +use base64::engine::Engine; use crate::coordinator::{DistributedExec, MetricsStore}; + use crate::execution_plans::{DistributedLeafExec, NetworkCoalesceExec}; use crate::metrics::DISTRIBUTED_DATAFUSION_TASK_ID_LABEL; use datafusion::common::{HashMap, Statistics, config_err}; @@ -252,7 +255,12 @@ pub async fn explain_analyze( .to_string()), Some(_) => { let executed = rewrite_distributed_plan_with_metrics(executed.clone(), format).await?; - Ok(display_plan_ascii(executed.as_ref(), true)) + let display_string = display_plan_ascii(executed.as_ref(), true); + if display_string.len() >= 10_000 { + Ok(Engine::encode(&STANDARD, display_string.as_bytes())) + } else { + Ok(display_string) + } } } }