You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/07/30 16:38:40 UTC
(datafusion) 01/01: Report peak MemoryPool reservation per query in benchmarks (#23985)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23985-541caab85a265c104888f1988fe54b08fa022f53
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 78c583cf3d076731b6568633e0cf5fe4c01f61ef
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Thu Jul 30 11:38:17 2026 -0500
Report peak MemoryPool reservation per query in benchmarks (#23985)
## Which issue does this PR close?
Part of #22758. Doesn't close an issue.
## Rationale for this change
DataFusion's `MemoryPool` deliberately accounts for only the "large"
allocations that scale with input size; intermediate batches flowing
between operators are assumed to be small and left untracked. The
`MemoryPool` docs therefore advise reserving "some overhead (e.g. 10%)"
on top of the configured limit.
Nothing measures what that overhead actually is.
`MemoryPool::reserved()` is a live value that has usually fallen back to
zero by the time a query finishes, so the peak is never observed, and
benchmarks report peak RSS with nothing to compare it against.
This adds the missing number so the two can be compared. It is
measurement only — nothing enforces a relationship between the pool's
accounting and actual allocation, and no threshold or CI check is added.
This is complementary to the accounting work discussed in #22898. That
issue proposes changing *what* gets accounted (Arrow `claim()` / builder
`with_pool()`); because `ArrowMemoryPool` grows a DataFusion
`MemoryReservation` against the pool it wraps, `reserved()` is where
both models converge. This PR just makes the peak of that value
observable per query, so the effect of any such change — or of a
regression — is visible as a number rather than inferred.
Concretely, that means the peak follows the accounting instead of being
pinned to the current set of manually tracked consumers: nothing claims
buffers today, but if #22898 lands in either form, those bytes show up
here without further changes. There is a test covering that path.
## What changes are included in this PR?
- `PeakRecordingPool`, a delegating `MemoryPool` wrapper recording the
high-water mark of `reserved()`. Every method delegates; wrapping does
not change how memory is granted, limited, or reported.
- `CommonOpt::runtime_env_builder` installs it around the pool it
already builds, so every benchmark run with a memory limit reports the
peak without further changes.
- `BenchQuery` gains `pool_peak_bytes`, reset per case by
`start_new_case`, so each query gets its own reading rather than
inheriting the high-water mark of the queries before it.
- `print_memory_stats` prints the run-wide peak after the existing
mimalloc line, leaving `mem_profile`'s output parser untouched.
Everything is confined to the `benchmarks` crate. No trait changes.
`arrow-buffer/pool` and `datafusion-execution/arrow_buffer_pool` are
enabled as **dev-dependencies only**, so the benchmark binaries build
with exactly the features they had before — confirmed absent from `cargo
tree --no-dev-dependencies`.
## Are these changes tested?
Yes — 5 unit tests plus a doctest, and verified end-to-end with `dfbench
nlj --query 1 -i 2 --memory-limit 512M`, which emits `"pool_peak_bytes":
262528` per query.
One of those tests builds an `ArrowMemoryPool` over the recording pool
and asserts an Arrow-side reservation both raises the peak and releases
on drop, pinning the interaction described above.
The pre-existing `test_runtime_env_builder_reads_env_var` still asserts
`MemoryLimit::Finite(2G)` *through* the wrapper, which is direct
evidence the delegation is transparent.
The wrapper is installed only when a memory limit is configured, since
without one there is no pool to wrap. `pool_peak_bytes` is omitted from
the results JSON in that case, leaving output byte-identical to before
for runs without a limit — confirmed `compare.py` parses both old and
new files.
## Are there any user-facing changes?
No. Confined to the benchmarks crate; adds one optional field to the
benchmark results JSON.
---------
Co-authored-by: Claude <no...@anthropic.com>
---
Cargo.lock | 2 +
benchmarks/Cargo.toml | 5 +
benchmarks/README.md | 8 +
benchmarks/src/bin/benchmark_runner.rs | 5 +-
benchmarks/src/bin/external_aggr.rs | 16 +-
benchmarks/src/clickbench.rs | 3 +-
benchmarks/src/dict.rs | 1 +
benchmarks/src/h2o.rs | 3 +-
benchmarks/src/hj.rs | 1 +
benchmarks/src/imdb/run.rs | 13 +-
benchmarks/src/nlj.rs | 1 +
benchmarks/src/smj.rs | 1 +
benchmarks/src/sort_pushdown.rs | 13 +-
benchmarks/src/sort_tpch.rs | 14 +-
benchmarks/src/sql_benchmark_runner.rs | 2 +-
benchmarks/src/tpcds/run.rs | 3 +-
benchmarks/src/tpch/run.rs | 3 +-
benchmarks/src/util/memory.rs | 30 ++-
benchmarks/src/util/memory_pool.rs | 381 +++++++++++++++++++++++++++++++++
benchmarks/src/util/mod.rs | 2 +
benchmarks/src/util/options.rs | 5 +-
benchmarks/src/util/run.rs | 140 +++++++++++-
22 files changed, 631 insertions(+), 21 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index a41734e064..451ff70b1c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1780,6 +1780,7 @@ name = "datafusion-benchmarks"
version = "54.1.0"
dependencies = [
"arrow",
+ "arrow-buffer",
"async-trait",
"bytes",
"clap",
@@ -1787,6 +1788,7 @@ dependencies = [
"datafusion",
"datafusion-common",
"datafusion-common-runtime",
+ "datafusion-execution",
"datafusion-proto",
"env_logger",
"futures",
diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml
index 5dae70761f..282b27e481 100644
--- a/benchmarks/Cargo.toml
+++ b/benchmarks/Cargo.toml
@@ -64,6 +64,11 @@ tokio = { workspace = true, features = ["rt-multi-thread", "parking_lot"] }
tokio-util = { version = "0.7.17" }
[dev-dependencies]
+# `pool`/`arrow_buffer_pool` are enabled only for tests, so the benchmark
+# binaries are built exactly as before. They let `memory_pool`'s tests cover
+# Arrow-side reservations reaching the pool via `ArrowMemoryPool`.
+arrow-buffer = { workspace = true, features = ["pool"] }
+datafusion-execution = { workspace = true, features = ["arrow_buffer_pool"] }
datafusion-proto = { workspace = true, features = ["parquet"] }
tempfile = { workspace = true }
diff --git a/benchmarks/README.md b/benchmarks/README.md
index 34c67e5151..b6a7705cf9 100644
--- a/benchmarks/README.md
+++ b/benchmarks/README.md
@@ -483,6 +483,14 @@ Your benchmark should create and use an instance of `BenchmarkRun` defined in `b
- Call its `start_new_case` method with a string that will appear in the "Query" column of the
compare output.
- Use `write_iter` to record elapsed times for the behavior you're benchmarking.
+- Call `set_memory_pool` with the `RuntimeEnv`'s memory pool (`ctx.runtime_env().memory_pool`),
+ and again for each new runtime if your benchmark builds one per query. Each case then reports a
+ `pool_peak_bytes` field: the peak `MemoryPool` reservation reached while running it, which is the
+ largest value across that case's iterations. The field is omitted when the benchmark runs without
+ `--memory-limit`, since no pool is installed to record. Comparing it against the peak RSS printed
+ by `print_memory_stats` shows how much of the run's memory the pool actually accounted for; the
+ pool only tracks the "large" allocations that scale with input size, so the two are expected to
+ differ.
- When all cases are done, call the `BenchmarkRun`'s `maybe_write_json` method, giving it the value
of the `--output` structopt field on `RunOpt`.
diff --git a/benchmarks/src/bin/benchmark_runner.rs b/benchmarks/src/bin/benchmark_runner.rs
index 71b3b1b9e0..c7a16086c9 100644
--- a/benchmarks/src/bin/benchmark_runner.rs
+++ b/benchmarks/src/bin/benchmark_runner.rs
@@ -299,6 +299,9 @@ async fn run_simple_benchmark(
let case_name = benchmark_case_name(benchmark);
+ // Each case gets its own `SessionContext`, so hand over its pool before the
+ // case starts.
+ run.set_memory_pool(&ctx.runtime_env().memory_pool);
run.start_new_case(&case_name);
for iteration in 0..config.common.iterations {
@@ -312,7 +315,7 @@ async fn run_simple_benchmark(
run.write_iter(elapsed, row_count);
}
- print_memory_stats();
+ print_memory_stats(&*ctx.runtime_env().memory_pool);
Ok(())
}
diff --git a/benchmarks/src/bin/external_aggr.rs b/benchmarks/src/bin/external_aggr.rs
index 42f25c2cb0..226a619192 100644
--- a/benchmarks/src/bin/external_aggr.rs
+++ b/benchmarks/src/bin/external_aggr.rs
@@ -39,7 +39,9 @@ use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::physical_plan::display::DisplayableExecutionPlan;
use datafusion::physical_plan::{collect, displayable};
use datafusion::prelude::*;
-use datafusion_benchmarks::util::{BenchmarkRun, CommonOpt, QueryResult};
+use datafusion_benchmarks::util::{
+ BenchmarkRun, CommonOpt, PeakRecordingPool, QueryResult,
+};
use datafusion_common::instant::Instant;
use datafusion_common::utils::get_available_parallelism;
use datafusion_common::{DEFAULT_PARQUET_EXTENSION, exec_err};
@@ -169,7 +171,7 @@ impl ExternalAggrConfig {
));
let query_results = self
- .benchmark_query(query_id, mem_limit, mem_pool_type)
+ .benchmark_query(query_id, mem_limit, mem_pool_type, &mut benchmark_run)
.await?;
for iter in query_results {
benchmark_run.write_iter(iter.elapsed, iter.row_count);
@@ -182,11 +184,15 @@ impl ExternalAggrConfig {
}
/// Benchmark query `query_id` in `AGGR_QUERIES`
+ ///
+ /// `benchmark_run` is handed this query's runtime, which is built here
+ /// because each query runs under its own memory limit.
async fn benchmark_query(
&self,
query_id: usize,
mem_limit: u64,
mem_pool_type: &str,
+ benchmark_run: &mut BenchmarkRun,
) -> Result<Vec<QueryResult>> {
let query_name =
format!("Q{query_id}({})", human_readable_size(mem_limit as usize));
@@ -198,6 +204,12 @@ impl ExternalAggrConfig {
return exec_err!("Invalid memory pool type: {}", mem_pool_type);
}
};
+ // This benchmark builds its pool directly rather than going through
+ // `CommonOpt::runtime_env_builder`, so it has to install the recorder
+ // itself to report a peak.
+ let memory_pool: Arc<dyn MemoryPool> =
+ Arc::new(PeakRecordingPool::new(memory_pool));
+ benchmark_run.set_memory_pool(&memory_pool);
let runtime_env = RuntimeEnvBuilder::new()
.with_memory_pool(memory_pool)
.build_arc()?;
diff --git a/benchmarks/src/clickbench.rs b/benchmarks/src/clickbench.rs
index 70aaeb7d2d..a2e65aa561 100644
--- a/benchmarks/src/clickbench.rs
+++ b/benchmarks/src/clickbench.rs
@@ -213,6 +213,7 @@ impl RunOpt {
self.register_hits(&ctx).await?;
let mut benchmark_run = BenchmarkRun::new();
+ benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool);
for query_id in query_range {
let query_path = get_query_path(&self.queries_path, query_id);
let Some(sql) = get_query_sql(&query_path)? else {
@@ -278,7 +279,7 @@ impl RunOpt {
println!("Query {query_id} avg time: {avg:.2} ms");
// Print memory usage stats using mimalloc (only when compiled with --features mimalloc_extended)
- print_memory_stats();
+ print_memory_stats(&*ctx.runtime_env().memory_pool);
Ok(query_results)
}
diff --git a/benchmarks/src/dict.rs b/benchmarks/src/dict.rs
index f8451715ea..e04b5f816a 100644
--- a/benchmarks/src/dict.rs
+++ b/benchmarks/src/dict.rs
@@ -333,6 +333,7 @@ impl RunOpt {
let rt = self.common.build_runtime()?;
let ctx = SessionContext::new_with_config_rt(config, rt);
let mut benchmark_run = BenchmarkRun::new();
+ benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool);
for query_id in query_range {
let query = &DICTIONARY_QUERIES[query_id - 1];
diff --git a/benchmarks/src/h2o.rs b/benchmarks/src/h2o.rs
index 8b6e04932c..feb4bf2fa1 100644
--- a/benchmarks/src/h2o.rs
+++ b/benchmarks/src/h2o.rs
@@ -109,6 +109,7 @@ impl RunOpt {
let iterations = self.common.iterations;
let mut benchmark_run = BenchmarkRun::new();
+ benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool);
for query_id in query_range {
benchmark_run.start_new_case(&format!("Query {query_id}"));
let sql = queries.get_query(query_id)?;
@@ -131,7 +132,7 @@ impl RunOpt {
println!("Query {query_id} avg time: {avg:.2} ms");
// Print memory usage stats using mimalloc (only when compiled with --features mimalloc_extended)
- print_memory_stats();
+ print_memory_stats(&*ctx.runtime_env().memory_pool);
if self.common.debug {
ctx.sql(sql)
diff --git a/benchmarks/src/hj.rs b/benchmarks/src/hj.rs
index 9a91e2714a..4f97b24d0f 100644
--- a/benchmarks/src/hj.rs
+++ b/benchmarks/src/hj.rs
@@ -564,6 +564,7 @@ impl RunOpt {
}
let mut benchmark_run = BenchmarkRun::new();
+ benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool);
for query_id in query_range {
let query_index = query_id - 1;
diff --git a/benchmarks/src/imdb/run.rs b/benchmarks/src/imdb/run.rs
index a8e2028887..5822bbcb0d 100644
--- a/benchmarks/src/imdb/run.rs
+++ b/benchmarks/src/imdb/run.rs
@@ -295,7 +295,7 @@ impl RunOpt {
let mut benchmark_run = BenchmarkRun::new();
for query_id in query_range {
benchmark_run.start_new_case(&format!("Query {query_id}"));
- let query_run = self.benchmark_query(query_id).await?;
+ let query_run = self.benchmark_query(query_id, &mut benchmark_run).await?;
for iter in query_run {
benchmark_run.write_iter(iter.elapsed, iter.row_count);
}
@@ -304,7 +304,13 @@ impl RunOpt {
Ok(())
}
- async fn benchmark_query(&self, query_id: usize) -> Result<Vec<QueryResult>> {
+ /// `benchmark_run` is handed this query's runtime, which is built here so
+ /// each query gets a pool of its own.
+ async fn benchmark_query(
+ &self,
+ query_id: usize,
+ benchmark_run: &mut BenchmarkRun,
+ ) -> Result<Vec<QueryResult>> {
let mut config = self
.common
.config()?
@@ -314,6 +320,7 @@ impl RunOpt {
self.hash_join_buffering_capacity;
let rt = self.common.build_runtime()?;
let ctx = SessionContext::new_with_config_rt(config, rt);
+ benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool);
// register tables
self.register_tables(&ctx).await?;
@@ -348,7 +355,7 @@ impl RunOpt {
println!("Query {query_id} avg time: {avg:.2} ms");
// Print memory usage stats using mimalloc (only when compiled with --features mimalloc_extended)
- print_memory_stats();
+ print_memory_stats(&*ctx.runtime_env().memory_pool);
Ok(query_results)
}
diff --git a/benchmarks/src/nlj.rs b/benchmarks/src/nlj.rs
index 361cc35ec2..485ee069d1 100644
--- a/benchmarks/src/nlj.rs
+++ b/benchmarks/src/nlj.rs
@@ -211,6 +211,7 @@ impl RunOpt {
let ctx = SessionContext::new_with_config_rt(config, rt);
let mut benchmark_run = BenchmarkRun::new();
+ benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool);
for query_id in query_range {
let query_index = query_id - 1; // Convert 1-based to 0-based index
diff --git a/benchmarks/src/smj.rs b/benchmarks/src/smj.rs
index 3d173b7116..9282f72c2f 100644
--- a/benchmarks/src/smj.rs
+++ b/benchmarks/src/smj.rs
@@ -550,6 +550,7 @@ impl RunOpt {
let ctx = SessionContext::new_with_config_rt(config, rt);
let mut benchmark_run = BenchmarkRun::new();
+ benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool);
for query_id in query_range {
let query_index = query_id - 1; // Convert 1-based to 0-based index
diff --git a/benchmarks/src/sort_pushdown.rs b/benchmarks/src/sort_pushdown.rs
index 86f1c0f5c1..77f889e702 100644
--- a/benchmarks/src/sort_pushdown.rs
+++ b/benchmarks/src/sort_pushdown.rs
@@ -137,7 +137,7 @@ impl RunOpt {
for query_id in query_ids {
benchmark_run.start_new_case(&format!("{query_id}"));
- let query_results = self.benchmark_query(query_id).await;
+ let query_results = self.benchmark_query(query_id, &mut benchmark_run).await;
match query_results {
Ok(query_results) => {
for iter in query_results {
@@ -156,7 +156,13 @@ impl RunOpt {
Ok(())
}
- async fn benchmark_query(&self, query_id: usize) -> Result<Vec<QueryResult>> {
+ /// `benchmark_run` is handed this query's runtime, which is built here so
+ /// each query gets a pool of its own.
+ async fn benchmark_query(
+ &self,
+ query_id: usize,
+ benchmark_run: &mut BenchmarkRun,
+ ) -> Result<Vec<QueryResult>> {
let sql = self.load_query(query_id)?;
let config = self.common.config()?;
@@ -168,6 +174,7 @@ impl RunOpt {
.with_default_features()
.build();
let ctx = SessionContext::from(state);
+ benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool);
self.register_tables(&ctx).await?;
@@ -191,7 +198,7 @@ impl RunOpt {
let avg = millis.iter().sum::<f64>() / millis.len() as f64;
println!("Query {query_id} avg time: {avg:.2} ms");
- print_memory_stats();
+ print_memory_stats(&*ctx.runtime_env().memory_pool);
Ok(query_results)
}
diff --git a/benchmarks/src/sort_tpch.rs b/benchmarks/src/sort_tpch.rs
index 2182d1a383..d5f81c04a3 100644
--- a/benchmarks/src/sort_tpch.rs
+++ b/benchmarks/src/sort_tpch.rs
@@ -187,7 +187,7 @@ impl RunOpt {
for query_id in query_range {
benchmark_run.start_new_case(&format!("{query_id}"));
- let query_results = self.benchmark_query(query_id).await;
+ let query_results = self.benchmark_query(query_id, &mut benchmark_run).await;
match query_results {
Ok(query_results) => {
for iter in query_results {
@@ -207,7 +207,14 @@ impl RunOpt {
}
/// Benchmark query `query_id` in `SORT_QUERIES`
- async fn benchmark_query(&self, query_id: usize) -> Result<Vec<QueryResult>> {
+ ///
+ /// `benchmark_run` is handed this query's runtime, which is built here so
+ /// each query gets a pool of its own.
+ async fn benchmark_query(
+ &self,
+ query_id: usize,
+ benchmark_run: &mut BenchmarkRun,
+ ) -> Result<Vec<QueryResult>> {
let config = self.common.config()?;
let rt = self.common.build_runtime()?;
let state = SessionStateBuilder::new()
@@ -216,6 +223,7 @@ impl RunOpt {
.with_default_features()
.build();
let ctx = SessionContext::from(state);
+ benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool);
// register tables
self.register_tables(&ctx).await?;
@@ -250,7 +258,7 @@ impl RunOpt {
println!("Query {query_id} avg time: {avg:.2} ms");
// Print memory usage stats using mimalloc (only when compiled with --features mimalloc_extended)
- print_memory_stats();
+ print_memory_stats(&*ctx.runtime_env().memory_pool);
Ok(query_results)
}
diff --git a/benchmarks/src/sql_benchmark_runner.rs b/benchmarks/src/sql_benchmark_runner.rs
index f1cb3ad2f7..420e780c64 100644
--- a/benchmarks/src/sql_benchmark_runner.rs
+++ b/benchmarks/src/sql_benchmark_runner.rs
@@ -110,7 +110,7 @@ fn run_criterion_benchmark(
match result {
Ok(()) => {
- print_memory_stats();
+ print_memory_stats(&*ctx.runtime_env().memory_pool);
Ok(())
}
Err(payload) => Err(panic_payload_to_error(payload.as_ref())),
diff --git a/benchmarks/src/tpcds/run.rs b/benchmarks/src/tpcds/run.rs
index 2e0274c935..3eaaf172c0 100644
--- a/benchmarks/src/tpcds/run.rs
+++ b/benchmarks/src/tpcds/run.rs
@@ -226,6 +226,7 @@ impl RunOpt {
self.hash_join_buffering_capacity;
let rt = self.common.build_runtime()?;
let ctx = SessionContext::new_with_config_rt(config, rt);
+ benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool);
// register tables
self.register_tables(&ctx).await?;
@@ -290,7 +291,7 @@ impl RunOpt {
println!("Query {query_id} avg time: {avg:.2} ms");
// Print memory stats using mimalloc (only when compiled with --features mimalloc_extended)
- print_memory_stats();
+ print_memory_stats(&*ctx.runtime_env().memory_pool);
Ok(query_results)
}
diff --git a/benchmarks/src/tpch/run.rs b/benchmarks/src/tpch/run.rs
index 422bcec9ea..47edfbac4b 100644
--- a/benchmarks/src/tpch/run.rs
+++ b/benchmarks/src/tpch/run.rs
@@ -137,6 +137,7 @@ impl RunOpt {
self.hash_join_buffering_capacity;
let rt = self.common.build_runtime()?;
let ctx = SessionContext::new_with_config_rt(config, rt);
+ benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool);
// register tables
self.register_tables(&ctx).await?;
let scale_factor = self.scale_factor()?;
@@ -208,7 +209,7 @@ impl RunOpt {
println!("Query {query_id} avg time: {avg:.2} ms");
// Print memory stats using mimalloc (only when compiled with --features mimalloc_extended)
- print_memory_stats();
+ print_memory_stats(&*ctx.runtime_env().memory_pool);
Ok(query_results)
}
diff --git a/benchmarks/src/util/memory.rs b/benchmarks/src/util/memory.rs
index 11b96ef227..2b186c79c3 100644
--- a/benchmarks/src/util/memory.rs
+++ b/benchmarks/src/util/memory.rs
@@ -15,8 +15,34 @@
// specific language governing permissions and limitations
// under the License.
-/// Print Peak RSS, Peak Commit, Page Faults based on mimalloc api
-pub fn print_memory_stats() {
+use datafusion::execution::memory_pool::MemoryPool;
+
+use super::PeakRecordingPool;
+
+/// Print Peak RSS, Peak Commit, Page Faults based on mimalloc api, followed by
+/// the peak reservation of `memory_pool` when a memory limit was configured.
+pub fn print_memory_stats(memory_pool: &dyn MemoryPool) {
+ print_allocator_stats();
+ print_pool_stats(memory_pool);
+}
+
+/// Print the peak reservation `memory_pool` has seen.
+///
+/// Prints nothing when the benchmark ran without a memory limit, since no
+/// [`PeakRecordingPool`] was installed to record. Comparing this against the
+/// peak RSS above shows how much of a run's memory the pool actually accounted
+/// for — DataFusion only tracks the "large" allocations that scale with input
+/// size, so the two are expected to differ.
+fn print_pool_stats(memory_pool: &dyn MemoryPool) {
+ if let Some(recorder) = PeakRecordingPool::from_pool(memory_pool) {
+ println!(
+ "Peak pool reserved: {}",
+ datafusion_common::human_readable_size(recorder.max_reserved())
+ );
+ }
+}
+
+fn print_allocator_stats() {
#[cfg(all(feature = "mimalloc", feature = "mimalloc_extended"))]
{
use datafusion_common::human_readable_size;
diff --git a/benchmarks/src/util/memory_pool.rs b/benchmarks/src/util/memory_pool.rs
new file mode 100644
index 0000000000..a3606ca0a7
--- /dev/null
+++ b/benchmarks/src/util/memory_pool.rs
@@ -0,0 +1,381 @@
+// 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.
+
+//! Records the peak [`MemoryPool`] reservation reached during a benchmark.
+//!
+//! DataFusion's [`MemoryPool`] deliberately accounts for only the "large"
+//! allocations that scale with input size; intermediate batches flowing between
+//! operators are assumed to be small and are left untracked. The [`MemoryPool`]
+//! documentation therefore advises reserving "some overhead (e.g. 10%)" on top
+//! of the configured limit.
+//!
+//! Nothing reports what that overhead actually is, because the peak reservation
+//! itself is never recorded — [`MemoryPool::reserved`] is a live value that has
+//! usually fallen back to zero by the time a query finishes. This module records
+//! the high-water mark so benchmarks can emit it alongside the peak RSS that
+//! [`print_memory_stats`] already prints, making the gap between the two
+//! measurable.
+//!
+//! This is measurement only: nothing here enforces a relationship between the
+//! two numbers.
+//!
+//! What lands in the peak is whatever the pool accounts for, so this follows
+//! the accounting rather than fixing it in place. Arrow-side reservations made
+//! through `ArrowMemoryPool` are included, because that adapter grows a
+//! DataFusion reservation against the pool it wraps; nothing claims buffers
+//! today, but the peak picks it up when something does.
+//!
+//! [`print_memory_stats`]: super::print_memory_stats
+
+use std::{
+ fmt::{Debug, Display, Formatter},
+ sync::{
+ Arc,
+ atomic::{AtomicUsize, Ordering},
+ },
+};
+
+use datafusion::execution::memory_pool::{
+ MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation,
+};
+use datafusion_common::Result;
+
+/// Wraps a [`MemoryPool`], recording the high-water mark of
+/// [`MemoryPool::reserved`] as reservations come and go.
+///
+/// Every method delegates to the wrapped pool, so wrapping does not change how
+/// memory is granted, limited, or reported. The one thing it does change is
+/// downcasting: `rt.memory_pool.downcast_ref::<FairSpillPool>()` now finds this
+/// wrapper instead of the pool it wraps. Nothing in the benchmarks relies on
+/// that, and [`Self::from_pool`] uses the same mechanism to find the recorder.
+///
+/// Both high-water marks are held per instance, so a benchmark that builds a
+/// fresh runtime per query gets a reading scoped to that query without any
+/// coordination.
+///
+/// # Example
+///
+/// ```
+/// # use std::sync::Arc;
+/// # use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool};
+/// # use datafusion_benchmarks::util::PeakRecordingPool;
+/// let recording = Arc::new(PeakRecordingPool::new(Arc::new(GreedyMemoryPool::new(1024))));
+/// let pool: Arc<dyn MemoryPool> = Arc::clone(&recording) as _;
+///
+/// let reservation = MemoryConsumer::new("example").register(&pool);
+/// reservation.try_grow(512)?;
+/// reservation.shrink(512);
+///
+/// // The pool is back to empty, but the high-water mark is retained.
+/// assert_eq!(pool.reserved(), 0);
+/// assert_eq!(recording.peak_reserved(), 512);
+///
+/// // The recorder can also be recovered from the pool it was installed as.
+/// assert_eq!(PeakRecordingPool::from_pool(&*pool).unwrap().peak_reserved(), 512);
+/// # Ok::<(), datafusion_common::DataFusionError>(())
+/// ```
+pub struct PeakRecordingPool {
+ inner: Arc<dyn MemoryPool>,
+ /// Running total of everything granted through this wrapper, kept so the
+ /// peak can be maintained without asking `inner` for its total.
+ reserved: AtomicUsize,
+ /// High-water mark since the last [`PeakRecordingPool::reset_peak`].
+ peak: AtomicUsize,
+ /// High-water mark since this pool was created. Never reset.
+ max: AtomicUsize,
+}
+
+impl PeakRecordingPool {
+ /// Wrap `inner`, recording its peak reservation from here on.
+ ///
+ /// `inner` is expected to be empty: the running total starts at zero, so
+ /// anything reserved before wrapping is not counted.
+ pub fn new(inner: Arc<dyn MemoryPool>) -> Self {
+ Self {
+ inner,
+ reserved: AtomicUsize::new(0),
+ peak: AtomicUsize::new(0),
+ max: AtomicUsize::new(0),
+ }
+ }
+
+ /// The recorder installed as `pool`, if there is one.
+ ///
+ /// Returns `None` whenever a benchmark runs without a memory limit, since
+ /// [`CommonOpt::runtime_env_builder`] only installs the wrapper alongside a
+ /// pool it has a limit for.
+ ///
+ /// [`CommonOpt::runtime_env_builder`]: super::CommonOpt::runtime_env_builder
+ pub fn from_pool(pool: &dyn MemoryPool) -> Option<&Self> {
+ pool.downcast_ref::<Self>()
+ }
+
+ /// Peak reservation, in bytes, since the last [`Self::reset_peak`].
+ pub fn peak_reserved(&self) -> usize {
+ self.peak.load(Ordering::Relaxed)
+ }
+
+ /// Peak reservation, in bytes, since this pool was created.
+ ///
+ /// Unlike [`Self::peak_reserved`] this is never reset, so it reports the
+ /// peak across every query that shared this pool.
+ pub fn max_reserved(&self) -> usize {
+ self.max.load(Ordering::Relaxed)
+ }
+
+ /// Reset the value returned by [`Self::peak_reserved`] to what is reserved
+ /// right now, so the next reading covers only what follows.
+ ///
+ /// [`BenchmarkRun::start_new_case`] calls this, giving each benchmark query
+ /// its own reading. Anything still held when a query starts — data the
+ /// benchmark loaded up front, say — stays in the reading, since the query
+ /// runs with those bytes reserved.
+ ///
+ /// [`BenchmarkRun::start_new_case`]: super::BenchmarkRun::start_new_case
+ pub fn reset_peak(&self) {
+ self.peak
+ .store(self.reserved.load(Ordering::Relaxed), Ordering::Relaxed);
+ }
+
+ /// Add `additional` granted bytes to the running total and publish it to
+ /// both high-water marks.
+ ///
+ /// Accumulating deltas rather than reading [`MemoryPool::reserved`] keeps
+ /// the wrapped pool's own bookkeeping off this path: `FairSpillPool` takes
+ /// its state lock to answer `reserved()`, which would double the lock
+ /// traffic of every accounted allocation in the benchmark being measured.
+ /// The total stays exact because the trait grants exactly what is asked
+ /// for — `grow` is infallible and `try_grow` either grants `additional` or
+ /// returns an error, leaving the reservation untouched.
+ fn record(&self, additional: usize) {
+ let reserved =
+ self.reserved.fetch_add(additional, Ordering::Relaxed) + additional;
+ self.peak.fetch_max(reserved, Ordering::Relaxed);
+ self.max.fetch_max(reserved, Ordering::Relaxed);
+ }
+}
+
+impl Debug for PeakRecordingPool {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ f.debug_struct("PeakRecordingPool")
+ .field("inner", &self.inner)
+ .field("peak", &self.peak_reserved())
+ .field("max", &self.max_reserved())
+ .finish()
+ }
+}
+
+impl Display for PeakRecordingPool {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ // Deferring to the wrapped pool keeps `SHOW ALL`-style output and error
+ // messages identical to running without the wrapper.
+ Display::fmt(&self.inner, f)
+ }
+}
+
+impl MemoryPool for PeakRecordingPool {
+ fn name(&self) -> &str {
+ self.inner.name()
+ }
+
+ fn register(&self, consumer: &MemoryConsumer) {
+ self.inner.register(consumer);
+ }
+
+ fn unregister(&self, consumer: &MemoryConsumer) {
+ self.inner.unregister(consumer);
+ }
+
+ fn grow(&self, reservation: &MemoryReservation, additional: usize) {
+ self.inner.grow(reservation, additional);
+ self.record(additional);
+ }
+
+ fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
+ self.inner.shrink(reservation, shrink);
+ self.reserved.fetch_sub(shrink, Ordering::Relaxed);
+ }
+
+ fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
+ self.inner.try_grow(reservation, additional)?;
+ self.record(additional);
+ Ok(())
+ }
+
+ fn reserved(&self) -> usize {
+ self.inner.reserved()
+ }
+
+ fn memory_limit(&self) -> MemoryLimit {
+ self.inner.memory_limit()
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use datafusion::execution::memory_pool::GreedyMemoryPool;
+
+ use super::*;
+
+ /// A recording pool over a `GreedyMemoryPool`, returned both as the
+ /// recorder (to read the marks) and as the pool reservations register with.
+ fn pool(limit: usize) -> (Arc<PeakRecordingPool>, Arc<dyn MemoryPool>) {
+ let recording = Arc::new(PeakRecordingPool::new(Arc::new(
+ GreedyMemoryPool::new(limit),
+ )));
+ let pool = Arc::clone(&recording) as Arc<dyn MemoryPool>;
+ (recording, pool)
+ }
+
+ #[test]
+ fn records_high_water_mark_across_reservations() {
+ let (recording, pool) = pool(1024);
+
+ let a = MemoryConsumer::new("a").register(&pool);
+ let b = MemoryConsumer::new("b").register(&pool);
+
+ a.try_grow(300).unwrap();
+ b.try_grow(400).unwrap();
+ // Peak of the sum, not the largest single reservation.
+ assert_eq!(recording.peak_reserved(), 700);
+
+ a.shrink(300);
+ b.try_grow(100).unwrap();
+
+ // Falling back below the peak leaves it untouched, and the later growth
+ // does not reach it.
+ assert_eq!(pool.reserved(), 500);
+ assert_eq!(recording.peak_reserved(), 700);
+ }
+
+ #[test]
+ fn failed_growth_does_not_move_the_peak() {
+ let (recording, pool) = pool(1024);
+
+ let reservation = MemoryConsumer::new("a").register(&pool);
+ reservation.try_grow(600).unwrap();
+ reservation
+ .try_grow(600)
+ .expect_err("should exceed the 1024 byte pool");
+
+ assert_eq!(recording.peak_reserved(), 600);
+ }
+
+ #[test]
+ fn reset_clears_the_window_but_not_the_run_maximum() {
+ let (recording, pool) = pool(1024);
+
+ let reservation = MemoryConsumer::new("a").register(&pool);
+ reservation.try_grow(800).unwrap();
+ reservation.shrink(800);
+
+ recording.reset_peak();
+ assert_eq!(recording.peak_reserved(), 0);
+ assert_eq!(recording.max_reserved(), 800);
+
+ reservation.try_grow(100).unwrap();
+ assert_eq!(recording.peak_reserved(), 100);
+ assert_eq!(recording.max_reserved(), 800);
+ }
+
+ #[test]
+ fn reset_keeps_what_is_still_reserved() {
+ let (recording, pool) = pool(1024);
+
+ // Something a benchmark loaded up front and holds across queries.
+ let held = MemoryConsumer::new("held").register(&pool);
+ held.try_grow(300).unwrap();
+
+ recording.reset_peak();
+ assert_eq!(recording.peak_reserved(), 300);
+
+ let query = MemoryConsumer::new("query").register(&pool);
+ query.try_grow(200).unwrap();
+ assert_eq!(recording.peak_reserved(), 500);
+ }
+
+ #[test]
+ fn marks_are_per_instance() {
+ let (one, one_pool) = pool(1024);
+ let (two, _two_pool) = pool(1024);
+
+ MemoryConsumer::new("a")
+ .register(&one_pool)
+ .try_grow(512)
+ .unwrap();
+
+ assert_eq!(one.peak_reserved(), 512);
+ assert_eq!(two.peak_reserved(), 0);
+ }
+
+ #[test]
+ fn is_recoverable_from_the_pool_it_is_installed_as() {
+ let (recording, pool) = pool(1024);
+
+ MemoryConsumer::new("a")
+ .register(&pool)
+ .try_grow(512)
+ .unwrap();
+
+ let found = PeakRecordingPool::from_pool(&*pool).expect("recorder installed");
+ assert_eq!(found.peak_reserved(), recording.peak_reserved());
+
+ // A pool with no recorder in front of it reports nothing.
+ let plain: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1024));
+ assert!(PeakRecordingPool::from_pool(&*plain).is_none());
+ }
+
+ #[test]
+ fn delegates_limit_and_name_to_the_wrapped_pool() {
+ let inner: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(4096));
+ let wrapped = PeakRecordingPool::new(Arc::clone(&inner));
+
+ assert_eq!(wrapped.name(), inner.name());
+ assert_eq!(wrapped.to_string(), inner.to_string());
+ assert!(matches!(wrapped.memory_limit(), MemoryLimit::Finite(4096)));
+ }
+
+ /// Arrow-side reservations reach the recorder too.
+ ///
+ /// [`ArrowMemoryPool`] implements Arrow's `MemoryPool` by growing a
+ /// DataFusion [`MemoryReservation`] against the pool it wraps, so a buffer
+ /// claimed through it lands in `grow` here. Nothing in DataFusion claims
+ /// buffers yet (see apache/datafusion#22898), but when something does, the
+ /// bytes show up in this peak without further changes — as long as the
+ /// adapter is built from the `RuntimeEnv`'s pool, which is the wrapped one.
+ /// This test pins that.
+ #[test]
+ fn records_reservations_arriving_through_the_arrow_adapter() {
+ use arrow_buffer::MemoryPool as ArrowMemoryPoolTrait;
+ use datafusion_execution::memory_pool::arrow::ArrowMemoryPool;
+
+ let (recording, pool) = pool(4096);
+
+ let arrow_pool =
+ ArrowMemoryPool::new(Arc::clone(&pool), MemoryConsumer::new("arrow"));
+ let reservation = arrow_pool.reserve(1024);
+
+ // The Arrow-side reservation is visible as DataFusion pool usage...
+ assert_eq!(pool.reserved(), 1024);
+ assert_eq!(recording.peak_reserved(), 1024);
+
+ // ...and dropping it releases the bytes while the peak is retained.
+ drop(reservation);
+ assert_eq!(pool.reserved(), 0);
+ assert_eq!(recording.peak_reserved(), 1024);
+ }
+}
diff --git a/benchmarks/src/util/mod.rs b/benchmarks/src/util/mod.rs
index 6dc11c0f42..43855ea468 100644
--- a/benchmarks/src/util/mod.rs
+++ b/benchmarks/src/util/mod.rs
@@ -18,9 +18,11 @@
//! Shared benchmark utilities
pub mod latency_object_store;
mod memory;
+mod memory_pool;
mod options;
mod run;
pub use memory::print_memory_stats;
+pub use memory_pool::PeakRecordingPool;
pub use options::CommonOpt;
pub use run::{BenchQuery, BenchmarkRun, QueryResult};
diff --git a/benchmarks/src/util/options.rs b/benchmarks/src/util/options.rs
index a3e6d2a4c5..c744d0bf31 100644
--- a/benchmarks/src/util/options.rs
+++ b/benchmarks/src/util/options.rs
@@ -30,7 +30,7 @@ use datafusion::{
use datafusion_common::{DataFusionError, Result};
use object_store::local::LocalFileSystem;
-use super::latency_object_store::LatencyObjectStore;
+use super::{latency_object_store::LatencyObjectStore, memory_pool::PeakRecordingPool};
// Common benchmark options (don't use doc comments otherwise this doc
// shows up in help files)
@@ -125,6 +125,9 @@ impl CommonOpt {
)));
}
};
+ // Record the peak reservation so benchmarks can report it next to
+ // peak RSS. Purely observational: every call is delegated.
+ let pool: Arc<dyn MemoryPool> = Arc::new(PeakRecordingPool::new(pool));
rt_builder = rt_builder
.with_memory_pool(pool)
.with_disk_manager_builder(DiskManagerBuilder::default());
diff --git a/benchmarks/src/util/run.rs b/benchmarks/src/util/run.rs
index df17674e62..6c63ceec64 100644
--- a/benchmarks/src/util/run.rs
+++ b/benchmarks/src/util/run.rs
@@ -15,6 +15,8 @@
// specific language governing permissions and limitations
// under the License.
+use super::memory_pool::PeakRecordingPool;
+use datafusion::execution::memory_pool::MemoryPool;
use datafusion::{DATAFUSION_VERSION, error::Result};
use datafusion_common::utils::get_available_parallelism;
use serde::{Serialize, Serializer};
@@ -22,6 +24,7 @@ use serde_json::Value;
use std::{
collections::HashMap,
path::Path,
+ sync::Arc,
time::{Duration, SystemTime},
};
@@ -91,6 +94,16 @@ pub struct BenchQuery {
#[serde(serialize_with = "serialize_start_time")]
start_time: SystemTime,
success: bool,
+ /// Peak [`MemoryPool`] reservation observed while running this query, in
+ /// bytes. Recorded for failed queries too, since a query that ran out of
+ /// memory is one whose peak is worth seeing.
+ ///
+ /// `None` (and omitted from the JSON) only when the benchmark ran without a
+ /// memory limit, since there is then no pool to record.
+ ///
+ /// [`MemoryPool`]: datafusion::execution::memory_pool::MemoryPool
+ #[serde(skip_serializing_if = "Option::is_none")]
+ pool_peak_bytes: Option<usize>,
}
/// Internal representation of a single benchmark query iteration result.
pub struct QueryResult {
@@ -102,6 +115,10 @@ pub struct BenchmarkRun {
context: RunContext,
queries: Vec<BenchQuery>,
current_case: Option<usize>,
+ /// The pool queries run against, when one was handed over with
+ /// [`BenchmarkRun::set_memory_pool`]. Only read through
+ /// [`BenchmarkRun::peak_recorder`].
+ memory_pool: Option<Arc<dyn MemoryPool>>,
}
impl Default for BenchmarkRun {
@@ -117,15 +134,44 @@ impl BenchmarkRun {
context: RunContext::new(),
queries: vec![],
current_case: None,
+ memory_pool: None,
}
}
+
+ /// Report the peak reservation of `memory_pool` alongside each query.
+ ///
+ /// Call this with the pool of the [`RuntimeEnv`] the queries run against.
+ /// Has no effect unless a [`PeakRecordingPool`] is installed, which
+ /// [`CommonOpt::runtime_env_builder`] does whenever a memory limit is
+ /// configured; without one `pool_peak_bytes` is omitted from the results.
+ ///
+ /// Benchmarks that build a runtime per query should call this each time, so
+ /// each query reports against the pool it actually ran on.
+ ///
+ /// [`RuntimeEnv`]: datafusion::execution::runtime_env::RuntimeEnv
+ /// [`CommonOpt::runtime_env_builder`]: super::CommonOpt::runtime_env_builder
+ pub fn set_memory_pool(&mut self, memory_pool: &Arc<dyn MemoryPool>) {
+ self.memory_pool = Some(Arc::clone(memory_pool));
+ }
+
+ /// The recorder in front of the pool set by [`Self::set_memory_pool`].
+ fn peak_recorder(&self) -> Option<&PeakRecordingPool> {
+ PeakRecordingPool::from_pool(self.memory_pool.as_deref()?)
+ }
+
/// begin a new case. iterations added after this will be included in the new case
pub fn start_new_case(&mut self, id: &str) {
+ // Give this query its own memory pool reading rather than inheriting
+ // the high-water mark of the queries that ran before it.
+ if let Some(recorder) = self.peak_recorder() {
+ recorder.reset_peak();
+ }
self.queries.push(BenchQuery {
query: id.to_owned(),
iterations: vec![],
start_time: SystemTime::now(),
success: true,
+ pool_peak_bytes: None,
});
if let Some(c) = self.current_case.as_mut() {
*c += 1;
@@ -135,10 +181,14 @@ impl BenchmarkRun {
}
/// Write a new iteration to the current case
pub fn write_iter(&mut self, elapsed: Duration, row_count: usize) {
+ // The peak is not reset between iterations, so this ends up holding the
+ // largest reservation seen across all of them.
+ let pool_peak_bytes = self.peak_recorder().map(PeakRecordingPool::peak_reserved);
if let Some(idx) = self.current_case {
self.queries[idx]
.iterations
- .push(QueryIter { elapsed, row_count })
+ .push(QueryIter { elapsed, row_count });
+ self.queries[idx].pool_peak_bytes = pool_peak_bytes;
} else {
panic!("no cases existed yet");
}
@@ -159,8 +209,12 @@ impl BenchmarkRun {
/// Mark current query
pub fn mark_failed(&mut self) {
+ // A query that failed under a memory limit wrote no iteration, so this
+ // is the only chance to record what it had reserved when it gave up.
+ let pool_peak_bytes = self.peak_recorder().map(PeakRecordingPool::peak_reserved);
if let Some(idx) = self.current_case {
self.queries[idx].success = false;
+ self.queries[idx].pool_peak_bytes = pool_peak_bytes;
} else {
unreachable!("Cannot mark failure: no current case");
}
@@ -182,3 +236,87 @@ impl BenchmarkRun {
Ok(())
}
}
+
+#[cfg(test)]
+mod tests {
+ use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryConsumer};
+
+ use super::*;
+
+ fn recording_pool(limit: usize) -> Arc<dyn MemoryPool> {
+ Arc::new(PeakRecordingPool::new(Arc::new(GreedyMemoryPool::new(
+ limit,
+ ))))
+ }
+
+ #[test]
+ fn each_case_reports_its_own_peak() {
+ let pool = recording_pool(1024);
+ let mut run = BenchmarkRun::new();
+ run.set_memory_pool(&pool);
+
+ run.start_new_case("q1");
+ let reservation = MemoryConsumer::new("q1").register(&pool);
+ reservation.try_grow(600).unwrap();
+ run.write_iter(Duration::from_millis(1), 1);
+ drop(reservation);
+
+ // The second case must not inherit the first case's high-water mark.
+ run.start_new_case("q2");
+ let reservation = MemoryConsumer::new("q2").register(&pool);
+ reservation.try_grow(100).unwrap();
+ run.write_iter(Duration::from_millis(1), 1);
+
+ assert_eq!(run.queries[0].pool_peak_bytes, Some(600));
+ assert_eq!(run.queries[1].pool_peak_bytes, Some(100));
+ }
+
+ #[test]
+ fn a_later_pool_replaces_an_earlier_one() {
+ let first = recording_pool(1024);
+ let mut run = BenchmarkRun::new();
+ run.set_memory_pool(&first);
+ MemoryConsumer::new("q1")
+ .register(&first)
+ .try_grow(600)
+ .unwrap();
+
+ // Benchmarks that build a runtime per query hand over the new pool
+ // before the next case; the reading follows it.
+ let second = recording_pool(1024);
+ run.set_memory_pool(&second);
+ run.start_new_case("q2");
+ MemoryConsumer::new("q2")
+ .register(&second)
+ .try_grow(100)
+ .unwrap();
+ run.write_iter(Duration::from_millis(1), 1);
+
+ assert_eq!(run.queries[0].pool_peak_bytes, Some(100));
+ }
+
+ #[test]
+ fn a_failed_query_still_reports_its_peak() {
+ let pool = recording_pool(1024);
+ let mut run = BenchmarkRun::new();
+ run.set_memory_pool(&pool);
+
+ run.start_new_case("q1");
+ let reservation = MemoryConsumer::new("q1").register(&pool);
+ reservation.try_grow(600).unwrap();
+ // No `write_iter`: the query failed before completing an iteration.
+ run.mark_failed();
+
+ assert_eq!(run.queries[0].pool_peak_bytes, Some(600));
+ }
+
+ #[test]
+ fn the_peak_is_omitted_without_a_recording_pool() {
+ let mut run = BenchmarkRun::new();
+ run.start_new_case("q1");
+ run.write_iter(Duration::from_millis(1), 1);
+
+ assert_eq!(run.queries[0].pool_peak_bytes, None);
+ assert!(!run.to_json().contains("pool_peak_bytes"));
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org