Addresss comments

Author: zhuqi-lucas

datafusion/physical-plan/src/sorts/sort.rs

@@ -24,7 +24,6 @@ use std::fmt;
 use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
 
-use crate::common::spawn_buffered;
 use crate::execution_plan::{Boundedness, CardinalityEffect, EmissionType};
 use crate::expressions::PhysicalSortExpr;
 use crate::limit::LimitStream;
@@ -687,23 +686,35 @@ impl ExternalSorter {
         let mut current_batches = Vec::new();
         let mut current_size = 0;
 
-        for batch in std::mem::take(&mut self.in_mem_batches) {
+        // Drain in_mem_batches using pop() to release memory earlier.
+        // This avoids holding onto the entire vector during iteration.
+        // Note:
+        // Now we use `sort_in_place_threshold_bytes` to determine, in future we can make it more dynamic.
+        while let Some(batch) = self.in_mem_batches.pop() {
             let batch_size = get_reserved_byte_for_record_batch(&batch);
+
+            // If adding this batch would exceed the memory threshold, merge current_batches.
             if current_size + batch_size > self.sort_in_place_threshold_bytes
                 && !current_batches.is_empty()
             {
+                // Merge accumulated batches into one.
                 let merged = concat_batches(&self.schema, &current_batches)?;
                 current_batches.clear();
+
+                // Update memory reservation.
                 self.reservation.try_shrink(current_size)?;
                 let merged_size = get_reserved_byte_for_record_batch(&merged);
                 self.reservation.try_grow(merged_size)?;
+
                 merged_batches.push(merged);
                 current_size = 0;
             }
+
             current_batches.push(batch);
             current_size += batch_size;
         }
 
+        // Merge any remaining batches after the loop.
         if !current_batches.is_empty() {
             let merged = concat_batches(&self.schema, &current_batches)?;
             self.reservation.try_shrink(current_size)?;
@@ -712,15 +723,19 @@ impl ExternalSorter {
             merged_batches.push(merged);
         }
 
+        // Create sorted streams directly without using spawn_buffered.
+        // This allows for sorting to happen inline and enables earlier batch drop.
         let streams = merged_batches
             .into_iter()
             .map(|batch| {
                 let metrics = self.metrics.baseline.intermediate();
                 let reservation = self
                     .reservation
                     .split(get_reserved_byte_for_record_batch(&batch));
+
+                // Sort the batch inline.
                 let input = self.sort_batch_stream(batch, metrics, reservation)?;
-                Ok(spawn_buffered(input, 1))
+                Ok(input)
             })
             .collect::<Result<_>>()?;