add benchmark utility to profile memory usage

Author: ding-young

Cargo.lock

@@ -45,6 +45,7 @@ mimalloc = { version = "0.1", optional = true, default-features = false }
 object_store = { workspace = true }
 parquet = { workspace = true, default-features = true }
 rand = { workspace = true }
+regex.workspace = true
 serde = { version = "1.0.219", features = ["derive"] }
 serde_json = { workspace = true }
 snmalloc-rs = { version = "0.3", optional = true }
@@ -53,5 +54,8 @@ test-utils = { path = "../test-utils/", version = "0.1.0" }
 tokio = { workspace = true, features = ["rt-multi-thread", "parking_lot"] }
 tokio-util = { version = "0.7.15" }
 
+[target.'cfg(target_os = "linux")'.dependencies]
+procfs = "0.17.0"
+
 [dev-dependencies]
 datafusion-proto = { workspace = true }

benchmarks/Cargo.toml

@@ -19,7 +19,7 @@ use std::fs;
 use std::io::ErrorKind;
 use std::path::{Path, PathBuf};
 
-use crate::util::{BenchmarkRun, CommonOpt, QueryResult};
+use crate::util::{print_memory_stats, BenchmarkRun, CommonOpt, QueryResult};
 use datafusion::logical_expr::{ExplainFormat, ExplainOption};
 use datafusion::{
     error::{DataFusionError, Result},
@@ -192,6 +192,10 @@ impl RunOpt {
         }
         let avg = millis.iter().sum::<f64>() / millis.len() as f64;
         println!("Query {query_id} avg time: {avg:.2} ms");
+
+        if self.common.memory_stat_enabled {
+            print_memory_stats();
+        }
         Ok(query_results)
     }
 

benchmarks/src/clickbench.rs

@@ -21,7 +21,7 @@ use std::sync::Arc;
 use super::{
     get_query_sql, get_tbl_tpch_table_schema, get_tpch_table_schema, TPCH_TABLES,
 };
-use crate::util::{BenchmarkRun, CommonOpt, QueryResult};
+use crate::util::{print_memory_stats, BenchmarkRun, CommonOpt, QueryResult};
 
 use arrow::record_batch::RecordBatch;
 use arrow::util::pretty::{self, pretty_format_batches};
@@ -184,6 +184,10 @@ impl RunOpt {
         let avg = millis.iter().sum::<f64>() / millis.len() as f64;
         println!("Query {query_id} avg time: {avg:.2} ms");
 
+        if self.common.memory_stat_enabled {
+            print_memory_stats();
+        }
+
         Ok(query_results)
     }
 

benchmarks/src/tpch/run.rs

@@ -0,0 +1,54 @@
+// 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.
+
+use datafusion::execution::memory_pool::human_readable_size;
+
+#[derive(Debug)]
+pub struct MemoryStats {
+    pub vm_rss_kb: Option<u64>,
+    pub vm_hwm_kb: Option<u64>,
+    pub vm_size_kb: Option<u64>,
+    pub vm_peak_kb: Option<u64>,
+}
+
+pub fn print_memory_stats() {
+    #[cfg(target_os = "linux")]
+    {
+        use procfs::process::Process;
+
+        let pid = std::process::id();
+        let process = Process::new(pid as i32).unwrap();
+        let statm = process.statm().unwrap();
+        let status = process.status().unwrap();
+        let page_size = procfs::page_size();
+
+        let resident_bytes = (statm.resident * page_size) as usize;
+        let vmpeak_bytes = status.vmpeak.map(|kb| (kb * 1024) as usize);
+        let vmhwm_bytes = status.vmhwm.map(|kb| (kb * 1024) as usize);
+
+        println!(
+            "VmPeak: {}, VmHWM: {}, RSS: {}",
+            vmpeak_bytes
+                .map(human_readable_size)
+                .unwrap_or_else(|| "N/A".to_string()),
+            vmhwm_bytes
+                .map(human_readable_size)
+                .unwrap_or_else(|| "N/A".to_string()),
+            human_readable_size(resident_bytes)
+        );
+    }
+}

benchmarks/src/util/memory.rs

@@ -16,8 +16,10 @@
 // under the License.
 
 //! Shared benchmark utilities
+mod memory;
 mod options;
 mod run;
 
+pub use memory::{print_memory_stats, MemoryStats};
 pub use options::CommonOpt;
 pub use run::{BenchQuery, BenchmarkRun, QueryResult};

benchmarks/src/util/mod.rs

@@ -61,6 +61,11 @@ pub struct CommonOpt {
     /// Activate debug mode to see more details
     #[structopt(short, long)]
     pub debug: bool,
+
+    /// Enable memory profiling to see VmPeak, VmHwm for running benchmark.
+    /// See more details in TODO
+    #[structopt(long)]
+    pub memory_stat_enabled: bool,
 }
 
 impl CommonOpt {