Skip to content

Commit 832a7eb

Browse files
committed
perf: skip measuring spilled rows when the largest fits the target
Measuring every row of a spilled batch made spilling aggregates with variable-width state, such as string keys, about 9% slower. Measure an upper bound on the row size once per batch instead, and only measure rows one by one when batch_size rows of that size could exceed the target. The rows of fixed-width columns are no longer measured one by one either.
1 parent c48d462 commit 832a7eb

2 files changed

Lines changed: 116 additions & 12 deletions

File tree

‎datafusion/physical-plan/src/sorts/stream.rs‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -306,6 +306,8 @@ pub(crate) struct IncrementalSortIterator {
306306
/// See [`Self::with_max_batch_bytes`].
307307
max_batch_bytes: Option<usize>,
308308
indices: Option<UInt32Array>,
309+
/// Row sizes of `batch`, measured with `indices` when `max_batch_bytes` is set
310+
row_sizes: Option<SpilledRowSizes>,
309311
cursor: usize,
310312
}
311313

@@ -322,6 +324,7 @@ impl IncrementalSortIterator {
322324
max_batch_bytes: None,
323325
cursor: 0,
324326
indices: None,
327+
row_sizes: None,
325328
}
326329
}
327330

@@ -335,15 +338,19 @@ impl IncrementalSortIterator {
335338

336339
/// Returns how many of the rows at `indices` fit in `max_batch_bytes`.
337340
fn rows_within(&self, indices: &[u32]) -> usize {
338-
let Some(max_batch_bytes) = self.max_batch_bytes else {
341+
let (Some(max_batch_bytes), Some(row_sizes)) =
342+
(self.max_batch_bytes, &self.row_sizes)
343+
else {
339344
return indices.len();
340345
};
341-
let row_sizes = SpilledRowSizes::new(&self.batch);
342346
if let Some(row_bytes) = row_sizes.fixed() {
343347
return (max_batch_bytes / row_bytes.max(1))
344348
.max(1)
345349
.min(indices.len());
346350
}
351+
if row_sizes.max_row().saturating_mul(indices.len()) <= max_batch_bytes {
352+
return indices.len();
353+
}
347354
let mut bytes = 0;
348355
for (rows, &row) in indices.iter().enumerate() {
349356
bytes += row_sizes.row(row as usize);
@@ -380,6 +387,9 @@ impl Iterator for IncrementalSortIterator {
380387
Err(e) => return Some(Err(e.into())),
381388
};
382389
self.indices = Some(indices);
390+
if self.max_batch_bytes.is_some() {
391+
self.row_sizes = Some(SpilledRowSizes::new(&self.batch));
392+
}
383393

384394
// Call again, this time it will hit the Some(indices) branch and return the first batch
385395
self.next()
@@ -404,6 +414,7 @@ impl Iterator for IncrementalSortIterator {
404414
let schema = self.batch.schema();
405415
let _ = mem::replace(&mut self.batch, RecordBatch::new_empty(schema));
406416
self.indices = None;
417+
self.row_sizes = None;
407418
}
408419

409420
// Return the new batch

‎datafusion/physical-plan/src/spill/spilled_row_sizes.rs‎

Lines changed: 103 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -31,27 +31,51 @@ use super::VIEW_SIZE_BYTES;
3131
/// variable-width values, and the elements of its nested values. Buffers that
3232
/// every row shares, such as validity bitmaps and dictionary values, are not
3333
/// counted.
34-
pub(crate) struct SpilledRowSizes(ColumnSize);
34+
pub(crate) struct SpilledRowSizes {
35+
/// Bytes that every row takes in the fixed-width columns
36+
fixed: usize,
37+
/// The other columns
38+
variable: Vec<ColumnSize>,
39+
/// See [`Self::max_row`]
40+
max_row: usize,
41+
}
3542

3643
impl SpilledRowSizes {
3744
pub(crate) fn new(batch: &RecordBatch) -> Self {
38-
Self(ColumnSize::Struct(
39-
batch
40-
.columns()
41-
.iter()
42-
.map(|column| ColumnSize::new(column.as_ref()))
43-
.collect(),
44-
))
45+
let mut fixed = 0;
46+
let mut variable = vec![];
47+
for column in batch.columns() {
48+
let column = ColumnSize::new(column.as_ref());
49+
match column.fixed() {
50+
Some(width) => fixed += width,
51+
None => variable.push(column),
52+
}
53+
}
54+
let max_row = variable.iter().fold(fixed, |bytes, column| {
55+
bytes.saturating_add(column.max_row())
56+
});
57+
Self {
58+
fixed,
59+
variable,
60+
max_row,
61+
}
4562
}
4663

4764
/// Returns the size in bytes of row `row`.
4865
pub(crate) fn row(&self, row: usize) -> usize {
49-
self.0.rows(row, row + 1)
66+
self.variable.iter().fold(self.fixed, |bytes, column| {
67+
bytes + column.rows(row, row + 1)
68+
})
5069
}
5170

5271
/// Returns the size in bytes of every row, if they all have the same size.
5372
pub(crate) fn fixed(&self) -> Option<usize> {
54-
self.0.fixed()
73+
self.variable.is_empty().then_some(self.fixed)
74+
}
75+
76+
/// Returns an upper bound on the size in bytes of any row.
77+
pub(crate) fn max_row(&self) -> usize {
78+
self.max_row
5579
}
5680
}
5781

@@ -129,6 +153,30 @@ impl ColumnSize {
129153
}
130154
}
131155

156+
/// Returns an upper bound on the size in bytes of a row.
157+
fn max_row(&self) -> usize {
158+
match self {
159+
Self::Fixed(width) => *width,
160+
Self::Bytes(offsets) => offsets.width() + offsets.max_len(),
161+
Self::Views(views) => {
162+
let max_len = views
163+
.iter()
164+
.map(|&view| view as u32)
165+
.filter(|&len| len > MAX_INLINE_VIEW_LEN)
166+
.max()
167+
.unwrap_or_default();
168+
VIEW_SIZE_BYTES + max_len as usize
169+
}
170+
Self::List(offsets, values) => offsets
171+
.width()
172+
.saturating_add(offsets.max_len().saturating_mul(values.max_row())),
173+
Self::FixedSizeList(size, values) => size.saturating_mul(values.max_row()),
174+
Self::Struct(fields) => fields
175+
.iter()
176+
.fold(0, |bytes, field| bytes.saturating_add(field.max_row())),
177+
}
178+
}
179+
132180
/// Returns the size in bytes of rows `start..end`.
133181
fn rows(&self, start: usize, end: usize) -> usize {
134182
let num_rows = end - start;
@@ -181,6 +229,21 @@ impl Offsets {
181229
Self::Large(offsets) => offsets[index] as usize,
182230
}
183231
}
232+
233+
/// Returns the largest distance between consecutive offsets.
234+
fn max_len(&self) -> usize {
235+
match self {
236+
Self::Small(offsets) => offsets
237+
.windows(2)
238+
.map(|pair| (pair[1] - pair[0]) as usize)
239+
.max(),
240+
Self::Large(offsets) => offsets
241+
.windows(2)
242+
.map(|pair| (pair[1] - pair[0]) as usize)
243+
.max(),
244+
}
245+
.unwrap_or_default()
246+
}
184247
}
185248

186249
impl From<&OffsetBuffer<i32>> for Offsets {
@@ -247,6 +310,36 @@ mod tests {
247310
);
248311
}
249312

313+
#[test]
314+
fn largest_row_bounds_every_row() {
315+
let ints: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3]));
316+
let strings: ArrayRef = Arc::new(StringArray::from(vec!["a", "bbbb", ""]));
317+
// 8 bytes of Int64, plus a 4 byte offset and the longest string
318+
assert_eq!(
319+
SpilledRowSizes::new(&batch(vec![ints, strings])).max_row(),
320+
16
321+
);
322+
323+
let long = "a string longer than twelve";
324+
let views: ArrayRef = Arc::new(StringViewArray::from(vec!["short", long]));
325+
assert_eq!(
326+
SpilledRowSizes::new(&batch(vec![views])).max_row(),
327+
16 + long.len()
328+
);
329+
330+
// An offset plus the most elements in a row, each of the largest size
331+
let list: ArrayRef =
332+
Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(vec![
333+
Some(vec![Some(1), Some(2), Some(3)]),
334+
Some(vec![]),
335+
Some(vec![Some(4)]),
336+
]));
337+
assert_eq!(
338+
SpilledRowSizes::new(&batch(vec![list])).max_row(),
339+
4 + 3 * 8
340+
);
341+
}
342+
250343
#[test]
251344
fn views_count_their_bytes_unless_inlined() {
252345
let long = "a string longer than twelve";

0 commit comments

Comments
 (0)