Report peak MemoryPool reservation per query in benchmarks

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 }
 

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

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

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()?;

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

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];

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)

benchmarks/src/hj.rs

@@ -518,6 +518,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/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)
     }

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
 

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
 

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

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

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())),

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

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

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;