Skip to content

Commit ea890b0

Browse files
committed
perf(sort-merge): cache current-row bytes in RowValues for SortPreservingMerge
Multi-column sorts serialize the sort key into arrow `Rows` and wrap it in `RowValues`. Every loser-tree comparison called `Rows::row(idx)` for each side, and each call walks `Arc<Rows>` -> offsets[idx] / offsets[idx+1] -> buffer slice. The offsets buffer is far too large to stay cache-resident, so those lookups typically cost an L2/L3 hit with a DRAM tail — repeated for the same idx on every compare of a stable loser-tree head. Cache the current row's `(ptr, len)` once per `Cursor::advance` (via the existing `CursorValues::set_offset` hook) and read it in `compare`, so the hot path resolves to two plain field loads plus a memcmp. The pointer is into the Arc-owned buffer heap and stays valid across struct moves (the cursor is written into a `Vec<Option<Cursor<..>>>` slot), so `Send`/`Sync` are implemented by hand with a SAFETY note. `eq` / `eq_to_previous` take arbitrary cross-batch indices and continue to index `Rows` directly. This is the first of a series splitting apart the SPM cursor-cache work for easier review; it carries the largest share of the win (multi-column sort-tpch). Follow-ups add the same pattern to the single-column string cursors and a null-wrapper fast path. sort_tpch10 benchmark on the full series showed the multi-column queries (Q4 +1.23x, Q8 +1.13x, Q9 +1.15x, Q5/Q6/Q11 ~+1.12x) driven by this cache. Tests: `test_row_values_cache_matches_rows_index` checks the cached ordering against per-row `Rows` indexing across every offset of a batch, including the cross-batch `eq` path; `test_row_values_single_row_batch` covers the up-front row-0 cache. Existing `sorts::*` merge tests (83) pass.
1 parent 2e3626e commit ea890b0

1 file changed

Lines changed: 122 additions & 4 deletions

File tree

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

Lines changed: 122 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -164,16 +164,35 @@ impl<T: CursorValues> Ord for Cursor<T> {
164164

165165
/// Implements [`CursorValues`] for [`Rows`]
166166
///
167-
/// Used for sorting when there are multiple columns in the sort key
167+
/// Used for sorting when there are multiple columns in the sort key.
168+
///
169+
/// Caches `(ptr, len)` for the current row's serialized bytes so the merge hot
170+
/// path compares two `&[u8]` slices directly rather than paying a
171+
/// `Rows::row(idx)` offset lookup for each side of each compare. The pointer
172+
/// is into `rows`'s Arc-owned buffer heap, so it stays valid even when this
173+
/// struct is moved (e.g. written into a `Vec<Option<Cursor<..>>>` slot).
168174
#[derive(Debug)]
169175
pub struct RowValues {
170176
rows: Arc<Rows>,
171177

178+
/// Number of rows — snapshot of `rows.num_rows()`. Read on every
179+
/// `Cursor::is_finished` / `advance` call.
180+
len: usize,
181+
/// Cached byte slice pointer for the current row.
182+
current_ptr: *const u8,
183+
/// Cached byte length for the current row.
184+
current_len: usize,
185+
172186
/// Tracks for the memory used by in the `Rows` of this
173187
/// cursor. Freed on drop
174188
_reservation: MemoryReservation,
175189
}
176190

191+
// SAFETY: `current_ptr` points into `rows`'s Arc-owned buffer heap. `Rows`
192+
// is `Send + Sync`; the referenced bytes are read-only after construction.
193+
unsafe impl Send for RowValues {}
194+
unsafe impl Sync for RowValues {}
195+
177196
impl RowValues {
178197
/// Create a new [`RowValues`] from `rows` and a `reservation`
179198
/// that tracks its memory. There must be at least one row
@@ -186,24 +205,45 @@ impl RowValues {
186205
reservation.size(),
187206
"memory reservation mismatch"
188207
);
189-
assert!(rows.num_rows() > 0);
208+
let len = rows.num_rows();
209+
assert!(len > 0);
210+
// Extract raw ptr + length while the temporary `Row` is still alive.
211+
// The pointer is into `rows`'s Arc buffer heap and stays valid.
212+
let (current_ptr, current_len) = {
213+
let row = rows.row(0);
214+
let bytes: &[u8] = row.as_ref();
215+
(bytes.as_ptr(), bytes.len())
216+
};
190217
Self {
191218
rows,
219+
len,
220+
current_ptr,
221+
current_len,
192222
_reservation: reservation,
193223
}
194224
}
225+
226+
#[inline(always)]
227+
fn current_slice(&self) -> &[u8] {
228+
// SAFETY: `set_offset` (or `new` for offset 0) populated `current_ptr`
229+
// / `current_len` from `rows.row(offset).as_ref()`, and the ptr is
230+
// into `rows`'s Arc heap that stays alive as long as `self` does.
231+
unsafe { std::slice::from_raw_parts(self.current_ptr, self.current_len) }
232+
}
195233
}
196234

