Add memory profiling support to DataFusion CLI and memory pool metrics

datafusion-cli/Cargo.toml

@@ -53,6 +53,7 @@ datafusion = { workspace = true, features = [
  "unicode_expressions",
  "compression",
 ] }
+datafusion-execution = { workspace = true }
 dirs = "6.0.0"
 env_logger = { workspace = true }
 futures = { workspace = true }

datafusion-cli/src/cli_context.rs

@@ -20,14 +20,38 @@ use std::sync::Arc;
 use datafusion::{
  dataframe::DataFrame,
  error::DataFusionError,
- execution::{context::SessionState, TaskContext},
+ execution::{context::SessionState, memory_pool::TrackedPool, TaskContext},
  logical_expr::LogicalPlan,
  prelude::SessionContext,
 };
 use object_store::ObjectStore;
 
 use crate::object_storage::{AwsOptions, GcpOptions};
 
+/// Registers table option extensions based on the provided URL scheme.
+///
+/// Supported schemes are:
+/// * `s3`, `oss`, `cos` - registers [`AwsOptions`]
+/// * `gs`, `gcs` - registers [`GcpOptions`]
+///
+/// Any other scheme is ignored.
+pub fn register_table_options_from_scheme(ctx: &SessionContext, scheme: &str) {
+ match scheme {
+ // For Amazon S3 or Alibaba Cloud OSS
+ "s3" | "oss" | "cos" => {
+ // Register AWS specific table options in the session context:
+ ctx.register_table_options_extension(AwsOptions::default())
+ }
+ // For Google Cloud Storage
+ "gs" | "gcs" => {
+ // Register GCP specific table options in the session context:
+ ctx.register_table_options_extension(GcpOptions::default())
+ }
+ // For unsupported schemes, do nothing:
+ _ => {}
+ }
+}
+
 #[async_trait::async_trait]
 /// The CLI session context trait provides a way to have a session context that can be used with datafusion's CLI code.
 pub trait CliSessionContext {
@@ -52,6 +76,19 @@ pub trait CliSessionContext {
  &self,
  plan: LogicalPlan,
  ) -> Result<DataFrame, DataFusionError>;
+
+ /// Return true if memory profiling is enabled.
+ fn memory_profiling(&self) -> bool {
+ false
+ }
+
+ /// Enable or disable memory profiling.
+ fn set_memory_profiling(&self, _enable: bool) {}
+
+ /// Return the tracked memory pool used for profiling, if any.
+ fn tracked_memory_pool(&self) -> Option<Arc<dyn TrackedPool>> {
+ None
+ }
 }
 
 #[async_trait::async_trait]
@@ -73,26 +110,82 @@ impl CliSessionContext for SessionContext {
  }
 
  fn register_table_options_extension_from_scheme(&self, scheme: &str) {
- match scheme {
- // For Amazon S3 or Alibaba Cloud OSS
- "s3" | "oss" | "cos" => {
- // Register AWS specific table options in the session context:
- self.register_table_options_extension(AwsOptions::default())
- }
- // For Google Cloud Storage
- "gs" | "gcs" => {
- // Register GCP specific table options in the session context:
- self.register_table_options_extension(GcpOptions::default())
- }
- // For unsupported schemes, do nothing:
- _ => {}
+ register_table_options_from_scheme(self, scheme);
+ }
+
+ async fn execute_logical_plan(
+ &self,
+ plan: LogicalPlan,
+ ) -> Result<DataFrame, DataFusionError> {
+ SessionContext::execute_logical_plan(self, plan).await
+ }
+}
+
+/// Session context used by the CLI with memory profiling support.
+pub struct ReplSessionContext {
+ ctx: SessionContext,
+ tracked_memory_pool: Option<Arc<dyn TrackedPool>>,
+}
+
+impl ReplSessionContext {
+ pub fn new(
+ ctx: SessionContext,
+ tracked_memory_pool: Option<Arc<dyn TrackedPool>>,
+ ) -> Self {
+ Self {
+ ctx,
+ tracked_memory_pool,
  }
  }
+}
+
+#[async_trait::async_trait]
+impl CliSessionContext for ReplSessionContext {
+ fn task_ctx(&self) -> Arc<TaskContext> {
+ self.ctx.task_ctx()
+ }
+
+ fn session_state(&self) -> SessionState {
+ self.ctx.state()
+ }
+
+ fn register_object_store(
+ &self,
+ url: &url::Url,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Option<Arc<dyn ObjectStore + 'static>> {
+ self.ctx.register_object_store(url, object_store)
+ }
+
+ fn register_table_options_extension_from_scheme(&self, scheme: &str) {
+ register_table_options_from_scheme(&self.ctx, scheme);
+ }
 
  async fn execute_logical_plan(
  &self,
  plan: LogicalPlan,
  ) -> Result<DataFrame, DataFusionError> {
- self.execute_logical_plan(plan).await
+ self.ctx.execute_logical_plan(plan).await
+ }
+
+ fn memory_profiling(&self) -> bool {
+ self.tracked_memory_pool
+ .as_ref()
+ .map(|pool| pool.tracking_enabled())
+ .unwrap_or(false)
+ }
+
+ fn set_memory_profiling(&self, enable: bool) {
+ if let Some(pool) = &self.tracked_memory_pool {
+ if enable {
+ pool.enable_tracking();
+ } else {
+ pool.disable_tracking();
+ }
+ }
+ }
+
+ fn tracked_memory_pool(&self) -> Option<Arc<dyn TrackedPool>> {
+ self.tracked_memory_pool.clone()
  }
 }

