[IGNORE] run bench for sort

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> {
@@ -293,7 +289,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
@@ -322,6 +318,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()
@@ -197,20 +207,22 @@ impl BatchBuilder {
         RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(Into::into)
     }
 
-    /// Drains the in_progress row indexes, and builds a new RecordBatch from them
+    /// Drains up to `n` in_progress row indexes, 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 fn build_record_batch(&mut self) -> Result<Option<RecordBatch>> {
-        if self.is_empty() {
+    pub fn build_record_batch(&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])
             })?;
 

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

@@ -44,8 +44,18 @@ pub trait CursorValues: Debug + Sync + Send {
     fn eq_to_previous(cursor: &Self, idx: usize) -> bool;
 
     /// Returns comparison of `l[l_idx]` and `r[r_idx]`
+    ///
+    /// Only called with the cursors' *current* offsets, so caching
+    /// implementations may serve it from their cache.
     fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering;
 
+    /// Returns comparison of `l[l_idx]` and `r[r_idx]` for *arbitrary*
+    /// indices, unlike [`Self::compare`]. Cold path: caching implementations
+    /// must override this to index directly.
+    fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering {
+        Self::compare(l, l_idx, r, r_idx)
+    }
+
     /// Notifies the values that the owning [`Cursor`] moved to `offset` (always
     /// `< len()`), so caching implementations can refresh the value(s) read by
     /// the hot comparisons. Default no-op (e.g. byte/row cursors don't benefit).
@@ -96,14 +106,19 @@ impl<T: CursorValues> Cursor<T> {
     /// Returns true if there are no more rows in this cursor
     #[inline]
     pub fn is_finished(&self) -> bool {
-        self.offset == self.values.len()
+        self.len() == 0
     }
 
     /// Advance the cursor, returning the previous row index
     #[inline]
     pub fn advance(&mut self) -> usize {
+        self.advance_n(1)
+    }
+
+    #[inline]
+    pub(crate) fn advance_n(&mut self, n: usize) -> usize {
         let t = self.offset;
-        self.offset += 1;
+        self.offset += n;
         // Refresh the cache for the new position. The guard keeps `set_offset`
         // in bounds; a finished cursor's stale cache is never read (it is taken
         // before the next comparison).
@@ -113,6 +128,18 @@ impl<T: CursorValues> Cursor<T> {
         t
     }
 
+    /// Compare this cursor's *last* row to `other`'s current row. Since the
+    /// values are sorted, `Less` means every remaining row of this cursor
+    /// sorts before all of `other`'s remaining rows.
+    pub fn compare_last_to(&self, other: &Self) -> Ordering {
+        T::compare_at(
+            &self.values,
+            self.values.len() - 1,
+            &other.values,
+            other.offset,
+        )
+    }
+
     pub fn is_eq_to_prev_one(&self, prev_cursor: Option<&Cursor<T>>) -> bool {
         if self.offset > 0 {
             self.is_eq_to_prev_row()
@@ -122,6 +149,10 @@ impl<T: CursorValues> Cursor<T> {
             false
         }
     }
+
+    pub(crate) fn len(&self) -> usize {
+        self.values.len() - self.offset
+    }
 }
 
 impl<T: CursorValues> PartialEq for Cursor<T> {
@@ -287,6 +318,11 @@ impl<T: ArrowNativeTypeOp> CursorValues for PrimitiveValues<T> {
         l.current.compare(r.current)
     }
 
+    fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering {
+        // Arbitrary indices (cold path), so index directly instead of using the cache.
+        l.values[l_idx].compare(r.values[r_idx])
+    }
+
     #[inline(always)]
     fn set_offset(&mut self, offset: usize) {
         // The caller (`Cursor::advance`) guarantees `offset < len`; inlined, that
@@ -516,6 +552,26 @@ impl<T: CursorValues> CursorValues for ArrayValues<T> {
         }
     }
 
+    fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering {
+        // Same null handling as `compare`, but delegates to the inner
+        // `compare_at` so arbitrary indices work with caching values.
+        match (l.is_null(l_idx), r.is_null(r_idx)) {
+            (true, true) => Ordering::Equal,
+            (true, false) => match l.options.nulls_first {
+                true => Ordering::Less,
+                false => Ordering::Greater,
+            },
+            (false, true) => match l.options.nulls_first {
+                true => Ordering::Greater,
+                false => Ordering::Less,
+            },
+            (false, false) => match l.options.descending {
+                true => T::compare_at(&r.values, r_idx, &l.values, l_idx),
+                false => T::compare_at(&l.values, l_idx, &r.values, r_idx),
+            },
+        }
+    }
+
     #[inline(always)]
     fn set_offset(&mut self, offset: usize) {
         // Forward to the wrapped values (e.g. caching `PrimitiveValues`).