You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/06/25 07:11:31 UTC
(datafusion) 01/01: Perf: cache primitive sort key in SortPreservingMerge to drop per-comparison bounds checks (#23162)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23162-a00f7499e139c29e13b8b94baba262435bef417b
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 284ae301c2a47b7480d67641982ddf772441b442
Author: Daniël Heres <da...@gmail.com>
AuthorDate: Thu Jun 25 09:11:09 2026 +0200
Perf: cache primitive sort key in SortPreservingMerge to drop per-comparison bounds checks (#23162)
## Which issue does this PR close?
<!-- No dedicated issue; this is a self-contained performance
improvement to the
sort-preserving merge hot path. Happy to file a tracking issue if
preferred. -->
- N/A (performance)
## Rationale for this change
Two inefficiencies in the hot path:
1. The single-column primitive/array cursor comparison was **not being
inlined**
(in profiles it appeared as a separate ~21% self-time symbol), and every
comparison
bounds-checked the underlying `ScalarBuffer` twice.
2. `maybe_poll_stream` was called on **every** output row, even though
it is a
no-op whenever the winner's cursor is still live (the common case).
```
┏━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ HEAD ┃ perf_spm-compare-cache ┃ Change ┃
┡━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ Q1 │ 165.13 / 166.08 ±0.75 / 166.98 ms │ 131.46 / 132.49 ±0.79 / 133.86 ms │ +1.25x faster │
│ Q2 │ 141.22 / 141.86 ±0.81 / 143.38 ms │ 118.06 / 120.08 ±1.30 / 121.91 ms │ +1.18x faster │
│ Q3 │ 652.56 / 657.09 ±2.91 / 661.10 ms │ 645.82 / 649.88 ±3.27 / 654.55 ms │ no change │
│ Q4 │ 195.87 / 199.80 ±6.83 / 213.43 ms │ 180.22 / 183.86 ±5.04 / 193.43 ms │ +1.09x faster │
│ Q5 │ 279.14 / 279.91 ±0.50 / 280.71 ms │ 260.15 / 260.70 ±0.48 / 261.54 ms │ +1.07x faster │
│ Q6 │ 292.50 / 293.43 ±0.71 / 294.56 ms │ 273.28 / 274.77 ±1.66 / 277.81 ms │ +1.07x faster │
│ Q7 │ 465.87 / 467.23 ±2.01 / 471.21 ms │ 445.83 / 449.21 ±3.47 / 455.33 ms │ no change │
│ Q8 │ 327.51 / 330.08 ±3.25 / 336.40 ms │ 319.12 / 325.35 ±6.20 / 336.30 ms │ no change │
│ Q9 │ 342.17 / 346.17 ±2.68 / 348.71 ms │ 336.40 / 348.61 ±13.37 / 368.15 ms │ no change │
│ Q10 │ 484.08 / 486.50 ±2.53 / 490.75 ms │ 474.15 / 490.16 ±13.58 / 507.56 ms │ no change │
│ Q11 │ 244.24 / 250.98 ±8.70 / 267.17 ms │ 229.07 / 237.38 ±8.10 / 249.40 ms │ +1.06x faster │
└───────┴───────────────────────────────────┴────────────────────────────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━┩
│ Total Time (HEAD) │ 3619.13ms │
│ Total Time (perf_spm-compare-cache) │ 3472.49ms │
│ Average Time (HEAD) │ 329.01ms │
│ Average Time (perf_spm-compare-cache) │ 315.68ms │
│ Queries Faster │ 6 │
│ Queries Slower │ 0 │
│ Queries with No Change │ 5 │
│ Queries with Failure │ 0 │
└───────────────────────────────────────┴───────────┘
```
## What changes are included in this PR?
- Inline the lightweight primitive/array cursor comparisons.
- Skip the per-row `maybe_poll_stream` call unless the winner's cursor
is
actually exhausted and needs a fresh `RecordBatch`.
- Cache the current (and previous) value of a primitive cursor,
refreshed once
per `advance()` via a new `CursorValues::set_offset` hook (default
no-op;
`ArrayValues` forwards to the inner cursor).
## Are these changes tested?
Existing tests
## Are there any user-facing changes?
No.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 4.8 <no...@anthropic.com>
---
datafusion/physical-plan/src/sorts/cursor.rs | 101 ++++++++++++++++++++++++---
datafusion/physical-plan/src/sorts/merge.rs | 14 +++-
2 files changed, 104 insertions(+), 11 deletions(-)
diff --git a/datafusion/physical-plan/src/sorts/cursor.rs b/datafusion/physical-plan/src/sorts/cursor.rs
index 288ec4cee1..8991922779 100644
--- a/datafusion/physical-plan/src/sorts/cursor.rs
+++ b/datafusion/physical-plan/src/sorts/cursor.rs
@@ -44,6 +44,14 @@ pub trait CursorValues {
/// Returns comparison of `l[l_idx]` and `r[r_idx]`
fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering;
+
+ /// 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).
+ #[inline]
+ fn set_offset(&mut self, offset: usize) {
+ let _ = offset;
+ }
}
/// A comparable cursor, used by sort operations
@@ -89,14 +97,22 @@ 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()
}
/// Advance the cursor, returning the previous row index
+ #[inline]
pub fn advance(&mut self) -> usize {
let t = self.offset;
self.offset += 1;
+ // 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).
+ if self.offset < self.values.len() {
+ self.values.set_offset(self.offset);
+ }
t
}
@@ -112,6 +128,7 @@ impl<T: CursorValues> Cursor<T> {
}
impl<T: CursorValues> PartialEq for Cursor<T> {
+ #[inline]
fn eq(&self, other: &Self) -> bool {
T::eq(&self.values, self.offset, &other.values, other.offset)
}
@@ -142,6 +159,7 @@ impl<T: CursorValues> PartialOrd for Cursor<T> {
}
impl<T: CursorValues> Ord for Cursor<T> {
+ #[inline]
fn cmp(&self, other: &Self) -> Ordering {
T::compare(&self.values, self.offset, &other.values, other.offset)
}
@@ -180,10 +198,14 @@ impl RowValues {
}
impl CursorValues for RowValues {
+ #[inline]
fn len(&self) -> usize {
self.rows.num_rows()
}
+ // No inline hint on purpose: for the heavyweight `Rows` byte comparison the
+ // compiler's own choice wins — both `#[inline]` and `#[inline(never)]`
+ // measurably regress the multi-column merge path.
fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool {
l.rows.row(l_idx) == r.rows.row(r_idx)
}
@@ -209,29 +231,74 @@ impl<T: ArrowPrimitiveType> CursorArray for PrimitiveArray<T> {
type Values = PrimitiveValues<T::Native>;
fn values(&self) -> Self::Values {
- PrimitiveValues(self.values().clone())
+ PrimitiveValues::new(self.values().clone())
}
}
+/// [`CursorValues`] for a primitive column.
+///
+/// Caches the value at the current (and previous) offset, refreshed once per
+/// [`Cursor::advance`] via [`CursorValues::set_offset`], so the hot loser-tree
+/// comparisons read a cached field instead of indexing the buffer each time.
#[derive(Debug)]
-pub struct PrimitiveValues<T: ArrowNativeTypeOp>(ScalarBuffer<T>);
+pub struct PrimitiveValues<T: ArrowNativeTypeOp> {
+ values: ScalarBuffer<T>,
+ /// Cached `values[offset]`.
+ current: T,
+ /// Cached `values[offset - 1]` (read by `eq_to_previous`, only past offset 0).
+ previous: T,
+ /// Current offset; used only to `debug_assert!` the cache is read in sync.
+ offset: usize,
+}
+
+impl<T: ArrowNativeTypeOp> PrimitiveValues<T> {
+ fn new(values: ScalarBuffer<T>) -> Self {
+ // Non-empty in practice; `unwrap_or_default` just avoids a panic.
+ let first = values.first().copied().unwrap_or_default();
+ Self {
+ values,
+ current: first,
+ previous: first,
+ offset: 0,
+ }
+ }
+}
impl<T: ArrowNativeTypeOp> CursorValues for PrimitiveValues<T> {
+ #[inline(always)]
fn len(&self) -> usize {
- self.0.len()
+ self.values.len()
}
+ #[inline(always)]
fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool {
- l.0[l_idx].is_eq(r.0[r_idx])
+ // Arbitrary indices (cross-batch comparison), so index directly.
+ l.values[l_idx].is_eq(r.values[r_idx])
}
+ #[inline(always)]
fn eq_to_previous(cursor: &Self, idx: usize) -> bool {
assert!(idx > 0);
- cursor.0[idx].is_eq(cursor.0[idx - 1])
+ debug_assert_eq!(idx, cursor.offset);
+ cursor.current.is_eq(cursor.previous)
}
+ #[inline(always)]
fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering {
- l.0[l_idx].compare(r.0[r_idx])
+ debug_assert_eq!(l_idx, l.offset);
+ debug_assert_eq!(r_idx, r.offset);
+ l.current.compare(r.current)
+ }
+
+ #[inline(always)]
+ fn set_offset(&mut self, offset: usize) {
+ // The caller (`Cursor::advance`) guarantees `offset < len`; inlined, that
+ // guard dominates the index below so its bounds check is elided — the
+ // length is checked once per row, not per comparison. The old `current`
+ // is `values[offset - 1]`, so it becomes `previous`.
+ self.previous = self.current;
+ self.current = self.values[offset];
+ self.offset = offset;
}
}
@@ -241,6 +308,7 @@ pub struct ByteArrayValues<T: OffsetSizeTrait> {
}
impl<T: OffsetSizeTrait> ByteArrayValues<T> {
+ #[inline]
fn value(&self, idx: usize) -> &[u8] {
assert!(idx < self.len());
// Safety: offsets are valid and checked bounds above
@@ -253,19 +321,23 @@ impl<T: OffsetSizeTrait> ByteArrayValues<T> {
}
impl<T: OffsetSizeTrait> CursorValues for ByteArrayValues<T> {
+ #[inline]
fn len(&self) -> usize {
self.offsets.len() - 1
}
+ #[inline]
fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool {
l.value(l_idx) == r.value(r_idx)
}
+ #[inline]
fn eq_to_previous(cursor: &Self, idx: usize) -> bool {
assert!(idx > 0);
cursor.value(idx) == cursor.value(idx - 1)
}
+ #[inline]
fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering {
l.value(l_idx).cmp(r.value(r_idx))
}
@@ -394,16 +466,19 @@ impl<T: CursorValues> ArrayValues<T> {
}
}
+ #[inline(always)]
fn is_null(&self, idx: usize) -> bool {
(idx < self.null_threshold) == self.options.nulls_first
}
}
impl<T: CursorValues> CursorValues for ArrayValues<T> {
+ #[inline(always)]
fn len(&self) -> usize {
self.values.len()
}
+ #[inline(always)]
fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool {
match (l.is_null(l_idx), r.is_null(r_idx)) {
(true, true) => true,
@@ -412,15 +487,19 @@ impl<T: CursorValues> CursorValues for ArrayValues<T> {
}
}
+ #[inline(always)]
fn eq_to_previous(cursor: &Self, idx: usize) -> bool {
assert!(idx > 0);
match (cursor.is_null(idx), cursor.is_null(idx - 1)) {
(true, true) => true,
- (false, false) => T::eq(&cursor.values, idx, &cursor.values, idx - 1),
+ // Delegate to inner `eq_to_previous` so a caching cursor can answer
+ // without indexing.
+ (false, false) => T::eq_to_previous(&cursor.values, idx),
_ => false,
}
}
+ #[inline(always)]
fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering {
match (l.is_null(l_idx), r.is_null(r_idx)) {
(true, true) => Ordering::Equal,
@@ -438,6 +517,12 @@ impl<T: CursorValues> CursorValues for ArrayValues<T> {
},
}
}
+
+ #[inline(always)]
+ fn set_offset(&mut self, offset: usize) {
+ // Forward to the wrapped values (e.g. caching `PrimitiveValues`).
+ self.values.set_offset(offset);
+ }
}
#[cfg(test)]
@@ -463,7 +548,7 @@ mod tests {
let reservation = consumer.register(&memory_pool);
let values = ArrayValues {
- values: PrimitiveValues(values),
+ values: PrimitiveValues::new(values),
null_threshold,
options,
_reservation: reservation,
diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs
index ad2b529790..4583d19e91 100644
--- a/datafusion/physical-plan/src/sorts/merge.rs
+++ b/datafusion/physical-plan/src/sorts/merge.rs
@@ -294,9 +294,17 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
// Adjust the loser tree if necessary, returning control if needed
if !self.loser_tree_adjusted {
let winner = self.loser_tree[0];
- if let Err(e) = ready!(self.maybe_poll_stream(cx, winner)) {
- self.done = true;
- return Poll::Ready(Some(Err(e)));
+ // Fast path: skip the `maybe_poll_stream` call (and its `Poll`
+ // plumbing) unless the winner's cursor is exhausted and needs a
+ // fresh batch — it is live for almost every row.
+ if self.cursors[winner].is_none() {
+ match ready!(self.maybe_poll_stream(cx, winner)) {
+ Ok(()) => {}
+ Err(e) => {
+ self.done = true;
+ return Poll::Ready(Some(Err(e)));
+ }
+ }
}
self.update_loser_tree();
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org