datafusion-cli/src/command.rs

@@ -46,6 +46,7 @@ pub enum Command {
  SearchFunctions(String),
  QuietMode(Option<bool>),
  OutputFormat(Option<String>),
+ MemoryProfiling,
 }
 
 pub enum OutputFormat {
@@ -110,6 +111,15 @@ impl Command {
  }
  Ok(())
  }
+ Self::MemoryProfiling => {
+ let enable = !ctx.memory_profiling();
+ ctx.set_memory_profiling(enable);
+ println!(
+ "Memory profiling {}",
+ if enable { "enabled" } else { "disabled" }
+ );
+ Ok(())
+ }
  Self::Quit => exec_err!("Unexpected quit, this should be handled outside"),
  Self::ListFunctions => display_all_functions(),
  Self::SearchFunctions(function) => {
@@ -142,11 +152,15 @@ impl Command {
  Self::OutputFormat(_) => {
  ("\\pset [NAME [VALUE]]", "set table output option\n(format)")
  }
+ Self::MemoryProfiling => (
+ "\\memory_profiling",
+ "toggle memory profiling (requires --top-memory-consumers N at startup for metrics)",
+ ),
  }
  }
 }
 
-const ALL_COMMANDS: [Command; 9] = [
+const ALL_COMMANDS: [Command; 10] = [
  Command::ListTables,
  Command::DescribeTableStmt(String::new()),
  Command::Quit,
@@ -156,6 +170,7 @@ const ALL_COMMANDS: [Command; 9] = [
  Command::SearchFunctions(String::new()),
  Command::QuietMode(None),
  Command::OutputFormat(None),
+ Command::MemoryProfiling,
 ];
 
 fn all_commands_info() -> RecordBatch {
@@ -206,6 +221,8 @@ impl FromStr for Command {
  Self::OutputFormat(Some(subcommand.to_string()))
  }
  ("pset", None) => Self::OutputFormat(None),
+ ("memory_profiling", None) => Self::MemoryProfiling,
+ ("memory_profiling", Some(_)) => return Err(()),
  _ => return Err(()),
  })
  }

datafusion-cli/src/exec.rs

@@ -19,6 +19,7 @@
 
 use crate::cli_context::CliSessionContext;
 use crate::helper::split_from_semicolon;
+use crate::memory_metrics::format_metrics;
 use crate::print_format::PrintFormat;
 use crate::{
  command::{Command, OutputFormat},
@@ -313,6 +314,12 @@ impl StatementExecutor {
  )?;
  reservation.free();
  }
+ if ctx.memory_profiling() {
+ if let Some(pool) = ctx.tracked_memory_pool() {
+ let metrics = pool.consumer_metrics();
+ println!("{}", format_metrics(&metrics));
+ }
+ }
 
  Ok(())
  }

datafusion-cli/src/lib.rs

@@ -30,6 +30,7 @@ pub mod exec;
 pub mod functions;
 pub mod helper;
 pub mod highlighter;
+pub mod memory_metrics;
 pub mod object_storage;
 pub mod pool_type;
 pub mod print_format;

datafusion-cli/src/main.rs

@@ -25,7 +25,8 @@ use std::sync::{Arc, LazyLock};
 use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::context::SessionConfig;
 use datafusion::execution::memory_pool::{
- FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool,
+ FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool, TrackedPool,
+ UnboundedMemoryPool,
 };
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::prelude::SessionContext;
@@ -39,6 +40,8 @@ use datafusion_cli::{
  DATAFUSION_CLI_VERSION,
 };
 
