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 <noreply@anthropic.com>

Author: adriangb

Cargo.lock

@@ -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 }
 

benchmarks/Cargo.toml

@@ -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`.
 

benchmarks/README.md

@@ -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(())
 }

benchmarks/src/bin/benchmark_runner.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()?;

benchmarks/src/bin/external_aggr.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)
     }

benchmarks/src/clickbench.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];

benchmarks/src/dict.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)

benchmarks/src/h2o.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;

benchmarks/src/hj.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)
     }