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,