yew1eb commented on code in PR #2406:
URL: https://github.com/apache/auron/pull/2406#discussion_r3627521403
##########
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:
Good points — all three addressed in `sort_batch_in_place`:
1. `limit` now goes straight to `sort_to_indices(Some(limit))`, so Arrow
uses partial sort for Top-N and only falls back to a full sort when `limit`
covers the batch (limit defaults to `usize::MAX` for a plain Sort, so no
regression there). The subsequent `.take(limit)` is gone.
2. The returned `UInt32Array` is handed directly to `take_batch` (which
takes `impl Into<PrimitiveArray>`) — the `Vec<u32>` copy in between is removed.
3. (deferred) Keeping the indices and only reordering the sort key +
projected output columns instead of the whole batch is a bigger refactor that
touches `prune`/`key_collector`; I left it for a follow-up rather than expand
this PR.
On (1): the partial-sort win only shows up under a real Top-N. I added a
`limit=10000` benchmark case (1M rows) — there Auron beats DataFusion ~3.6-4.5x
(e.g. 0.007s vs 0.026s at repeat=1), vs. the ~2.3x *slower* figure for a full
sort. So passing the limit through is the meaningful change.
##########
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:
Done — on the fast path the key rows from `prune` are already in final
order, so I append them sequentially (`for row_idx in 0..keys.num_rows()`) and
skip the identity `Vec<u32>` allocation. The non-fast path still builds a
sorted-index `Vec` because it needs the indices for `take_batch` too.
--
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]
