diff --git a/.gitignore b/.gitignore index 47cd74cc..8e01bad0 100644 --- a/.gitignore +++ b/.gitignore @@ -105,4 +105,6 @@ bench/url_files/* # Tags tags +# Local scratch / per-machine files +notes .claude/settings.local.json diff --git a/Cargo.lock b/Cargo.lock index ce00ab47..6e236ec1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2658,8 +2658,10 @@ dependencies = [ "object_store_opendal", "opendal", "parking_lot", + "prost", "rustls", "serde", + "serde_json", "tokio", "tokio-metrics", "tokio-stream", diff --git a/Cargo.toml b/Cargo.toml index ca83770f..53d1ffde 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,6 +42,7 @@ prost = "0.14" ratatui = { optional = true, version = "0.30" } ratatui-textarea = { features = ["search"], optional = true, version = "0.8" } serde = { features = ["derive"], version = "1.0.197" } +serde_json = "1.0" strum = { features = ["derive"], version = "0.26.2" } tokio = { features = [ "macros", diff --git a/crates/datafusion-app/Cargo.toml b/crates/datafusion-app/Cargo.toml index 45124803..f9effe07 100644 --- a/crates/datafusion-app/Cargo.toml +++ b/crates/datafusion-app/Cargo.toml @@ -44,10 +44,12 @@ opendal = { features = [ "services-huggingface", ], git = "https://github.com/apache/opendal", optional = true, rev = "24aaff9c62b1" } parking_lot = "0.12.3" +prost = { optional = true, version = "0.14" } rustls = { default-features = false, features = [ "aws-lc-rs", ], optional = true, version = "0.23" } serde = { features = ["derive"], version = "1.0.197" } +serde_json = { optional = true, version = "1.0" } tokio = { features = ["macros", "rt-multi-thread"], version = "1.36.0" } tokio-metrics = { features = [ "metrics-rs-integration", @@ -62,6 +64,7 @@ vortex-datafusion = { optional = true, version = "0.78" } [dev-dependencies] criterion = { features = ["async_tokio"], version = "0.5.1" } +tempfile = "3" [features] clickhouse = [ @@ -70,7 +73,13 @@ clickhouse = [ ] default = ["functions-parquet"] deltalake = ["dep:deltalake"] -flightsql = ["dep:arrow-flight", "dep:base64", "dep:tonic"] +flightsql = [ + "dep:arrow-flight", + "dep:base64", + "dep:prost", + "dep:serde_json", + "dep:tonic", +] functions-arrow = ["dep:datafusion-functions-arrow"] functions-json = ["dep:datafusion-functions-json"] functions-parquet = ["dep:datafusion-functions-parquet"] diff --git a/crates/datafusion-app/src/flightsql.rs b/crates/datafusion-app/src/flightsql.rs index e79dfd0b..a1ff3349 100644 --- a/crates/datafusion-app/src/flightsql.rs +++ b/crates/datafusion-app/src/flightsql.rs @@ -483,4 +483,102 @@ impl FlightSQLContext { )) } } + + /// Get raw metrics batch without reconstruction (for --analyze-raw) + pub async fn analyze_query_raw( + &self, + query: &str, + ) -> Result<(String, datafusion::arrow::array::RecordBatch)> { + self.fetch_analyze_batches(query).await + } + + /// Reconstruct ExecutionStats from metrics (for --analyze) + pub async fn analyze_query(&self, query: &str) -> Result { + let (query_str, metrics_batch) = self.fetch_analyze_batches(query).await?; + + // Reconstruct ExecutionStats from metrics table + let stats = crate::stats::ExecutionStats::from_metrics_table(metrics_batch, query_str)?; + + Ok(stats) + } + + /// Shared logic to fetch analyze batch and query from server + async fn fetch_analyze_batches( + &self, + query: &str, + ) -> Result<(String, datafusion::arrow::array::RecordBatch)> { + use crate::stats::{ + is_compatible_protocol_version, PROTOCOL_VERSION_METADATA_KEY, QUERY_ID_METADATA_KEY, + }; + use arrow_flight::utils::flight_data_to_batches; + use arrow_flight::{Action, FlightData}; + + // Validate that query contains only a single statement + let dialect = datafusion::sql::sqlparser::dialect::GenericDialect {}; + let statements = DFParser::parse_sql_with_dialect(query, &dialect)?; + if statements.len() != 1 { + return Err(eyre::eyre!("Only a single SQL statement can be analyzed")); + } + + // 1. Create JSON request and encode as Action body + let request = crate::stats::AnalyzeQueryRequest::with_sql(query); + let request_body = serde_json::to_vec(&request) + .map_err(|e| eyre::eyre!("Failed to serialize request: {}", e))?; + + let action = Action { + r#type: "analyze_query".to_string(), + body: request_body.into(), + }; + + // 2. Call do_action on the FlightSQL service + let mut client = self.client.lock().await; + let client = client + .as_mut() + .ok_or_else(|| eyre::eyre!("No FlightSQL client configured"))?; + + let mut stream = client + .do_action(action.into_request()) + .await + .map_err(|e| eyre::eyre!("do_action failed: {}", e))?; + + // 3. Collect all Result messages, decoding each body as FlightData + let mut all_flight_data = Vec::new(); + while let Some(result) = stream.next().await { + let result = result.map_err(|e| eyre::eyre!("Stream error: {}", e))?; + let flight_data = ::decode(result.body.as_ref()) + .map_err(|e| eyre::eyre!("Failed to decode FlightData: {}", e))?; + all_flight_data.push(flight_data); + } + + // 4. Decode the metrics batches. The framing is a schema message + // followed by any number of data (and dictionary) messages; a server + // may legitimately split a large metrics table into several batches. + let metrics_batches = flight_data_to_batches(&all_flight_data) + .map_err(|e| eyre::eyre!("Failed to decode metrics batch: {}", e))?; + + if metrics_batches.is_empty() { + return Err(eyre::eyre!("No metrics batch found in response")); + } + + let schema = metrics_batches[0].schema(); + if let Some(version) = schema.metadata().get(PROTOCOL_VERSION_METADATA_KEY) { + if !is_compatible_protocol_version(version) { + return Err(eyre::eyre!( + "Server analyze protocol version {} is not compatible with client version {}", + version, + crate::stats::ANALYZE_PROTOCOL_VERSION + )); + } + } + if let Some(query_id) = schema.metadata().get(QUERY_ID_METADATA_KEY) { + debug!("Analyze response query_id: {query_id}"); + } + + let metrics_batch = + datafusion::arrow::compute::concat_batches(&schema, &metrics_batches) + .map_err(|e| eyre::eyre!("Failed to concatenate metrics batches: {}", e))?; + + // Use the original query (client retains it, server doesn't send it back) + Ok((query.to_string(), metrics_batch)) + } } diff --git a/crates/datafusion-app/src/local.rs b/crates/datafusion-app/src/local.rs index d576ae91..86594d7a 100644 --- a/crates/datafusion-app/src/local.rs +++ b/crates/datafusion-app/src/local.rs @@ -665,7 +665,7 @@ impl ExecutionContext { physical_plan, ) } else { - Err(eyre::eyre!("Only a single statement can be benchmarked")) + Err(eyre::eyre!("Only a single statement can be analyzed")) } } diff --git a/crates/datafusion-app/src/stats.rs b/crates/datafusion-app/src/stats.rs index a4842a41..9eb9a071 100644 --- a/crates/datafusion-app/src/stats.rs +++ b/crates/datafusion-app/src/stats.rs @@ -16,8 +16,12 @@ // under the License. use datafusion::{ + arrow::{ + array::{Array, ArrayRef, Int32Array, RecordBatch, StringArray, UInt64Array}, + datatypes::{DataType, Field, Schema, SchemaRef}, + }, datasource::{ - physical_plan::{FileScanConfig, ParquetSource}, + physical_plan::{ArrowSource, CsvSource, FileScanConfig, JsonSource, ParquetSource}, source::DataSourceExec, }, physical_plan::{ @@ -27,25 +31,334 @@ use datafusion::{ CrossJoinExec, HashJoinExec, NestedLoopJoinExec, SortMergeJoinExec, SymmetricHashJoinExec, }, - metrics::MetricValue, + limit::{GlobalLimitExec, LocalLimitExec}, + metrics::{MetricValue, MetricsSet}, projection::ProjectionExec, sorts::{sort::SortExec, sort_preserving_merge::SortPreservingMergeExec}, - visit_execution_plan, ExecutionPlan, ExecutionPlanVisitor, + union::UnionExec, + windows::{BoundedWindowAggExec, WindowAggExec}, + ExecutionPlan, }, }; use itertools::Itertools; -use std::{sync::Arc, time::Duration}; +use log::debug; +use serde::{Deserialize, Serialize}; +use std::{collections::HashMap, sync::Arc, time::Duration}; + +/// Version of the analyze protocol implemented by this crate. Carried in the +/// metrics batch schema metadata under [`PROTOCOL_VERSION_METADATA_KEY`]. +pub const ANALYZE_PROTOCOL_VERSION: &str = "0.1"; + +/// Schema metadata key holding the protocol version of the metrics batch +pub const PROTOCOL_VERSION_METADATA_KEY: &str = "analyze.protocol_version"; + +/// Schema metadata key holding an opaque per-request correlation id +pub const QUERY_ID_METADATA_KEY: &str = "analyze.query_id"; + +/// Request structure for the analyze_query action +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AnalyzeQueryRequest { + /// SQL query to analyze (currently the only supported format) + pub sql: Option, + /// Protocol version the client speaks. Servers reject unknown major versions. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub protocol_version: Option, + // Future extensibility fields (not yet implemented): + // /// Substrait query plan (binary or JSON) + // pub substrait: Option>, + // /// Serialized logical plan + // pub logical_plan: Option, + // /// Serialized physical plan + // pub physical_plan: Option, +} + +impl AnalyzeQueryRequest { + /// Create a new request with a SQL query + pub fn with_sql(sql: impl Into) -> Self { + Self { + sql: Some(sql.into()), + protocol_version: Some(ANALYZE_PROTOCOL_VERSION.to_string()), + } + } + + /// Get the SQL query, returning an error if not present + pub fn sql(&self) -> color_eyre::Result<&str> { + self.sql + .as_deref() + .ok_or_else(|| color_eyre::eyre::eyre!("sql field is required")) + } +} + +/// Returns true when `version` is compatible with the protocol version this +/// crate implements (same major version) +pub fn is_compatible_protocol_version(version: &str) -> bool { + let major = |v: &str| v.split('.').next().map(str::to_string); + major(version) == major(ANALYZE_PROTOCOL_VERSION) +} + +/// Operator categories used to group metrics. This enum is the single source +/// of truth for classification, wire encoding, and display ordering. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum OperatorCategory { + Io, + Projection, + Filter, + Sort, + Aggregate, + Join, + Window, + Distinct, + Limit, + Union, + Exchange, + Other, +} + +impl OperatorCategory { + /// Compute categories in display order (everything except `Io`) + pub const COMPUTE: [OperatorCategory; 11] = [ + OperatorCategory::Projection, + OperatorCategory::Filter, + OperatorCategory::Sort, + OperatorCategory::Aggregate, + OperatorCategory::Join, + OperatorCategory::Window, + OperatorCategory::Distinct, + OperatorCategory::Limit, + OperatorCategory::Union, + OperatorCategory::Exchange, + OperatorCategory::Other, + ]; + + pub fn as_str(&self) -> &'static str { + match self { + OperatorCategory::Io => "io", + OperatorCategory::Projection => "projection", + OperatorCategory::Filter => "filter", + OperatorCategory::Sort => "sort", + OperatorCategory::Aggregate => "aggregate", + OperatorCategory::Join => "join", + OperatorCategory::Window => "window", + OperatorCategory::Distinct => "distinct", + OperatorCategory::Limit => "limit", + OperatorCategory::Union => "union", + OperatorCategory::Exchange => "exchange", + OperatorCategory::Other => "other", + } + } + + /// Display label for the category + fn label(&self) -> &'static str { + match self { + OperatorCategory::Io => "IO", + OperatorCategory::Projection => "Projection", + OperatorCategory::Filter => "Filter", + OperatorCategory::Sort => "Sort", + OperatorCategory::Aggregate => "Aggregate", + OperatorCategory::Join => "Join", + OperatorCategory::Window => "Window", + OperatorCategory::Distinct => "Distinct", + OperatorCategory::Limit => "Limit", + OperatorCategory::Union => "Union", + OperatorCategory::Exchange => "Exchange", + OperatorCategory::Other => "Other", + } + } + + fn from_str(s: &str) -> Option { + match s { + "io" => Some(OperatorCategory::Io), + "projection" => Some(OperatorCategory::Projection), + "filter" => Some(OperatorCategory::Filter), + "sort" => Some(OperatorCategory::Sort), + "aggregate" => Some(OperatorCategory::Aggregate), + "join" => Some(OperatorCategory::Join), + "window" => Some(OperatorCategory::Window), + "distinct" => Some(OperatorCategory::Distinct), + "limit" => Some(OperatorCategory::Limit), + "union" => Some(OperatorCategory::Union), + "exchange" => Some(OperatorCategory::Exchange), + "other" => Some(OperatorCategory::Other), + _ => None, + } + } + + /// Classify a plan node. Downcasts are used where the operator type is + /// public; a name-based fallback covers the remaining cases so operators + /// from newer DataFusion versions still land in a sensible category. + fn classify(plan: &dyn ExecutionPlan) -> Self { + if io_format(plan).is_some() { + return OperatorCategory::Io; + } + if plan.downcast_ref::().is_some() { + return OperatorCategory::Projection; + } + if plan.downcast_ref::().is_some() { + return OperatorCategory::Filter; + } + if plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + { + return OperatorCategory::Sort; + } + if plan.downcast_ref::().is_some() { + return OperatorCategory::Aggregate; + } + if plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + { + return OperatorCategory::Join; + } + if plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + { + return OperatorCategory::Window; + } + if plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + { + return OperatorCategory::Limit; + } + if plan.downcast_ref::().is_some() { + return OperatorCategory::Union; + } + // Name-based fallback for operators without a public type + let name = plan.name(); + if name.contains("Window") { + return OperatorCategory::Window; + } + if name.contains("Distinct") || name.contains("Deduplicate") { + return OperatorCategory::Distinct; + } + // Network-boundary operators from distributed engines (e.g. + // datafusion-distributed's NetworkShuffleExec/NetworkCoalesceExec, + // Ballista's ShuffleWriterExec/ShuffleReaderExec/ExchangeExec, or a + // RemoteExec-style fragment shipper). Local RepartitionExec is NOT an + // exchange: the category is reserved for network boundaries. + if name.contains("Network") + || name.contains("Exchange") + || name.contains("Shuffle") + || name.contains("RemoteExec") + { + return OperatorCategory::Exchange; + } + OperatorCategory::Other + } +} + +impl std::fmt::Display for OperatorCategory { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.as_str()) + } +} + +/// Represents the file format type for I/O operations +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum IOFormatType { + Csv, + Parquet, + Arrow, + Json, +} + +impl IOFormatType { + fn namespace_prefix(&self) -> &'static str { + match self { + IOFormatType::Csv => "io.csv", + IOFormatType::Parquet => "io.parquet", + IOFormatType::Arrow => "io.arrow", + IOFormatType::Json => "io.json", + } + } + + fn from_namespace(namespace: &str) -> Option { + match namespace { + "csv" => Some(IOFormatType::Csv), + "parquet" => Some(IOFormatType::Parquet), + "arrow" => Some(IOFormatType::Arrow), + "json" => Some(IOFormatType::Json), + _ => None, + } + } + + fn label(&self) -> &'static str { + match self { + IOFormatType::Csv => "CSV", + IOFormatType::Parquet => "Parquet", + IOFormatType::Arrow => "Arrow", + IOFormatType::Json => "JSON", + } + } +} + +/// Determine the I/O format of a plan node, if it is a file scan. In +/// DataFusion 51 all file scans are `DataSourceExec` nodes wrapping a +/// `FileScanConfig` that carries the format-specific `FileSource`. +fn io_format(plan: &dyn ExecutionPlan) -> Option { + let data_source_exec = plan.downcast_ref::()?; + let file_scan_config = data_source_exec + .data_source() + .downcast_ref::()?; + let source = file_scan_config.file_source(); + if source.downcast_ref::().is_some() { + Some(IOFormatType::Parquet) + } else if source.downcast_ref::().is_some() { + Some(IOFormatType::Csv) + } else if source.downcast_ref::().is_some() { + Some(IOFormatType::Json) + } else if source.downcast_ref::().is_some() { + Some(IOFormatType::Arrow) + } else { + None + } +} + +/// Identity and category of a single node in the execution plan. Node ids are +/// assigned by pre-order traversal with the root as 0, so a plan's DAG can be +/// reconstructed from `(node_id, parent_node_id)` pairs even when the same +/// operator type appears multiple times. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PlanNodeInfo { + node_id: i32, + parent_node_id: Option, + name: String, + category: OperatorCategory, +} + +impl PlanNodeInfo { + pub fn node_id(&self) -> i32 { + self.node_id + } + + pub fn parent_node_id(&self) -> Option { + self.parent_node_id + } + + pub fn name(&self) -> &str { + &self.name + } + + pub fn category(&self) -> OperatorCategory { + self.category + } +} #[derive(Clone, Debug)] pub struct ExecutionStats { query: String, rows: usize, - batches: i32, + batches: usize, bytes: usize, durations: ExecutionDurationStats, io: Option, compute: Option, - plan: Arc, + /// The executed physical plan. `None` when the stats were reconstructed + /// from a metrics table (e.g. received over FlightSQL). + plan: Option>, + nodes: Vec, } impl ExecutionStats { @@ -53,56 +366,85 @@ impl ExecutionStats { query: String, durations: ExecutionDurationStats, rows: usize, - batches: i32, + batches: usize, bytes: usize, plan: Arc, ) -> color_eyre::Result { + let collected = collect_node_stats(&plan); Ok(Self { query, durations, rows, batches, bytes, - plan, + plan: Some(plan), io: None, compute: None, + nodes: collected.nodes, }) } + /// Collect I/O and compute metrics from the executed plan. No-op when the + /// stats were reconstructed from a metrics table (no plan available). pub fn collect_stats(&mut self) { - if let Some(io) = collect_plan_io_stats(Arc::clone(&self.plan)) { - self.io = Some(io) - } - if let Some(compute) = collect_plan_compute_stats(Arc::clone(&self.plan)) { - self.compute = Some(compute) + let Some(plan) = &self.plan else { return }; + let collected = collect_node_stats(plan); + self.nodes = collected.nodes; + if !collected.io_nodes.is_empty() { + self.io = Some(ExecutionIOStats { + nodes: collected.io_nodes, + }); } + self.compute = Some(ExecutionComputeStats { + elapsed_compute: collected.elapsed_compute, + computes: collected.computes, + }); } - pub fn rows_selectivity(&self) -> f64 { - let maybe_io_output_rows = self.io.as_ref().and_then(|io| io.parquet_output_rows); - if let Some(io_output_rows) = maybe_io_output_rows { - self.rows as f64 / io_output_rows as f64 - } else { - 0.0 - } + pub fn nodes(&self) -> &[PlanNodeInfo] { + &self.nodes } - pub fn bytes_selectivity(&self) -> f64 { - let maybe_io_output_bytes = self.io.as_ref().and_then(|io| io.bytes_scanned.clone()); - if let Some(io_output_bytes) = maybe_io_output_bytes { - self.bytes as f64 / io_output_bytes.as_usize() as f64 - } else { - 0.0 - } + /// Fraction of scanned rows that made it into the query result. `None` + /// when no scan output-row metrics are available. + pub fn rows_selectivity(&self) -> Option { + let scan_rows: u64 = self + .io + .as_ref()? + .nodes + .iter() + .filter_map(|n| n.output_rows) + .sum(); + (scan_rows > 0).then(|| self.rows as f64 / scan_rows as f64) } - pub fn selectivity_efficiency(&self) -> f64 { - if let Some(io) = &self.io { - io.parquet_rg_pruned_stats_ratio() / self.rows_selectivity() - } else { - 0.0 - } + /// Ratio of result bytes (in-memory Arrow size) to bytes scanned. `None` + /// when no bytes-scanned metrics are available. + pub fn bytes_selectivity(&self) -> Option { + let scanned: u64 = self + .io + .as_ref()? + .nodes + .iter() + .filter_map(|n| n.bytes_scanned) + .sum(); + (scanned > 0).then(|| self.bytes as f64 / scanned as f64) } + + /// Ratio of the Parquet row-group pruning rate to row selectivity. Higher + /// values mean pruning removed more data relative to how selective the + /// query was. `None` when either input is unavailable. + pub fn selectivity_efficiency(&self) -> Option { + let matched_ratio = self.io.as_ref()?.parquet_rg_matched_ratio()?; + let selectivity = self.rows_selectivity()?; + (selectivity > 0.0).then(|| matched_ratio / selectivity) + } +} + +fn fmt_opt_ratio(value: Option) -> String { + value + .map(|v| format!("{:.2}", v)) + .unwrap_or_else(|| "N/A".to_string()) } impl std::fmt::Display for ExecutionStats { @@ -124,14 +466,18 @@ impl std::fmt::Display for ExecutionStats { writeln!( f, "{:<20} {:<20} {:<20}", - format!("{} ({:.2})", self.rows, self.rows_selectivity()), - format!("{} ({:.2})", self.bytes, self.bytes_selectivity()), + format!("{} ({})", self.rows, fmt_opt_ratio(self.rows_selectivity())), + format!( + "{} ({})", + self.bytes, + fmt_opt_ratio(self.bytes_selectivity()) + ), self.batches, )?; writeln!(f)?; writeln!(f, "{}", self.durations)?; writeln!(f, "{:<20}", "Parquet Efficiency (Pruning / Selectivity)")?; - writeln!(f, "{:<20.2}", self.selectivity_efficiency())?; + writeln!(f, "{:<20}", fmt_opt_ratio(self.selectivity_efficiency()))?; writeln!(f)?; if let Some(io_stats) = &self.io { writeln!(f, "{}", io_stats)?; @@ -143,7 +489,7 @@ impl std::fmt::Display for ExecutionStats { } } -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, Eq)] pub struct ExecutionDurationStats { parsing: Duration, logical_planning: Duration, @@ -189,228 +535,174 @@ impl std::fmt::Display for ExecutionDurationStats { } } -#[derive(Clone, Debug)] -pub struct ExecutionIOStats { - bytes_scanned: Option, - time_opening: Option, - time_scanning: Option, - parquet_output_rows: Option, - parquet_pruned_page_index: Option, - parquet_matched_page_index: Option, - parquet_rg_pruned_stats: Option, - parquet_rg_matched_stats: Option, - parquet_rg_pruned_bloom_filter: Option, - parquet_rg_matched_bloom_filter: Option, +/// I/O metrics for a single scan node. Values are plain integers so they +/// round-trip losslessly through the metrics table. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct IONodeStats { + node_id: Option, + operator_name: String, + format: IOFormatType, + bytes_scanned: Option, + time_opening_ns: Option, + time_scanning_ns: Option, + // Parquet-specific pruning metrics + output_rows: Option, + rg_pruned: Option, + rg_matched: Option, + bloom_pruned: Option, + bloom_matched: Option, + page_index_pruned: Option, + page_index_matched: Option, } -impl ExecutionIOStats { - fn parquet_rg_pruned_stats_ratio(&self) -> f64 { - if let (Some(pruned), Some(matched)) = ( - self.parquet_rg_matched_stats.as_ref(), - self.parquet_rg_pruned_stats.as_ref(), - ) { - let pruned = pruned.as_usize() as f64; - let matched = matched.as_usize() as f64; - matched / (pruned + matched) - } else { - 0.0 - } +impl IONodeStats { + fn matched_ratio(pruned: Option, matched: Option) -> Option { + let (pruned, matched) = (pruned?, matched?); + let total = pruned + matched; + (total > 0).then(|| matched as f64 / total as f64) } - fn parquet_rg_pruned_bloom_filter_ratio(&self) -> f64 { - if let (Some(pruned), Some(matched)) = ( - self.parquet_rg_matched_bloom_filter.as_ref(), - self.parquet_rg_pruned_bloom_filter.as_ref(), - ) { - let pruned = pruned.as_usize() as f64; - let matched = matched.as_usize() as f64; - matched / (pruned + matched) - } else { - 0.0 - } + fn rg_matched_ratio(&self) -> Option { + Self::matched_ratio(self.rg_pruned, self.rg_matched) } - fn parquet_rg_pruned_page_index_ratio(&self) -> f64 { - if let (Some(pruned), Some(matched)) = ( - self.parquet_matched_page_index.as_ref(), - self.parquet_pruned_page_index.as_ref(), - ) { - let pruned = pruned.as_usize() as f64; - let matched = matched.as_usize() as f64; - matched / (pruned + matched) - } else { - 0.0 - } + fn bloom_matched_ratio(&self) -> Option { + Self::matched_ratio(self.bloom_pruned, self.bloom_matched) } - fn row_group_count(&self) -> usize { - if let (Some(pruned), Some(matched)) = ( - self.parquet_rg_matched_stats.as_ref(), - self.parquet_rg_pruned_stats.as_ref(), - ) { - let pruned = pruned.as_usize(); - let matched = matched.as_usize(); - pruned + matched - } else { - 0 + fn page_index_matched_ratio(&self) -> Option { + Self::matched_ratio(self.page_index_pruned, self.page_index_matched) + } + + fn row_group_count(&self) -> Option { + match (self.rg_pruned, self.rg_matched) { + (Some(p), Some(m)) => Some(p + m), + _ => None, } } } -impl std::fmt::Display for ExecutionIOStats { +impl std::fmt::Display for IONodeStats { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let node = self + .node_id + .map(|id| format!("node {id}")) + .unwrap_or_else(|| "node ?".to_string()); writeln!( f, - "======================= IO Summary ========================" + "{} ({}) [{}]", + self.operator_name, + node, + self.format.label() )?; writeln!( f, "{:<20} {:<20} {:<20}", "Bytes Scanned", "Time Opening", "Time Scanning" )?; - writeln!( - f, - "{:<20} {:<20} {:<20}", - self.bytes_scanned - .as_ref() - .map(|m| m.to_string()) - .unwrap_or("None".to_string()), - self.time_opening - .as_ref() - .map(|m| m.to_string()) - .unwrap_or("None".to_string()), - self.time_scanning - .as_ref() - .map(|m| m.to_string()) - .unwrap_or("None".to_string()) - )?; - writeln!(f)?; - writeln!( - f, - "Parquet Pruning Stats (Output Rows: {}, Row Groups: {} [{}ms per row group])", - self.parquet_output_rows - .as_ref() - .map(|m| m.to_string()) - .unwrap_or("None".to_string()), - self.row_group_count(), - self.time_scanning - .as_ref() - .map(|ts| format!( - "{:.2}", - (ts.as_usize() / 1_000_000) as f64 / self.row_group_count() as f64 - )) + let fmt_opt_u64 = |v: Option| v.map(|v| v.to_string()).unwrap_or("None".to_string()); + let fmt_opt_ns = |v: Option| { + v.map(|v| format!("{:?}", Duration::from_nanos(v))) .unwrap_or("None".to_string()) - )?; + }; writeln!( f, "{:<20} {:<20} {:<20}", - "Matched RG Stats %", "Matched RG Bloom %", "Matched Page Index %" - )?; - writeln!( - f, - "{:<20.2} {:<20.2} {:<20.2}", - self.parquet_rg_pruned_stats_ratio(), - self.parquet_rg_pruned_bloom_filter_ratio(), - self.parquet_rg_pruned_page_index_ratio() + fmt_opt_u64(self.bytes_scanned), + fmt_opt_ns(self.time_opening_ns), + fmt_opt_ns(self.time_scanning_ns), )?; + if self.format == IOFormatType::Parquet { + let scan_time_per_rg = match (self.time_scanning_ns, self.row_group_count()) { + (Some(ns), Some(rgs)) if rgs > 0 => { + format!( + "{:.2}ms scan time per row group", + ns as f64 / 1e6 / rgs as f64 + ) + } + _ => "N/A".to_string(), + }; + writeln!( + f, + "Parquet Pruning Stats (Output Rows: {}, Row Groups: {} [{}])", + fmt_opt_u64(self.output_rows), + self.row_group_count() + .map(|v| v.to_string()) + .unwrap_or("None".to_string()), + scan_time_per_rg, + )?; + writeln!( + f, + "{:<20} {:<20} {:<20}", + "Matched RG Stats", "Matched RG Bloom", "Matched Page Index" + )?; + writeln!( + f, + "{:<20} {:<20} {:<20}", + fmt_opt_ratio(self.rg_matched_ratio()), + fmt_opt_ratio(self.bloom_matched_ratio()), + fmt_opt_ratio(self.page_index_matched_ratio()), + )?; + } Ok(()) } } -/// Visitor to collect IO metrics from an execution plan -/// -/// IO metrics are collected from nodes that perform IO operations, such as -/// `CsvExec`, `ParquetExec`, and `ArrowExec`. -struct PlanIOVisitor { - bytes_scanned: Option, - time_opening: Option, - time_scanning: Option, - parquet_output_rows: Option, - parquet_pruned_page_index: Option, - parquet_matched_page_index: Option, - parquet_rg_pruned_stats: Option, - parquet_rg_matched_stats: Option, - parquet_rg_pruned_bloom_filter: Option, - parquet_rg_matched_bloom_filter: Option, -} - -impl PlanIOVisitor { - fn new() -> Self { - Self { - bytes_scanned: None, - time_opening: None, - time_scanning: None, - parquet_output_rows: None, - parquet_pruned_page_index: None, - parquet_matched_page_index: None, - parquet_rg_pruned_stats: None, - parquet_rg_matched_stats: None, - parquet_rg_pruned_bloom_filter: None, - parquet_rg_matched_bloom_filter: None, - } - } - - fn collect_io_metrics(&mut self, plan: &dyn ExecutionPlan) { - let io_metrics = plan.metrics(); - if let Some(metrics) = io_metrics { - self.bytes_scanned = metrics.sum_by_name("bytes_scanned"); - self.time_opening = metrics.sum_by_name("time_elapsed_opening"); - self.time_scanning = metrics.sum_by_name("time_elapsed_scanning_total"); - - if let Some(data_source_exec) = plan.downcast_ref::() { - if data_source_exec - .data_source() - .downcast_ref::() - .is_some_and(|config| { - config - .file_source() - .downcast_ref::() - .is_some() - }) - { - self.parquet_output_rows = metrics.output_rows(); - self.parquet_rg_pruned_stats = - metrics.sum_by_name("row_groups_pruned_statistics"); - self.parquet_rg_matched_stats = - metrics.sum_by_name("row_groups_matched_statistics"); - } - } - } - } +/// I/O metrics for all scan nodes in the plan +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ExecutionIOStats { + nodes: Vec, } -impl From for ExecutionIOStats { - fn from(value: PlanIOVisitor) -> Self { - Self { - bytes_scanned: value.bytes_scanned, - time_opening: value.time_opening, - time_scanning: value.time_scanning, - parquet_output_rows: value.parquet_output_rows, - parquet_pruned_page_index: value.parquet_pruned_page_index, - parquet_matched_page_index: value.parquet_matched_page_index, - parquet_rg_pruned_stats: value.parquet_rg_pruned_stats, - parquet_rg_matched_stats: value.parquet_rg_matched_stats, - parquet_rg_pruned_bloom_filter: value.parquet_rg_pruned_bloom_filter, - parquet_rg_matched_bloom_filter: value.parquet_rg_matched_bloom_filter, +impl ExecutionIOStats { + pub fn nodes(&self) -> &[IONodeStats] { + &self.nodes + } + + /// Aggregate Parquet row-group matched ratio across all scan nodes + fn parquet_rg_matched_ratio(&self) -> Option { + let mut pruned = 0u64; + let mut matched = 0u64; + let mut any = false; + for node in &self.nodes { + if let (Some(p), Some(m)) = (node.rg_pruned, node.rg_matched) { + pruned += p; + matched += m; + any = true; + } + } + if !any { + return None; } + IONodeStats::matched_ratio(Some(pruned), Some(matched)) } } -impl ExecutionPlanVisitor for PlanIOVisitor { - type Error = datafusion::common::DataFusionError; - - fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> color_eyre::Result { - if is_io_plan(plan) { - self.collect_io_metrics(plan); +impl std::fmt::Display for ExecutionIOStats { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!( + f, + "======================= IO Summary ========================" + )?; + for (i, node) in self.nodes.iter().enumerate() { + if i > 0 { + writeln!(f)?; + } + write!(f, "{}", node)?; } - Ok(true) + Ok(()) } } -#[derive(Clone, Debug)] +/// Per-partition elapsed-compute values for a single plan node. +/// +/// `elapsed_computes` is sorted ascending; the wire `partition_id` is the rank +/// in this ordering, not the physical partition number. +#[derive(Clone, Debug, PartialEq, Eq)] pub struct PartitionsComputeStats { + node_id: Option, name: String, - /// Sorted elapsed compute times + category: OperatorCategory, elapsed_computes: Vec, } @@ -433,58 +725,60 @@ impl PartitionsComputeStats { } } -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, Eq)] pub struct ExecutionComputeStats { elapsed_compute: Option, - projection_compute: Option>, - filter_compute: Option>, - sort_compute: Option>, - join_compute: Option>, - aggregate_compute: Option>, - other_compute: Option>, + computes: Vec, } impl ExecutionComputeStats { - fn display_compute( + fn display_category( &self, f: &mut std::fmt::Formatter<'_>, - compute: &Option>, - label: &str, + category: OperatorCategory, ) -> std::fmt::Result { - if let (Some(filter_compute), Some(elapsed_compute)) = (compute, &self.elapsed_compute) { - let partitions = filter_compute.iter().fold(0, |acc, c| acc + c.partitions()); - writeln!( - f, - "{label} Stats ({} nodes, {} partitions)", - filter_compute.len(), - partitions - )?; + let nodes: Vec<&PartitionsComputeStats> = self + .computes + .iter() + .filter(|c| c.category == category) + .collect(); + if nodes.is_empty() { + return writeln!(f, "{}: No data", category.label()); + } + let partitions = nodes.iter().fold(0, |acc, c| acc + c.partitions()); + writeln!( + f, + "{}: {} nodes, {} partitions", + category.label(), + nodes.len(), + partitions + )?; + writeln!( + f, + "{:<30} {:<16} {:<16} {:<16} {:<16} {:<16}", + "Node(Partitions)", "Min", "Median", "Mean", "Max", "Total (%)" + )?; + nodes.iter().try_for_each(|node| { + let (min, median, mean, max, total) = node.summary_stats(); + let total = match &self.elapsed_compute { + Some(elapsed) if *elapsed > 0 => format!( + "{} ({:.2}%)", + total, + (total as f32 / *elapsed as f32) * 100.0 + ), + _ => total.to_string(), + }; writeln!( f, "{:<30} {:<16} {:<16} {:<16} {:<16} {:<16}", - "Node(Partitions)", "Min", "Median", "Mean", "Max", "Total (%)" - )?; - filter_compute.iter().try_for_each(|node| { - let (min, median, mean, max, total) = node.summary_stats(); - let total = format!( - "{} ({:.2}%)", - total, - (total as f32 / *elapsed_compute as f32) * 100.0 - ); - writeln!( - f, - "{:<30} {:<16} {:<16} {:<16} {:<16} {:<16}", - format!("{}({})", node.name, node.elapsed_computes.len()), - min, - median, - mean, - max, - total, - ) - }) - } else { - writeln!(f, "No {label} Stats") - } + format!("{}({})", node.name, node.elapsed_computes.len()), + min, + median, + mean, + max, + total, + ) + }) } } @@ -504,253 +798,594 @@ impl std::fmt::Display for ExecutionComputeStats { .unwrap_or("None".to_string()), )?; writeln!(f)?; - self.display_compute(f, &self.projection_compute, "Projection")?; - writeln!(f)?; - self.display_compute(f, &self.filter_compute, "Filter")?; - writeln!(f)?; - self.display_compute(f, &self.sort_compute, "Sort")?; - writeln!(f)?; - self.display_compute(f, &self.join_compute, "Join")?; - writeln!(f)?; - self.display_compute(f, &self.aggregate_compute, "Aggregate")?; - writeln!(f)?; - self.display_compute(f, &self.other_compute, "Other")?; - writeln!(f) + + for category in OperatorCategory::COMPUTE { + self.display_category(f, category)?; + writeln!(f)?; + } + Ok(()) } } +/// Result of walking an execution plan once: node identity plus per-node +/// I/O and compute metrics #[derive(Default)] -pub struct PlanComputeVisitor { +struct CollectedPlanStats { + nodes: Vec, + io_nodes: Vec, + computes: Vec, elapsed_compute: Option, - filter_computes: Vec, - sort_computes: Vec, - projection_computes: Vec, - join_computes: Vec, - aggregate_computes: Vec, - other_computes: Vec, -} - -impl PlanComputeVisitor { - fn add_elapsed_compute(&mut self, node_elapsed_compute: Option) { - match (self.elapsed_compute, node_elapsed_compute) { - (Some(agg_elapsed_compute), Some(node_elapsed_compute)) => { - self.elapsed_compute = Some(agg_elapsed_compute + node_elapsed_compute) +} + +fn collect_node_stats(plan: &Arc) -> CollectedPlanStats { + let mut collected = CollectedPlanStats::default(); + let mut next_id = 0i32; + walk_plan(plan, None, &mut next_id, &mut collected); + collected +} + +fn walk_plan( + plan: &Arc, + parent_node_id: Option, + next_id: &mut i32, + out: &mut CollectedPlanStats, +) { + let node_id = *next_id; + *next_id += 1; + + let category = OperatorCategory::classify(plan.as_ref()); + out.nodes.push(PlanNodeInfo { + node_id, + parent_node_id, + name: plan.name().to_string(), + category, + }); + + if let Some(metrics) = plan.metrics() { + if category == OperatorCategory::Io { + if let Some(format) = io_format(plan.as_ref()) { + out.io_nodes.push(collect_io_node_stats( + node_id, + plan.name(), + format, + &metrics, + )); } - (Some(_), None) | (None, None) => {} - (None, Some(node_elapsed_compute)) => self.elapsed_compute = Some(node_elapsed_compute), - } - } - - fn collect_compute_metrics(&mut self, plan: &dyn ExecutionPlan) { - let compute_metrics = plan.metrics(); - if let Some(metrics) = compute_metrics { - self.add_elapsed_compute(metrics.elapsed_compute()); - } - self.collect_filter_metrics(plan); - self.collect_sort_metrics(plan); - self.collect_projection_metrics(plan); - self.collect_join_metrics(plan); - self.collect_aggregate_metrics(plan); - self.collect_other_metrics(plan); - } - - // TODO: Refactor to have a single function that takes predicate and collector - fn collect_filter_metrics(&mut self, plan: &dyn ExecutionPlan) { - if is_filter_plan(plan) { - if let Some(metrics) = plan.metrics() { - let sorted_computes: Vec = metrics - .iter() - .filter_map(|m| match m.value() { - MetricValue::ElapsedCompute(t) => Some(t.value()), - _ => None, - }) - .sorted() - .collect(); - let p = PartitionsComputeStats { - name: plan.name().to_string(), - elapsed_computes: sorted_computes, - }; - self.filter_computes.push(p) + } else { + if let Some(node_elapsed) = metrics.elapsed_compute() { + out.elapsed_compute = Some(out.elapsed_compute.unwrap_or(0) + node_elapsed); } - } - } - - fn collect_sort_metrics(&mut self, plan: &dyn ExecutionPlan) { - if is_sort_plan(plan) { - if let Some(metrics) = plan.metrics() { - let sorted_computes: Vec = metrics - .iter() - .filter_map(|m| match m.value() { - MetricValue::ElapsedCompute(t) => Some(t.value()), - _ => None, - }) - .sorted() - .collect(); - let p = PartitionsComputeStats { + let sorted_computes: Vec = metrics + .iter() + .filter_map(|m| match m.value() { + MetricValue::ElapsedCompute(t) => Some(t.value()), + _ => None, + }) + .sorted() + .collect(); + if !sorted_computes.is_empty() { + out.computes.push(PartitionsComputeStats { + node_id: Some(node_id), name: plan.name().to_string(), + category, elapsed_computes: sorted_computes, - }; - self.sort_computes.push(p) + }); } } } - fn collect_projection_metrics(&mut self, plan: &dyn ExecutionPlan) { - if is_projection_plan(plan) { - if let Some(metrics) = plan.metrics() { - let sorted_computes: Vec = metrics - .iter() - .filter_map(|m| match m.value() { - MetricValue::ElapsedCompute(t) => Some(t.value()), - _ => None, - }) - .sorted() - .collect(); - let p = PartitionsComputeStats { - name: plan.name().to_string(), - elapsed_computes: sorted_computes, - }; - self.projection_computes.push(p) - } - } + for child in plan.children() { + walk_plan(child, Some(node_id), next_id, out); } +} - fn collect_join_metrics(&mut self, plan: &dyn ExecutionPlan) { - if is_join_plan(plan) { - if let Some(metrics) = plan.metrics() { - let sorted_computes: Vec = metrics - .iter() - .filter_map(|m| match m.value() { - MetricValue::ElapsedCompute(t) => Some(t.value()), - _ => None, - }) - .sorted() - .collect(); - let p = PartitionsComputeStats { - name: plan.name().to_string(), - elapsed_computes: sorted_computes, - }; - self.join_computes.push(p) - } - } - } +/// Sum a named metric across partitions, as a plain integer +fn metric_u64(metrics: &MetricsSet, name: &str) -> Option { + metrics.sum_by_name(name).map(|v| v.as_usize() as u64) +} - fn collect_aggregate_metrics(&mut self, plan: &dyn ExecutionPlan) { - if is_aggregate_plan(plan) { - if let Some(metrics) = plan.metrics() { - let sorted_computes: Vec = metrics - .iter() - .filter_map(|m| match m.value() { - MetricValue::ElapsedCompute(t) => Some(t.value()), - _ => None, - }) - .sorted() - .collect(); - let p = PartitionsComputeStats { - name: plan.name().to_string(), - elapsed_computes: sorted_computes, - }; - self.aggregate_computes.push(p) - } +/// Sum a named pruning metric across partitions, returning `(pruned, matched)`. +/// `PruningMetrics` cannot be read via `as_usize` (it always returns 0). +fn pruning_counts(metrics: &MetricsSet, name: &str) -> Option<(u64, u64)> { + metrics.sum_by_name(name).and_then(|v| match v { + MetricValue::PruningMetrics { + pruning_metrics, .. + } => Some(( + pruning_metrics.pruned() as u64, + pruning_metrics.matched() as u64, + )), + _ => None, + }) +} + +fn collect_io_node_stats( + node_id: i32, + operator_name: &str, + format: IOFormatType, + metrics: &MetricsSet, +) -> IONodeStats { + let mut stats = IONodeStats { + node_id: Some(node_id), + operator_name: operator_name.to_string(), + format, + bytes_scanned: metric_u64(metrics, "bytes_scanned"), + time_opening_ns: metric_u64(metrics, "time_elapsed_opening"), + time_scanning_ns: metric_u64(metrics, "time_elapsed_scanning_total"), + output_rows: None, + rg_pruned: None, + rg_matched: None, + bloom_pruned: None, + bloom_matched: None, + page_index_pruned: None, + page_index_matched: None, + }; + if format == IOFormatType::Parquet { + stats.output_rows = metrics.output_rows().map(|v| v as u64); + if let Some((pruned, matched)) = pruning_counts(metrics, "row_groups_pruned_statistics") { + stats.rg_pruned = Some(pruned); + stats.rg_matched = Some(matched); + } + if let Some((pruned, matched)) = pruning_counts(metrics, "row_groups_pruned_bloom_filter") { + stats.bloom_pruned = Some(pruned); + stats.bloom_matched = Some(matched); + } + if let Some((pruned, matched)) = pruning_counts(metrics, "page_index_rows_pruned") { + stats.page_index_pruned = Some(pruned); + stats.page_index_matched = Some(matched); } } + stats +} - fn collect_other_metrics(&mut self, plan: &dyn ExecutionPlan) { - if !is_filter_plan(plan) - && !is_sort_plan(plan) - && !is_projection_plan(plan) - && !is_aggregate_plan(plan) - && !is_join_plan(plan) - { - if let Some(metrics) = plan.metrics() { - let sorted_computes: Vec = metrics - .iter() - .filter_map(|m| match m.value() { - MetricValue::ElapsedCompute(t) => Some(t.value()), - _ => None, - }) - .sorted() - .collect(); - let p = PartitionsComputeStats { - name: plan.name().to_string(), - elapsed_computes: sorted_computes, - }; - self.other_computes.push(p) - } - } +pub fn collect_plan_io_stats(plan: Arc) -> Option { + let collected = collect_node_stats(&plan); + if collected.io_nodes.is_empty() { + None + } else { + Some(ExecutionIOStats { + nodes: collected.io_nodes, + }) } } -fn is_filter_plan(plan: &dyn ExecutionPlan) -> bool { - plan.downcast_ref::().is_some() +pub fn collect_plan_compute_stats(plan: Arc) -> Option { + let collected = collect_node_stats(&plan); + Some(ExecutionComputeStats { + elapsed_compute: collected.elapsed_compute, + computes: collected.computes, + }) } -fn is_sort_plan(plan: &dyn ExecutionPlan) -> bool { - plan.downcast_ref::().is_some() - || plan.downcast_ref::().is_some() +/// Standard Arrow schema for analyze metrics. The schema metadata carries the +/// protocol version under [`PROTOCOL_VERSION_METADATA_KEY`]. +pub fn analyze_metrics_schema() -> SchemaRef { + let fields = vec![ + Field::new("metric_name", DataType::Utf8, false), + Field::new("value", DataType::UInt64, false), + Field::new("value_type", DataType::Utf8, false), + Field::new("operator_name", DataType::Utf8, true), + Field::new("partition_id", DataType::Int32, true), + Field::new("operator_category", DataType::Utf8, true), + Field::new("node_id", DataType::Int32, true), + Field::new("parent_node_id", DataType::Int32, true), + ]; + let metadata = HashMap::from([( + PROTOCOL_VERSION_METADATA_KEY.to_string(), + ANALYZE_PROTOCOL_VERSION.to_string(), + )]); + Arc::new(Schema::new_with_metadata(fields, metadata)) } -fn is_projection_plan(plan: &dyn ExecutionPlan) -> bool { - plan.downcast_ref::().is_some() +/// A single row of the metrics table +struct MetricRow<'a> { + metric_name: &'a str, + value: u64, + value_type: &'a str, + operator_name: Option<&'a str>, + partition_id: Option, + operator_category: Option, + node_id: Option, + parent_node_id: Option, } -fn is_join_plan(plan: &dyn ExecutionPlan) -> bool { - plan.downcast_ref::().is_some() - || plan.downcast_ref::().is_some() - || plan.downcast_ref::().is_some() - || plan.downcast_ref::().is_some() - || plan.downcast_ref::().is_some() +impl<'a> MetricRow<'a> { + /// A query- or stage-level metric, not attached to any operator + fn query_level(metric_name: &'a str, value: u64, value_type: &'a str) -> Self { + Self { + metric_name, + value, + value_type, + operator_name: None, + partition_id: None, + operator_category: None, + node_id: None, + parent_node_id: None, + } + } } -fn is_aggregate_plan(plan: &dyn ExecutionPlan) -> bool { - plan.downcast_ref::().is_some() +/// Helper to build metrics table rows +struct MetricsTableBuilder { + metric_names: Vec, + values: Vec, + value_types: Vec, + operator_names: Vec>, + partition_ids: Vec>, + operator_categories: Vec>, + node_ids: Vec>, + parent_node_ids: Vec>, } -impl From for ExecutionComputeStats { - fn from(value: PlanComputeVisitor) -> Self { +impl MetricsTableBuilder { + fn new() -> Self { Self { - elapsed_compute: value.elapsed_compute, - filter_compute: Some(value.filter_computes), - sort_compute: Some(value.sort_computes), - projection_compute: Some(value.projection_computes), - join_compute: Some(value.join_computes), - aggregate_compute: Some(value.aggregate_computes), - other_compute: Some(value.other_computes), + metric_names: Vec::new(), + values: Vec::new(), + value_types: Vec::new(), + operator_names: Vec::new(), + partition_ids: Vec::new(), + operator_categories: Vec::new(), + node_ids: Vec::new(), + parent_node_ids: Vec::new(), } } -} -impl ExecutionPlanVisitor for PlanComputeVisitor { - type Error = datafusion::common::DataFusionError; + fn add(&mut self, row: MetricRow<'_>) { + self.metric_names.push(row.metric_name.to_string()); + self.values.push(row.value); + self.value_types.push(row.value_type.to_string()); + self.operator_names + .push(row.operator_name.map(String::from)); + self.partition_ids.push(row.partition_id); + self.operator_categories + .push(row.operator_category.map(|c| c.as_str().to_string())); + self.node_ids.push(row.node_id); + self.parent_node_ids.push(row.parent_node_id); + } - fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> color_eyre::Result { - if !is_io_plan(plan) { - self.collect_compute_metrics(plan); - } - Ok(true) + fn build(self, schema: SchemaRef) -> color_eyre::Result { + let metric_names_array: ArrayRef = Arc::new(StringArray::from(self.metric_names)); + let values_array: ArrayRef = Arc::new(UInt64Array::from(self.values)); + let value_types_array: ArrayRef = Arc::new(StringArray::from(self.value_types)); + let operator_names_array: ArrayRef = Arc::new(StringArray::from(self.operator_names)); + let partition_ids_array: ArrayRef = Arc::new(Int32Array::from(self.partition_ids)); + let operator_categories_array: ArrayRef = + Arc::new(StringArray::from(self.operator_categories)); + let node_ids_array: ArrayRef = Arc::new(Int32Array::from(self.node_ids)); + let parent_node_ids_array: ArrayRef = Arc::new(Int32Array::from(self.parent_node_ids)); + + Ok(RecordBatch::try_new( + schema, + vec![ + metric_names_array, + values_array, + value_types_array, + operator_names_array, + partition_ids_array, + operator_categories_array, + node_ids_array, + parent_node_ids_array, + ], + )?) } } -fn is_io_plan(plan: &dyn ExecutionPlan) -> bool { - let io_plans = ["CsvExec", "ParquetExec", "ArrowExec"]; - io_plans.contains(&plan.name()) -} +impl ExecutionStats { + /// Serialize ExecutionStats to metrics table format + pub fn to_metrics_table(&self) -> color_eyre::Result { + let schema = analyze_metrics_schema(); + let mut rows = MetricsTableBuilder::new(); -pub fn collect_plan_io_stats(plan: Arc) -> Option { - let mut visitor = PlanIOVisitor::new(); - if visit_execution_plan(plan.as_ref(), &mut visitor).is_ok() { - Some(visitor.into()) - } else { - None + rows.add(MetricRow::query_level( + "query.rows", + self.rows as u64, + "count", + )); + rows.add(MetricRow::query_level( + "query.batches", + self.batches as u64, + "count", + )); + rows.add(MetricRow::query_level( + "query.bytes", + self.bytes as u64, + "bytes", + )); + + rows.add(MetricRow::query_level( + "stage.parsing", + self.durations.parsing.as_nanos() as u64, + "duration_ns", + )); + rows.add(MetricRow::query_level( + "stage.logical_planning", + self.durations.logical_planning.as_nanos() as u64, + "duration_ns", + )); + rows.add(MetricRow::query_level( + "stage.physical_planning", + self.durations.physical_planning.as_nanos() as u64, + "duration_ns", + )); + rows.add(MetricRow::query_level( + "stage.execution", + self.durations.execution.as_nanos() as u64, + "duration_ns", + )); + rows.add(MetricRow::query_level( + "stage.total", + self.durations.total.as_nanos() as u64, + "duration_ns", + )); + + // Map node_id -> parent_node_id for operator-level rows + let parents: HashMap> = self + .nodes + .iter() + .map(|n| (n.node_id, n.parent_node_id)) + .collect(); + let parent_of = + |node_id: Option| node_id.and_then(|id| parents.get(&id).copied()).flatten(); + + if let Some(io) = &self.io { + for node in &io.nodes { + let namespace = node.format.namespace_prefix(); + let parent_node_id = parent_of(node.node_id); + let mut add_io_metric = |suffix: &str, value: Option, value_type: &str| { + if let Some(value) = value { + rows.add(MetricRow { + metric_name: &format!("{namespace}.{suffix}"), + value, + value_type, + operator_name: Some(&node.operator_name), + partition_id: None, + operator_category: Some(OperatorCategory::Io), + node_id: node.node_id, + parent_node_id, + }); + } + }; + add_io_metric("bytes_scanned", node.bytes_scanned, "bytes"); + add_io_metric("time_opening", node.time_opening_ns, "duration_ns"); + add_io_metric("time_scanning", node.time_scanning_ns, "duration_ns"); + if node.format == IOFormatType::Parquet { + add_io_metric("output_rows", node.output_rows, "count"); + add_io_metric("rg_pruned", node.rg_pruned, "count"); + add_io_metric("rg_matched", node.rg_matched, "count"); + add_io_metric("bloom_pruned", node.bloom_pruned, "count"); + add_io_metric("bloom_matched", node.bloom_matched, "count"); + add_io_metric("page_index_pruned", node.page_index_pruned, "count"); + add_io_metric("page_index_matched", node.page_index_matched, "count"); + } + } + } + + if let Some(compute) = &self.compute { + if let Some(elapsed) = compute.elapsed_compute { + rows.add(MetricRow::query_level( + "compute.elapsed_compute", + elapsed as u64, + "duration_ns", + )); + } + + for node in &compute.computes { + let parent_node_id = parent_of(node.node_id); + for (partition_id, elapsed) in node.elapsed_computes.iter().enumerate() { + rows.add(MetricRow { + metric_name: "compute.elapsed_compute", + value: *elapsed as u64, + value_type: "duration_ns", + operator_name: Some(&node.name), + partition_id: Some(partition_id as i32), + operator_category: Some(node.category), + node_id: node.node_id, + parent_node_id, + }); + } + } + } + + rows.build(schema) } -} -pub fn collect_plan_compute_stats(plan: Arc) -> Option { - let mut visitor = PlanComputeVisitor::default(); - if visit_execution_plan(plan.as_ref(), &mut visitor).is_ok() { - Some(visitor.into()) - } else { - None + /// Deserialize ExecutionStats from a metrics table + pub fn from_metrics_table(batch: RecordBatch, query: String) -> color_eyre::Result { + let column_string = |idx: usize, name: &str| -> color_eyre::Result<&StringArray> { + batch + .column(idx) + .as_any() + .downcast_ref::() + .ok_or_else(|| color_eyre::eyre::eyre!("Invalid {name} column type")) + }; + let column_i32 = |idx: usize, name: &str| -> color_eyre::Result<&Int32Array> { + batch + .column(idx) + .as_any() + .downcast_ref::() + .ok_or_else(|| color_eyre::eyre::eyre!("Invalid {name} column type")) + }; + + let metric_names = column_string(0, "metric_name")?; + let values = batch + .column(1) + .as_any() + .downcast_ref::() + .ok_or_else(|| color_eyre::eyre::eyre!("Invalid value column type"))?; + let operator_names = column_string(3, "operator_name")?; + let partition_ids = column_i32(4, "partition_id")?; + let operator_categories = column_string(5, "operator_category")?; + let node_ids = column_i32(6, "node_id")?; + let parent_node_ids = column_i32(7, "parent_node_id")?; + + let opt_str = |arr: &StringArray, idx: usize| -> Option { + (!arr.is_null(idx)).then(|| arr.value(idx).to_string()) + }; + let opt_i32 = |arr: &Int32Array, idx: usize| -> Option { + (!arr.is_null(idx)).then(|| arr.value(idx)) + }; + + let mut rows = 0usize; + let mut batches = 0usize; + let mut bytes = 0usize; + let mut parsing = Duration::ZERO; + let mut logical_planning = Duration::ZERO; + let mut physical_planning = Duration::ZERO; + let mut execution = Duration::ZERO; + let mut total = Duration::ZERO; + let mut elapsed_compute: Option = None; + + // Operator identity reconstructed from the table + let mut node_map: HashMap = HashMap::new(); + // (node_id, operator_name) -> per-partition compute values + type ComputeKey = (Option, String); + let mut compute_map: HashMap)> = + HashMap::new(); + // (node_id, operator_name) -> (format, metric suffix -> value) + let mut io_map: HashMap)> = HashMap::new(); + let mut io_order: Vec = Vec::new(); + + for row_idx in 0..batch.num_rows() { + let metric_name = metric_names.value(row_idx); + let value = values.value(row_idx); + let operator_name = opt_str(operator_names, row_idx); + let partition_id = opt_i32(partition_ids, row_idx); + let category = + opt_str(operator_categories, row_idx).and_then(|c| OperatorCategory::from_str(&c)); + let node_id = opt_i32(node_ids, row_idx); + let parent_node_id = opt_i32(parent_node_ids, row_idx); + + if let (Some(id), Some(name), Some(cat)) = (node_id, &operator_name, category) { + node_map.entry(id).or_insert_with(|| PlanNodeInfo { + node_id: id, + parent_node_id, + name: name.clone(), + category: cat, + }); + } + + match (metric_name, category) { + ("query.rows", None) => rows = value as usize, + ("query.batches", None) => batches = value as usize, + ("query.bytes", None) => bytes = value as usize, + ("stage.parsing", None) => parsing = Duration::from_nanos(value), + ("stage.logical_planning", None) => logical_planning = Duration::from_nanos(value), + ("stage.physical_planning", None) => { + physical_planning = Duration::from_nanos(value) + } + ("stage.execution", None) => execution = Duration::from_nanos(value), + ("stage.total", None) => total = Duration::from_nanos(value), + ("compute.elapsed_compute", None) => elapsed_compute = Some(value as usize), + ("compute.elapsed_compute", Some(cat)) => { + let key = ( + node_id, + operator_name + .clone() + .unwrap_or_else(|| "Unknown".to_string()), + ); + compute_map + .entry(key) + .or_insert_with(|| (cat, Vec::new())) + .1 + .push((partition_id.unwrap_or(0), value)); + } + (name, Some(OperatorCategory::Io)) => { + // io.. + let mut parts = name.splitn(3, '.'); + match (parts.next(), parts.next(), parts.next()) { + (Some("io"), Some(fmt), Some(suffix)) => { + if let Some(format) = IOFormatType::from_namespace(fmt) { + let key = ( + node_id, + operator_name + .clone() + .unwrap_or_else(|| "Unknown".to_string()), + ); + let entry = io_map.entry(key.clone()).or_insert_with(|| { + io_order.push(key); + (format, HashMap::new()) + }); + entry.1.insert(suffix.to_string(), value); + } else { + debug!("Unknown io namespace in metric: {}", name); + } + } + _ => debug!("Malformed io metric name: {}", name), + } + } + (name, category) => { + debug!("Unknown metric: {} (category: {:?})", name, category); + } + } + } + + let io_nodes: Vec = io_order + .into_iter() + .filter_map(|key| { + let (format, metrics) = io_map.remove(&key)?; + let (node_id, operator_name) = key; + let get = |suffix: &str| metrics.get(suffix).copied(); + Some(IONodeStats { + node_id, + operator_name, + format, + bytes_scanned: get("bytes_scanned"), + time_opening_ns: get("time_opening"), + time_scanning_ns: get("time_scanning"), + output_rows: get("output_rows"), + rg_pruned: get("rg_pruned"), + rg_matched: get("rg_matched"), + bloom_pruned: get("bloom_pruned"), + bloom_matched: get("bloom_matched"), + page_index_pruned: get("page_index_pruned"), + page_index_matched: get("page_index_matched"), + }) + }) + .collect(); + + let computes: Vec = compute_map + .into_iter() + .map(|((node_id, name), (category, mut partitions))| { + partitions.sort_by_key(|(pid, _)| *pid); + PartitionsComputeStats { + node_id, + name, + category, + elapsed_computes: partitions.iter().map(|(_, v)| *v as usize).collect(), + } + }) + .sorted_by_key(|c| (c.node_id, c.name.clone())) + .collect(); + + let io = (!io_nodes.is_empty()).then_some(ExecutionIOStats { nodes: io_nodes }); + let compute = + (elapsed_compute.is_some() || !computes.is_empty()).then_some(ExecutionComputeStats { + elapsed_compute, + computes, + }); + + let nodes: Vec = node_map + .into_values() + .sorted_by_key(|n| n.node_id) + .collect(); + + Ok(ExecutionStats { + query, + rows, + batches, + bytes, + durations: ExecutionDurationStats::new( + parsing, + logical_planning, + physical_planning, + execution, + total, + ), + io, + compute, + plan: None, + nodes, + }) } } @@ -762,3 +1397,298 @@ pub fn print_io_summary(plan: Arc) { println!("No IO metrics found"); } } + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::Int64Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::parquet::arrow::ArrowWriter; + use datafusion::parquet::file::properties::WriterProperties; + use datafusion::physical_plan::collect; + use datafusion::prelude::{ParquetReadOptions, SessionContext}; + + /// Execute `sql` to completion and return fully collected ExecutionStats, + /// mirroring the analyze path in `ExecutionContext::analyze_query` + async fn analyze(ctx: &SessionContext, sql: &str) -> ExecutionStats { + let df = ctx.sql(sql).await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + let bytes: usize = batches.iter().map(|b| b.get_array_memory_size()).sum(); + let durations = ExecutionDurationStats::new( + Duration::from_nanos(1), + Duration::from_nanos(2), + Duration::from_nanos(3), + Duration::from_nanos(4), + Duration::from_nanos(10), + ); + let mut stats = + ExecutionStats::try_new(sql.to_string(), durations, rows, batches.len(), bytes, plan) + .unwrap(); + stats.collect_stats(); + stats + } + + fn metric_names(batch: &RecordBatch) -> Vec { + let names = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()) + .map(|i| names.value(i).to_string()) + .collect() + } + + fn categories(batch: &RecordBatch) -> Vec> { + let cats = batch + .column(5) + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()) + .map(|i| (!cats.is_null(i)).then(|| cats.value(i).to_string())) + .collect() + } + + #[test] + fn test_protocol_version_compatibility() { + assert!(is_compatible_protocol_version(ANALYZE_PROTOCOL_VERSION)); + assert!(is_compatible_protocol_version("0.9")); + assert!(!is_compatible_protocol_version("99.0")); + } + + #[tokio::test] + async fn test_node_ids_are_unique_and_preorder() { + let ctx = SessionContext::new(); + let stats = analyze( + &ctx, + "SELECT column2, COUNT(*) FROM (VALUES (1,'a'),(2,'b'),(3,'a')) AS t(column1, column2) GROUP BY column2", + ) + .await; + + let nodes = stats.nodes(); + assert!(!nodes.is_empty()); + // Pre-order assignment: ids are 0..n in push order, root first + for (i, node) in nodes.iter().enumerate() { + assert_eq!(node.node_id(), i as i32); + } + assert_eq!(nodes[0].parent_node_id(), None, "root has no parent"); + // Every non-root parent id refers to an existing, earlier node + for node in &nodes[1..] { + let parent = node.parent_node_id().expect("non-root node has a parent"); + assert!(parent >= 0 && parent < node.node_id()); + } + } + + #[tokio::test] + async fn test_duplicate_operators_stay_distinct() { + let ctx = SessionContext::new(); + let stats = analyze( + &ctx, + "SELECT column2, COUNT(*) FROM (VALUES (1,'a'),(2,'b'),(3,'a')) AS t(column1, column2) GROUP BY column2", + ) + .await; + + // A GROUP BY plans two AggregateExec nodes (partial + final) + let compute = stats.compute.as_ref().unwrap(); + let agg_node_ids: Vec> = compute + .computes + .iter() + .filter(|c| c.name == "AggregateExec") + .map(|c| c.node_id) + .collect(); + assert!( + agg_node_ids.len() >= 2, + "expected two AggregateExec nodes, got {:?}", + agg_node_ids + ); + let distinct: std::collections::HashSet<_> = agg_node_ids.iter().collect(); + assert_eq!( + distinct.len(), + agg_node_ids.len(), + "AggregateExec nodes must have distinct node ids" + ); + + // And they survive the wire round trip as distinct nodes + let batch = stats.to_metrics_table().unwrap(); + let roundtripped = ExecutionStats::from_metrics_table(batch, stats.query.clone()).unwrap(); + let rt_agg: Vec> = roundtripped + .compute + .as_ref() + .unwrap() + .computes + .iter() + .filter(|c| c.name == "AggregateExec") + .map(|c| c.node_id) + .collect(); + assert_eq!(agg_node_ids, rt_agg); + } + + #[tokio::test] + async fn test_metrics_table_round_trip() { + let ctx = SessionContext::new(); + let stats = analyze( + &ctx, + "SELECT column2, COUNT(*) FROM (VALUES (1,'a'),(2,'b'),(3,'a')) AS t(column1, column2) GROUP BY column2", + ) + .await; + + let batch = stats.to_metrics_table().unwrap(); + assert_eq!( + batch.schema().metadata().get(PROTOCOL_VERSION_METADATA_KEY), + Some(&ANALYZE_PROTOCOL_VERSION.to_string()) + ); + + let roundtripped = ExecutionStats::from_metrics_table(batch, stats.query.clone()).unwrap(); + assert_eq!(roundtripped.query, stats.query); + assert_eq!(roundtripped.rows, stats.rows); + assert_eq!(roundtripped.batches, stats.batches); + assert_eq!(roundtripped.bytes, stats.bytes); + assert_eq!(roundtripped.durations, stats.durations); + assert_eq!(roundtripped.io, stats.io); + assert_eq!(roundtripped.compute, stats.compute); + assert!(roundtripped.plan.is_none()); + // Every node referenced by metrics is reconstructed with identical identity + for node in roundtripped.nodes() { + let original = stats + .nodes() + .iter() + .find(|n| n.node_id() == node.node_id()) + .expect("reconstructed node exists in original plan"); + assert_eq!(node, original); + } + } + + #[tokio::test] + async fn test_compute_categories_emitted_on_wire() { + let ctx = SessionContext::new(); + let stats = analyze( + &ctx, + "SELECT v, ROW_NUMBER() OVER (ORDER BY v) AS rn FROM \ + (SELECT v FROM (VALUES (1),(2),(3)) t(v) \ + UNION ALL SELECT v FROM (VALUES (4),(5)) u(v)) \ + LIMIT 3", + ) + .await; + + let batch = stats.to_metrics_table().unwrap(); + let cats: std::collections::HashSet = + categories(&batch).into_iter().flatten().collect(); + for expected in ["window", "limit", "union"] { + assert!( + cats.contains(expected), + "expected category {expected} on the wire, got {:?}", + cats + ); + } + + // And the round trip keeps them + let roundtripped = ExecutionStats::from_metrics_table(batch, stats.query.clone()).unwrap(); + assert_eq!(roundtripped.compute, stats.compute); + } + + #[tokio::test] + async fn test_parquet_io_and_pruning_metrics() { + // 1000 sorted values in 10 row groups of 100; `v > 950` prunes 9 of + // 10 row groups via statistics + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from((0..1000).collect::>())) as ArrayRef], + ) + .unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("t.parquet"); + let file = std::fs::File::create(&path).unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_size(100) + .build(); + let mut writer = ArrowWriter::try_new(file, schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let ctx = SessionContext::new(); + ctx.register_parquet("t", path.to_str().unwrap(), ParquetReadOptions::default()) + .await + .unwrap(); + + let stats = analyze(&ctx, "SELECT v FROM t WHERE v > 950").await; + assert_eq!(stats.rows, 49); + + let io = stats.io.as_ref().expect("io stats collected"); + assert_eq!(io.nodes().len(), 1); + let node = &io.nodes()[0]; + assert_eq!(node.format, IOFormatType::Parquet); + assert_eq!(node.operator_name, "DataSourceExec"); + assert!(node.bytes_scanned.unwrap() > 0); + assert_eq!(node.rg_pruned, Some(9)); + assert_eq!(node.rg_matched, Some(1)); + assert!(node.output_rows.unwrap() > 0); + + // The scan node is classified as io, not compute + let compute = stats.compute.as_ref().unwrap(); + assert!( + compute.computes.iter().all(|c| c.name != "DataSourceExec"), + "scan must not appear as a compute node" + ); + + // Selectivity helpers have data to work with + assert!(stats.rows_selectivity().is_some()); + assert!(stats.bytes_selectivity().is_some()); + assert!(stats.selectivity_efficiency().is_some()); + + // Round trip preserves I/O values exactly + let table = stats.to_metrics_table().unwrap(); + let names = metric_names(&table); + for expected in [ + "io.parquet.bytes_scanned", + "io.parquet.time_opening", + "io.parquet.time_scanning", + "io.parquet.output_rows", + "io.parquet.rg_pruned", + "io.parquet.rg_matched", + ] { + assert!( + names.iter().any(|n| n == expected), + "expected {expected} in metrics table" + ); + } + let roundtripped = ExecutionStats::from_metrics_table(table, stats.query.clone()).unwrap(); + assert_eq!(roundtripped.io, stats.io); + } + + #[tokio::test] + async fn test_multiple_scans_reported_separately() { + // Two parquet scans in one query (self join) must produce two + // distinct I/O nodes rather than one overwritten aggregate + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from((0..10).collect::>())) as ArrayRef], + ) + .unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("t.parquet"); + let file = std::fs::File::create(&path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let ctx = SessionContext::new(); + ctx.register_parquet("t", path.to_str().unwrap(), ParquetReadOptions::default()) + .await + .unwrap(); + + let stats = analyze(&ctx, "SELECT a.v FROM t a JOIN t b ON a.v = b.v").await; + let io = stats.io.as_ref().expect("io stats collected"); + assert_eq!(io.nodes().len(), 2, "each scan is a separate io node"); + let ids: std::collections::HashSet<_> = io.nodes().iter().map(|n| n.node_id).collect(); + assert_eq!(ids.len(), 2, "scan nodes have distinct node ids"); + + let table = stats.to_metrics_table().unwrap(); + let roundtripped = ExecutionStats::from_metrics_table(table, stats.query.clone()).unwrap(); + assert_eq!(roundtripped.io, stats.io); + } +} diff --git a/data/test_io_formats.arrow b/data/test_io_formats.arrow new file mode 100644 index 00000000..97e70d66 Binary files /dev/null and b/data/test_io_formats.arrow differ diff --git a/data/test_io_formats.csv b/data/test_io_formats.csv new file mode 100644 index 00000000..e492cfdd --- /dev/null +++ b/data/test_io_formats.csv @@ -0,0 +1,11 @@ +id,name,value,category +1,apple,100,fruit +2,banana,150,fruit +3,carrot,75,vegetable +4,date,200,fruit +5,eggplant,120,vegetable +6,fig,90,fruit +7,garlic,45,vegetable +8,honeydew,180,fruit +9,iceberg,60,vegetable +10,jalapeño,30,vegetable diff --git a/data/test_io_formats.json b/data/test_io_formats.json new file mode 100644 index 00000000..ce9d7a66 --- /dev/null +++ b/data/test_io_formats.json @@ -0,0 +1,10 @@ +{"id":1,"name":"apple","value":100,"category":"fruit"} +{"id":2,"name":"banana","value":150,"category":"fruit"} +{"id":3,"name":"carrot","value":75,"category":"vegetable"} +{"id":4,"name":"date","value":200,"category":"fruit"} +{"id":5,"name":"eggplant","value":120,"category":"vegetable"} +{"id":6,"name":"fig","value":90,"category":"fruit"} +{"id":7,"name":"garlic","value":45,"category":"vegetable"} +{"id":8,"name":"honeydew","value":180,"category":"fruit"} +{"id":9,"name":"iceberg","value":60,"category":"vegetable"} +{"id":10,"name":"jalapeño","value":30,"category":"vegetable"} diff --git a/data/test_io_formats.parquet b/data/test_io_formats.parquet new file mode 100644 index 00000000..c40309b3 Binary files /dev/null and b/data/test_io_formats.parquet differ diff --git a/docs/arrow_flight_analyze_protocol.md b/docs/arrow_flight_analyze_protocol.md new file mode 100644 index 00000000..a949d60c --- /dev/null +++ b/docs/arrow_flight_analyze_protocol.md @@ -0,0 +1,512 @@ +# Arrow Flight Analyze Protocol Specification + +**Version**: 0.1 +**Status**: Experimental + +This document specifies a protocol extension for Apache Arrow Flight services to provide detailed query execution metrics. The protocol is modeled on Apache DataFusion's execution metrics and is intended to generalize to other engines; see [Relationship to DataFusion](#relationship-to-datafusion). + +## Overview + +The Arrow Flight Analyze Protocol enables clients to retrieve detailed execution metrics for queries through a custom Arrow Flight action. This provides: + +- Query execution timing breakdown (parsing, planning, execution) +- I/O statistics (bytes scanned, file operations) +- Format-specific metrics (Parquet pruning, etc.) +- Per-operator compute time by partition +- Execution plan hierarchy, via stable node ids, for reconstructing query plan structure +- Extensible metric model for custom execution plan nodes + +### Protocol Scope + +This protocol is an **Apache Arrow Flight** extension, not specific to Flight SQL. While it naturally pairs with Flight SQL for SQL query analysis, any Arrow Flight service can implement the `analyze_query` action to provide execution metrics. + +The examples in this specification use SQL for illustration, but the protocol works with any query representation that the Flight service supports. + +### Related Work + +The protocol composes ideas from existing systems rather than inventing new ones; the gap it fills is that **Arrow Flight has no standardized, in-band, machine-readable way to return execution telemetry**: + +- **`EXPLAIN ANALYZE`** (Postgres, DataFusion, DuckDB, ...): rich per-operator metrics, but delivered as engine-specific text or JSON out of band from the transport. Postgres's `EXPLAIN (ANALYZE, FORMAT JSON)` is the closest precedent for structured output; this protocol replaces the format-specific document with a flat Arrow relation. +- **Trino/Presto query stats and Spark's SQL metrics APIs**: execution metrics over a separate REST channel, with stable plan-node ids identifying operators. The node-id approach here follows that precedent. +- **ClickHouse `system.query_log`**: flat metric rows queryable as a table — the same "metrics are data" shape this protocol uses, but post-hoc rather than in-band. +- **OpenMetrics/Prometheus**: the namespaced metric-name convention (`io.parquet.bytes_scanned`) follows their naming discipline. The `value_type` field is a deliberately minimal unit model; adopters needing richer semantics (temporality, exemplars) should look at the OpenTelemetry metrics data model, which this protocol does not attempt to replicate. +- **Substrait**: a future request/response field may carry Substrait plans; Substrait plan-relation ids would then be the natural cross-engine operator identity, complementing the per-execution `node_id` used here. +- **Distributed DataFusion engines** ([DataFusion Distributed](https://datafusion-contrib.github.io/datafusion-distributed/), [Apache DataFusion Ballista](https://github.com/apache/datafusion-ballista), and production Arrow Flight remote-execution systems): a distributed plan is still one plan tree, with network-boundary operators between stages; after execution each worker's per-node metrics are folded back into the coordinator's copy of the plan before display. The [`distributed.*` namespace](#distributed-execution-metrics) and the unified-tree rule follow that model. Notably, systems that ship per-node metrics back today correlate them with the coordinator's tree *positionally* (DFS pre-order) — Ballista's executors report each task's metric sets to the scheduler keyed by pre-order node position within a stage, and its `GetJobMetrics` RPC (per-job metrics retained on the scheduler after completion) is a working precedent for the two-phase analyze in [Future Work](#future-work). The `node_id`/`parent_node_id` columns make exactly that correlation explicit on the wire. + +### Relationship to DataFusion + +The reference implementation is built on Apache DataFusion, and the standard metric set below is derived from DataFusion's `MetricSet` values. Non-DataFusion implementers should treat the metric tables as a **DataFusion mapping** of the abstract categories (query, stage, io, compute): implement the metrics that have equivalents in their engine, keep the namespaces and `value_type` conventions, and omit the rest. Clients are required to tolerate missing optional metrics (see [Client Metric Handling](#client-metric-handling)). + +## Action Specification + +### Action Type + +**Action Name**: `"analyze_query"` + +**Purpose**: Execute a SQL query with metrics collection enabled and return detailed execution statistics. + +### Request Format + +**Request Body**: JSON-encoded query request structure. JSON is used for the request because bodies are small, debuggability matters more than throughput here, and it avoids coupling the extension to a protobuf schema registry. (Flight SQL precedent would be protobuf `Any`; adopters that need it can layer that in a future version.) + +The request body should be a JSON object with the following structure: + +```json +{ + "sql": "SELECT * FROM table WHERE id > 100", + "protocol_version": "0.1" +} +``` + +**Current Fields**: +- `sql` (string, required): The SQL query to analyze. Must contain exactly one SQL statement. Multiple statements (e.g., separated by semicolons) are not supported and will result in an error. +- `protocol_version` (string, optional): The protocol version the client implements. Servers MUST reject a request whose major version they do not support with `invalid_argument`. When absent, the server assumes the client speaks the server's version. + +**Future Extensibility**: +The protocol is designed to be extensible. Additional query representation fields may be supported in the future: +- `substrait` (bytes): Substrait query plan (binary or JSON) +- `logical_plan` (string): Serialized logical plan +- `physical_plan` (string): Serialized physical plan + +Servers should ignore unknown fields and clients should only send one query representation field at a time. + +**Request Encoding**: The JSON object should be serialized to UTF-8 bytes in the `Action.body` field. + +### Response Format + +The response is a stream of `arrow_flight::Result` messages. Each `Result.body` contains one serialized `FlightData` message (protobuf encoding). Concatenated, the `FlightData` messages form a standard Arrow IPC stream: one schema message followed by **one or more** record batch messages (and dictionary batches if needed). Clients MUST NOT assume the metrics table arrives as a single batch; servers MAY split large tables. + +**Response Metadata** (Arrow schema metadata on the metrics batch schema): + +| Key | Required | Description | +|-----|----------|-------------| +| `analyze.protocol_version` | yes | Protocol version the server implements (e.g., `"0.1"`) | +| `analyze.query_id` | recommended | Opaque correlation id for this request (e.g., a UUID). Lets clients issuing concurrent analyze calls match responses to requests. | + +**Note**: The query text is NOT echoed in the response. The client is responsible for retaining the original query and correlating it with the response (using `analyze.query_id` when concurrent requests are in flight). + +#### Metrics Batch + +**Purpose**: A flat Arrow table where each row represents a single metric observation. + +**Schema**: +| Column | Type | Nullable | Description | +|--------|------|----------|-------------| +| metric_name | Utf8 | false | Namespaced metric name (e.g., "query.rows", "stage.parsing", "io.parquet.bytes_scanned") | +| value | UInt64 | false | Numeric value of the metric | +| value_type | Utf8 | false | Unit of value: "duration_ns", "bytes", or "count" | +| operator_name | Utf8 | true | Execution plan node display name (e.g., "FilterExec", "DataSourceExec"). A label, NOT an identifier — plans routinely contain multiple nodes with the same name. | +| partition_id | Int32 | true | Partition rank for per-partition metrics (see Compute Metrics) | +| operator_category | Utf8 | true | Category: "filter", "sort", "projection", "join", "aggregate", "window", "distinct", "limit", "union", "exchange", "io", "other" — see [Operator Categories](#operator-categories) | +| node_id | Int32 | true | Stable id of the plan node this metric belongs to (NULL for query/stage-level metrics) | +| parent_node_id | Int32 | true | `node_id` of the node's parent (NULL for the root node and for query/stage-level metrics) | + +**Canonical row key**: `(metric_name, node_id, partition_id)`. The `operator_name`, `operator_category`, and the namespace prefix of `metric_name` are denormalized presentation hints; they carry no identity. + +**Value types**: Durations are always nanoseconds (`duration_ns`); sizes are bytes; everything else is a plain count. There is deliberately no fractional/ratio value type in v0.1 — ratios (selectivity, pruning effectiveness) are derived by clients from count metrics. + +**Cardinality**: One row per metric, per node, per partition. For large plans this is `O(operators × partitions × metrics)`; see [Operational Considerations](#operational-considerations). + +### Execution Plan Hierarchy + +Every node in the execution plan is assigned a stable integer `node_id` by **pre-order traversal, with the root as 0**. The `(node_id, parent_node_id)` pairs on operator-level metric rows fully describe the plan tree, even when the same operator type appears multiple times (e.g., partial and final `AggregateExec`, repeated `RepartitionExec` nodes, self-joins). + +NULL rules: +- **Query/stage-level rows** (`operator_name = NULL`): both `node_id` and `parent_node_id` are NULL. +- **Root operator rows**: `node_id = 0`, `parent_node_id = NULL`. +- **All other operator rows**: both fields are non-NULL. + +**Example**: For a plan `ProjectionExec -> FilterExec -> DataSourceExec` (root first): +- ProjectionExec: `node_id = 0`, `parent_node_id = NULL` +- FilterExec: `node_id = 1`, `parent_node_id = 0` +- DataSourceExec: `node_id = 2`, `parent_node_id = 1` + +## Operator Categories + +The `operator_category` column groups plan nodes by their function so clients can organize metrics without recognizing every engine-specific operator name. Categories describe what the **physical operator** does, not the logical intent of the query — e.g. an engine that plans `SELECT DISTINCT` as a grouped aggregation reports those nodes as `aggregate`, not `distinct`. + +| Category | Meaning | Representative operators (DataFusion) | +|----------|---------|----------------------------------------| +| `io` | Scan/data-source nodes reading external data. Report format-specific metrics under `io.{format}.*` and are excluded from the compute breakdown (see [Format-Specific I/O Metrics](#format-specific-io-metrics)) | `DataSourceExec` | +| `projection` | Column selection and expression evaluation | `ProjectionExec` | +| `filter` | Row-level predicate evaluation | `FilterExec` | +| `sort` | Ordering, including order-preserving merges and top-N variants | `SortExec` (incl. TopK), `SortPreservingMergeExec` | +| `aggregate` | Grouped or scalar aggregation, at any phase (partial/final) | `AggregateExec` | +| `join` | All join algorithms, including cross joins | `HashJoinExec`, `SortMergeJoinExec`, `NestedLoopJoinExec`, `CrossJoinExec`, `SymmetricHashJoinExec` | +| `window` | Window-function evaluation | `WindowAggExec`, `BoundedWindowAggExec` | +| `distinct` | Dedicated deduplication operators (engines without one report deduplication under the operator that implements it, typically `aggregate`) | — | +| `limit` | Row-count truncation | `GlobalLimitExec`, `LocalLimitExec` | +| `union` | Concatenating/interleaving multiple inputs without join semantics | `UnionExec`, `InterleaveExec` | +| `exchange` | Network-boundary data movement between workers in a distributed plan (see [Distributed Execution Metrics](#distributed-execution-metrics)). **Not** for process-local repartitioning | `NetworkShuffleExec`, `ShuffleWriterExec`/`ShuffleReaderExec`, `RemoteExec`-style nodes | +| `other` | Everything else: local repartition/coalesce, empty relations, unnest, table functions, custom nodes | `RepartitionExec`, `CoalesceBatchesExec`, `UnnestExec` | + +**Classification rules**: +- Classify by the operator's **dominant function** — the work its metrics measure. A hybrid operator goes to the category of its dominant cost (e.g. a fused sort-with-limit is `sort`). +- Category assignment for ambiguous operators is implementation-defined; `other` is always valid, and clients MUST handle it (and unknown categories) gracefully per [Client Metric Handling](#client-metric-handling). +- The category is a denormalized presentation hint carried on each of the node's metric rows; all rows for one `node_id` MUST carry the same category. + +## Metric Namespaces + +Metric names use a hierarchical namespace structure to prevent collisions and provide clear semantic grouping: + +**Format**: `{namespace}.{metric_name}` + +**Standard Namespaces**: +- `query.*` - Query-level metrics (rows, batches, bytes) +- `stage.*` - Execution stage durations (parsing, logical_planning, physical_planning, execution, total) +- `io.parquet.*` - Parquet-specific I/O metrics +- `io.csv.*` - CSV-specific I/O metrics +- `io.json.*` - JSON-specific I/O metrics +- `io.arrow.*` - Arrow IPC-specific I/O metrics +- `compute.*` - Compute metrics (elapsed_compute with operator breakdown) +- `index.*` - Reserved for index-related metrics (future: index_hits, index_scans) +- `distributed.*` - Distributed execution metrics (see [Distributed Execution Metrics](#distributed-execution-metrics)) + +**Important**: There is no generic `io.*` namespace. Each file format reports its own complete set of I/O metrics under its specific namespace (e.g., `io.parquet.*`, `io.csv.*`). This prevents mixing aggregated and raw data. Only namespaced metric names are valid on the wire. + +## Standard Metrics + +### Query-Level Metrics + +These metrics have `operator_name = NULL`, `partition_id = NULL`, `operator_category = NULL`, `node_id = NULL`, `parent_node_id = NULL`: + +| Metric Name | Value Type | Description | +|-------------|------------|-------------| +| `query.rows` | count | Total number of output rows | +| `query.batches` | count | Total number of output batches | +| `query.bytes` | bytes | Total in-memory Arrow size of the output batches (as reported by `RecordBatch::get_array_memory_size` in the reference implementation). This is NOT a serialized/wire size; servers implementing this metric MUST use the in-memory definition. | + +### Duration Metrics + +Timing breakdown for query execution phases. All have `operator_name = NULL`, `partition_id = NULL`, `operator_category = NULL`, `node_id = NULL`, `parent_node_id = NULL`: + +| Metric Name | Value Type | Description | +|-------------|------------|-------------| +| `stage.parsing` | duration_ns | Query parsing time in nanoseconds | +| `stage.logical_planning` | duration_ns | Logical plan creation time | +| `stage.physical_planning` | duration_ns | Physical plan creation time | +| `stage.execution` | duration_ns | Query execution wall-clock time | +| `stage.total` | duration_ns | Total wall-clock time for the request | + +**Timing semantics**: `stage.*` metrics are wall-clock durations of non-overlapping request phases; they are additive and sum to approximately `stage.total`. `compute.elapsed_compute` and `io.*.time_*` metrics are per-operator CPU/IO time summed across partitions that execute **concurrently**; under pipelined, parallel execution they routinely exceed `stage.execution` and MUST NOT be treated as additive with the stage timers or with each other. + +### Format-Specific I/O Metrics + +Each I/O metric row carries: +- `operator_name`: The scan operator's display name (in DataFusion 51+, file scans are `"DataSourceExec"`) +- `operator_category = "io"` +- `partition_id = NULL` (values are summed across partitions) +- `node_id` / `parent_node_id`: identity of the scan node + +A query with multiple scans (joins, unions) produces a separate set of I/O rows per scan node, distinguished by `node_id`. A mixed-format query reports each scan under its own format namespace. + +**Common I/O Metrics** (each format provides these under its own namespace): + +| Metric Pattern | Value Type | Description | +|----------------|------------|-------------| +| `io.{format}.bytes_scanned` | bytes | Total bytes read from storage (where available) | +| `io.{format}.time_opening` | duration_ns | Time spent opening files | +| `io.{format}.time_scanning` | duration_ns | Time spent reading/scanning data | + +#### Parquet Metrics + +| Metric Name | Value Type | Description | +|-------------|------------|-------------| +| `io.parquet.bytes_scanned` | bytes | Total bytes read from Parquet files | +| `io.parquet.time_opening` | duration_ns | Time spent opening Parquet files | +| `io.parquet.time_scanning` | duration_ns | Time spent reading/scanning Parquet data | +| `io.parquet.output_rows` | count | Number of rows produced by the scan node | +| `io.parquet.rg_pruned` | count | Row groups pruned by statistics | +| `io.parquet.rg_matched` | count | Row groups matched (not pruned) by statistics | +| `io.parquet.bloom_pruned` | count | Row groups pruned by bloom filters | +| `io.parquet.bloom_matched` | count | Row groups matched by bloom filters | +| `io.parquet.page_index_pruned` | count | Rows pruned by the page index | +| `io.parquet.page_index_matched` | count | Rows matched by the page index | + +#### CSV / JSON / Arrow IPC Metrics + +The reference implementation emits the common I/O metrics (`time_opening`, `time_scanning`, and `bytes_scanned` where the engine reports it) under `io.csv.*`, `io.json.*`, and `io.arrow.*`. Format-specific names such as `io.csv.rows_parsed`, `io.csv.parse_errors`, `io.json.invalid_rows`, or `io.arrow.dictionary_hits` are **reserved** for future use; they are listed here so implementations do not repurpose them, but v0.1 does not define or emit them. + +#### Other Formats + +Additional formats can define their own metrics under their namespace (e.g., `io.orc.stripe_pruned`). Use `io.{format}.*` for custom format metrics. + +### Compute Metrics + +Metrics for CPU-intensive operators. The `compute.elapsed_compute` metric appears in two forms: + +#### Aggregate Compute Time + +Total compute time across all (non-scan) operators: +- `metric_name = "compute.elapsed_compute"` +- All other columns NULL + +#### Per-Operator, Per-Partition Compute Time + +Detailed breakdown by operator and partition: +- `metric_name = "compute.elapsed_compute"` +- `operator_name`: Display name of the operator (e.g., "FilterExec") +- `partition_id`: **Rank** of the value among the node's partitions when sorted ascending (0-based). It preserves the per-partition distribution for skew analysis but is not guaranteed to be the physical partition number. +- `operator_category`: One of `filter`, `sort`, `projection`, `join`, `aggregate`, `window`, `distinct`, `limit`, `union`, `exchange`, `other` (see [Operator Categories](#operator-categories)) +- `node_id` / `parent_node_id`: identity of the operator node + +**Note**: Scan nodes report under `io` and are excluded from the compute breakdown and from the aggregate compute total. Exchange (network boundary) nodes report their local CPU work here under category `exchange`; their network-side metrics are reported separately under `distributed.*`. + +### Distributed Execution Metrics + +Distributed engines execute parts of the plan on remote workers, connected by network-boundary **exchange operators** — shuffle/coalesce nodes in a stage/task scheduler (e.g. DataFusion Distributed's `NetworkShuffleExec` and `NetworkCoalesceExec`), or a `RemoteExec`-style node in a single-hop plan-fragment shipper. The `distributed.*` namespace carries the network-side metrics of those boundaries. The reference implementation is a single-node engine and does not emit these metrics; the definitions below are derived from surveyed distributed DataFusion systems so that distributed adopters map onto real metric sets. Servers implementing distributed execution SHOULD include `"distributed"` in the `namespaces` list of `analyze_query_capabilities`. + +**One table, one tree.** A distributed query still produces a single metrics table over the **unified** plan: servers MUST fold each worker's sub-plan metrics back into the coordinator's plan tree before emitting, and assign `node_id` by pre-order traversal of that unified tree. In stage-based engines (Ballista-style), stages are stitched back together at their exchange boundaries (e.g. a shuffle reader leaf reconnects to the producer stage's shuffle writer root). Remote operators are therefore indistinguishable from local ones in the response, apart from the exchange nodes between them. + +**Trees, not DAGs.** Stage graphs can be DAGs: a broadcast stage may feed several consumers (e.g. a broadcast join build side read by N exchange nodes). Servers MUST emit each producer node exactly once — attach the shared subtree under exactly one of its consumer exchange nodes (implementation-defined; e.g. the first in pre-order) so its metrics are not double-counted. The `distributed.consumers` metric on the subtree's root records the fan-out; an explicit plan-edge representation is future work. + +**Adaptive execution.** Engines that re-plan stages mid-flight (adaptive query execution: runtime join selection, stage elimination, stage re-execution) MUST emit the plan as it **finally executed**. Node ids are assigned after the query completes, over that final plan; retried stages report the winning attempt's metrics only. + +**Aggregation across tasks.** When multiple tasks (workers) execute the same stage in parallel, each plan node's metric values are aggregated across tasks before emission — sums for counts, bytes, and durations — mirroring how I/O metrics are summed across partitions. Per-partition `compute.elapsed_compute` rows rank values over the union of all tasks' partitions. Engines that cannot attribute values to individual partitions exactly (a known hazard for subtrees beneath partition-coalescing operators in task-parallel engines) SHOULD emit an aggregated row (`partition_id = NULL`) rather than approximate or inflated per-partition rows. A per-task or per-worker breakdown requires a task identity the v0.1 schema does not carry; see [Future Work](#future-work). + +#### Exchange-Node Metrics + +Each exchange operator reports with `operator_category = "exchange"`, `partition_id = NULL`, and its own `node_id`/`parent_node_id`. A boundary may be represented by a single receiving operator (fragment shippers, network shuffle readers) or by a sender/receiver pair (a shuffle writer at a producer stage's root plus a shuffle reader leaf in the consumer); both sides are `exchange` nodes and each reports the metrics for its own side. + +**Receiver-side metrics**: + +| Metric Name | Value Type | Description | +|-------------|------------|-------------| +| `distributed.bytes_transferred` | bytes | Serialized bytes received across the boundary (wire size — NOT the in-memory size used by `query.bytes`) | +| `distributed.messages` | count | Data messages (e.g. Arrow Flight `FlightData`) received | +| `distributed.rpc_calls` | count | RPCs issued to serve the boundary | +| `distributed.rpc_retries` | count | Failed attempts that were retried (connection acquisition, endpoint failover, fetch retries) | +| `distributed.tasks` | count | Remote tasks feeding this boundary | +| `distributed.time_to_first_batch` | duration_ns | Time from issuing the request to receiving the first record batch | +| `distributed.time_transferring` | duration_ns | Time spent receiving data across the boundary | +| `distributed.network_latency_sum` | duration_ns | Sum of per-message network latencies (clients derive the average using `_count`) | +| `distributed.network_latency_count` | count | Number of latency observations contributing to `_sum` | + +**Sender-side metrics** (on writer-style exchange nodes, where they exist): + +| Metric Name | Value Type | Description | +|-------------|------------|-------------| +| `distributed.bytes_written` | bytes | Serialized bytes written for downstream consumers (e.g. shuffle files) | +| `distributed.time_writing` | duration_ns | Time spent serializing and writing exchange output | +| `distributed.consumers` | count | Downstream consumers of this node's output (> 1 when a broadcast stage feeds several exchanges; see "Trees, not DAGs") | + +Byte accounting is deliberately split: engines often source the two sides differently (received bytes from the reader's metrics, written bytes from shuffle file statistics), and with broadcast or partial fetches they are not equal, so there is no single "transferred" number spanning both sides. + +`distributed.network_latency_min`, `_max`, `_p50`, `_p95`, and `_p99` are **reserved** for pre-aggregated latency distribution values (each a `duration_ns` computed server-side, e.g. from a quantile sketch); v0.1 does not define or emit them, consistent with the no-server-computed-ratios rule (`_sum`/`_count` are the required pair). `distributed.time_queued` is likewise **reserved** for time spent waiting on admission control or backpressure permits before fetching — engines with such governors MUST NOT fold that wait into `distributed.network_latency_sum`. + +Exchange operators MAY additionally report `compute.elapsed_compute` (category `exchange`) for their local CPU work, like any other operator. + +#### Query-Level Distributed Metrics + +These rows have all identity columns NULL (like other query-level metrics): + +| Metric Name | Value Type | Description | +|-------------|------------|-------------| +| `distributed.stages` | count | Number of distributed stages in the plan | +| `distributed.tasks` | count | Total tasks scheduled across all stages | +| `distributed.plan_bytes_sent` | bytes | Serialized plan fragments/tasks shipped to workers | +| `distributed.time_distributing_plan` | duration_ns | Time spent serializing and shipping plan fragments to workers | + +`distributed.tasks` follows the `compute.elapsed_compute` precedent: the query-level row is the total, exchange-node rows are per-boundary. + +**Timing semantics.** All `distributed.*` durations are wall-clock observations at a boundary that overlap the execution of upstream stages and each other (`time_to_first_batch` overlaps `time_transferring`); like the `compute.*` and `io.*` timers, they MUST NOT be treated as additive with the `stage.*` timers or with each other. + +## Example Response + +The client retains the original query: `"SELECT * FROM table WHERE id > 100"` + +Plan: `ProjectionExec (0) -> FilterExec (1) -> DataSourceExec (2)` + +``` +metric_name value value_type operator_name partition_id operator_category node_id parent_node_id +------------------------------ ------------- ------------ ---------------- ------------- ------------------ -------- -------------- +query.rows 1000 count NULL NULL NULL NULL NULL +query.batches 10 count NULL NULL NULL NULL NULL +query.bytes 50000 bytes NULL NULL NULL NULL NULL +stage.parsing 12000000 duration_ns NULL NULL NULL NULL NULL +stage.logical_planning 45000000 duration_ns NULL NULL NULL NULL NULL +stage.physical_planning 78000000 duration_ns NULL NULL NULL NULL NULL +stage.execution 234000000 duration_ns NULL NULL NULL NULL NULL +stage.total 369000000 duration_ns NULL NULL NULL NULL NULL +io.parquet.bytes_scanned 1000000 bytes DataSourceExec NULL io 2 1 +io.parquet.time_opening 50000000 duration_ns DataSourceExec NULL io 2 1 +io.parquet.time_scanning 150000000 duration_ns DataSourceExec NULL io 2 1 +io.parquet.output_rows 10000 count DataSourceExec NULL io 2 1 +io.parquet.rg_pruned 16 count DataSourceExec NULL io 2 1 +io.parquet.rg_matched 4 count DataSourceExec NULL io 2 1 +compute.elapsed_compute 12345678 duration_ns NULL NULL NULL NULL NULL +compute.elapsed_compute 1200 duration_ns ProjectionExec 0 projection 0 NULL +compute.elapsed_compute 1250 duration_ns ProjectionExec 1 projection 0 NULL +compute.elapsed_compute 1400 duration_ns FilterExec 0 filter 1 0 +compute.elapsed_compute 1500 duration_ns FilterExec 1 filter 1 0 +``` + +## Capability Discovery + +**Action Name**: `"analyze_query_capabilities"` + +Servers implementing `analyze_query` SHOULD also implement this action so clients can probe support and version before issuing (potentially expensive) analyze calls. The response is a single `arrow_flight::Result` whose body is a UTF-8 JSON object: + +```json +{ + "protocol_version": "0.1", + "request_formats": ["sql"], + "namespaces": ["query", "stage", "compute", "io.parquet", "io.csv", "io.arrow", "io.json"] +} +``` + +Clients MUST ignore unknown fields. Servers not implementing the action return `unimplemented`, which clients should treat the same as an `unimplemented` response to `analyze_query` itself. + +## Implementation Guide + +### Server Implementation + +To implement this protocol in an Arrow Flight service: + +1. **Register Action Handler** + - Implement `do_action` or `do_action_fallback` to recognize action type `"analyze_query"` + +2. **Parse and Validate Request** + - Decode `Action.body` as JSON + - Reject unsupported `protocol_version` major versions with `invalid_argument` + +3. **Execute Query** + - Run the query to completion with execution plan metrics collection enabled + +4. **Collect Metrics** + - In a distributed engine, first gather each worker's per-node metrics and fold them into the coordinator's plan tree (see [Distributed Execution Metrics](#distributed-execution-metrics)); all subsequent steps operate on the unified tree + - Assign each plan node a `node_id` by pre-order traversal (root = 0) + - Traverse the plan extracting metrics from each operator; emit one row per metric value with the node's `node_id`/`parent_node_id` + +5. **Build Response** + - Create the metrics RecordBatch with the 8-field schema + - Set `analyze.protocol_version` (and ideally `analyze.query_id`) in the schema metadata + - Encode as FlightData using `batches_to_flight_data()` or equivalent + - Serialize each FlightData to bytes (protobuf encoding) and wrap in `arrow_flight::Result { body }` + - Stream Result messages to the client + +### Client Implementation + +To consume this protocol: + +1. **Send Request**: JSON body with `sql` and `protocol_version`, action type `"analyze_query"`. +2. **Receive Stream**: collect all `arrow_flight::Result` messages. +3. **Decode**: decode each `Result.body` as a `FlightData` protobuf message; feed the sequence to an IPC decoder (e.g., `flight_data_to_batches`). Handle any number of data batches — concatenate them into one table. +4. **Validate**: check `analyze.protocol_version` in the schema metadata for major-version compatibility. +5. **Reconstruct**: parse the metrics table; rebuild the plan tree from `(node_id, parent_node_id)`; group per-partition compute rows by `node_id` (never by `operator_name` — names are not identities). + +### Error Handling + +**Server Behavior**: +Any error during request parsing, query execution, metrics collection, or response serialization results in complete failure. No partial metrics are returned. + +**Error Codes**: +- `Status::unimplemented` - Server doesn't support the analyze protocol +- `Status::invalid_argument` - Invalid SQL, malformed request, unsupported protocol version, or multiple SQL statements provided +- `Status::internal` - Query execution, metrics collection, or serialization failure + +**Client Handling**: +- Handle `unimplemented` gracefully with a clear user message +- Retry transient errors as appropriate — but note that a retried analyze re-executes the query (see below) +- Any error response means no metrics were collected + +## Operational Considerations + +**Re-execution cost.** `analyze_query` executes the query to completion and discards the results. Analyzing a query therefore costs a full extra execution, and the analyzed run is a *different execution* from the one the user experienced (caches may be warm, data may have changed). There is currently no fused "results plus metrics" mode; a two-phase design (execute normally, then fetch metrics for a statement id) and a streaming `analyze_query_live` variant are candidate future extensions. + +**Resource controls.** Because the action runs an arbitrary query server-side, deployments should apply the same admission control, timeouts, and cost caps to `analyze_query` as to regular query execution. The protocol itself defines no limits. + +**Payload size.** The metrics table is `O(operators × partitions × metrics)` rows. A scan with thousands of partitions produces a correspondingly large batch. Servers MAY split the table across multiple record batches (clients must handle this); future versions may add sampling or aggregation options. + +**Distributed collection.** In a distributed engine the response cannot be built until every worker's metrics have arrived at the coordinator, adding a synchronization step after the result stream completes (surveyed systems use a dedicated metrics channel, piggyback metrics on the final data message, or retain per-job metrics on the scheduler for post-hoc retrieval). Per [Error Handling](#error-handling), if any worker's metrics cannot be retrieved the action fails — no partial metrics are returned. This also applies to *silent* loss inside the engine's own metric transport: surveyed systems have failure modes where a single non-serializable metric value drops an entire task's metric set, or where only successfully completed stages report at all — an `analyze_query` implementation MUST surface such gaps as `internal` errors rather than emit a silently incomplete table. Engines should also note that result caches must be bypassed for `analyze_query`: a cached response carries no live metrics. + +**Authorization and audit.** The SQL travels inside an `Action` body. Proxies or middleware that authorize/audit only `DoGet`/`GetFlightInfo` tickets will not see it — route `do_action` through the same authorization and audit path as query execution. Also note that analyze output discloses plan internals (operator structure, partition counts, pruning effectiveness); deployments may want to gate the action separately from query execution. + +## Extensibility + +### Design Principles + +The protocol is designed to be: + +1. **Format-Agnostic**: Any file format (Parquet, CSV, JSON, ORC, Avro, etc.) can add metrics using the namespace conventions +2. **Execution-Agnostic**: Custom execution plan nodes can emit metrics in standard categories +3. **Forward-Compatible**: Clients display unknown metrics that carry a recognized category rather than rejecting them +4. **Language-Agnostic**: A flat Arrow table works in any language with an Arrow implementation +5. **Type-Safe**: Proper Arrow types prevent parsing ambiguities + +### Adding Custom Metrics + +Servers can add custom metrics as long as they follow the 8-field schema: + +**Custom Format Metrics**: +``` +metric_name: "io.{format}.{metric_name}" +operator_name: (scan node display name) +operator_category: "io" +node_id / parent_node_id: (scan node identity) +``` + +**Custom Compute Operators**: +- Choose the closest standard `operator_category` per the classification rules in [Operator Categories](#operator-categories) +- Use the `compute.elapsed_compute` metric name with operator name, partition rank, and node identity + +**Custom Query-Level Metrics**: +- Add new namespaced metric names with appropriate value types +- Use NULL for operator_name, partition_id, operator_category, node_id, and parent_node_id + +### Client Metric Handling + +Clients should handle metrics according to these principles: + +1. **Display all metrics** with recognized `operator_category` values, even if `metric_name` is unknown + - This enables forward compatibility with new server metrics + - Allows debugging of custom or experimental operators + +2. **Categorize and group** metrics by operator_category for presentation + +3. **Optionally validate** metric names against a known set + - Strict mode: Warn or error on unknown metric_name + - Permissive mode (default): Display all metrics + - Configuration option: `unknown_metrics_policy: "allow" | "warn" | "error"` + +4. **Gracefully handle** metrics with unknown operator_category + - Display in an "other"/"unknown" section; log for debugging + +5. **Validate required metrics** exist: + - Query-level: `query.rows`, `query.batches`, `query.bytes` + - Stage durations: `stage.parsing`, `stage.logical_planning`, `stage.physical_planning`, `stage.execution`, `stage.total` + +6. **Do not fail** on missing optional metrics (format-specific, compute per-partition, etc.) + +## Future Work + +Tracked candidates for later protocol versions, roughly in priority order: + +- Two-phase analyze (execute normally → fetch metrics by statement id) and/or a results-then-metrics mode, to avoid double execution +- `analyze_query_live`: streaming incremental metrics for long-running queries +- Per-operator/per-partition start and end timestamps for timeline (flamegraph) reconstruction +- Partition skew summary metrics +- A fractional value column (e.g., Float64) if server-computed ratios prove necessary +- Substrait request support, with Substrait plan-relation ids as cross-engine operator identity +- Response served via `Ticket` + standard `DoGet` instead of `Result`-wrapped FlightData, gaining standard streaming/dictionary handling +- `index.*` namespace definitions +- Per-task/per-worker breakdown of distributed metrics (DataFusion Distributed's "PerTask" format). Note that in surveyed systems worker identity is a set of strings (service, cluster, host/URL) rather than an integer, so this likely needs a nullable Utf8 identity column rather than reusing `partition_id` +- A text value column for non-numeric per-node diagnostics (e.g. a remote worker's end-of-execution dynamic-filter snapshot, which surveyed systems transport alongside numeric metrics) +- An explicit plan-edge representation for DAG-shaped distributed plans (broadcast stages feeding multiple consumers), replacing the emit-once/`distributed.consumers` convention +- A `memory.*` namespace for spill and memory metrics (`spilled_bytes`, `spilled_rows`, `spill_count`, peak memory) — first-class in DataFusion's metric model and in Ballista's metric transport, but not yet mapped by this protocol +- Config/resource context (memory limits, target partitions) in response metadata +- A UDTF form (`SELECT * FROM analyze('...')`) so metrics can be queried directly + +## References + +- [Apache Arrow Flight SQL Protocol](https://arrow.apache.org/docs/format/FlightSql.html) +- [Apache Arrow IPC Format](https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format) +- [DataFusion Execution Plans](https://docs.rs/datafusion/latest/datafusion/physical_plan/trait.ExecutionPlan.html) +- [DataFusion Distributed](https://datafusion-contrib.github.io/datafusion-distributed/) +- [Apache DataFusion Ballista](https://github.com/apache/datafusion-ballista) +- [OpenTelemetry Metrics Data Model](https://opentelemetry.io/docs/specs/otel/metrics/data-model/) +- [Substrait](https://substrait.io/) + +## License + +This specification is provided under the Apache License 2.0, consistent with the Apache Arrow project. diff --git a/docs/cli.md b/docs/cli.md index 11a7d4a7..5fa967e2 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -186,12 +186,86 @@ The output from `EXPLAIN ANALYZE` provides a wealth of information on a queries To help with this the `--analyze` flag can used to generate a summary of the underlying `ExecutionPlan` `MetricSet`s. The summary presents the information in a way that is hopefully easier to understand and easier to draw conclusions on a query's performance. -This feature is still in it's early stages and is expected to evolve. Once it has gone through enough real world testing and it has been confirmed the metrics make sense documentation will be added on the exact calculations - until then the source will need to be inspected to see the calculations. +**Important**: The analyze feature only supports a single SQL statement. If you provide multiple statements (e.g., separated by semicolons) or multiple files/commands, an error will be returned. + +**Note**: Analyze executes the query to completion and discards the results, so analyzing an expensive query costs a full extra execution. + +**Derived ratios** shown in the formatted output: +- *Output Rows (%)*: `query.rows / sum of scan output_rows` (row selectivity) +- *Output Bytes (%)*: `query.bytes / sum of bytes_scanned` +- *Parquet Efficiency*: row-group matched ratio (`rg_matched / (rg_pruned + rg_matched)`) divided by row selectivity +- Ratios display as `N/A` when the underlying metrics are unavailable + +### Local Analyze ```sh -dft -c "SELECT ..." --analyze +# Analyze a query locally +dft -c "SELECT * FROM table WHERE id > 100" --analyze + +# Analyze from a file +dft -f query.sql --analyze + +# Warm up caches or create tables before analyzing (same as --bench) +dft -c "SELECT * FROM t WHERE id > 100" --analyze --run-before "CREATE EXTERNAL TABLE t STORED AS PARQUET LOCATION 'data/t.parquet'" ``` +### FlightSQL Analyze + +The `--analyze` flag also works with FlightSQL execution, providing identical output to local analyze: + +```sh +# Analyze query on FlightSQL server +dft -c "SELECT * FROM table WHERE id > 100" --analyze --flightsql + +# Analyze from a file via FlightSQL +dft -f query.sql --analyze --flightsql +``` + +**Requirements:** +- The Arrow Flight service must support the `"analyze_query"` custom action +- See the [Arrow Flight Analyze Protocol Specification](arrow_flight_analyze_protocol.md) for implementation details +- Servers without analyze support will return an "unimplemented" error + +**How it works:** +1. Client sends a `do_action("analyze_query")` request with the SQL query +2. Server executes the query with metrics collection enabled +3. Server serializes execution statistics to a single flat Arrow metrics table (one row per metric, with `node_id`/`parent_node_id` identifying plan nodes) +4. Client deserializes and reconstructs the full execution statistics +5. Output is formatted identically to local analyze + +### Analyze Output + +Both local and FlightSQL analyze produce identical output including: + +- **Execution Summary**: Output rows/bytes, batch counts, selectivity ratios +- **Timing Breakdown**: Parsing, logical planning, physical planning, execution, total time +- **I/O Statistics**: Bytes scanned, file opening/scanning times +- **Parquet Metrics** (when applicable): + - Row group pruning effectiveness (statistics, bloom filters, page index) + - Per-row-group timing +- **Compute Statistics**: Per-operator elapsed compute time by partition + - Breakdown by operator category (projection, filter, sort, aggregate, join, window, distinct, limit, union, other) + - Min/median/mean/max timing per operator + +### Raw Metrics Mode + +For debugging or custom analysis, use `--analyze-raw` to print the raw metrics table without formatting: + +```sh +# Local raw metrics +dft -c "SELECT ..." --analyze-raw + +# FlightSQL raw metrics +dft -c "SELECT ..." --analyze-raw --flightsql + +# Write the raw metrics table to a file (Arrow/CSV/JSON/Parquet, by extension) +dft -c "SELECT ..." --analyze-raw --output metrics.parquet +``` + +This outputs a single flat Arrow metrics table with columns `(metric_name, value, value_type, operator_name, partition_id, operator_category, node_id, parent_node_id)`. The query text is not included; retain it client-side. See the [Arrow Flight Analyze Protocol Specification](arrow_flight_analyze_protocol.md) for the full schema and semantics. + +Writing raw metrics to a file makes a simple pipeline for tracking query performance over time: append metrics batches to a file on a schedule and query them with `dft` itself. + ## Generate TPC-H Data Generate TPC-H data into your configured DB path diff --git a/docs/flightsql_server.md b/docs/flightsql_server.md index 3f6186eb..ddca789b 100644 --- a/docs/flightsql_server.md +++ b/docs/flightsql_server.md @@ -34,6 +34,10 @@ The server implements the FlightSQL protocol, providing: - **SQL information** - Query server capabilities and version information via `CommandGetSqlInfo` - **Type metadata** - Get XDBC/ODBC type information via `CommandGetXdbcTypeInfo` for understanding supported data types +### Custom Actions +- **Query Analysis** - Get detailed execution metrics via custom `"analyze_query"` action + - See [Arrow Flight Analyze Protocol](arrow_flight_analyze_protocol.md) for the complete specification + ## Client Connections (TODO - Test this) You can connect to the server using any FlightSQL-compatible client: @@ -75,6 +79,7 @@ Available metrics include: - `get_flight_info_xdbc_type_info_latency_ms` - Type info metadata latency - `do_action_create_prepared_statement_latency_ms` - Prepared statement creation latency - `do_action_close_prepared_statement_latency_ms` - Prepared statement cleanup latency + - `do_action_analyze_query_latency_ms` - Analyze query action latency - `get_flight_info_prepared_statement_latency_ms` - Prepared statement flight info latency - `do_get_prepared_statement_latency_ms` - Prepared statement execution latency - Active prepared statements (`prepared_statements_active` gauge) @@ -96,3 +101,57 @@ target_partitions = 8 ``` See the [Config Reference](config.md) for all available options. + +## Query Analysis Support + +The `dft` FlightSQL server implements the [Arrow Flight Analyze Protocol](arrow_flight_analyze_protocol.md), which provides detailed query execution metrics through a custom action. + +### Quick Start + +The analyze protocol is automatically enabled when you start the FlightSQL server: + +```sh +dft serve-flightsql +``` + +Clients can then use the `--analyze` flag: + +```sh +# Analyze query via FlightSQL +dft -c "SELECT * FROM table WHERE id > 100" --analyze --flightsql + +# View raw metrics table +dft -c "SELECT ..." --analyze-raw --flightsql +``` + +### What's Provided + +The server collects and returns: + +- **Execution timing**: Parsing, planning, execution, and total time +- **I/O statistics**: Bytes scanned, file operations, format-specific metrics (per scan node) +- **Parquet metrics**: Row group pruning, bloom filters, page index effectiveness +- **Compute breakdown**: Per-operator, per-partition CPU time by category (projection, filter, sort, aggregate, join, window, distinct, limit, union, other) +- **Plan hierarchy**: Stable `node_id`/`parent_node_id` for every operator so the plan tree can be reconstructed + +**Note**: `analyze_query` executes the query to completion server-side and discards the results, so an analyze call costs a full query execution. + +### Implementation Details + +The `dft` server implementation: + +1. Handles the `"analyze_query"` and `"analyze_query_capabilities"` actions in `do_action_fallback()` +2. Validates the client's `protocol_version` if provided +3. Executes queries using `ExecutionContext::analyze_query()` and collects metrics from DataFusion execution plan `MetricSet`s +4. Serializes to a single flat Arrow metrics batch following the protocol spec, with `analyze.protocol_version` and a per-request `analyze.query_id` in the schema metadata +5. Returns a FlightData stream with metrics encoded as rows + +### For Other Implementers + +If you're implementing an Arrow Flight service and want to support the analyze protocol, see the complete [Arrow Flight Analyze Protocol Specification](arrow_flight_analyze_protocol.md). + +The protocol is designed to be: +- Modeled on DataFusion, intended to generalize to other engines +- Format-agnostic (supports Parquet, CSV, JSON, ORC, custom formats) +- Extensible (servers can add custom metrics) +- Language-agnostic (simple flat table format) diff --git a/src/args.rs b/src/args.rs index 3799d609..0e8cddb9 100644 --- a/src/args.rs +++ b/src/args.rs @@ -103,10 +103,16 @@ pub struct DftArgs { #[clap( long, - help = "Print a summary of the query's execution plan and statistics" + help = "Print a summary of the query's execution plan and statistics. Only a single SQL statement is allowed. Works with both local and FlightSQL execution (requires server support)." )] pub analyze: bool, + #[clap( + long, + help = "Print raw execution metrics as Arrow table without formatting. Only a single SQL statement is allowed. Useful for debugging and custom analysis. Implies --analyze." + )] + pub analyze_raw: bool, + #[clap(long, help = "Run the provided query before running the benchmark")] pub run_before: Option, diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 4276875e..ad48dd28 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -25,6 +25,7 @@ use color_eyre::eyre::eyre; use color_eyre::Result; use datafusion::arrow::array::{RecordBatch, RecordBatchWriter}; use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::ipc::writer::FileWriter as ArrowIpcWriter; use datafusion::arrow::util::pretty::pretty_format_batches; use datafusion::arrow::{csv, json}; use datafusion::sql::parser::DFParser; @@ -226,7 +227,7 @@ impl CliApp { self.args.commands.is_empty(), self.args.flightsql, self.args.bench, - self.args.analyze, + self.args.analyze || self.args.analyze_raw, ) { // Error cases (_, _, true, _, _) => Err(eyre!( @@ -236,12 +237,15 @@ impl CliApp { Err(eyre!("Cannot benchmark without a command or file")) } (true, true, _, _, _) => Err(eyre!("No files or commands provided to execute")), - (false, false, _, false, _) => Err(eyre!( + (false, false, _, false, false) => Err(eyre!( "Cannot execute both files and commands at the same time" )), (_, _, false, true, true) => Err(eyre!( "The `benchmark` and `analyze` flags are mutually exclusive" )), + (false, false, _, _, true) => { + Err(eyre!("Analyze requires exactly one command or file")) + } // Execution cases (false, true, _, false, false) => self.execute_files(&self.args.files).await, @@ -252,8 +256,14 @@ impl CliApp { (true, false, _, true, false) => self.benchmark_commands(&self.args.commands).await, // Analyze cases - (false, true, _, false, true) => self.analyze_files(&self.args.files).await, - (true, false, _, false, true) => self.analyze_commands(&self.args.commands).await, + (false, true, _, false, true) => { + self.validate_analyze_target()?; + self.analyze_files(&self.args.files).await + } + (true, false, _, false, true) => { + self.validate_analyze_target()?; + self.analyze_commands(&self.args.commands).await + } } #[cfg(feature = "flightsql")] match ( @@ -261,22 +271,22 @@ impl CliApp { self.args.commands.is_empty(), self.args.flightsql, self.args.bench, - self.args.analyze, + self.args.analyze || self.args.analyze_raw, ) { // Error cases (true, true, _, _, _) => Err(eyre!("No files or commands provided to execute")), (false, false, false, true, _) => { Err(eyre!("Cannot benchmark without a command or file")) } - (false, false, _, _, _) => Err(eyre!( + (false, false, _, _, false) => Err(eyre!( "Cannot execute both files and commands at the same time" )), (_, _, _, true, true) => Err(eyre!( "The `benchmark` and `analyze` flags are mutually exclusive" )), - (_, _, true, false, true) => Err(eyre!( - "The `analyze` flag is not currently supported with FlightSQL" - )), + (false, false, _, _, true) => { + Err(eyre!("Analyze requires exactly one command or file")) + } // Execution cases (true, false, false, false, false) => self.execute_commands(&self.args.commands).await, @@ -301,11 +311,36 @@ impl CliApp { (true, false, false, true, false) => self.benchmark_commands(&self.args.commands).await, // Analyze cases - (true, false, false, false, true) => self.analyze_commands(&self.args.commands).await, - (false, true, false, false, true) => self.analyze_files(&self.args.files).await, + (true, false, false, false, true) => { + self.validate_analyze_target()?; + self.analyze_commands(&self.args.commands).await + } + (false, true, false, false, true) => { + self.validate_analyze_target()?; + self.analyze_files(&self.args.files).await + } + (true, false, true, false, true) => { + self.validate_analyze_target()?; + self.flightsql_analyze_commands(&self.args.commands).await + } + (false, true, true, false, true) => { + self.validate_analyze_target()?; + self.flightsql_analyze_files(&self.args.files).await + } } } + /// Analyze accepts exactly one file or exactly one command + fn validate_analyze_target(&self) -> Result<()> { + if self.args.files.len() > 1 { + return Err(eyre!("Analyze requires exactly one file")); + } + if self.args.commands.len() > 1 { + return Err(eyre!("Analyze requires exactly one command")); + } + Ok(()) + } + async fn execute_files(&self, files: &[PathBuf]) -> Result<()> { info!("Executing files: {:?}", files); for file in files { @@ -332,6 +367,12 @@ impl CliApp { } async fn analyze_files(&self, files: &[PathBuf]) -> Result<()> { + if let Some(run_before_query) = &self.args.run_before { + self.app_execution + .execution_ctx() + .execute_sql_and_discard_results(run_before_query) + .await?; + } info!("Analyzing files: {:?}", files); for file in files { let query = std::fs::read_to_string(file)?; @@ -476,6 +517,12 @@ impl CliApp { } async fn analyze_commands(&self, commands: &[String]) -> color_eyre::Result<()> { + if let Some(run_before_query) = &self.args.run_before { + self.app_execution + .execution_ctx() + .execute_sql_and_discard_results(run_before_query) + .await?; + } info!("Analyzing commands: {:?}", commands); for command in commands { self.analyze_from_string(command).await?; @@ -529,6 +576,68 @@ impl CliApp { Ok(()) } + #[cfg(feature = "flightsql")] + async fn flightsql_analyze_commands(&self, commands: &[String]) -> Result<()> { + info!("Analyzing FlightSQL commands: {:?}", commands); + for command in commands { + self.flightsql_analyze_from_string(command).await?; + } + Ok(()) + } + + #[cfg(feature = "flightsql")] + async fn flightsql_analyze_files(&self, files: &[PathBuf]) -> Result<()> { + info!("Analyzing FlightSQL files: {:?}", files); + for file in files { + let sql = std::fs::read_to_string(file)?; + self.flightsql_analyze_from_string(&sql).await?; + } + Ok(()) + } + + #[cfg(feature = "flightsql")] + async fn flightsql_analyze_from_string(&self, sql: &str) -> Result<()> { + if self.args.analyze_raw { + // Raw mode: print metrics table directly + let (query_str, metrics_batch) = self + .app_execution + .flightsql_ctx() + .analyze_query_raw(sql) + .await?; + + if let Some(output_path) = &self.args.output { + // Write metrics batch to file + let schema = metrics_batch.schema(); + let mut writer = path_to_writer(output_path, schema)?; + writer.write(&metrics_batch)?; + writer.close().await?; + } else { + // Print to stdout + println!("==================== Query ===================="); + println!("{}", query_str); + println!("\n==================== Metrics ===================="); + self.print_batch(&metrics_batch)?; + } + } else { + // Normal mode: reconstruct and display ExecutionStats + let stats = self + .app_execution + .flightsql_ctx() + .analyze_query(sql) + .await?; + + // Display using existing ExecutionStats::Display implementation + println!("{}", stats); + } + Ok(()) + } + + fn print_batch(&self, batch: &datafusion::arrow::array::RecordBatch) -> Result<()> { + use datafusion::arrow::util::pretty::print_batches; + print_batches(std::slice::from_ref(batch))?; + Ok(()) + } + async fn exec_from_string(&self, sql: &str) -> Result<()> { let dialect = datafusion::sql::sqlparser::dialect::GenericDialect {}; let statements = DFParser::parse_sql_with_dialect(sql, &dialect)?; @@ -607,7 +716,23 @@ impl CliApp { .analyze_query(sql) .await?; stats.collect_stats(); - println!("{}", stats); + if self.args.analyze_raw { + // Raw mode: print or write the metrics table directly + let metrics_batch = stats.to_metrics_table()?; + if let Some(output_path) = &self.args.output { + let schema = metrics_batch.schema(); + let mut writer = path_to_writer(output_path, schema)?; + writer.write(&metrics_batch)?; + writer.close().await?; + } else { + println!("==================== Query ===================="); + println!("{}", sql); + println!("\n==================== Metrics ===================="); + self.print_batch(&metrics_batch)?; + } + } else { + println!("{}", stats); + } Ok(()) } @@ -880,6 +1005,7 @@ enum AnyWriter { Csv(csv::writer::Writer), Json(json::writer::LineDelimitedWriter), Parquet(ArrowWriter), + Arrow(ArrowIpcWriter), #[cfg(feature = "vortex")] Vortex(VortexFileWriter), } @@ -890,12 +1016,13 @@ impl AnyWriter { AnyWriter::Csv(w) => Ok(w.write(batch)?), AnyWriter::Json(w) => Ok(w.write(batch)?), AnyWriter::Parquet(w) => Ok(w.write(batch)?), + AnyWriter::Arrow(w) => Ok(w.write(batch)?), #[cfg(feature = "vortex")] AnyWriter::Vortex(w) => Ok(w.write(batch)?), } } - async fn close(self) -> Result<()> { + async fn close(mut self) -> Result<()> { match self { AnyWriter::Csv(w) => Ok(w.close()?), AnyWriter::Json(w) => Ok(w.close()?), @@ -903,6 +1030,10 @@ impl AnyWriter { w.close()?; Ok(()) } + AnyWriter::Arrow(ref mut w) => { + w.finish()?; + Ok(()) + } #[cfg(feature = "vortex")] AnyWriter::Vortex(w) => w.close().await, } @@ -923,6 +1054,10 @@ fn path_to_writer(path: &Path, schema: SchemaRef) -> Result { let writer = ArrowWriter::try_new(file, schema, Some(props))?; Ok(AnyWriter::Parquet(writer)) } + "arrow" | "ipc" => { + let writer = ArrowIpcWriter::try_new(file, &schema)?; + Ok(AnyWriter::Arrow(writer)) + } #[cfg(feature = "vortex")] "vortex" => Ok(AnyWriter::Vortex(VortexFileWriter::new( file, schema, path, @@ -930,11 +1065,11 @@ fn path_to_writer(path: &Path, schema: SchemaRef) -> Result { _ => { #[cfg(feature = "vortex")] return Err(eyre!( - "Only 'csv', 'parquet', 'json', and 'vortex' file types can be output" + "Only 'csv', 'parquet', 'json', 'arrow', and 'vortex' file types can be output" )); #[cfg(not(feature = "vortex"))] return Err(eyre!( - "Only 'csv', 'parquet', and 'json' file types can be output" + "Only 'csv', 'parquet', 'json', and 'arrow' file types can be output" )); } }; diff --git a/src/server/flightsql/service.rs b/src/server/flightsql/service.rs index 1910500f..6c56f3fc 100644 --- a/src/server/flightsql/service.rs +++ b/src/server/flightsql/service.rs @@ -839,6 +839,151 @@ impl FlightSqlService for FlightSqlServiceImpl { res } + async fn do_action_fallback( + &self, + request: Request, + ) -> Result::DoActionStream>, Status> { + use arrow_flight::utils::batches_to_flight_data; + let action = request.into_inner(); + counter!("requests", "endpoint" => "do_action_fallback").increment(1); + let start = Timestamp::now(); + + match action.r#type.as_str() { + "analyze_query" => { + // 1. Parse JSON request body + let request: datafusion_app::stats::AnalyzeQueryRequest = + serde_json::from_slice(&action.body).map_err(|e| { + Status::invalid_argument(format!("Invalid JSON request: {}", e)) + })?; + + // 2. Validate protocol version if the client sent one + if let Some(version) = &request.protocol_version { + if !datafusion_app::stats::is_compatible_protocol_version(version) { + return Err(Status::invalid_argument(format!( + "Unsupported analyze protocol version: {} (server implements {})", + version, + datafusion_app::stats::ANALYZE_PROTOCOL_VERSION + ))); + } + } + + // 3. Extract SQL query (only supported format for now) + let sql = request + .sql() + .map_err(|e| Status::invalid_argument(e.to_string()))?; + + info!("Analyzing query via do_action: {}", sql); + + // 4. Execute analyze_query on ExecutionContext + let mut stats = self + .execution + .analyze_query(sql) + .await + .map_err(|e| Status::internal(format!("Analyze failed: {}", e)))?; + + stats.collect_stats(); // Collect IO and compute metrics from plan + + // 5. Convert ExecutionStats to metrics table format + let metrics_batch = stats.to_metrics_table().map_err(|e| { + Status::internal(format!("Metrics serialization failed: {}", e)) + })?; + + // 6. Attach a correlation id so clients can match concurrent + // responses to requests. The protocol version is already in + // the schema metadata from `analyze_metrics_schema`. + let query_id = uuid::Uuid::new_v4(); + let mut metadata = metrics_batch.schema().metadata().clone(); + metadata.insert( + datafusion_app::stats::QUERY_ID_METADATA_KEY.to_string(), + query_id.to_string(), + ); + let schema_with_meta = Arc::new( + metrics_batch + .schema() + .as_ref() + .clone() + .with_metadata(metadata), + ); + let metrics_batch = datafusion::arrow::array::RecordBatch::try_new( + Arc::clone(&schema_with_meta), + metrics_batch.columns().to_vec(), + ) + .map_err(|e| Status::internal(format!("Failed to attach metadata: {}", e)))?; + + // 7. Encode metrics batch as FlightData + let flight_data = batches_to_flight_data(&schema_with_meta, vec![metrics_batch]) + .map_err(|e| { + Status::internal(format!("Failed to encode metrics batch: {}", e)) + })?; + + // 8. Convert FlightData to arrow_flight::Result messages + // Note: Query is NOT included in response; clients must retain the original request + let results: Vec = flight_data + .into_iter() + .map(|fd| { + // Serialize FlightData to bytes + let bytes = fd.encode_to_vec(); + arrow_flight::Result { body: bytes.into() } + }) + .collect(); + + // 9. Create stream of Result messages + let stream = futures::stream::iter(results.into_iter().map(Ok)).boxed(); + + // Record metrics + let duration = Timestamp::now() - start; + histogram!("do_action_analyze_query_latency_ms") + .record(duration.get_milliseconds() as f64); + + let ctx = self.execution.session_ctx(); + let req = ObservabilityRequestDetails { + request_id: None, + path: "/do_action/analyze_query".to_string(), + sql: Some(sql.to_string()), + start_ms: start.as_millisecond(), + duration_ms: duration.get_milliseconds(), + rows: None, + status: 0, + }; + if let Err(e) = self + .execution + .observability() + .try_record_request(ctx, req) + .await + { + error!("Error recording request: {}", e); + } + + Ok(Response::new(stream)) + } + "analyze_query_capabilities" => { + let capabilities = serde_json::json!({ + "protocol_version": datafusion_app::stats::ANALYZE_PROTOCOL_VERSION, + "request_formats": ["sql"], + "namespaces": [ + "query", + "stage", + "compute", + "io.parquet", + "io.csv", + "io.arrow", + "io.json", + ], + }); + let body = serde_json::to_vec(&capabilities).map_err(|e| { + Status::internal(format!("Failed to serialize capabilities: {}", e)) + })?; + let result = arrow_flight::Result { body: body.into() }; + let stream = futures::stream::iter([Ok(result)]).boxed(); + Ok(Response::new(stream)) + } + _ => Err(Status::unimplemented(format!( + "Unknown action: {}", + action.r#type + ))), + } + } + async fn register_sql_info(&self, _id: i32, _result: &SqlInfo) {} } diff --git a/tests/cli_cases/analyze.rs b/tests/cli_cases/analyze.rs new file mode 100644 index 00000000..054f07db --- /dev/null +++ b/tests/cli_cases/analyze.rs @@ -0,0 +1,182 @@ +// 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. + +//! Tests for local (non-FlightSQL) `--analyze` and `--analyze-raw`. +//! +//! The `data/test_io_formats.{arrow,csv,json,parquet}` files at the repo root +//! are fixtures for these and the FlightSQL analyze tests. + +use assert_cmd::Command; + +use super::sql_in_file; + +#[test] +fn test_analyze_command() { + let assert = Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT 1 + 2") + .arg("--analyze") + .assert() + .success(); + + let output = String::from_utf8_lossy(&assert.get_output().stdout).to_string(); + assert!(output.contains("Query"), "Should contain Query section"); + assert!( + output.contains("Execution Summary"), + "Should contain Execution Summary" + ); + assert!( + output.contains("Parsing"), + "Should contain timing breakdown" + ); + assert!( + output.contains("Compute Summary"), + "Should contain Compute Summary" + ); +} + +#[test] +fn test_analyze_file() { + let file = sql_in_file("SELECT 1 + 1"); + let assert = Command::cargo_bin("dft") + .unwrap() + .arg("-f") + .arg(file.path()) + .arg("--analyze") + .assert() + .success(); + + let output = String::from_utf8_lossy(&assert.get_output().stdout).to_string(); + assert!( + output.contains("Execution Summary"), + "Should contain Execution Summary" + ); +} + +#[test] +fn test_analyze_raw_command() { + let assert = Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT 1 + 2") + .arg("--analyze-raw") + .assert() + .success(); + + // Raw mode prints the metrics table, not the formatted summary + let output = String::from_utf8_lossy(&assert.get_output().stdout).to_string(); + assert!( + output.contains("metric_name"), + "Should contain metric_name column" + ); + assert!(output.contains("node_id"), "Should contain node_id column"); + assert!( + output.contains("parent_node_id"), + "Should contain parent_node_id column" + ); + assert!( + output.contains("query.rows"), + "Should contain query.rows metric" + ); + assert!( + output.contains("stage.execution"), + "Should contain stage.execution metric" + ); + assert!( + !output.contains("Execution Summary"), + "Raw mode should not print the formatted summary" + ); +} + +#[test] +fn test_analyze_raw_parquet_io_metrics() { + let assert = Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT * FROM parquet_scan WHERE value > 50") + .arg("--analyze-raw") + .arg("--run-before") + .arg("CREATE EXTERNAL TABLE parquet_scan STORED AS PARQUET LOCATION 'data/test_io_formats.parquet'") + .assert() + .success(); + + let output = String::from_utf8_lossy(&assert.get_output().stdout).to_string(); + assert!( + output.contains("io.parquet.bytes_scanned"), + "Should contain io.parquet.bytes_scanned for a Parquet scan" + ); + assert!( + output.contains("io.parquet.rg_pruned"), + "Should contain io.parquet.rg_pruned for a Parquet scan" + ); + assert!( + output.contains("io.parquet.rg_matched"), + "Should contain io.parquet.rg_matched for a Parquet scan" + ); +} + +#[test] +fn test_analyze_raw_output_to_file() { + let dir = tempfile::tempdir().unwrap(); + let out = dir.path().join("metrics.csv"); + + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT 1") + .arg("--analyze-raw") + .arg("--output") + .arg(&out) + .assert() + .success(); + + let contents = std::fs::read_to_string(&out).unwrap(); + assert!( + contents.contains("metric_name"), + "Output file should contain the metrics table header" + ); + assert!( + contents.contains("query.rows"), + "Output file should contain query.rows metric" + ); +} + +#[test] +fn test_analyze_multiple_commands_fails() { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT 1") + .arg("-c") + .arg("SELECT 2") + .arg("--analyze") + .assert() + .failure(); +} + +#[test] +fn test_analyze_and_bench_mutually_exclusive() { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT 1") + .arg("--analyze") + .arg("--bench") + .assert() + .failure(); +} diff --git a/tests/cli_cases/mod.rs b/tests/cli_cases/mod.rs index b3658c99..9306f091 100644 --- a/tests/cli_cases/mod.rs +++ b/tests/cli_cases/mod.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +mod analyze; mod basic; mod bench; mod config; diff --git a/tests/extension_cases/flightsql.rs b/tests/extension_cases/flightsql.rs index 141802ba..f7df89e0 100644 --- a/tests/extension_cases/flightsql.rs +++ b/tests/extension_cases/flightsql.rs @@ -47,7 +47,9 @@ pub async fn test_execute_with_no_flightsql_server() { #[tokio::test] pub async fn test_execute() { - let test_server = TestFlightSqlServiceImpl::new(); + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; let assert = tokio::task::spawn_blocking(|| { @@ -76,7 +78,9 @@ pub async fn test_execute() { #[tokio::test] pub async fn test_invalid_sql_command() { - let test_server = TestFlightSqlServiceImpl::new(); + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; let assert = tokio::task::spawn_blocking(|| { @@ -92,9 +96,9 @@ pub async fn test_invalid_sql_command() { .await .unwrap(); - // arrow-flight surfaces server-side failures as a Tonic status with the underlying - // SQL error in the message - let expected = r##"Expected: an SQL statement"##; + // arrow-flight surfaces server-side failures as a Tonic status; the real + // FlightSQL service maps SQL parse failures to an internal error + let expected = r##"error parsing SQL query"##; assert.stderr(contains_str(expected)); fixture.shutdown_and_wait().await; } @@ -136,7 +140,9 @@ pub async fn test_execute_multiple_commands() { #[tokio::test] pub async fn test_command_in_file() { - let test_server = TestFlightSqlServiceImpl::new(); + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; let file = sql_in_file("SELECT 1 + 1"); let assert = tokio::task::spawn_blocking(move || { @@ -163,7 +169,9 @@ pub async fn test_command_in_file() { #[tokio::test] pub async fn test_invalid_sql_command_in_file() { - let test_server = TestFlightSqlServiceImpl::new(); + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; let file = sql_in_file("SELEC 1"); let assert = tokio::task::spawn_blocking(move || { @@ -178,9 +186,9 @@ pub async fn test_invalid_sql_command_in_file() { .await .unwrap(); - // arrow-flight surfaces server-side failures as a Tonic status with the underlying - // SQL error in the message - let expected = r##"Expected: an SQL statement"##; + // arrow-flight surfaces server-side failures as a Tonic status; the real + // FlightSQL service maps SQL parse failures to an internal error + let expected = r##"error parsing SQL query"##; assert.stderr(contains_str(expected)); fixture.shutdown_and_wait().await; } @@ -292,7 +300,9 @@ SELECT 1 #[tokio::test] pub async fn test_bench_files() { - let test_server = TestFlightSqlServiceImpl::new(); + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; let file = sql_in_file(r#"SELECT 1 + 1;"#); let assert = tokio::task::spawn_blocking(move || { @@ -383,7 +393,9 @@ SELECT 1 #[tokio::test] pub async fn test_bench_files_and_save() { - let test_server = TestFlightSqlServiceImpl::new(); + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; let file = sql_in_file(r#"SELECT 1 + 1;"#); @@ -1358,6 +1370,948 @@ pub async fn test_prepared_statement_complex_query() { fixture.shutdown_and_wait().await; } +// ============================================================================ +// FlightSQL Analyze Tests +// ============================================================================ + +#[tokio::test] +pub async fn test_analyze_command() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + let assert = tokio::task::spawn_blocking(|| { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT 1 + 2") + .arg("--analyze") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .success() + }) + .await + .unwrap(); + + // Verify output contains expected sections + let output = String::from_utf8_lossy(&assert.get_output().stdout); + assert!(output.contains("Query"), "Should contain Query section"); + assert!( + output.contains("Execution Summary"), + "Should contain Execution Summary" + ); + assert!( + output.contains("Output Rows"), + "Should contain Output Rows metric" + ); + assert!( + output.contains("Parsing"), + "Should contain timing breakdown" + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_raw_command() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + let assert = tokio::task::spawn_blocking(|| { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT 1 + 2") + .arg("--analyze-raw") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .success() + }) + .await + .unwrap(); + + // Verify raw mode outputs query string and metrics table + let output = String::from_utf8_lossy(&assert.get_output().stdout); + assert!(output.contains("Query"), "Should contain Query header"); + assert!(output.contains("Metrics"), "Should contain Metrics header"); + assert!( + output.contains("SELECT 1 + 2"), + "Should contain the SQL query" + ); + assert!( + output.contains("metric_name"), + "Should contain metric_name column" + ); + assert!(output.contains("value"), "Should contain value column"); + assert!( + output.contains("value_type"), + "Should contain value_type column" + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_file() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + let file = sql_in_file("SELECT 1 + 1"); + + let assert = tokio::task::spawn_blocking(move || { + Command::cargo_bin("dft") + .unwrap() + .arg("-f") + .arg(file.path()) + .arg("--analyze") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .success() + }) + .await + .unwrap(); + + // Verify output contains expected sections + let output = String::from_utf8_lossy(&assert.get_output().stdout); + assert!( + output.contains("Execution Summary"), + "Should contain Execution Summary" + ); + assert!( + output.contains("Output Rows"), + "Should contain Output Rows metric" + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_raw_file() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + let file = sql_in_file("SELECT 1 + 1"); + + let assert = tokio::task::spawn_blocking(move || { + Command::cargo_bin("dft") + .unwrap() + .arg("-f") + .arg(file.path()) + .arg("--analyze-raw") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .success() + }) + .await + .unwrap(); + + // Verify raw mode outputs query string and metrics table + let output = String::from_utf8_lossy(&assert.get_output().stdout); + assert!(output.contains("Query"), "Should contain Query header"); + assert!(output.contains("Metrics"), "Should contain Metrics header"); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_multiple_commands() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + let assert = tokio::task::spawn_blocking(|| { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT 1") + .arg("-c") + .arg("SELECT 2") + .arg("--analyze") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .failure() + }) + .await + .unwrap(); + + // Verify error message about requiring exactly one command + let stderr = String::from_utf8_lossy(&assert.get_output().stderr); + assert!( + stderr.contains("Analyze requires exactly one command"), + "Should contain error about requiring one command" + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_invalid_sql() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + let assert = tokio::task::spawn_blocking(|| { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELEC 1") + .arg("--analyze") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .failure() + }) + .await + .unwrap(); + + // Verify error is reported + let stderr = String::from_utf8_lossy(&assert.get_output().stderr); + assert!( + stderr.contains("Error") || stderr.contains("error"), + "Should contain error message" + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_with_timing_metrics() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + let assert = tokio::task::spawn_blocking(|| { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT 1") + .arg("--analyze") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .success() + }) + .await + .unwrap(); + + // Verify all timing metrics are present + let output = String::from_utf8_lossy(&assert.get_output().stdout); + assert!(output.contains("Parsing"), "Should contain Parsing time"); + assert!( + output.contains("Logical Planning"), + "Should contain Logical Planning time" + ); + assert!( + output.contains("Physical Planning"), + "Should contain Physical Planning time" + ); + assert!( + output.contains("Execution"), + "Should contain Execution time" + ); + assert!(output.contains("Total"), "Should contain Total time"); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_raw_metrics_schema() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + let assert = tokio::task::spawn_blocking(|| { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT 1") + .arg("--analyze-raw") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .success() + }) + .await + .unwrap(); + + // Verify raw metrics table has all expected columns (8-field schema) + let output = String::from_utf8_lossy(&assert.get_output().stdout); + assert!( + output.contains("metric_name"), + "Should contain metric_name column" + ); + assert!(output.contains("value"), "Should contain value column"); + assert!( + output.contains("value_type"), + "Should contain value_type column" + ); + assert!( + output.contains("operator_name"), + "Should contain operator_name column" + ); + assert!( + output.contains("partition_id"), + "Should contain partition_id column" + ); + assert!( + output.contains("operator_category"), + "Should contain operator_category column" + ); + assert!(output.contains("node_id"), "Should contain node_id column"); + assert!( + output.contains("parent_node_id"), + "Should contain parent_node_id column" + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_raw_duration_metrics() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + let assert = tokio::task::spawn_blocking(|| { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT 1") + .arg("--analyze-raw") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .success() + }) + .await + .unwrap(); + + // Verify duration metrics are in the raw output + let output = String::from_utf8_lossy(&assert.get_output().stdout); + assert!(output.contains("parsing"), "Should contain parsing metric"); + assert!( + output.contains("logical_planning"), + "Should contain logical_planning metric" + ); + assert!( + output.contains("physical_planning"), + "Should contain physical_planning metric" + ); + assert!( + output.contains("execution"), + "Should contain execution metric" + ); + assert!(output.contains("total"), "Should contain total metric"); + assert!( + output.contains("duration_ns"), + "Should contain duration_ns value_type" + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_output_metrics() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + let assert = tokio::task::spawn_blocking(|| { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT 1") + .arg("--analyze-raw") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .success() + }) + .await + .unwrap(); + + // Verify output metrics (rows, batches, bytes) are present + let output = String::from_utf8_lossy(&assert.get_output().stdout); + assert!(output.contains("rows"), "Should contain rows metric"); + assert!(output.contains("batches"), "Should contain batches metric"); + assert!(output.contains("bytes"), "Should contain bytes metric"); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_multiple_statements_in_single_command() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + let assert = tokio::task::spawn_blocking(|| { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg("SELECT 1; SELECT 2") + .arg("--analyze") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .failure() + }) + .await + .unwrap(); + + // Verify error message about single statement requirement + let stderr = String::from_utf8_lossy(&assert.get_output().stderr); + assert!( + stderr.contains("Only a single SQL statement can be analyzed"), + "Should contain error about single statement requirement" + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_file_with_multiple_statements() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + let file = sql_in_file("SELECT 1; SELECT 2"); + + let assert = tokio::task::spawn_blocking(move || { + Command::cargo_bin("dft") + .unwrap() + .arg("-f") + .arg(file.path()) + .arg("--analyze") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .failure() + }) + .await + .unwrap(); + + // Verify error message about single statement requirement + let stderr = String::from_utf8_lossy(&assert.get_output().stderr); + assert!( + stderr.contains("Only a single SQL statement can be analyzed"), + "Should contain error about single statement requirement" + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_multiple_files() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + let file1 = sql_in_file("SELECT 1"); + let file2 = sql_in_file("SELECT 2"); + + let assert = tokio::task::spawn_blocking(move || { + Command::cargo_bin("dft") + .unwrap() + .arg("-f") + .arg(file1.path()) + .arg("-f") + .arg(file2.path()) + .arg("--analyze") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .failure() + }) + .await + .unwrap(); + + // Verify error message about requiring exactly one file + let stderr = String::from_utf8_lossy(&assert.get_output().stderr); + assert!( + stderr.contains("Analyze requires exactly one file"), + "Should contain error about requiring one file" + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_operator_hierarchy() { + let ctx = ExecutionContext::test(); + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + // Run a complex query with multiple operators to test operator hierarchy + // (parent-child relationships). This query has: + // - VALUES clause (scan) + // - Filter (WHERE) + // - Aggregate (GROUP BY) + // - Sort (ORDER BY) + // - Limit + // - Projection (SELECT columns) + let query = r#" + SELECT + category, + count, + total + FROM ( + SELECT + column2 as category, + COUNT(*) as count, + SUM(column3) as total + FROM (VALUES + (1, 'a', 100), + (2, 'b', 200), + (3, 'a', 300), + (4, 'b', 400), + (5, 'a', 500) + ) AS t(column1, column2, column3) + WHERE column1 > 1 + GROUP BY column2 + ) AS subquery + ORDER BY count DESC + LIMIT 2 + "#; + + let assert = tokio::task::spawn_blocking(move || { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg(query) + .arg("--analyze-raw") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .success() + }) + .await + .unwrap(); + + let output = String::from_utf8_lossy(&assert.get_output().stdout); + + // Verify the schema includes hierarchy fields + assert!(output.contains("node_id"), "Should contain node_id column"); + assert!( + output.contains("parent_node_id"), + "Should contain parent_node_id column" + ); + + // Verify we have a complex execution plan with multiple operators + // For this query we expect operators like: + // - Projection (root - final SELECT columns) + // - Limit + // - Sort + // - Aggregate (GROUP BY) + // - Filter (WHERE clause) + // - MemoryExec or similar scan operator (leaf) + + // Check that we have metrics from multiple operator types + assert!( + output.contains("ProjectionExec") || output.contains("projection"), + "Should have projection operator" + ); + assert!( + output.contains("AggregateExec") || output.contains("aggregate"), + "Should have aggregate operator" + ); + assert!( + output.contains("FilterExec") + || output.contains("filter") + || output.contains("CoalesceBatchesExec"), + "Should have filter or coalesce operator" + ); + + // For a complex query with multiple operators, verify that operator names are present + // This confirms that the hierarchy was collected and included in the output + assert!( + output.contains("Exec"), + "Should contain operator names in output" + ); + + // A GROUP BY plans two AggregateExec nodes (partial + final). With + // node-id based identity they must appear as distinct nodes, not be + // merged under one operator name. Extract the node_id column (7th field + // in the pretty-printed table) for every AggregateExec row. + let agg_node_ids: std::collections::HashSet = output + .lines() + .filter(|l| l.contains("AggregateExec")) + .filter_map(|l| l.split('|').map(str::trim).nth(7).map(str::to_string)) + .collect(); + assert!( + agg_node_ids.len() >= 2, + "Partial and final AggregateExec should have distinct node ids, got: {:?}", + agg_node_ids + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_operator_hierarchy_with_csv() { + use datafusion::prelude::CsvReadOptions; + + // Create an ExecutionContext and register the test CSV file as a table + let mut ctx = ExecutionContext::test(); + + // Register the aggregate_test_100.csv file + ctx.session_ctx() + .register_csv( + "test_data", + "data/aggregate_test_100.csv", + CsvReadOptions::new(), + ) + .await + .expect("Failed to register CSV table"); + + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + // Run a complex query that will create a rich operator hierarchy with CSV scan + // This query has: CsvExec -> Filter -> Aggregate -> Sort -> Limit -> Projection + let analyze_query = r#" + SELECT + c1, + COUNT(*) as count, + AVG(c2) as avg_c2, + SUM(c3) as sum_c3 + FROM test_data + WHERE c2 > 2 + GROUP BY c1 + HAVING COUNT(*) > 5 + ORDER BY count DESC + LIMIT 10 + "#; + + let assert = tokio::task::spawn_blocking(move || { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg(analyze_query) + .arg("--analyze-raw") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .success() + }) + .await + .unwrap(); + + let output = String::from_utf8_lossy(&assert.get_output().stdout); + + // Verify the schema includes hierarchy fields + assert!(output.contains("node_id"), "Should contain node_id column"); + assert!( + output.contains("parent_node_id"), + "Should contain parent_node_id column" + ); + + // Verify we have expected operators from the complex query + // Note: CSV scan operators don't report compute metrics, so we check other operators + assert!( + output.contains("AggregateExec"), + "Should have AggregateExec operator for GROUP BY" + ); + assert!( + output.contains("FilterExec"), + "Should have FilterExec operator for WHERE clause" + ); + assert!( + output.contains("SortExec"), + "Should have SortExec operator for ORDER BY" + ); + assert!( + output.contains("ProjectionExec"), + "Should have ProjectionExec operator for SELECT" + ); + + // Verify we have compute metrics + assert!( + output.contains("compute.") || output.contains("elapsed_compute"), + "Should contain compute metrics" + ); + + // Verify we have I/O metrics with correct CSV namespace + assert!( + output.contains("io.csv."), + "Should contain io.csv.* metrics for a CSV scan" + ); + assert!( + output.contains("io.csv.time_scanning"), + "Should contain io.csv.time_scanning metric" + ); + + // Verify stage metrics with namespacing + assert!( + output.contains("stage.") || output.contains("parsing"), + "Should contain stage metrics" + ); + + // Verify query-level metrics with namespacing + assert!( + output.contains("query.rows") || output.contains("rows"), + "Should contain query-level metrics" + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_io_metrics_parquet_namespace() { + // Test that Parquet files report metrics under io.parquet.* namespace + let mut ctx = ExecutionContext::test(); + + // Register the test Parquet file + ctx.session_ctx() + .sql("CREATE EXTERNAL TABLE test_parquet STORED AS PARQUET LOCATION 'data/test_io_formats.parquet'") + .await + .expect("Failed to register Parquet table") + .collect() + .await + .expect("Failed to execute CREATE TABLE"); + + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + let analyze_query = r#" + SELECT + category, + COUNT(*) as count, + SUM(value) as total + FROM test_parquet + WHERE value > 50 + GROUP BY category + ORDER BY count DESC + "#; + + let assert = tokio::task::spawn_blocking(move || { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg(analyze_query) + .arg("--analyze-raw") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .success() + }) + .await + .unwrap(); + + let output = String::from_utf8_lossy(&assert.get_output().stdout); + + // Verify schema + assert!(output.contains("node_id"), "Should contain node_id column"); + assert!( + output.contains("parent_node_id"), + "Should contain parent_node_id column" + ); + + // Verify Parquet-specific I/O metrics are present with the right namespace + assert!( + output.contains("io.parquet.bytes_scanned"), + "Should contain io.parquet.bytes_scanned for a Parquet scan" + ); + assert!( + output.contains("io.parquet.rg_pruned"), + "Should contain io.parquet.rg_pruned for a Parquet scan" + ); + assert!( + output.contains("io.parquet.rg_matched"), + "Should contain io.parquet.rg_matched for a Parquet scan" + ); + // Verify no other format namespaces are used + assert!( + !output.contains("io.csv."), + "Should NOT use io.csv.* namespace for Parquet files" + ); + assert!( + !output.contains("io.json."), + "Should NOT use io.json.* namespace for Parquet files" + ); + + // Verify compute metrics are always present + assert!( + output.contains("compute.") || output.contains("elapsed_compute"), + "Should contain compute metrics" + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_io_metrics_json_namespace() { + // Test that JSON files report metrics under io.json.* namespace + let mut ctx = ExecutionContext::test(); + + // Register the test JSON file + ctx.session_ctx() + .sql("CREATE EXTERNAL TABLE test_json STORED AS JSON LOCATION 'data/test_io_formats.json'") + .await + .expect("Failed to register JSON table") + .collect() + .await + .expect("Failed to execute CREATE TABLE"); + + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + let analyze_query = r#" + SELECT + category, + COUNT(*) as count + FROM test_json + WHERE value > 75 + GROUP BY category + "#; + + let assert = tokio::task::spawn_blocking(move || { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg(analyze_query) + .arg("--analyze-raw") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .success() + }) + .await + .unwrap(); + + let output = String::from_utf8_lossy(&assert.get_output().stdout); + + // Verify schema + assert!(output.contains("node_id"), "Should contain node_id column"); + assert!( + output.contains("parent_node_id"), + "Should contain parent_node_id column" + ); + + // Verify JSON-specific I/O metrics are present with the right namespace + assert!( + output.contains("io.json."), + "Should contain io.json.* metrics for a JSON scan" + ); + // Verify no other format namespaces are used + assert!( + !output.contains("io.csv."), + "Should NOT use io.csv.* namespace for JSON files" + ); + assert!( + !output.contains("io.parquet."), + "Should NOT use io.parquet.* namespace for JSON files" + ); + + // Verify compute metrics are always present + assert!( + output.contains("compute.") || output.contains("elapsed_compute"), + "Should contain compute metrics" + ); + + fixture.shutdown_and_wait().await; +} + +#[tokio::test] +pub async fn test_analyze_io_metrics_arrow_namespace() { + // Test that Arrow IPC files report metrics under io.arrow.* namespace + let mut ctx = ExecutionContext::test(); + + // Register the test Arrow IPC file + ctx.session_ctx() + .sql("CREATE EXTERNAL TABLE test_arrow STORED AS ARROW LOCATION 'data/test_io_formats.arrow'") + .await + .expect("Failed to register Arrow table") + .collect() + .await + .expect("Failed to execute CREATE TABLE"); + + let exec = AppExecution::new(ctx); + let test_server = FlightSqlServiceImpl::new(exec); + let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await; + + let analyze_query = r#" + SELECT + name, + category, + value + FROM test_arrow + WHERE value > 100 + ORDER BY value DESC + "#; + + let assert = tokio::task::spawn_blocking(move || { + Command::cargo_bin("dft") + .unwrap() + .arg("-c") + .arg(analyze_query) + .arg("--analyze-raw") + .arg("--flightsql") + .timeout(Duration::from_secs(5)) + .assert() + .success() + }) + .await + .unwrap(); + + let output = String::from_utf8_lossy(&assert.get_output().stdout); + + // Verify schema + assert!(output.contains("node_id"), "Should contain node_id column"); + assert!( + output.contains("parent_node_id"), + "Should contain parent_node_id column" + ); + + // Verify Arrow-specific I/O metrics are present with the right namespace + assert!( + output.contains("io.arrow."), + "Should contain io.arrow.* metrics for an Arrow scan" + ); + // Verify no other format namespaces are used + assert!( + !output.contains("io.csv."), + "Should NOT use io.csv.* namespace for Arrow files" + ); + assert!( + !output.contains("io.parquet."), + "Should NOT use io.parquet.* namespace for Arrow files" + ); + assert!( + !output.contains("io.json."), + "Should NOT use io.json.* namespace for Arrow files" + ); + + // Verify compute metrics are always present + assert!( + output.contains("compute.") || output.contains("elapsed_compute"), + "Should contain compute metrics" + ); + + fixture.shutdown_and_wait().await; +} + #[tokio::test] pub async fn test_execute_with_headers_file() { let test_server = TestFlightSqlServiceImpl::new();