Repository navigation
Perf: Support automatically concat_batches for sort which will improve performance #15375
Description
Is your feature request related to a problem or challenge?
We should investigate and improve the sort code to support concat_batches for more cases besides the following case:
// If less than sort_in_place_threshold_bytes, concatenate and sort in place
if self.reservation.size() < self.sort_in_place_threshold_bytes {
// Concatenate memory batches together and sort
let batch = concat_batches(&self.schema, &self.in_mem_batches)?;
self.in_mem_batches.clear();
self.reservation
.try_resize(get_reserved_byte_for_record_batch(&batch))?;
let reservation = self.reservation.take();
return self.sort_batch_stream(batch, metrics, reservation);
}See details about the performance improvement:
Describe the solution you'd like
No response
Describe alternatives you've considered
No response
Additional context
No response
Activity
zhuqi-lucas commentedon Mar 24, 2025
take
zhuqi-lucas commentedon Mar 24, 2025
cc @alamb
I did some POC testing, the performance improvement is very promising, even 30% improvement for some queries:
2010YOUY01 commentedon Mar 24, 2025
This is a great observation, and the POC optimization has a high ROI. Here are some additional thoughts:
This is just intuition, but I think the sorting phase should be faster than the sort-preserving merge phase. Because the sorting implementation is much simpler, it can rely on the existing optimized arrow row format and also quick sort implementation from the standard library. On the contrary, merging phase we have an in-house implementation of a loser tree heap, I think it's a bit complex so maybe it is also hard to optimize manually.
As a result, perhaps it's preferred to put more work to be done during sorting, and leave minimal work to be done during merging (however, merging is still needed inside partial sort, because it can keep all the cores busy when partial sort is all finished)
Example
There is a sort query to run in 4 partitions, and each partition will process 100 input batches
Implementation
The POC will copy the batches with
cocnat()than do a bigger sort, an alternative to try to avoid copies is: first sort all elements' indices (2-level index consists of(batch_idx, row_idx)), and get a permutation array.Use the interleave kernel to construct the final result https://docs.rs/arrow/latest/arrow/compute/kernels/interleave/fn.interleave.html
zhuqi-lucas commentedon Mar 25, 2025
Thank you @2010YOUY01 for review and good suggestion, i will improve my POC code and add more testing.
Dandandan commentedon Apr 12, 2025
Really nice observation! I think we should drive this further.
Some further observations I saw when looking at the current implementation on master for the in memory merging part:
spawn_bufferedin the code for streaming merge () sorting can be done on the spot here and batches dropped as soon as possible.zhuqi-lucas commentedon Apr 12, 2025
Thank you @Dandandan , addressed your comments. And we can make it as the first version. And in future we may can improve it as described by @2010YOUY01 :
#15375 (comment)
And i think current implementation is also reasonable because the sort_in_place_threshold_bytes is a already used config, we can first reuse it to concat batch and it's safe.
Dandandan commentedon Apr 12, 2025
@2010YOUY01 that sound like a very promising future direction. I might try something experimenting on this soon if none beats me to it.