+use datafusion_cli::cli_context::ReplSessionContext;
+
 use clap::Parser;
 use datafusion::common::config_err;
 use datafusion::config::ConfigOptions;
@@ -123,7 +126,7 @@ struct Args {
  #[clap(
  long,
  help = "The number of top memory consumers to display when query fails due to memory exhaustion. To disable memory consumer tracking, set this value to 0",
- default_value = "3"
+ default_value = "5"
  )]
  top_memory_consumers: usize,
 
@@ -174,28 +177,30 @@ async fn main_inner() -> Result<()> {
  let session_config = get_session_config(&args)?;
 
  let mut rt_builder = RuntimeEnvBuilder::new();
- // set memory pool size
- if let Some(memory_limit) = args.memory_limit {
- // set memory pool type
- let pool: Arc<dyn MemoryPool> = match args.mem_pool_type {
- PoolType::Fair if args.top_memory_consumers == 0 => {
- Arc::new(FairSpillPool::new(memory_limit))
- }
- PoolType::Fair => Arc::new(TrackConsumersPool::new(
- FairSpillPool::new(memory_limit),
- NonZeroUsize::new(args.top_memory_consumers).unwrap(),
- )),
- PoolType::Greedy if args.top_memory_consumers == 0 => {
- Arc::new(GreedyMemoryPool::new(memory_limit))
+
+ // set memory pool type
+ let base_memory_pool: Arc<dyn MemoryPool> =
+ if let Some(memory_limit) = args.memory_limit {
+ match args.mem_pool_type {
+ PoolType::Fair => Arc::new(FairSpillPool::new(memory_limit)),
+ PoolType::Greedy => Arc::new(GreedyMemoryPool::new(memory_limit)),
  }
- PoolType::Greedy => Arc::new(TrackConsumersPool::new(
- GreedyMemoryPool::new(memory_limit),
- NonZeroUsize::new(args.top_memory_consumers).unwrap(),
- )),
+ } else {
+ Arc::new(UnboundedMemoryPool::default())
  };
 
- rt_builder = rt_builder.with_memory_pool(pool)
- }
+ let tracked_pool: Option<Arc<dyn TrackedPool>> = if args.top_memory_consumers > 0 {
+ let tracked = Arc::new(TrackConsumersPool::new(
+ base_memory_pool.clone(),
+ NonZeroUsize::new(args.top_memory_consumers).unwrap(),
+ ));
+ tracked.disable_tracking();
+ rt_builder = rt_builder.with_memory_pool(tracked.clone());
+ Some(tracked as Arc<dyn TrackedPool>)
+ } else {
+ rt_builder = rt_builder.with_memory_pool(base_memory_pool);
+ None
+ };
 
  // set disk limit
  if let Some(disk_limit) = args.disk_limit {
@@ -208,24 +213,25 @@ async fn main_inner() -> Result<()> {
  let runtime_env = rt_builder.build_arc()?;
 
  // enable dynamic file query
- let ctx = SessionContext::new_with_config_rt(session_config, runtime_env)
+ let session_ctx = SessionContext::new_with_config_rt(session_config, runtime_env)
  .enable_url_table();
-ctx.refresh_catalogs().await?;
+session_ctx.refresh_catalogs().await?;
  // install dynamic catalog provider that can register required object stores
-ctx.register_catalog_list(Arc::new(DynamicObjectStoreCatalog::new(
-ctx.state().catalog_list().clone(),
-ctx.state_weak_ref(),
+session_ctx.register_catalog_list(Arc::new(DynamicObjectStoreCatalog::new(
+session_ctx.state().catalog_list().clone(),
+session_ctx.state_weak_ref(),
  )));
  // register `parquet_metadata` table function to get metadata from parquet files
- ctx.register_udtf("parquet_metadata", Arc::new(ParquetMetadataFunc {}));
-
+ session_ctx.register_udtf("parquet_metadata", Arc::new(ParquetMetadataFunc {}));
  // register `metadata_cache` table function to get the contents of the file metadata cache
-ctx.register_udtf(
+session_ctx.register_udtf(
  "metadata_cache",
  Arc::new(MetadataCacheFunc::new(
-ctx.task_ctx().runtime_env().cache_manager.clone(),
+session_ctx.task_ctx().runtime_env().cache_manager.clone(),
  )),
  );
+ // wrap the SessionContext in a REPL context (adds profiling, top consumers, etc.)
+ let ctx = ReplSessionContext::new(session_ctx, tracked_pool);
 
  let mut print_options = PrintOptions {
  format: args.format,