feat: Improve datafusion-cli memory usage and considering reserve mem…

datafusion-cli/examples/cli-session-context.rs

@@ -88,6 +88,7 @@ pub async fn main() {
         format: datafusion_cli::print_format::PrintFormat::Automatic,
         quiet: false,
         maxrows: datafusion_cli::print_options::MaxRows::Unlimited,
+        stop_after_max_rows: false,
         color: true,
     };
 

datafusion-cli/src/exec.rs

@@ -17,11 +17,6 @@
 
 //! Execution functions
 
-use std::collections::HashMap;
-use std::fs::File;
-use std::io::prelude::*;
-use std::io::BufReader;
-
 use crate::cli_context::CliSessionContext;
 use crate::helper::split_from_semicolon;
 use crate::print_format::PrintFormat;
@@ -31,6 +26,11 @@ use crate::{
     object_storage::get_object_store,
     print_options::{MaxRows, PrintOptions},
 };
+use futures::StreamExt;
+use std::collections::HashMap;
+use std::fs::File;
+use std::io::prelude::*;
+use std::io::BufReader;
 
 use datafusion::common::instant::Instant;
 use datafusion::common::{plan_datafusion_err, plan_err};
@@ -39,10 +39,12 @@ use datafusion::datasource::listing::ListingTableUrl;
 use datafusion::error::{DataFusionError, Result};
 use datafusion::logical_expr::{DdlStatement, LogicalPlan};
 use datafusion::physical_plan::execution_plan::EmissionType;
-use datafusion::physical_plan::{collect, execute_stream, ExecutionPlanProperties};
+use datafusion::physical_plan::{execute_stream, ExecutionPlanProperties};
 use datafusion::sql::parser::{DFParser, Statement};
 use datafusion::sql::sqlparser::dialect::dialect_from_str;
 
+use datafusion::execution::memory_pool::MemoryConsumer;
+use datafusion::physical_plan::spill::get_record_batch_memory_size;
 use datafusion::sql::sqlparser;
 use rustyline::error::ReadlineError;
 use rustyline::Editor;
@@ -235,6 +237,10 @@ pub(super) async fn exec_and_print(
         let df = ctx.execute_logical_plan(plan).await?;
         let physical_plan = df.create_physical_plan().await?;
 
+        // Track memory usage for the query result if it's bounded
+        let mut reservation =
+            MemoryConsumer::new("DataFusion-Cli").register(task_ctx.memory_pool());
+
         if physical_plan.boundedness().is_unbounded() {
             if physical_plan.pipeline_behavior() == EmissionType::Final {
                 return plan_err!(
@@ -249,8 +255,21 @@ pub(super) async fn exec_and_print(
         } else {
             // Bounded stream; collected results are printed after all input consumed.
             let schema = physical_plan.schema();
-            let results = collect(physical_plan, task_ctx.clone()).await?;
+            let mut stream = execute_stream(physical_plan, task_ctx.clone())?;
+            let mut results = vec![];
+            while let Some(batch) = stream.next().await {
+                let batch = batch?;
+                reservation.try_grow(get_record_batch_memory_size(&batch))?;
+                results.push(batch);
+                if let MaxRows::Limited(max_rows) = print_options.maxrows {
+                    // Stop collecting results if the number of rows exceeds the limit
+                    if results.len() >= max_rows && print_options.stop_after_max_rows {
+                        break;
+                    }
+                }
+            }
             adjusted.into_inner().print_batches(schema, &results, now)?;
+            reservation.free();
         }
     }
 

datafusion-cli/src/main.rs

@@ -121,6 +121,13 @@ struct Args {
     )]
     maxrows: MaxRows,
 
+    #[clap(
+        short,
+        long,
+        help = "Whether to stop early when max rows is reached, this will help reduce the memory usage when the result is large"
+    )]
+    stop_after_max_rows: bool,
+
     #[clap(long, help = "Enables console syntax highlighting")]
     color: bool,
 }
@@ -186,6 +193,7 @@ async fn main_inner() -> Result<()> {
         quiet: args.quiet,
         maxrows: args.maxrows,
         color: args.color,
+        stop_after_max_rows: args.stop_after_max_rows,
     };
 
     let commands = args.command;

datafusion-cli/src/print_options.rs

@@ -72,6 +72,7 @@ pub struct PrintOptions {
     pub quiet: bool,
     pub maxrows: MaxRows,
     pub color: bool,
+    pub stop_after_max_rows: bool,
 }
 
 // Returns the query execution details formatted

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

@@ -54,6 +54,9 @@ OPTIONS:
             The max number of rows to display for 'Table' format
             [possible values: numbers(0/10/...), inf(no limit)] [default: 40]
 
+        --stop-after-max-rows
+            Enable to stop early when maxrows is reached, default to false.
+
         --mem-pool-type <MEM_POOL_TYPE>
             Specify the memory pool type 'greedy' or 'fair', default to 'greedy'