197235
impl CursorValues for RowValues {
198236
#[inline]
199237
fn len(&self) -> usize {
200-
self.rows.num_rows()
238+
self.len
201239
}
202240

203241
// No inline hint on purpose: for the heavyweight `Rows` byte comparison the
204242
// compiler's own choice wins — both `#[inline]` and `#[inline(never)]`
205243
// measurably regress the multi-column merge path.
206244
fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool {
245+
// Arbitrary indices (cross-batch); can't use the cache which only
246+
// holds the current offset.
207247
l.rows.row(l_idx) == r.rows.row(r_idx)
208248
}
209249

@@ -213,7 +253,26 @@ impl CursorValues for RowValues {
213253
}
214254

215255
fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering {
216-
l.rows.row(l_idx).cmp(&r.rows.row(r_idx))
256+
// Merge callers always compare at current offsets; the cache is up
257+
// to date. (Debug-only: verify the invariant.)
258+
debug_assert!(l_idx < l.len && r_idx < r.len);
259+
let _ = (l_idx, r_idx);
260+
l.current_slice().cmp(r.current_slice())
261+
}
262+
263+
#[inline(always)]
264+
fn set_offset(&mut self, offset: usize) {
265+
// Refresh the cached byte-slice for the new row. Caller guarantees
266+
// `offset < len`. `Rows::row(idx).as_ref()` returns `&[u8]` into the
267+
// Arc-owned buffer heap, so the pointer we stow stays valid after
268+
// the temporary `Row` drops.
269+
let (ptr, len) = {
270+
let row = self.rows.row(offset);
271+
let bytes: &[u8] = row.as_ref();
272+
(bytes.as_ptr(), bytes.len())
273+
};
274+
self.current_ptr = ptr;
275+
self.current_len = len;
217276
}
218277
}
219278

@@ -555,6 +614,65 @@ mod tests {
555614
Cursor::new(values)
556615
}
557616

617+
/// Builds a `RowValues` cursor from a single string column, so tests can
618+
/// drive the multi-column `Rows` path with concrete data.
619+
fn new_row_values(strings: &[&str]) -> Cursor<RowValues> {
620+
use arrow::array::{ArrayRef, StringArray};
621+
use arrow::datatypes::DataType;
622+
use arrow::row::{RowConverter, SortField};
623+
624+
let array: ArrayRef = Arc::new(StringArray::from(strings.to_vec()));
625+
let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
626+
let rows = converter.convert_columns(&[array]).unwrap();
627+
628+
let memory_pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1_000_000));
629+
let consumer = MemoryConsumer::new("test");
630+
let mut reservation = consumer.register(&memory_pool);
631+
reservation.grow(rows.size());
632+
633+
Cursor::new(RowValues::new(Arc::new(rows), reservation))
634+
}
635+
636+
/// The `(current_ptr, current_len)` cache refreshed on `advance` must yield
637+
/// exactly the same ordering as indexing the underlying `Rows` per compare.
638+
/// Drives the cache across every offset of a multi-row batch.
639+
#[test]
640+
fn test_row_values_cache_matches_rows_index() {
641+
// Deliberately unsorted so the comparisons exercise <, >, and ==.
642+
let a = new_row_values(&["banana", "apple", "cherry", "apple"]);
643+
let b = new_row_values(&["apricot", "apple", "blueberry", "date"]);
644+
645+
// Reference comparison straight off the arrow `Rows`, no cache.
646+
let expected: Vec<Ordering> = (0..4)
647+
.map(|i| a.values.rows.row(i).cmp(&b.values.rows.row(i)))
648+
.collect();
649+
650+
let mut a = a;
651+
let mut b = b;
652+
let mut got = Vec::with_capacity(4);
653+
for _ in 0..4 {
654+
got.push(a.cmp(&b));
655+
a.advance();
656+
b.advance();
657+
}
658+
assert_eq!(got, expected);
659+
660+
// "apple" appears at index 1 in both and index 3 in `a`: the cross-batch
661+
// `eq` path (arbitrary indices, bypasses the cache) must still hold.
662+
assert!(RowValues::eq(&a.values, 1, &b.values, 1));
663+
assert!(RowValues::eq(&a.values, 3, &b.values, 1));
664+
assert!(!RowValues::eq(&a.values, 0, &b.values, 0));
665+
}
666+
667+
/// A single-row `Rows` batch: `new` caches row 0 up front, and it must not
668+
/// index past the end.
669+
#[test]
670+
fn test_row_values_single_row_batch() {
671+
let cursor = new_row_values(&["solo"]);
672+
assert_eq!(cursor.values.len(), 1);
673+
assert!(!cursor.is_finished());
674+
}
675+
558676
#[test]
559677
fn test_primitive_nulls_first() {
560678
let options = SortOptions {

0 commit comments

Comments
 (0)