add top-memory-consumers option in cli

Author: ding-young

datafusion-cli/src/main.rs

@@ -17,13 +17,16 @@
 
 use std::collections::HashMap;
 use std::env;
+use std::num::NonZeroUsize;
 use std::path::Path;
 use std::process::ExitCode;
 use std::sync::{Arc, LazyLock};
 
 use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::context::SessionConfig;
-use datafusion::execution::memory_pool::{FairSpillPool, GreedyMemoryPool, MemoryPool};
+use datafusion::execution::memory_pool::{
+    FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool,
+};
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::execution::DiskManager;
 use datafusion::prelude::SessionContext;
@@ -118,6 +121,13 @@ struct Args {
     )]
     mem_pool_type: PoolType,
 
+    #[clap(
+        long,
+        help = "The number of top memory consumers to display when query fails due to memory exhaustion. If you don't want to track MemoryConsumers, set this value to 0",
+        default_value = "3"
+    )]
+    top_memory_consumers: usize,
+
     #[clap(
         long,
         help = "The max number of rows to display for 'Table' format\n[possible values: numbers(0/10/...), inf(no limit)]",
@@ -169,9 +179,22 @@ async fn main_inner() -> Result<()> {
     if let Some(memory_limit) = args.memory_limit {
         // set memory pool type
         let pool: Arc<dyn MemoryPool> = match args.mem_pool_type {
-            PoolType::Fair => Arc::new(FairSpillPool::new(memory_limit)),
-            PoolType::Greedy => Arc::new(GreedyMemoryPool::new(memory_limit)),
+            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))
+            }
+            PoolType::Greedy => Arc::new(TrackConsumersPool::new(
+                GreedyMemoryPool::new(memory_limit),
+                NonZeroUsize::new(args.top_memory_consumers).unwrap(),
+            )),
         };
+
         rt_builder = rt_builder.with_memory_pool(pool)
     }
 

docs/source/user-guide/cli/usage.md

@@ -57,6 +57,9 @@ OPTIONS:
         --mem-pool-type <MEM_POOL_TYPE>
             Specify the memory pool type 'greedy' or 'fair', default to 'greedy'
 
+        --top_memory_consumers <TOP_MEMORY_CONSUMERS>
+            The number of top memory consumers to display when query fails due to memory exhaustion. If you don't want to track MemoryConsumers, set this value to 0
+
     -d, --disk-limit <DISK_LIMIT>
             Available disk space for spilling queries (e.g. '10g'), default to None (uses DataFusion's default value of '100g')