From 98319dcabc80efee5d1e885b3e5b2fe926651ac8 Mon Sep 17 00:00:00 2001 From: yew1eb Date: Mon, 22 Jun 2026 23:53:42 +0800 Subject: [PATCH 1/2] [AURON #2405] Add primitive fast path for single-column integer sort --- .../datafusion-ext-plans/src/sort_exec.rs | 226 ++++++++++++++++-- 1 file changed, 206 insertions(+), 20 deletions(-) diff --git a/native-engine/datafusion-ext-plans/src/sort_exec.rs b/native-engine/datafusion-ext-plans/src/sort_exec.rs index 8b56b6ebe..bf2eca65c 100644 --- a/native-engine/datafusion-ext-plans/src/sort_exec.rs +++ b/native-engine/datafusion-ext-plans/src/sort_exec.rs @@ -28,7 +28,8 @@ use std::{ use arrow::{ array::ArrayRef, - datatypes::{Schema, SchemaRef}, + compute::{SortOptions, sort_to_indices}, + datatypes::{DataType, Schema, SchemaRef}, record_batch::{RecordBatch, RecordBatchOptions}, row::{RowConverter, Rows, SortField}, }; @@ -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::>() + } else 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(), + ) + }) + .take(self.limit) + .collect::>() } else { - (0..keys.num_rows() as u32).sorted_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(), - ) - }) - } - .take(self.limit) - .collect::>(); + (0..keys.num_rows() as u32) + .sorted_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(), + ) + }) + .take(self.limit) + .collect::>() + }; // build keys let mut key_collector = InMemRowsKeyCollector::default(); @@ -673,7 +697,13 @@ impl ExternalSorter { // build batch let sorted_batch = if !self.prune_sort_keys_from_batch.is_all_pruned() { - take_batch(batch, sorted_indices)? + if is_fast_path { + // batch is already sorted and pruned by `sort_batch_in_place`; + // avoid an extra `take_batch` copy. + batch + } else { + take_batch(batch, sorted_indices)? + } } else { create_zero_column_batch(batch.num_rows()) }; @@ -1087,6 +1117,41 @@ fn create_zero_column_batch(num_rows: usize) -> RecordBatch { .expect("failed to create empty RecordBatch") } +struct PrimitiveSortKey { + sort_expr: PhysicalSortExpr, +} + +impl PrimitiveSortKey { + fn try_new(input_schema: &SchemaRef, exprs: &[PhysicalSortExpr]) -> Option { + if exprs.len() != 1 { + return None; + } + let expr = &exprs[0]; + // must be a direct column reference + expr.expr.as_any().downcast_ref::()?; + let data_type = expr.expr.data_type(input_schema).ok()?; + if !matches!( + data_type, + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Date32 + | DataType::Date64 + | DataType::Timestamp(_, _) + ) { + return None; + } + Some(Self { + sort_expr: expr.clone(), + }) + } +} + struct PruneSortKeysFromBatch { input_projection: Vec, sort_row_converter: Arc>, @@ -1095,6 +1160,7 @@ struct PruneSortKeysFromBatch { restored_col_mappers: Vec, restored_schema: SchemaRef, pruned_schema: SchemaRef, + primitive_fast_path: Option, } #[derive(Clone, Copy)] @@ -1165,6 +1231,7 @@ impl PruneSortKeysFromBatch { restored_col_mappers, pruned_schema, restored_schema, + primitive_fast_path: PrimitiveSortKey::try_new(&input_schema, exprs), }) } @@ -1172,6 +1239,39 @@ impl PruneSortKeysFromBatch { 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 { + // 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)?; + let indices = indices + .values() + .iter() + .copied() + .take(limit) + .collect::>(); + take_batch(batch, indices) + } + fn pruned_schema(&self) -> SchemaRef { self.pruned_schema.clone() } @@ -1584,7 +1684,7 @@ mod fuzztest { use std::{sync::Arc, time::Instant}; use arrow::{ - array::{ArrayRef, StringArray, UInt32Array}, + array::{ArrayRef, Int64Array, StringArray, UInt32Array}, compute::{SortOptions, concat_batches}, record_batch::RecordBatch, }; @@ -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 = (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); + let session_ctx = + SessionContext::new_with_config(SessionConfig::new().with_batch_size(10000)); + let task_ctx = session_ctx.task_ctx(); + let num_rows = 1_000_000; + let batch = build_repeated_i64_batch(num_rows, repeat, 42); + let schema = batch.schema(); + let sort_exprs = vec![PhysicalSortExpr { + expr: Arc::new(Column::new("l_orderkey", 0)), + options: SortOptions::default(), + }]; + + let input = Arc::new(TestMemoryExec::try_new( + &[vec![batch]], + schema.clone(), + None, + )?); + let sort: Arc = if use_auron { + Arc::new(SortExec::new(input, sort_exprs.clone(), None, 0)) + } else { + Arc::new(datafusion::physical_plan::sorts::sort::SortExec::new( + LexOrdering::new(sort_exprs.iter().cloned()).expect("invalid sort exprs"), + input, + )) + }; + let start = Instant::now(); + let output = datafusion::physical_plan::collect(sort.clone(), task_ctx.clone()).await?; + let elapsed = start.elapsed().as_secs_f64(); + let total_rows: usize = output.iter().map(|b| b.num_rows()).sum(); + Ok((total_rows, elapsed)) + } + + #[tokio::test] + #[ignore = "manual benchmark"] + async fn bench_native_sort_varying_repeat_in_mem() -> Result<()> { + for repeat in [1usize, 4, 20, 100] { + let (_, elapsed_auron) = bench_sort_repeat(repeat, 1_000_000_000, true).await?; + let (_, elapsed_df) = bench_sort_repeat(repeat, 1_000_000_000, false).await?; + eprintln!( + "[sort in-mem] repeat={:>3}, auron={:.3}s, datafusion={:.3}s, speedup(df/auron)={:.2}x", + repeat, + elapsed_auron, + elapsed_df, + elapsed_auron / elapsed_df + ); + } + Ok(()) + } + + #[tokio::test] + #[ignore = "manual benchmark"] + async fn bench_native_sort_varying_repeat_external() -> Result<()> { + for repeat in [1usize, 4, 20, 100] { + let (_, elapsed_auron) = bench_sort_repeat(repeat, 2_000_000, true).await?; + let (_, elapsed_df) = bench_sort_repeat(repeat, 2_000_000, false).await?; + eprintln!( + "[sort external] repeat={:>3}, auron={:.3}s, datafusion={:.3}s, speedup(df/auron)={:.2}x", + repeat, + elapsed_auron, + elapsed_df, + elapsed_auron / elapsed_df + ); + } + Ok(()) + } + #[tokio::test] async fn fuzztest_in_mem_sorting() -> Result<()> { let time_start = Instant::now(); From 2148ad73153c22b37b0e302edd745defe21bdad5 Mon Sep 17 00:00:00 2001 From: yew1eb Date: Wed, 22 Jul 2026 12:42:15 +0800 Subject: [PATCH 2/2] Address review: pass limit to sort_to_indices, drop Vec copies, fix spill asymmetry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address PR #2406 review feedback from lyne7-sc: - sort_batch_in_place: pass `limit` to `sort_to_indices` (Arrow uses partial sort for Top-N; full sort when limit covers the batch) instead of sorting fully then .take(limit). - sort_batch_in_place: hand the returned UInt32Array directly to take_batch instead of copying through an intermediate Vec. - fast path: append key rows sequentially (rows are already in final order) instead of allocating an identity Vec. - benchmark: the external/spill group compared Auron spilling (2 MB MemManager limit) against DataFusion sorting fully in memory (DF does not honor the limit) — not comparable. Report Auron spill cost on its own; the in-mem group remains the fair Auron-vs-DataFusion comparison. - benchmark: add a Top-N (limit=10000) case. Top-N turns the tables: auron beats DataFusion ~3.6-4.5x, validating the partial-sort path. --- .../datafusion-ext-plans/src/sort_exec.rs | 137 ++++++++++++------ 1 file changed, 95 insertions(+), 42 deletions(-) diff --git a/native-engine/datafusion-ext-plans/src/sort_exec.rs b/native-engine/datafusion-ext-plans/src/sort_exec.rs index bf2eca65c..967ff6f44 100644 --- a/native-engine/datafusion-ext-plans/src/sort_exec.rs +++ b/native-engine/datafusion-ext-plans/src/sort_exec.rs @@ -661,37 +661,54 @@ impl ExternalSorter { // NOTE: we use stable merge sort for longer keys due to less comparison let (keys, batch) = self.prune_sort_keys_from_batch.prune(batch)?; + // On the fast path the batch is already sorted (and limited) above, so + // the key rows from `prune` are in final order — we can append them + // sequentially and skip allocating an identity `Vec`. The non-fast + // path needs an actual sorted-index vector for both key appending and + // the `take_batch` below. 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::>() + None } else 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(), - ) - }) - .take(self.limit) - .collect::>() + Some( + (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(), + ) + }) + .take(self.limit) + .collect::>(), + ) } else { - (0..keys.num_rows() as u32) - .sorted_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(), - ) - }) - .take(self.limit) - .collect::>() + Some( + (0..keys.num_rows() as u32) + .sorted_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(), + ) + }) + .take(self.limit) + .collect::>(), + ) }; // build keys let mut key_collector = InMemRowsKeyCollector::default(); key_collector.reserve(keys.num_rows(), keys.size()); - for &row_idx in &sorted_indices { - key_collector.add_key(keys.row(row_idx as usize).as_ref()); + match sorted_indices.as_deref() { + Some(indices) => { + for &row_idx in indices { + key_collector.add_key(keys.row(row_idx as usize).as_ref()); + } + } + // fast path: rows are already in final order, append sequentially. + None => { + for row_idx in 0..keys.num_rows() as u32 { + key_collector.add_key(keys.row(row_idx as usize).as_ref()); + } + } } key_collector.freeze(); @@ -702,7 +719,7 @@ impl ExternalSorter { // avoid an extra `take_batch` copy. batch } else { - take_batch(batch, sorted_indices)? + take_batch(batch, sorted_indices.expect("non-fast-path indices"))? } } else { create_zero_column_batch(batch.num_rows()) @@ -1262,13 +1279,13 @@ impl PruneSortKeysFromBatch { descending: key.sort_expr.options.descending, nulls_first: key.sort_expr.options.nulls_first, }; - let indices = sort_to_indices(&array, Some(options), None)?; - let indices = indices - .values() - .iter() - .copied() - .take(limit) - .collect::>(); + // Pass `limit` straight to `sort_to_indices` so Arrow uses partial + // sorting (it only emits the top-`limit` indices); when `limit` + // covers the whole batch Arrow internally falls back to a full sort, + // so this is never worse than `None`. The returned `UInt32Array` is + // handed directly to `take_batch` (which accepts a `PrimitiveArray`) + // instead of being copied through an intermediate `Vec`. + let indices = sort_to_indices(&array, Some(options), Some(limit))?; take_batch(batch, indices) } @@ -1718,6 +1735,18 @@ mod fuzztest { } async fn bench_sort_repeat(repeat: usize, mem: usize, use_auron: bool) -> Result<(usize, f64)> { + bench_sort_repeat_with_limit(repeat, mem, use_auron, None).await + } + + /// Variant of `bench_sort_repeat` that can pass a Top-N `limit` to the + /// Auron `SortExec`, so the `sort_to_indices(limit)` partial-sort path + /// can be exercised (vs. a full sort when `limit` is `None`). + async fn bench_sort_repeat_with_limit( + repeat: usize, + mem: usize, + use_auron: bool, + limit: Option, + ) -> Result<(usize, f64)> { MemManager::init(mem); let session_ctx = SessionContext::new_with_config(SessionConfig::new().with_batch_size(10000)); @@ -1736,12 +1765,15 @@ mod fuzztest { None, )?); let sort: Arc = if use_auron { - Arc::new(SortExec::new(input, sort_exprs.clone(), None, 0)) + Arc::new(SortExec::new(input, sort_exprs.clone(), limit, 0)) } else { - Arc::new(datafusion::physical_plan::sorts::sort::SortExec::new( - LexOrdering::new(sort_exprs.iter().cloned()).expect("invalid sort exprs"), - input, - )) + Arc::new( + datafusion::physical_plan::sorts::sort::SortExec::new( + LexOrdering::new(sort_exprs.iter().cloned()).expect("invalid sort exprs"), + input, + ) + .with_fetch(limit), + ) }; let start = Instant::now(); let output = datafusion::physical_plan::collect(sort.clone(), task_ctx.clone()).await?; @@ -1767,18 +1799,39 @@ mod fuzztest { Ok(()) } + #[tokio::test] + #[ignore = "manual benchmark"] + async fn bench_native_sort_topn_limit_in_mem() -> Result<()> { + // Top-N scenario: limit=10000 << num_rows=1M, so the primitive fast + // path's `sort_to_indices(Some(limit))` partial sort should be cheaper + // than a full sort. Compares Auron vs DataFusion, both with limit. + let limit = 10_000usize; + for repeat in [1usize, 100] { + let (_, elapsed_auron) = + bench_sort_repeat_with_limit(repeat, 1_000_000_000, true, Some(limit)).await?; + let (_, elapsed_df) = + bench_sort_repeat_with_limit(repeat, 1_000_000_000, false, Some(limit)).await?; + eprintln!( + "[sort top-N limit={limit}] repeat={repeat:>3}, auron={elapsed_auron:.3}s, datafusion={elapsed_df:.3}s, speedup(df/auron)={:.2}x", + elapsed_auron / elapsed_df + ); + } + Ok(()) + } + #[tokio::test] #[ignore = "manual benchmark"] async fn bench_native_sort_varying_repeat_external() -> Result<()> { + // NOTE: this measures Auron's external (spilling) sort in isolation. + // DataFusion's native SortExec has no spilling path, so under the same + // 2 MB `MemManager` limit Auron spills while DataFusion keeps sorting + // fully in memory — the two are not comparable. We therefore report + // Auron's spill cost on its own; use the in-mem group for a fair + // Auron-vs-DataFusion comparison (both stay in memory at 1 GB). for repeat in [1usize, 4, 20, 100] { let (_, elapsed_auron) = bench_sort_repeat(repeat, 2_000_000, true).await?; - let (_, elapsed_df) = bench_sort_repeat(repeat, 2_000_000, false).await?; eprintln!( - "[sort external] repeat={:>3}, auron={:.3}s, datafusion={:.3}s, speedup(df/auron)={:.2}x", - repeat, - elapsed_auron, - elapsed_df, - elapsed_auron / elapsed_df + "[sort external/spill, auron-only] repeat={repeat:>3}, auron={elapsed_auron:.3}s" ); } Ok(())