[IGNORE] run bench for sort

Cargo.toml

@@ -159,6 +159,7 @@ datafusion-session = { path = "datafusion/session", version = "54.0.0" }
 datafusion-spark = { path = "datafusion/spark", version = "54.0.0" }
 datafusion-sql = { path = "datafusion/sql", version = "54.0.0" }
 datafusion-substrait = { path = "datafusion/substrait", version = "54.0.0" }
+genawaiter = { version = "0.99.1", default-features = false }
 
 doc-comment = "0.3"
 env_logger = "0.11"

datafusion/physical-plan/Cargo.toml

@@ -78,6 +78,7 @@ datafusion-physical-expr-common = { workspace = true }
 datafusion-proto-common = { workspace = true, optional = true }
 datafusion-proto-models = { workspace = true, optional = true }
 futures = { workspace = true }
+genawaiter = { workspace = true, features = ["futures03"] }
 half = { workspace = true }
 hashbrown = { workspace = true }
 indexmap = { workspace = true }

datafusion/physical-plan/src/coalesce/mod.rs

@@ -68,6 +68,18 @@ impl LimitedBatchCoalescer {
         }
     }
 
+    pub fn with_biggest_coalesce_batch_size(
+        self,
+        biggest_coalesce_batch_size: Option<usize>,
+    ) -> Self {
+        Self {
+            inner: self
+                .inner
+                .with_biggest_coalesce_batch_size(biggest_coalesce_batch_size),
+            ..self
+        }
+    }
+
     /// Return the schema of the output batches
     pub fn schema(&self) -> SchemaRef {
         self.inner.schema()

datafusion/physical-plan/src/coalesce_batches.rs

@@ -200,16 +200,12 @@ impl ExecutionPlan for CoalesceBatchesExec {
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
-        Ok(Box::pin(CoalesceBatchesStream {
-            input: self.input.execute(partition, context)?,
-            coalescer: LimitedBatchCoalescer::new(
-                self.input.schema(),
-                self.target_batch_size,
-                self.fetch,
-            ),
-            baseline_metrics: BaselineMetrics::new(&self.metrics, partition),
-            completed: false,
-        }))
+        Ok(Box::pin(CoalesceBatchesStream::new(
+            self.input.execute(partition, context)?,
+            self.target_batch_size,
+            self.fetch,
+            BaselineMetrics::new(&self.metrics, partition),
+        )))
     }
 
     fn metrics(&self) -> Option<MetricsSet> {
@@ -287,7 +283,7 @@ impl ExecutionPlan for CoalesceBatchesExec {
 }
 
 /// Stream for [`CoalesceBatchesExec`]. See [`CoalesceBatchesExec`] for more details.
-struct CoalesceBatchesStream {
+pub(crate) struct CoalesceBatchesStream {
     /// The input plan
     input: SendableRecordBatchStream,
     /// Buffer for combining batches
@@ -316,6 +312,36 @@ impl Stream for CoalesceBatchesStream {
 }
 
 impl CoalesceBatchesStream {
+    pub(crate) fn new(
+        input: SendableRecordBatchStream,
+        target_batch_size: usize,
+        fetch: Option<usize>,
+        baseline_metrics: BaselineMetrics,
+    ) -> Self {
+        CoalesceBatchesStream {
+            coalescer: LimitedBatchCoalescer::new(
+                input.schema(),
+                target_batch_size,
+                fetch,
+            ),
+            input,
+            baseline_metrics,
+            completed: false,
+        }
+    }
+
+    pub(crate) fn with_biggest_coalesce_batch_size(
+        self,
+        biggest_coalesce_batch_size: Option<usize>,
+    ) -> Self {
+        Self {
+            coalescer: self
+                .coalescer
+                .with_biggest_coalesce_batch_size(biggest_coalesce_batch_size),
+            ..self
+        }
+    }
+
     fn poll_next_inner(
         self: &mut Pin<&mut Self>,
         cx: &mut Context<'_>,

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

@@ -114,6 +114,16 @@ impl BatchBuilder {
         self.indices.push((cursor.batch_idx, row_idx));
     }
 
+    /// Append the next `n` rows from `stream_idx`
+    pub fn push_n_rows(&mut self, stream_idx: usize, n: usize) {
+        let cursor = &mut self.cursors[stream_idx];
+        let row_idx = cursor.row_idx;
+
+        cursor.row_idx += n;
+        self.indices
+            .extend((0..n).map(|i| (cursor.batch_idx, row_idx + i)));
+    }
+
     /// Returns the number of in-progress rows in this [`BatchBuilder`]
     pub fn len(&self) -> usize {
         self.indices.len()
@@ -205,12 +215,28 @@ impl BatchBuilder {
     ///
     /// Returns `None` if no pending rows
     pub fn build_record_batch(&mut self) -> Result<Option<RecordBatch>> {
-        if self.is_empty() {
+        self.build_record_batch_with_up_to_n_rows(self.indices.len())
+    }
+
+    /// Drains the in_progress row indexes up to n rows, and builds a new RecordBatch from them
+    ///
+    /// Will then drop any batches for which all rows have been yielded to the output.
+    /// If an offset overflow occurs (e.g. string/list offsets exceed i32::MAX),
+    /// retries with progressively fewer rows until it succeeds.
+    ///
+    /// Returns `None` if no pending rows
+    pub(crate) fn build_record_batch_with_up_to_n_rows(
+        &mut self,
+        n: usize,
+    ) -> Result<Option<RecordBatch>> {
+        if self.is_empty() || n == 0 {
             return Ok(None);
         }
 
+        let batch_size = self.indices.len().min(n);
+
         let (rows_to_emit, columns) =
-            retry_interleave(self.indices.len(), self.indices.len(), |rows_to_emit| {
+            retry_interleave(batch_size, batch_size, |rows_to_emit| {
                 self.try_interleave_columns(&self.indices[..rows_to_emit])
             })?;