lyne7-sc commented on code in PR #2406:
URL: https://github.com/apache/auron/pull/2406#discussion_r3627098109
##########
native-engine/datafusion-ext-plans/src/sort_exec.rs:
##########
@@ -1165,13 +1231,47 @@ impl PruneSortKeysFromBatch {
restored_col_mappers,
pruned_schema,
restored_schema,
+ primitive_fast_path: PrimitiveSortKey::try_new(&input_schema,
exprs),
})
}
fn is_all_pruned(&self) -> bool {
self.pruned_schema.fields().is_empty()
}
+ /// Fast-path batch in-place sort for a single-column primitive key.
+ ///
+ /// Sorts the batch by the primitive sort key using Arrow `sort_to_indices`
+ /// (which compares the raw values, e.g. `i64`, instead of `RowConverter`
+ /// byte encoding), then reorders the batch by the resulting indices and
+ /// applies `limit` rows. The caller can then treat the returned batch as
+ /// already sorted and pruned. No-op precondition: `primitive_fast_path`
+ /// must be `Some`; the caller checks `is_fast_path` before calling.
+ fn sort_batch_in_place(&self, batch: RecordBatch, limit: usize) ->
Result<RecordBatch> {
+ // safety/correctness: only reached when is_fast_path == true
+ let key = self
+ .primitive_fast_path
+ .as_ref()
+ .expect("sort_batch_in_place requires a primitive fast path");
+ let array = key
+ .sort_expr
+ .expr
+ .evaluate(&batch)
+ .and_then(|cv| cv.into_array(batch.num_rows()))?;
+ let options = SortOptions {
+ descending: key.sort_expr.options.descending,
+ nulls_first: key.sort_expr.options.nulls_first,
+ };
+ let indices = sort_to_indices(&array, Some(options), None)?;
Review Comment:
I wonder if we could keep the sorted indices around a little longer instead
of materializing the fully reordered input batch here. That way, we could use
the indices to reorder only the sort key and projected output columns, instead
of every column in the input batch.
A few related optimizations may be possible:
- We could pass `limit` directly to `sort_to_indices`, so Arrow can use
partial sorting.
- The returned `UInt32Array` could be passed directly to `take_batch`,
avoiding the intermediate copy into a `Vec<u32>`.
##########
native-engine/datafusion-ext-plans/src/sort_exec.rs:
##########
@@ -642,26 +643,49 @@ impl ExternalSorter {
self.mem_total_size
.fetch_add(batch.get_batch_mem_size(), SeqCst);
- // sort keys
+ // For a single-column primitive key (e.g. TPC-H
`lineitem.l_orderkey`),
+ // skip the RowConverter byte-comparison sort and use Arrow
+ // `sort_to_indices` on the raw array instead (~6x faster on the batch
+ // in-place sort). The batch is sorted and limited here; the rest of
the
+ // pipeline (encode for K-way merge, output) is unchanged.
+ let is_fast_path = self
+ .prune_sort_keys_from_batch
+ .primitive_fast_path
+ .is_some();
+ let batch = if is_fast_path {
+ self.prune_sort_keys_from_batch
+ .sort_batch_in_place(batch, self.limit)?
+ } else {
+ batch
+ };
+
// NOTE: we use stable merge sort for longer keys due to less
comparison
let (keys, batch) = self.prune_sort_keys_from_batch.prune(batch)?;
- let sorted_indices = if keys.size() / keys.num_rows() <= 8 {
- (0..keys.num_rows() as u32).sorted_unstable_by_key(|&row_idx|
unsafe {
- // safety: bypass boundary and lifetime checking
- std::mem::transmute::<_, &'static [u8]>(
- keys.row_unchecked(row_idx as usize).as_ref(),
- )
- })
+ let sorted_indices = if is_fast_path {
+ // the batch was already sorted (and limited) above, so the key
rows
+ // produced by `prune` are in final order — use identity indices.
+ (0..keys.num_rows() as u32).collect::<Vec<_>>()
Review Comment:
nit: Since the batch is already sorted on the fast path, could we append the
key rows sequentially instead of allocating this identity `Vec<u32>`?
##########
native-engine/datafusion-ext-plans/src/sort_exec.rs:
##########
@@ -1598,6 +1698,92 @@ mod fuzztest {
use crate::sort_exec::SortExec;
+ /// Benchmark helper: build a single Int64 column where each value is
+ /// repeated `repeat` times, shuffled, to simulate TPC-H
+ /// lineitem.l_orderkey distribution.
+ fn build_repeated_i64_batch(num_rows: usize, repeat: usize, seed: u64) ->
RecordBatch {
+ use rand::{Rng, SeedableRng};
+ let unique_keys = (num_rows + repeat - 1) / repeat;
+ let mut values: Vec<i64> = (0..unique_keys)
+ .flat_map(|v| std::iter::repeat(v as i64).take(repeat))
+ .take(num_rows)
+ .collect();
+ let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
+ rand::seq::SliceRandom::shuffle(values.as_mut_slice(), &mut rng);
+ let schema = Arc::new(arrow::datatypes::Schema::new(vec![
+ arrow::datatypes::Field::new("l_orderkey",
arrow::datatypes::DataType::Int64, false),
+ ]));
+ RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(values))
as ArrayRef])
+ .expect("failed to create benchmark batch")
+ }
+
+ async fn bench_sort_repeat(repeat: usize, mem: usize, use_auron: bool) ->
Result<(usize, f64)> {
+ MemManager::init(mem);
Review Comment:
Just a question: memory limit looks Auron-specific, so could this be
comparing Auron spilling with DataFusion sorting in memory?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
