Remove SortKeyCursor (~5-10% faster merge)

datafusion/core/src/physical_plan/sorts/cursor.rs

@@ -18,24 +18,15 @@
 use arrow::row::{Row, Rows};
 use std::cmp::Ordering;
 
-/// A `SortKeyCursor` is created from a `RecordBatch`, and a set of
-/// `PhysicalExpr` that when evaluated on the `RecordBatch` yield the sort keys.
-///
-/// Additionally it maintains a row cursor that can be advanced through the rows
-/// of the provided `RecordBatch`
-///
-/// `SortKeyCursor::compare` can then be used to compare the sort key pointed to
-/// by this row cursor, with that of another `SortKeyCursor`. A cursor stores
-/// a row comparator for each other cursor that it is compared to.
-pub struct SortKeyCursor {
-    stream_idx: usize,
+/// A [`Cursor`] for [`Rows`]
+pub struct RowCursor {
     cur_row: usize,
     num_rows: usize,
 
     rows: Rows,
 }
 
-impl std::fmt::Debug for SortKeyCursor {
+impl std::fmt::Debug for RowCursor {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         f.debug_struct("SortKeyCursor")
             .field("cur_row", &self.cur_row)
@@ -44,70 +35,39 @@ impl std::fmt::Debug for SortKeyCursor {
     }
 }
 
-impl SortKeyCursor {
+impl RowCursor {
     /// Create a new SortKeyCursor
-    pub fn new(stream_idx: usize, rows: Rows) -> Self {
+    pub fn new(rows: Rows) -> Self {
         Self {
-            stream_idx,
             cur_row: 0,
             num_rows: rows.num_rows(),
             rows,
         }
     }
 
-    #[inline(always)]
-    /// Return the stream index of this cursor
-    pub fn stream_idx(&self) -> usize {
-        self.stream_idx
-    }
-
-    #[inline(always)]
-    /// Return true if the stream is finished
-    pub fn is_finished(&self) -> bool {
-        self.num_rows == self.cur_row
-    }
-
-    #[inline(always)]
-    /// Returns the cursor's current row, and advances the cursor to the next row
-    pub fn advance(&mut self) -> usize {
-        assert!(!self.is_finished());
-        let t = self.cur_row;
-        self.cur_row += 1;
-        t
-    }
-
     /// Returns the current row
     fn current(&self) -> Row<'_> {
         self.rows.row(self.cur_row)
     }
 }
 
-impl PartialEq for SortKeyCursor {
+impl PartialEq for RowCursor {
     fn eq(&self, other: &Self) -> bool {
         self.current() == other.current()
     }
 }
 
-impl Eq for SortKeyCursor {}
+impl Eq for RowCursor {}
 
-impl PartialOrd for SortKeyCursor {
+impl PartialOrd for RowCursor {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
         Some(self.cmp(other))
     }
 }
 
-impl Ord for SortKeyCursor {
+impl Ord for RowCursor {
     fn cmp(&self, other: &Self) -> Ordering {
-        // Order finished cursors greater (last)
-        match (self.is_finished(), other.is_finished()) {
-            (true, true) => Ordering::Equal,
-            (_, true) => Ordering::Less,
-            (true, _) => Ordering::Greater,
-            _ => self
-                .current()
-                .cmp(&other.current())
-                .then_with(|| self.stream_idx.cmp(&other.stream_idx)),
-        }
+        self.current().cmp(&other.current())
     }
 }
 
@@ -117,17 +77,19 @@ pub trait Cursor: Ord {
     fn is_finished(&self) -> bool;
 
     /// Advance the cursor, returning the previous row index
-    ///
-    /// Returns `None` if [`Self::is_finished`]
-    fn advance(&mut self) -> Option<usize>;
+    fn advance(&mut self) -> usize;
 }
 
-impl Cursor for SortKeyCursor {
+impl Cursor for RowCursor {
+    #[inline]
     fn is_finished(&self) -> bool {
-        self.is_finished()
+        self.num_rows == self.cur_row
     }
 
-    fn advance(&mut self) -> Option<usize> {
-        (!self.is_finished()).then(|| self.advance())
+    #[inline]
+    fn advance(&mut self) -> usize {
+        let t = self.cur_row;
+        self.cur_row += 1;
+        t
     }
 }

datafusion/core/src/physical_plan/sorts/merge.rs

@@ -19,7 +19,7 @@ use crate::common::Result;
 use crate::physical_plan::metrics::MemTrackingMetrics;
 use crate::physical_plan::sorts::builder::BatchBuilder;
 use crate::physical_plan::sorts::cursor::Cursor;
-use crate::physical_plan::sorts::stream::{PartitionedStream, SortKeyCursorStream};
+use crate::physical_plan::sorts::stream::{PartitionedStream, RowCursorStream};
 use crate::physical_plan::{
     PhysicalSortExpr, RecordBatchStream, SendableRecordBatchStream,
 };
@@ -37,7 +37,7 @@ pub(crate) fn streaming_merge(
     tracking_metrics: MemTrackingMetrics,
     batch_size: usize,
 ) -> Result<SendableRecordBatchStream> {
-    let streams = SortKeyCursorStream::try_new(schema.as_ref(), expressions, streams)?;
+    let streams = RowCursorStream::try_new(schema.as_ref(), expressions, streams)?;
 
     Ok(Box::pin(SortPreservingMergeStream::new(
         Box::new(streams),
@@ -119,11 +119,7 @@ impl<C: Cursor> SortPreservingMergeStream<C> {
         cx: &mut Context<'_>,
         idx: usize,
     ) -> Poll<Result<()>> {
-        if self.cursors[idx]
-            .as_ref()
-            .map(|cursor| !cursor.is_finished())
-            .unwrap_or(false)
-        {
+        if self.cursors[idx].is_some() {
             // Cursor is not finished - don't need a new RecordBatch yet
             return Poll::Ready(Ok(()));
         }
@@ -176,8 +172,7 @@ impl<C: Cursor> SortPreservingMergeStream<C> {
             }
 
             let stream_idx = self.loser_tree[0];
-            let cursor = self.cursors[stream_idx].as_mut();
-            if cursor.and_then(Cursor::advance).is_some() {
+            if self.advance(stream_idx) {
                 self.loser_tree_adjusted = false;
                 self.in_progress.push_row(stream_idx);
                 if self.in_progress.len() < self.batch_size {
@@ -189,13 +184,27 @@ impl<C: Cursor> SortPreservingMergeStream<C> {
         }
     }
 
+    fn advance(&mut self, stream_idx: usize) -> bool {
+        let slot = &mut self.cursors[stream_idx];
+        match slot.as_mut() {
+            Some(c) => {
+                c.advance();
+                if c.is_finished() {
+                    *slot = None;
+                }
+                true
+            }
+            None => false,
+        }
+    }
+
     /// Returns `true` if the cursor at index `a` is greater than at index `b`
     #[inline]
     fn is_gt(&self, a: usize, b: usize) -> bool {
         match (&self.cursors[a], &self.cursors[b]) {
             (None, _) => true,
             (_, None) => false,
-            (Some(a), Some(b)) => b < a,
+            (Some(ac), Some(bc)) => ac.cmp(bc).then_with(|| a.cmp(&b)).is_gt(),
         }
     }
 

datafusion/core/src/physical_plan/sorts/mod.rs

@@ -25,6 +25,5 @@ pub mod sort;
 pub mod sort_preserving_merge;
 mod stream;
 
-pub use cursor::SortKeyCursor;
 pub use index::RowIndex;
 pub(crate) use merge::streaming_merge;

datafusion/core/src/physical_plan/sorts/stream.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::common::Result;
-use crate::physical_plan::sorts::cursor::SortKeyCursor;
+use crate::physical_plan::sorts::cursor::RowCursor;
 use crate::physical_plan::SendableRecordBatchStream;
 use crate::physical_plan::{PhysicalExpr, PhysicalSortExpr};
 use arrow::datatypes::Schema;
@@ -73,9 +73,9 @@ impl FusedStreams {
 }
 
 /// A [`PartitionedStream`] that wraps a set of [`SendableRecordBatchStream`]
-/// and computes [`SortKeyCursor`] based on the provided [`PhysicalSortExpr`]
+/// and computes [`RowCursor`] based on the provided [`PhysicalSortExpr`]
 #[derive(Debug)]
-pub(crate) struct SortKeyCursorStream {
+pub(crate) struct RowCursorStream {
     /// Converter to convert output of physical expressions
     converter: RowConverter,
     /// The physical expressions to sort by
@@ -84,7 +84,7 @@ pub(crate) struct SortKeyCursorStream {
     streams: FusedStreams,
 }
 
-impl SortKeyCursorStream {
+impl RowCursorStream {
     pub(crate) fn try_new(
         schema: &Schema,
         expressions: &[PhysicalSortExpr],
@@ -107,24 +107,20 @@ impl SortKeyCursorStream {
         })
     }
 
-    fn convert_batch(
-        &mut self,
-        batch: &RecordBatch,
-        stream_idx: usize,
-    ) -> Result<SortKeyCursor> {
+    fn convert_batch(&mut self, batch: &RecordBatch) -> Result<RowCursor> {
         let cols = self
             .column_expressions
             .iter()
             .map(|expr| Ok(expr.evaluate(batch)?.into_array(batch.num_rows())))
             .collect::<Result<Vec<_>>>()?;
 
         let rows = self.converter.convert_columns(&cols)?;
-        Ok(SortKeyCursor::new(stream_idx, rows))
+        Ok(RowCursor::new(rows))
     }
 }
 
-impl PartitionedStream for SortKeyCursorStream {
-    type Output = Result<(SortKeyCursor, RecordBatch)>;
+impl PartitionedStream for RowCursorStream {
+    type Output = Result<(RowCursor, RecordBatch)>;
 
     fn partitions(&self) -> usize {
         self.streams.0.len()
@@ -137,7 +133,7 @@ impl PartitionedStream for SortKeyCursorStream {
     ) -> Poll<Option<Self::Output>> {
         Poll::Ready(ready!(self.streams.poll_next(cx, stream_idx)).map(|r| {
             r.and_then(|batch| {
-                let cursor = self.convert_batch(&batch, stream_idx)?;
+                let cursor = self.convert_batch(&batch)?;
                 Ok((cursor, batch))
             })
         }))