diff --git a/benchmarks/README.md b/benchmarks/README.md index 489f2d485a9f..da56568a1d00 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -964,6 +964,40 @@ Several queries are included to test sort merge joins under various workloads. ./bench.sh run smj ``` + +## Memory-Limited Join + +One join workload through a fixed memory budget (300MB by default), in each configuration a user +can pick today. `HashJoinExec` cannot spill its build side, so some rows are *expected* to fail +with `Resources exhausted`: that matrix is the recorded baseline for external hash join work, and +the benchmark prints each row's outcome next to it and calls out any row that flipped. + +| row | role | baseline | +| --- | --- | --- | +| 1 | default settings — the planner picks `HashJoinExec` | fails | +| 2 | the only workaround — `prefer_hash_join=false`, so the sorts (which spill) carry the join | completes | +| 3a | where the ceiling is — the build side filtered to a size that fits | completes | +| 3b | just past the ceiling — the same filter, slightly larger | fails | +| 4 | what the workaround costs when it isn't needed — row 3a forced through `SortMergeJoinExec` | completes, ~3.5x slower than 3a | +| 5 | control, not a way to run the join — hash *aggregation* through the same budget | completes | + +Row 4 divided by row 3a is reported as the "SMJ tax". Both inputs are the same generated relation +(20M rows, written on the first run and cached under `--path`), so there is no smaller side for the +planner to swap in. Spill metrics are reported per operator, and `-q` runs a single row. + +### Example Run + +```bash +# Data is generated on the first run +./bench.sh run join_mem + +# Or directly, e.g. only the ceiling row, under a different budget +cargo run --release --bin dfbench -- join-mem --query 3b --memory-limit 512M +``` + +The same matrix is asserted at test scale in +`datafusion/core/tests/memory_limit/join_failure_matrix.rs`. + ## Cancellation Test performance of cancelling queries. diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh index df9b7f6c94f1..2a9cb1cb44d9 100755 --- a/benchmarks/bench.sh +++ b/benchmarks/bench.sh @@ -157,6 +157,8 @@ cancellation: How long cancelling a query takes nlj: Benchmark for simple nested loop joins, testing various join scenarios hj: Benchmark for simple hash joins, testing various join scenarios smj: Benchmark for simple sort merge joins, testing various join scenarios +join_mem: One join workload through a fixed memory budget (default 300M), in each configuration a user can pick today. + Rows failing with 'Resources exhausted' are the recorded baseline (hash join cannot spill), not a broken run dict: Benchmark for dictionary-encoded group-by scenarios compile_profile: Compile and execute TPC-H across selected Cargo profiles, reporting timing and binary size @@ -374,6 +376,10 @@ main() { # smj uses range() function, no data generation needed echo "SMJ benchmark does not require data generation" ;; + join_mem) + # join_mem generates its own parquet file on first run + echo "join_mem benchmark generates its data on first run" + ;; dict) # dict generates in-memory data, no data generation needed echo "DICT benchmark does not require data generation" @@ -604,6 +610,9 @@ main() { smj) run_smj ;; + join_mem) + run_join_mem + ;; dict) run_dict ;; @@ -1577,6 +1586,15 @@ run_smj() { debug_run $CARGO_COMMAND --bin dfbench -- smj --iterations 5 -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG} } +# Runs the memory-limited join benchmark (the join failure matrix) +run_join_mem() { + JOIN_MEM_DIR="${DATA_DIR}/join_mem" + RESULTS_FILE="${RESULTS_DIR}/join_mem.json" + echo "RESULTS_FILE: ${RESULTS_FILE}" + echo "Running join_mem benchmark..." + debug_run $CARGO_COMMAND --bin dfbench -- join-mem --iterations 3 --path "${JOIN_MEM_DIR}" -o "${RESULTS_FILE}" ${QUERY_ARG} +} + # Runs the dict benchmark run_dict() { RESULTS_FILE="${RESULTS_DIR}/dict.json" diff --git a/benchmarks/src/bin/dfbench.rs b/benchmarks/src/bin/dfbench.rs index 29cc8d63d2d8..48ab4aee4032 100644 --- a/benchmarks/src/bin/dfbench.rs +++ b/benchmarks/src/bin/dfbench.rs @@ -32,8 +32,8 @@ static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc; static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc; use datafusion_benchmarks::{ - cancellation, clickbench, dict, h2o, hj, imdb, nlj, smj, sort_tpch, statistics, - tpcds, tpch, + cancellation, clickbench, dict, h2o, hj, imdb, join_mem, nlj, smj, sort_tpch, + statistics, tpcds, tpch, }; #[derive(Debug, Parser)] @@ -51,6 +51,7 @@ enum Options { H2o(h2o::RunOpt), HJ(hj::RunOpt), Imdb(imdb::RunOpt), + JoinMem(join_mem::RunOpt), Nlj(nlj::RunOpt), Smj(smj::RunOpt), Statistics(statistics::RunOpt), @@ -73,6 +74,7 @@ pub async fn main() -> Result<()> { Options::H2o(opt) => opt.run().await, Options::HJ(opt) => opt.run().await, Options::Imdb(opt) => Box::pin(opt.run()).await, + Options::JoinMem(opt) => opt.run().await, Options::Nlj(opt) => opt.run().await, Options::Smj(opt) => opt.run().await, Options::Statistics(opt) => opt.run().await, diff --git a/benchmarks/src/join_mem.rs b/benchmarks/src/join_mem.rs new file mode 100644 index 000000000000..cfa6fefe78f2 --- /dev/null +++ b/benchmarks/src/join_mem.rs @@ -0,0 +1,534 @@ +// 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. + +//! Memory-limited join benchmark: one join workload through a fixed memory +//! budget, in each configuration a user can pick today. +//! +//! `HashJoinExec` cannot spill its build side, so some rows are *expected* to +//! fail with `ResourcesExhausted`. The benchmark keeps that matrix +//! reproducible while external hash join is built, and reports when a row +//! flips. Row 4 divided by row 3a is the "SMJ tax": what the only workaround +//! costs on a join that would have fit in memory. +//! +//! Both inputs read the same generated file, so there is no smaller side for +//! the planner to swap in: the failure is not one a better build side avoids. +//! +//! The rows and their recorded outcomes are in `benchmarks/README.md`; +//! `datafusion/core/tests/memory_limit/join_failure_matrix.rs` asserts the same +//! matrix at test scale. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; + +use clap::Args; +use datafusion::error::Result; +use datafusion::physical_plan::{ExecutionPlan, displayable, execute_stream}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::instant::Instant; +use datafusion_common::{DataFusionError, exec_err, human_readable_size}; +use futures::StreamExt; + +use crate::util::{BenchmarkRun, CommonOpt, QueryResult}; + +/// Budget the matrix was recorded under, used when none is configured. +const DEFAULT_MEMORY_LIMIT: usize = 300 * 1024 * 1024; + +/// Partition count it was recorded under: the per-partition build side is what +/// does or does not fit in the budget. +const DEFAULT_PARTITIONS: usize = 4; + +/// Row 4 / row 3a when the matrix was recorded. +const RECORDED_SMJ_TAX: f64 = 3.5; + +/// Run the memory-limited join benchmark +/// +/// Runs the join failure matrix under a fixed memory budget: each row is a +/// configuration a user can choose today, and its recorded outcome is the +/// baseline external hash join has to improve on. +#[derive(Debug, Args, Clone)] +#[command(verbatim_doc_comment)] +pub struct RunOpt { + /// Matrix row to run (1, 2, 3a, 3b, 4 or 5). If not specified, runs all rows + #[arg(short = 'q', long = "query")] + query: Option, + + /// Common options (iterations, memory limit, memory pool type, partitions, etc.) + #[command(flatten)] + common: CommonOpt, + + /// Directory holding the generated parquet file, generated on first use. + /// Defaults to `datafusion-join-mem` under the system temp dir + #[arg(short = 'p', long = "path")] + path: Option, + + /// Rows in the generated table. Both join inputs read it + #[arg(long = "rows", default_value = "20000000")] + rows: usize, + + /// Build-side row cap that fits in the budget (rows 3a and 4) + #[arg(long = "fit-rows", default_value = "10000000")] + fit_rows: usize, + + /// Build-side row cap just past the budget (row 3b) + #[arg(long = "over-rows", default_value = "12000000")] + over_rows: usize, + + /// If present, write results json here + #[arg(short = 'o', long = "output")] + output_path: Option, +} + +/// Which query a row runs. The two filtered variants bracket the point where +/// the build side stops fitting in the budget. +#[derive(Debug, Clone, Copy)] +enum MatrixQuery { + Join, + JoinFits, + JoinOverflows, + /// Hash aggregation over the same data, as a control + Control, +} + +/// One row of the failure matrix. +#[derive(Debug, Clone, Copy)] +struct MatrixRow { + label: &'static str, + /// What the row is here to show + role: &'static str, + query: MatrixQuery, + prefer_hash_join: bool, + /// Operator the row is about; the plan is checked to contain it + operator: Option<&'static str>, + /// Whether the row completed when the matrix was recorded + completes: bool, +} + +const MATRIX: &[MatrixRow] = &[ + MatrixRow { + label: "1", + role: "default settings - the planner picks HashJoinExec", + query: MatrixQuery::Join, + prefer_hash_join: true, + operator: Some("HashJoinExec"), + completes: false, + }, + MatrixRow { + label: "2", + role: "the only workaround - prefer_hash_join=false, sorts spill", + query: MatrixQuery::Join, + prefer_hash_join: false, + operator: Some("SortMergeJoinExec"), + completes: true, + }, + MatrixRow { + label: "3a", + role: "where the ceiling is - build side filtered to a fitting size", + query: MatrixQuery::JoinFits, + prefer_hash_join: true, + operator: Some("HashJoinExec"), + completes: true, + }, + MatrixRow { + label: "3b", + role: "just past the ceiling - the same filter, slightly larger", + query: MatrixQuery::JoinOverflows, + prefer_hash_join: true, + operator: Some("HashJoinExec"), + completes: false, + }, + MatrixRow { + label: "4", + role: "what the workaround costs - row 3a forced through SMJ", + query: MatrixQuery::JoinFits, + prefer_hash_join: false, + operator: Some("SortMergeJoinExec"), + completes: true, + }, + MatrixRow { + label: "5", + role: "control - hash aggregation through the same budget", + query: MatrixQuery::Control, + prefer_hash_join: true, + operator: None, + completes: true, + }, +]; + +/// Spill metrics of one operator of an executed plan. +struct OperatorSpill { + operator: String, + spill_count: usize, + spilled_bytes: usize, +} + +/// What one matrix row did. +struct RowResult { + iterations: Vec, + /// The allocation that failed, when the row ran out of budget + error: Option, + /// Spill metrics of the last executed plan + spills: Vec, +} + +impl RowResult { + fn completes(&self) -> bool { + self.error.is_none() + } + + fn mean_elapsed(&self) -> Option { + let total: Duration = self.iterations.iter().map(|iter| iter.elapsed).sum(); + (!self.iterations.is_empty()).then(|| total / self.iterations.len() as u32) + } +} + +fn outcome(completes: bool) -> &'static str { + if completes { "completes" } else { "exhausted" } +} + +impl RunOpt { + pub async fn run(self) -> Result<()> { + let common = self.common_with_defaults(); + println!("Running memory-limited join benchmark: {self:#?}\n"); + + let rows = match &self.query { + None => MATRIX.iter().collect::>(), + Some(label) => match MATRIX.iter().find(|row| row.label == *label) { + Some(row) => vec![row], + None => return exec_err!("Matrix row {label} not found"), + }, + }; + let data = self.ensure_data(&common).await?; + + let mut benchmark_run = BenchmarkRun::new(); + let mut results = Vec::with_capacity(rows.len()); + + for row in rows { + let sql = self.sql(row); + let ctx = self.context(&common, row, &data).await?; + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); + benchmark_run.start_new_case(&format!("row {} ({})", row.label, row.role)); + + println!("--- row {}: {}\n{sql}", row.label, row.role); + let result = self.run_row(&ctx, row, &sql, common.iterations).await?; + + for iter in &result.iterations { + benchmark_run.write_iter(iter.elapsed, iter.row_count); + } + if !result.completes() { + benchmark_run.mark_failed(); + } + results.push((row, result)); + } + + benchmark_run.maybe_write_json(self.output_path.as_ref())?; + self.report(&common, &results); + Ok(()) + } + + /// Fill in the settings the matrix was recorded under, for whatever the + /// caller left unset. + fn common_with_defaults(&self) -> CommonOpt { + let mut common = self.common.clone(); + // Leave the env var path in `runtime_env_builder` alone if that is the + // one carrying the limit. + if common.memory_limit.is_none() + && std::env::var("DATAFUSION_RUNTIME_MEMORY_LIMIT").is_err() + { + common.memory_limit = Some(DEFAULT_MEMORY_LIMIT); + } + common.partitions = common.partitions.or(Some(DEFAULT_PARTITIONS)); + common + } + + fn sql(&self, row: &MatrixRow) -> String { + let join_with_build_limit = |limit| { + format!( + "SELECT count(*) FROM t_probe p \ + JOIN (SELECT * FROM t_build WHERE k <= {limit}) b ON p.k = b.k" + ) + }; + match row.query { + MatrixQuery::Join => { + "SELECT count(*) FROM t_probe p JOIN t_build b ON p.k = b.k".to_string() + } + MatrixQuery::JoinFits => join_with_build_limit(self.fit_rows), + MatrixQuery::JoinOverflows => join_with_build_limit(self.over_rows), + MatrixQuery::Control => { + "SELECT count(DISTINCT payload) FROM t_build".to_string() + } + } + } + + /// A context with its own budgeted runtime, so each row starts from an + /// empty pool. + async fn context( + &self, + common: &CommonOpt, + row: &MatrixRow, + data: &Path, + ) -> Result { + let mut config = common.config()?; + config.options_mut().optimizer.prefer_hash_join = row.prefer_hash_join; + let ctx = SessionContext::new_with_config_rt(config, common.build_runtime()?); + + let path = data.to_str().ok_or_else(|| { + DataFusionError::Execution(format!("non-UTF-8 data path {}", data.display())) + })?; + for table in ["t_build", "t_probe"] { + ctx.register_parquet(table, path, Default::default()) + .await?; + } + Ok(ctx) + } + + /// Run one row. Running out of budget is a result, not an error: it is what + /// several rows are here to record. + async fn run_row( + &self, + ctx: &SessionContext, + row: &MatrixRow, + sql: &str, + iterations: usize, + ) -> Result { + let mut result = RowResult { + iterations: vec![], + error: None, + spills: vec![], + }; + + for i in 0..iterations { + let plan = ctx.sql(sql).await?.create_physical_plan().await?; + if let Some(operator) = row.operator { + let plan_display = displayable(plan.as_ref()).indent(true).to_string(); + if !plan_display.contains(operator) { + return exec_err!( + "row {} is about {operator}, but its plan does not use it:\n{plan_display}", + row.label + ); + } + } + + let start = Instant::now(); + let executed = drain(Arc::clone(&plan), ctx).await; + let elapsed = start.elapsed(); + result.spills = collect_spills(plan.as_ref()); + + match executed { + Ok(row_count) => { + println!( + "row {} iteration {i} returned {row_count} rows in {elapsed:?}", + row.label + ); + result.iterations.push(QueryResult { elapsed, row_count }); + } + // Anything but exhaustion is a real failure. + Err(e) + if !matches!( + e.find_root(), + DataFusionError::ResourcesExhausted(_) + ) => + { + return Err(e); + } + Err(e) => { + println!("row {} iteration {i} failed in {elapsed:?}", row.label); + println!(" {}", e.find_root()); + result.error = Some(failed_allocation(&e.find_root().to_string())); + return Ok(result); + } + } + } + + Ok(result) + } + + /// Generate the table once and cache it on disk. + /// + /// Generated with no memory limit on purpose: writing the file under the + /// benchmark's budget is a different fight than the one under test. + async fn ensure_data(&self, common: &CommonOpt) -> Result { + let dir = self + .path + .clone() + .unwrap_or_else(|| std::env::temp_dir().join("datafusion-join-mem")); + std::fs::create_dir_all(&dir)?; + + let file = dir.join(format!("join_mem_{}_rows.parquet", self.rows)); + if file.exists() { + println!("Using existing data file {}", file.display()); + return Ok(file); + } + + println!("Generating {} rows into {}", self.rows, file.display()); + let start = Instant::now(); + let ctx = + SessionContext::new_with_config(common.update_config(SessionConfig::new())); + // Write under a temporary name, so an interrupted run leaves no + // truncated file behind to be picked up as cached data. + let partial = file.with_extension("parquet.partial"); + ctx.sql(&format!( + "COPY (SELECT v AS k, \ + concat('payload-', v, '-', repeat('x', 24)) AS payload \ + FROM generate_series(1, {}) AS t(v)) \ + TO '{}' STORED AS PARQUET", + self.rows, + partial.display() + )) + .await? + .collect() + .await?; + std::fs::rename(&partial, &file)?; + println!("Generated in {:?}", start.elapsed()); + + Ok(file) + } + + /// Print the matrix: what each row did, next to what it did when recorded. + fn report(&self, common: &CommonOpt, results: &[(&MatrixRow, RowResult)]) { + let budget = common + .memory_limit + .map(human_readable_size) + .unwrap_or_else(|| "unlimited".to_string()); + println!( + "\nJoin failure matrix: {budget} {} pool, {} partitions, {} rows\n", + common.mem_pool_type, + common.partitions.unwrap_or(DEFAULT_PARTITIONS), + self.rows + ); + println!( + "{:<4} {:<11} {:<11} {:<10} role", + "row", "baseline", "actual", "mean" + ); + + for (row, result) in results { + let mean = result + .mean_elapsed() + .map(|mean| format!("{:.3}s", mean.as_secs_f64())) + .unwrap_or_else(|| "-".to_string()); + println!( + "{:<4} {:<11} {:<11} {mean:<10} {}", + row.label, + outcome(row.completes), + outcome(result.completes()), + row.role + ); + if let Some(error) = &result.error { + println!(" error: {error}"); + } + for spill in &result.spills { + println!( + " spilled: {} spill_count={} spilled_bytes={}", + spill.operator, + spill.spill_count, + human_readable_size(spill.spilled_bytes) + ); + } + } + + if let Some(tax) = smj_tax(results) { + println!( + "\nSMJ tax (row 4 / row 3a): {tax:.1}x on a join that would have fit \ + (recorded: {RECORDED_SMJ_TAX:.1}x)" + ); + } + + let flipped: Vec<_> = results + .iter() + .filter(|(row, result)| result.completes() != row.completes) + .map(|(row, result)| { + format!("row {}: now {}", row.label, outcome(result.completes())) + }) + .collect(); + match flipped.is_empty() { + true => println!("\nEvery row matched its recorded baseline."), + false => println!("\nFlipped from the baseline: {}", flipped.join(", ")), + } + } +} + +/// Execute `plan`, dropping each batch so the budget is spent on the join +/// rather than on holding results. +async fn drain(plan: Arc, ctx: &SessionContext) -> Result { + let mut stream = execute_stream(plan, ctx.task_ctx())?; + let mut row_count = 0; + while let Some(batch) = stream.next().await { + row_count += batch?.num_rows(); + } + Ok(row_count) +} + +/// Spill metrics of every operator in `plan` that spilled, in plan order and +/// merged by operator name. +fn collect_spills(plan: &dyn ExecutionPlan) -> Vec { + let metrics = plan.metrics(); + let spill_count = metrics.as_ref().and_then(|m| m.spill_count()).unwrap_or(0); + let spilled_bytes = metrics + .as_ref() + .and_then(|m| m.spilled_bytes()) + .unwrap_or(0); + + let mut spills = vec![]; + if spill_count > 0 || spilled_bytes > 0 { + spills.push(OperatorSpill { + operator: plan.name().to_string(), + spill_count, + spilled_bytes, + }); + } + for child in plan.children() { + for spill in collect_spills(child.as_ref()) { + match spills + .iter_mut() + .find(|held| held.operator == spill.operator) + { + Some(held) => { + held.spill_count += spill.spill_count; + held.spilled_bytes += spill.spilled_bytes; + } + None => spills.push(spill), + } + } + } + spills +} + +/// The line of a pool error naming the allocation that failed, which the pool +/// prints after its list of top consumers. +fn failed_allocation(error: &str) -> String { + error + .lines() + .rev() + .find(|line| !line.trim().is_empty()) + .unwrap_or(error) + .trim_start_matches("Error: ") + .to_string() +} + +/// Row 4 divided by row 3a: the same join, once as the planner would run it and +/// once forced through the workaround. `None` unless both rows completed. +fn smj_tax(results: &[(&MatrixRow, RowResult)]) -> Option { + let mean = |label: &str| { + results + .iter() + .find(|(row, _)| row.label == label) + .and_then(|(_, result)| result.mean_elapsed()) + .map(|mean| mean.as_secs_f64()) + }; + let (fits, forced) = (mean("3a")?, mean("4")?); + (fits > 0.0).then_some(forced / fits) +} diff --git a/benchmarks/src/lib.rs b/benchmarks/src/lib.rs index 0b3783421f84..a3974e191be4 100644 --- a/benchmarks/src/lib.rs +++ b/benchmarks/src/lib.rs @@ -22,6 +22,7 @@ pub mod dict; pub mod h2o; pub mod hj; pub mod imdb; +pub mod join_mem; pub mod nlj; pub mod smj; pub mod sort_pushdown; diff --git a/datafusion/core/tests/memory_limit/budgeted_env.rs b/datafusion/core/tests/memory_limit/budgeted_env.rs new file mode 100644 index 000000000000..867d0e4918e1 --- /dev/null +++ b/datafusion/core/tests/memory_limit/budgeted_env.rs @@ -0,0 +1,316 @@ +// 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. + +//! Running SQL under a fixed memory budget, and saying what happened. +//! +//! [`TestCase`] in this module's parent covers queries over one of the built-in +//! [`Scenario`] tables, and asserts on the error text. These helpers cover the +//! other shape: arbitrary SQL against a context the test set up itself, where +//! *whether* the query completed — and what it had to spill to get there — is +//! the thing under test, not the message it failed with. +//! +//! ```text +//! let ctx = BudgetedEnv::new(16 * 1024 * 1024).build_ctx(); +//! let outcome = run_under_budget(&ctx, "SELECT ...").await; +//! println!("{}", outcome.summary()); +//! outcome.assert_completed().assert_spilled("SortExec"); +//! ``` +//! +//! [`Scenario`]: super::Scenario + +use std::num::NonZeroUsize; +use std::sync::Arc; +use std::time::Duration; + +use datafusion::physical_plan::{ExecutionPlan, displayable, execute_stream}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::DataFusionError; +use datafusion_common::human_readable_size; +use datafusion_common::instant::Instant; +use datafusion_execution::disk_manager::DiskManagerBuilder; +use datafusion_execution::memory_pool::{FairSpillPool, MemoryPool, TrackConsumersPool}; +use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; +use futures::StreamExt; + +/// How many consumers an exhaustion error lists. Higher than the pool's default +/// of 3, so that a failure names every partition of the operator that hit the +/// budget rather than cutting off at the largest few. +const TRACKED_CONSUMERS: usize = 8; + +/// A [`RuntimeEnv`] with a fixed memory budget and spilling available. +/// +/// The budget is enforced by a [`FairSpillPool`], the pool +/// `datafusion-cli --mem-pool-type fair` installs: spilling consumers get an +/// equal share of what is left after the unspillable ones. +#[derive(Debug, Clone)] +pub struct BudgetedEnv { + budget: usize, + config: SessionConfig, +} + +impl BudgetedEnv { + /// A budget of `budget` bytes. + pub fn new(budget: usize) -> Self { + Self { + budget, + config: SessionConfig::new(), + } + } + + pub fn with_config(mut self, config: SessionConfig) -> Self { + self.config = config; + self + } + + /// The memory pool this budget is enforced by. + pub fn build_pool(&self) -> Arc { + let tracked = + NonZeroUsize::new(TRACKED_CONSUMERS).expect("non-zero tracked consumers"); + Arc::new(TrackConsumersPool::new( + FairSpillPool::new(self.budget), + tracked, + )) + } + + pub fn build_runtime(&self) -> Arc { + RuntimeEnvBuilder::new() + .with_memory_pool(self.build_pool()) + // Operators that can spill, may: whether a query needs the disk to + // finish is part of what these tests report. + .with_disk_manager_builder(DiskManagerBuilder::default()) + .build_arc() + .expect("building a budgeted runtime") + } + + /// A context whose queries run against this budget. + pub fn build_ctx(&self) -> SessionContext { + SessionContext::new_with_config_rt(self.config.clone(), self.build_runtime()) + } +} + +/// Spill metrics of one operator of an executed plan, summed over its +/// partitions. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OperatorSpill { + pub operator: String, + pub spill_count: usize, + pub spilled_bytes: usize, +} + +/// What one query did under a budget. +pub struct QueryOutcome { + /// The executed plan, for metrics and for asserting on operators + pub plan: Arc, + pub elapsed: Duration, + /// Rows returned, or rows returned before the failure + pub row_count: usize, + /// Why the query stopped, if it did not finish + pub error: Option, + /// Every operator that spilled, in plan order + pub spills: Vec, +} + +impl QueryOutcome { + pub fn completed(&self) -> bool { + self.error.is_none() + } + + /// Whether the query failed because it ran out of budget, as opposed to any + /// other error. + pub fn exhausted_budget(&self) -> bool { + matches!( + self.error.as_ref().map(DataFusionError::find_root), + Some(DataFusionError::ResourcesExhausted(_)) + ) + } + + /// Spills recorded by every operator named `operator`. + pub fn spill_count_of(&self, operator: &str) -> usize { + self.spills + .iter() + .filter(|spill| spill.operator == operator) + .map(|spill| spill.spill_count) + .sum() + } + + /// One line saying what the query did: worth printing from a test, since + /// what a budget does to a query is a report as much as an assertion. + pub fn summary(&self) -> String { + let outcome = match &self.error { + None => format!("completed, {} row(s)", self.row_count), + Some(error) => format!("failed: {}", error.find_root()), + }; + let spills = if self.spills.is_empty() { + "nothing spilled".to_string() + } else { + self.spills + .iter() + .map(|spill| { + format!( + "{} spilled {} in {} event(s)", + spill.operator, + human_readable_size(spill.spilled_bytes), + spill.spill_count + ) + }) + .collect::>() + .join(", ") + }; + format!("{outcome} in {:?}; {spills}", self.elapsed) + } + + /// The plan, as it would be displayed. Handy in assertion messages. + pub fn plan_display(&self) -> String { + displayable(self.plan.as_ref()).indent(true).to_string() + } + + /// Assert the query ran to completion under the budget. + pub fn assert_completed(&self) -> &Self { + assert!( + self.completed(), + "expected the query to complete under the budget, but it failed with: {}\n{}", + self.error.as_ref().expect("checked above"), + self.plan_display(), + ); + self + } + + /// Assert the query failed because the budget ran out. + pub fn assert_exhausted_budget(&self) -> &Self { + match &self.error { + None => panic!( + "expected the query to run out of budget, but it completed\n{}", + self.plan_display() + ), + Some(error) => assert!( + self.exhausted_budget(), + "expected a ResourcesExhausted failure, got: {error}" + ), + } + self + } + + /// Assert the plan contains `operator`, so the test is measuring the + /// operator it means to. + pub fn assert_operator(&self, operator: &str) -> &Self { + let plan = self.plan_display(); + assert!( + plan.contains(operator), + "expected the plan to use {operator}, got:\n{plan}" + ); + self + } + + /// Assert `operator` spilled at least once. + pub fn assert_spilled(&self, operator: &str) -> &Self { + self.assert_operator(operator); + assert!( + self.spill_count_of(operator) > 0, + "expected {operator} to spill, but it did not. Spills: {:?}", + self.spills + ); + self + } + + /// Assert `operator` did not spill. + pub fn assert_did_not_spill(&self, operator: &str) -> &Self { + self.assert_operator(operator); + assert_eq!( + self.spill_count_of(operator), + 0, + "expected {operator} not to spill. Spills: {:?}", + self.spills + ); + self + } +} + +/// Run `sql` on `ctx` and report what it did. +/// +/// Batches are dropped as they arrive, so the budget is spent on the query +/// rather than on holding its output. An execution failure — running out of +/// budget, most of all — is part of the outcome; only a planning failure +/// panics, since that means the test asked for something it cannot run. +pub async fn run_under_budget(ctx: &SessionContext, sql: &str) -> QueryOutcome { + let plan = ctx + .sql(sql) + .await + .unwrap_or_else(|e| panic!("planning `{sql}`: {e}")) + .create_physical_plan() + .await + .unwrap_or_else(|e| panic!("planning `{sql}`: {e}")); + + let start = Instant::now(); + let mut row_count = 0; + let mut error = None; + match execute_stream(Arc::clone(&plan), ctx.task_ctx()) { + Err(e) => error = Some(e), + Ok(mut stream) => { + while let Some(batch) = stream.next().await { + match batch { + Ok(batch) => row_count += batch.num_rows(), + Err(e) => { + error = Some(e); + break; + } + } + } + } + } + let elapsed = start.elapsed(); + + let mut spills = vec![]; + collect_spills(plan.as_ref(), &mut spills); + + QueryOutcome { + plan, + elapsed, + row_count, + error, + spills, + } +} + +/// Collect the spill metrics of every operator in `plan` that spilled, merging +/// repeats of the same operator name. +fn collect_spills(plan: &dyn ExecutionPlan, spills: &mut Vec) { + let metrics = plan.metrics(); + let spill_count = metrics.as_ref().and_then(|m| m.spill_count()).unwrap_or(0); + let spilled_bytes = metrics + .as_ref() + .and_then(|m| m.spilled_bytes()) + .unwrap_or(0); + + if spill_count > 0 || spilled_bytes > 0 { + let name = plan.name(); + match spills.iter_mut().find(|spill| spill.operator == name) { + Some(spill) => { + spill.spill_count += spill_count; + spill.spilled_bytes += spilled_bytes; + } + None => spills.push(OperatorSpill { + operator: name.to_string(), + spill_count, + spilled_bytes, + }), + } + } + + for child in plan.children() { + collect_spills(child.as_ref(), spills); + } +} diff --git a/datafusion/core/tests/memory_limit/join_failure_matrix.rs b/datafusion/core/tests/memory_limit/join_failure_matrix.rs new file mode 100644 index 000000000000..0036dab6dd3b --- /dev/null +++ b/datafusion/core/tests/memory_limit/join_failure_matrix.rs @@ -0,0 +1,229 @@ +// 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. + +//! What one join query does under a fixed memory budget, in each configuration +//! a user can choose today. +//! +//! `HashJoinExec` cannot spill, so a build side that does not fit the budget +//! fails outright; the only way to run such a join today is +//! `prefer_hash_join=false`, which hands the work to `SortMergeJoinExec`, whose +//! sorts *can* spill. These tests record that matrix, so that the rows flip +//! deliberately rather than quietly: +//! +//! | row | role | today | +//! |-----|------|-------| +//! | 1 | default settings — the planner picks `HashJoinExec` | fails | +//! | 2 | the only workaround — `prefer_hash_join=false` | completes, sorts spill | +//! | 3a | where the ceiling is — build side filtered to a fitting size | completes | +//! | 3b | just past the ceiling — the same filter, slightly larger | fails | +//! | 4 | what the workaround costs when it isn't needed — row 3a forced through SMJ | completes | +//! | 5 | control — hash *aggregation* through the same budget | completes, spills | +//! +//! Rows 1, 3b are the ones external hash join is meant to turn into +//! `completes`; when that happens, these tests are the ones to update. +//! +//! Both inputs are the same relation, so there is no smaller side for the +//! planner to swap in: the failure is not one a better build side can avoid. +//! Row 4 divided by row 3a is the "SMJ tax" — the cost of the workaround on a +//! join that would have fit. Its *magnitude* (~3.5x at the recorded scale) is +//! measured by the `join_mem` benchmark; timings are not asserted here. + +use std::sync::LazyLock; + +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use tempfile::TempDir; + +use crate::memory_limit::budgeted_env::{BudgetedEnv, run_under_budget}; + +/// Budget every row runs under. +const BUDGET: usize = 16 * 1024 * 1024; + +/// Rows in each join input. Sized so the build-side hash tables cannot fit in +/// [`BUDGET`]: `count(*)` projects the key alone, ~8 B/row of batches, and the +/// hash map on top of it costs ~19 B/row and is asked for in one allocation +/// once every build batch is already held. +const ROWS: usize = 2_000_000; + +/// Build-side cap that comfortably fits the budget (rows 3a and 4). +const FIT_ROWS: usize = 200_000; + +/// Build-side cap that comfortably exceeds it (row 3b). +const OVER_ROWS: usize = 1_000_000; + +const HASH_JOIN: &str = "HashJoinExec"; +const SORT_MERGE_JOIN: &str = "SortMergeJoinExec"; +const SORT: &str = "SortExec"; +const AGGREGATE: &str = "AggregateExec"; + +/// The join, over the whole build side. +const JOIN: &str = "SELECT count(*) FROM t_probe p JOIN t_build b ON p.k = b.k"; + +/// The join, with the build side filtered down to `limit` rows. +fn join_with_build_limit(limit: usize) -> String { + format!( + "SELECT count(*) FROM t_probe p \ + JOIN (SELECT * FROM t_build WHERE k <= {limit}) b ON p.k = b.k" + ) +} + +/// The control: a hash aggregation over the same data, through the same budget. +const CONTROL: &str = "SELECT count(DISTINCT payload) FROM t_build"; + +/// The relation both join inputs read, generated once per test binary. +/// +/// A parquet file rather than `generate_series` directly, because a series is a +/// sorted source with no statistics: the sort-merge rows would skip their sorts +/// entirely, and the planner could not tell which side of the filtered join is +/// smaller. Both are exactly what the matrix is about. +/// +/// Generated with no memory limit, on a runtime of its own: writing the file is +/// not the thing under test, and a test's own runtime cannot block on this. +static DATA: LazyLock = LazyLock::new(|| { + std::thread::spawn(|| { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("building the data generation runtime"); + runtime.block_on(async { + let dir = TempDir::new().expect("creating the data directory"); + let ctx = SessionContext::new(); + ctx.sql(&format!( + "COPY (SELECT v AS k, \ + concat('payload-', v, '-', repeat('x', 24)) AS payload \ + FROM generate_series(1, {ROWS}) AS t(v)) \ + TO '{}' STORED AS PARQUET", + data_path(&dir) + )) + .await + .expect("generating the join inputs") + .collect() + .await + .expect("generating the join inputs"); + dir + }) + }) + .join() + .expect("generating the join inputs") +}); + +fn data_path(dir: &TempDir) -> String { + dir.path() + .join("t.parquet") + .to_str() + .expect("temp dir path is not UTF-8") + .to_string() +} + +/// A context with the budget installed, both join inputs registered, and the +/// join algorithm pinned. +/// +/// The sort spill reservation is lowered from its 10 MB default because the +/// sort-merge rows plan one sort per partition per side, and at this budget the +/// defaults would reserve most of the pool before any data is read. +async fn budgeted_ctx(prefer_hash_join: bool) -> SessionContext { + let mut config = SessionConfig::new() + .with_target_partitions(2) + .with_sort_spill_reservation_bytes(1024 * 1024); + config.options_mut().optimizer.prefer_hash_join = prefer_hash_join; + + let ctx = BudgetedEnv::new(BUDGET).with_config(config).build_ctx(); + + let path = data_path(&DATA); + for table in ["t_build", "t_probe"] { + ctx.register_parquet(table, &path, ParquetReadOptions::default()) + .await + .expect("registering an input table"); + } + + ctx +} + +/// Row 1: at default settings the planner picks `HashJoinExec`, which has no +/// way to spill its build side, and the query fails. +#[tokio::test] +async fn row_1_hash_join_at_default_settings_exhausts_the_budget() { + let ctx = budgeted_ctx(true).await; + let outcome = run_under_budget(&ctx, JOIN).await; + println!("row 1: {}", outcome.summary()); + + outcome.assert_operator(HASH_JOIN).assert_exhausted_budget(); +} + +/// Row 2: the only workaround. `prefer_hash_join=false` plans a +/// `SortMergeJoinExec`, whose sorts spill, and the same query completes. The +/// join operator itself spills nothing — the sorts carry all of it. +#[tokio::test] +async fn row_2_sort_merge_join_workaround_completes() { + let ctx = budgeted_ctx(false).await; + let outcome = run_under_budget(&ctx, JOIN).await; + println!("row 2: {}", outcome.summary()); + + outcome + .assert_operator(SORT_MERGE_JOIN) + .assert_completed() + .assert_spilled(SORT) + .assert_did_not_spill(SORT_MERGE_JOIN); +} + +/// Row 3a: where the ceiling is. The same join with a build side small enough +/// to fit runs on the hash join without trouble. +#[tokio::test] +async fn row_3a_hash_join_below_the_ceiling_completes() { + let ctx = budgeted_ctx(true).await; + let outcome = run_under_budget(&ctx, &join_with_build_limit(FIT_ROWS)).await; + println!("row 3a: {}", outcome.summary()); + + outcome.assert_operator(HASH_JOIN).assert_completed(); +} + +/// Row 3b: just past it. The same query with a larger build side fails, at the +/// hash table build. +#[tokio::test] +async fn row_3b_hash_join_above_the_ceiling_exhausts_the_budget() { + let ctx = budgeted_ctx(true).await; + let outcome = run_under_budget(&ctx, &join_with_build_limit(OVER_ROWS)).await; + println!("row 3b: {}", outcome.summary()); + + outcome.assert_operator(HASH_JOIN).assert_exhausted_budget(); +} + +/// Row 4: what the workaround costs when it isn't needed. Row 3a's join — one +/// that fits in memory — still completes when forced through the sort-merge +/// path, but pays for sorting both inputs to get there. +#[tokio::test] +async fn row_4_fitting_join_forced_through_sort_merge_completes() { + let ctx = budgeted_ctx(false).await; + let outcome = run_under_budget(&ctx, &join_with_build_limit(FIT_ROWS)).await; + println!("row 4: {}", outcome.summary()); + + outcome + .assert_operator(SORT_MERGE_JOIN) + .assert_completed() + .assert_spilled(SORT); +} + +/// Row 5: the control. It is not a way to run the join; it shows the budget +/// itself is workable, by pushing more data than the join ever holds through +/// the same pool in an operator that can spill. +#[tokio::test] +async fn row_5_control_hash_aggregation_completes_through_the_same_budget() { + let ctx = budgeted_ctx(true).await; + let outcome = run_under_budget(&ctx, CONTROL).await; + println!("row 5: {}", outcome.summary()); + + outcome.assert_completed().assert_spilled(AGGREGATE); +} diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 02856ebb5ab0..4ca9ecd5609c 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -20,6 +20,8 @@ use std::num::NonZeroUsize; use std::sync::{Arc, LazyLock}; +mod budgeted_env; +mod join_failure_matrix; #[cfg(feature = "extended_tests")] mod memory_limit_validation; mod nlj_spill_unmatched;