gabotechs commented on code in PR #23975: URL: https://github.com/apache/datafusion/pull/23975#discussion_r3713745292
########## benchmarks/src/statistics.rs: ########## @@ -0,0 +1,656 @@ +// 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. + +//! Reports planning statistics alongside runtime metrics for benchmark queries. + +use std::collections::{BTreeMap, HashMap}; +use std::fs; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::sync::{Arc, LazyLock}; + +use clap::Args; +use datafusion::error::{DataFusionError, Result}; +use datafusion::physical_plan::metrics::MetricValue; +use datafusion::physical_plan::operator_statistics::StatisticsRegistry; +use datafusion::physical_plan::{ExecutionPlan, collect}; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use datafusion_common::stats::Precision; +use regex::Regex; +use serde::{Deserialize, Serialize}; + +/// Generate reports that compare planning statistics with runtime metrics. +#[derive(Debug, Args)] +#[command(verbatim_doc_comment)] +pub struct RunOpt { + /// Query filename stem. If not specified, runs every `.sql` file. + #[arg(short, long)] + query: Option<String>, + + /// Branch whose results should be compared. Defaults to the previous run on this branch. + #[arg(long)] + compare: Option<String>, + + /// Path to Parquet data. Top-level files and directories are registered as tables. + #[arg(required = true, short = 'p', long)] + path: PathBuf, + + /// Path to a SQL file or directory of SQL query files. + #[arg(required = true, short = 'Q', long = "query_path")] + query_path: PathBuf, +} + +impl RunOpt { + pub async fn run(self) -> Result<()> { + let mut config = SessionConfig::from_env()?.with_collect_statistics(true); + config.options_mut().optimizer.prefer_hash_join = true; + let ctx = SessionContext::new_with_config(config); + register_parquet_files(&ctx, &self.path).await?; + + let branch = current_branch_name(); + let result_path = self.report_path(&branch); + let comparison_branch = self + .compare + .as_deref() + .map_or(branch.as_str(), |branch| branch); + let comparison_path = self.report_path(comparison_branch); + let previous = load_comparison_report(&comparison_path)?; + backup_previous_report(&result_path)?; + let comparison_description = self.compare.as_ref().map_or_else( + || format!("previous run on branch '{branch}'"), + |branch| format!("branch '{branch}'"), + ); + + let mut reports = vec![]; + let mut successful_reports = vec![]; + for query_path in query_files(&self.query_path, self.query.as_deref())? { + let query = query_path + .file_stem() + .expect("query file has a filename") + .to_string_lossy() + .to_string(); + let sql = fs::read_to_string(query_path)?; + for (statement, sql) in sql + .split(';') + .filter(|sql| !sql.trim().is_empty()) + .enumerate() + { + let statement = statement + 1; + let report = match self.report_query(&ctx, sql).await { + Ok(operators) => QueryReport { + query: query.clone(), + statement, + operators, + success: true, + error: None, + }, + Err(error) => QueryReport { + query: query.clone(), + statement, + operators: vec![], + success: false, + error: Some(error.to_string()), + }, + }; + print_query_report(&report, previous.as_deref()); + if report.success { + successful_reports.push(report.clone()); + store_report(&result_path, &successful_reports)?; + } + reports.push(report); + } + } + print_q_error_summary( + &reports, + previous.as_deref(), + &branch, + &comparison_description, + ); + Ok(()) + } + + fn report_path(&self, branch: &str) -> PathBuf { + let report_name = self.query.as_ref().map_or_else( + || "statistics.json".to_string(), + |query| format!("statistics-{query}.json"), + ); + PathBuf::from("target/dfbench/statistics") + .join(normalize_branch_name(branch)) + .join(report_name) + } + + async fn report_query( + &self, + ctx: &SessionContext, + sql: &str, + ) -> Result<Vec<OperatorReport>> { + let dataframe = ctx.sql(sql).await?; + let (state, logical_plan) = dataframe.into_parts(); + let logical_plan = state.optimize(&logical_plan)?; + let physical_plan = state.create_physical_plan(&logical_plan).await?; + + let statistics = capture_statistics(physical_plan.as_ref())?; + collect(Arc::clone(&physical_plan), state.task_ctx()).await?; + + let mut report = Vec::with_capacity(statistics.len()); + append_runtime_metrics(physical_plan.as_ref(), &statistics, &mut report); + Ok(report) + } +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +struct QueryReport { + query: String, + statement: usize, + operators: Vec<OperatorReport>, + #[serde(default = "default_success")] + success: bool, + #[serde(default)] + error: Option<String>, +} + +fn default_success() -> bool { + true +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +struct OperatorReport { + path: String, + name: String, + estimated_rows: StatisticValue, + estimated_bytes: StatisticValue, + runtime_output_rows: Option<usize>, + runtime_output_bytes: Option<usize>, + q_error: Option<QError>, +} + +#[derive(Debug, Serialize)] +struct CapturedStatistics { + path: String, + name: String, + estimated_rows: StatisticValue, + estimated_bytes: StatisticValue, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)] +#[serde(tag = "precision", content = "value", rename_all = "snake_case")] +enum StatisticValue { + Exact(usize), + Inexact(usize), + Absent, +} + +/// Symmetric multiplicative error between an estimated and actual row count. +/// +/// See [How Good Are Query Optimizers, Really?](https://vldb.org/pvldb/vol9/p204-leis.pdf) +/// for the q-error definition and motivation. +/// +/// The q-error metric is infinite when exactly one value is zero. A zero +/// estimate for an actual zero result is reported separately as an exact +/// estimate and is excluded from q-error percentiles. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(tag = "kind", content = "value", rename_all = "snake_case")] +enum QError { + Finite(f64), + Infinite, + ExactZero, +} + +fn capture_statistics(plan: &dyn ExecutionPlan) -> Result<Vec<CapturedStatistics>> { + let statistics_context = StatisticsRegistry::default_with_builtin_providers(); + let mut result = vec![]; + capture_statistics_inner(plan, &statistics_context, "0", &mut result)?; + Ok(result) +} + +fn capture_statistics_inner( + plan: &dyn ExecutionPlan, + statistics_context: &StatisticsRegistry, + path: &str, + result: &mut Vec<CapturedStatistics>, +) -> Result<()> { + let statistics = statistics_context.compute_base(plan)?; + result.push(CapturedStatistics { + path: path.to_string(), + name: plan.name().to_string(), + estimated_rows: statistic_value(statistics.num_rows), + estimated_bytes: statistic_value(statistics.total_byte_size), + }); + for (child_index, child) in plan.children().iter().enumerate() { + capture_statistics_inner( + child.as_ref(), + statistics_context, + &format!("{path}.{child_index}"), + result, + )?; + } + Ok(()) +} + +fn append_runtime_metrics( + plan: &dyn ExecutionPlan, + statistics: &[CapturedStatistics], + report: &mut Vec<OperatorReport>, +) { + let statistics_entry = &statistics[report.len()]; + let (runtime_output_rows, runtime_output_bytes) = + plan.metrics().map_or((None, None), |m| { + ( + m.output_rows(), + m.sum(|metric| matches!(metric.value(), MetricValue::OutputBytes(_))) + .map(|metric| metric.as_usize()), + ) + }); + report.push(OperatorReport { + path: statistics_entry.path.clone(), + name: statistics_entry.name.clone(), + estimated_rows: statistics_entry.estimated_rows.clone(), + estimated_bytes: statistics_entry.estimated_bytes.clone(), + runtime_output_rows, + runtime_output_bytes, + q_error: q_error(&statistics_entry.estimated_rows, runtime_output_rows), + }); + for child in plan.children() { + append_runtime_metrics(child.as_ref(), statistics, report); + } +} + +fn q_error(estimated: &StatisticValue, actual: Option<usize>) -> Option<QError> { + let estimated = match estimated { + StatisticValue::Exact(value) | StatisticValue::Inexact(value) => *value, + StatisticValue::Absent => return None, + }; + let actual = actual?; + if estimated == 0 && actual == 0 { + Some(QError::ExactZero) + } else if estimated == 0 || actual == 0 { + Some(QError::Infinite) + } else { + Some(QError::Finite( + estimated.max(actual) as f64 / estimated.min(actual) as f64, + )) + } +} + +fn statistic_value(value: Precision<usize>) -> StatisticValue { + match value { + Precision::Exact(value) => StatisticValue::Exact(value), + Precision::Inexact(value) => StatisticValue::Inexact(value), + Precision::Absent => StatisticValue::Absent, + } +} + +fn load_comparison_report(comparison_path: &Path) -> Result<Option<Vec<QueryReport>>> { + if comparison_path.exists() { + let previous = fs::read_to_string(comparison_path)?; + let previous: Vec<QueryReport> = serde_json::from_str(&previous) + .map_err(|error| DataFusionError::External(Box::new(error)))?; + Ok(Some(previous)) + } else { + eprintln!( + "No comparison report found at {}", + comparison_path.display() + ); + Ok(None) + } +} + +fn backup_previous_report(result_path: &Path) -> Result<()> { + if result_path.exists() { + let previous_path = result_path.with_extension("previous.json"); + fs::copy(result_path, previous_path)?; + } + Ok(()) +} + +fn store_report(result_path: &Path, report: &[QueryReport]) -> Result<()> { + let parent = result_path.parent().expect("result path has a parent"); + fs::create_dir_all(parent)?; + let temporary_path = result_path.with_extension("tmp"); + fs::write(&temporary_path, serialize_report(report)?)?; + fs::rename(temporary_path, result_path)?; + Ok(()) +} + +fn current_branch_name() -> String { + Command::new("git") + .args(["rev-parse", "--abbrev-ref", "HEAD"]) + .output() + .ok() + .filter(|output| output.status.success()) + .and_then(|output| String::from_utf8(output.stdout).ok()) + .map(|branch| branch.trim().to_string()) + .filter(|branch| branch != "HEAD") + .unwrap_or_else(|| "detached".to_string()) +} + +fn normalize_branch_name(branch: &str) -> String { + branch + .chars() + .map(|character| match character { + 'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_' | '.' => character, + _ => '_', + }) + .collect() +} + +fn print_query_report(query: &QueryReport, previous: Option<&[QueryReport]>) { + if !query.success { + println!( + "=== {} (statement {}) FAILED ===\n{}", + query.query, + query.statement, + query.error.as_deref().unwrap_or("unknown error") + ); + return; + } + let previous = previous.map(operator_reports); + println!("=== {} ===", query.query); + for operator in &query.operators { + let identifier = OperatorIdentifier::new(query, operator); + let previous_q_error = previous + .as_ref() + .and_then(|reports| reports.get(&identifier)) + .and_then(|operator| operator.q_error.as_ref()); + let depth = operator.path.matches('.').count(); + println!( + "{:indent$}{}: rows={} vs {}, q-error: previous={}, current={}, change={}", + "", + operator.name, + display_statistic(&operator.estimated_rows), + display_option(operator.runtime_output_rows), + display_q_error(previous_q_error), + display_q_error(operator.q_error.as_ref()), + display_improvement(previous_q_error, operator.q_error.as_ref()), + indent = depth * 2, + ); + } +} + +fn display_option(value: Option<usize>) -> String { + value.map_or_else(|| "?".to_string(), |value| value.to_string()) +} + +fn display_statistic(value: &StatisticValue) -> String { + match value { + StatisticValue::Exact(value) => format!("Exact({value})"), + StatisticValue::Inexact(value) => format!("Inexact({value})"), + StatisticValue::Absent => "Absent".to_string(), + } +} + +fn display_q_error(value: Option<&QError>) -> String { + match value { + Some(QError::Finite(value)) => format!("{value:.2}x"), + Some(QError::Infinite) => "infinite".to_string(), + Some(QError::ExactZero) => "exact zero".to_string(), + None => "?".to_string(), + } +} + +fn display_improvement(previous: Option<&QError>, current: Option<&QError>) -> String { + match (previous, current) { + (Some(QError::Finite(previous)), Some(QError::Finite(current))) => { + let improvement = (previous - current) / previous * 100.0; + if improvement >= 20.0 { + format!("✅ {improvement:.1}%") + } else if improvement <= -20.0 { + format!("❌ {improvement:.1}%") + } else { + format!("{improvement:.1}%") + } + } + (Some(QError::Infinite), Some(QError::Infinite)) => "0.0%".to_string(), + (Some(QError::ExactZero), Some(QError::ExactZero)) => "0.0%".to_string(), + (Some(QError::Infinite | QError::Finite(_)), Some(QError::ExactZero)) => { + "✅ exact zero".to_string() + } + (Some(QError::ExactZero), Some(QError::Infinite | QError::Finite(_))) => { + "❌ no longer exact zero".to_string() + } + (Some(QError::Infinite), Some(QError::Finite(_))) => "✅ resolved".to_string(), + (Some(QError::Finite(_)), Some(QError::Infinite)) => "❌ infinite".to_string(), + _ => "?".to_string(), + } +} + +fn print_q_error_summary( + reports: &[QueryReport], + previous: Option<&[QueryReport]>, + branch: &str, + comparison_description: &str, +) { + let q_errors = sorted_finite_q_errors(reports); + let previous_q_errors = previous.map(sorted_finite_q_errors); + let q_error_counts = q_error_counts(reports); + let total = reports + .iter() + .map(|query| query.operators.len()) + .sum::<usize>(); + + println!("=== q-error summary ==="); + println!("current: branch '{branch}'"); + println!("comparison: {comparison_description}"); + let failed = reports.iter().filter(|query| !query.success).count(); + println!( + "queries: {} succeeded, {failed} failed", + reports.len() - failed + ); + println!( + "evaluated operators: {}/{}", + q_error_counts.evaluated, total + ); + println!( + "finite q-errors: {}, exact-zero estimates: {}, infinite q-errors: {}", + q_error_counts.finite, q_error_counts.exact_zero, q_error_counts.infinite + ); + for percentile in [50, 75, 95, 99] { + let current = percentile_q_error(&q_errors, percentile); + let previous = previous_q_errors + .as_ref() + .and_then(|q_errors| percentile_q_error(q_errors, percentile)); + println!( + "p{percentile}: previous={}, current={}, change={}", + display_q_error(previous), + display_q_error(current), + display_improvement(previous, current), + ); + } +} + +struct QErrorCounts { + evaluated: usize, + finite: usize, + exact_zero: usize, + infinite: usize, +} + +fn q_error_counts(reports: &[QueryReport]) -> QErrorCounts { + let mut counts = QErrorCounts { + evaluated: 0, + finite: 0, + exact_zero: 0, + infinite: 0, + }; + for q_error in reports + .iter() + .flat_map(|query| &query.operators) + .filter_map(|operator| operator.q_error.as_ref()) + { + counts.evaluated += 1; + match q_error { + QError::Finite(_) => counts.finite += 1, + QError::ExactZero => counts.exact_zero += 1, + QError::Infinite => counts.infinite += 1, + } + } + counts +} + +fn sorted_finite_q_errors(reports: &[QueryReport]) -> Vec<&QError> { + let mut q_errors = reports + .iter() + .flat_map(|query| &query.operators) + .filter_map(|operator| match operator.q_error.as_ref() { + Some(QError::Finite(_)) => operator.q_error.as_ref(), + Some(QError::ExactZero | QError::Infinite) | None => None, + }) + .collect::<Vec<_>>(); + q_errors.sort_by(|left, right| q_error_value(left).total_cmp(&q_error_value(right))); + q_errors +} + +fn percentile_q_error<'a>( + q_errors: &[&'a QError], + percentile: usize, +) -> Option<&'a QError> { + if q_errors.is_empty() { + return None; + } + let index = (percentile * q_errors.len()) + .div_ceil(100) + .saturating_sub(1); + q_errors.get(index).copied() +} + +fn q_error_value(q_error: &QError) -> f64 { + match q_error { + QError::Finite(value) => *value, + QError::Infinite => f64::INFINITY, + QError::ExactZero => 1.0, + } +} + +fn operator_reports( + reports: &[QueryReport], +) -> HashMap<OperatorIdentifier<'_>, &OperatorReport> { + reports + .iter() + .flat_map(|query| { + query + .operators + .iter() + .map(move |operator| (OperatorIdentifier::new(query, operator), operator)) + }) + .collect() +} + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +struct OperatorIdentifier<'a> { + query: &'a str, + statement: usize, + path: &'a str, + name: &'a str, +} + +impl<'a> OperatorIdentifier<'a> { + fn new(query: &'a QueryReport, operator: &'a OperatorReport) -> Self { + Self { + query: query.query.as_str(), + statement: query.statement, + path: operator.path.as_str(), + name: operator.name.as_str(), + } + } +} + +fn serialize_report(report: &[QueryReport]) -> Result<String> { + serde_json::to_string_pretty(report) + .map_err(|error| DataFusionError::External(Box::new(error))) +} + +fn query_files(path: &Path, query: Option<&str>) -> Result<Vec<PathBuf>> { Review Comment: Added some tests here https://github.com/apache/datafusion/pull/23975/commits/6366a30fae76699e650aeb87296947d49a2504a5. I want to be careful in not introducing too many tests, as those also require maintenance, and this is just a benchmarking tool that is not in a production path. Let me know if you think we should be covering some more. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
