diff --git a/native-engine/datafusion-ext-plans/src/sort_exec.rs b/native-engine/datafusion-ext-plans/src/sort_exec.rs index 8b56b6ebe..967ff6f44 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,38 +643,84 @@ 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(), - ) - }) + // 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 { + None + } else if keys.size() / keys.num_rows() <= 8 { + 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(); // 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.expect("non-fast-path indices"))? + } } else { create_zero_column_batch(batch.num_rows()) }; @@ -1087,6 +1134,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 +1177,7 @@ struct PruneSortKeysFromBatch { restored_col_mappers: Vec, restored_schema: SchemaRef, pruned_schema: SchemaRef, + primitive_fast_path: Option, } #[derive(Clone, Copy)] @@ -1165,6 +1248,7 @@ impl PruneSortKeysFromBatch { restored_col_mappers, pruned_schema, restored_schema, + primitive_fast_path: PrimitiveSortKey::try_new(&input_schema, exprs), }) } @@ -1172,6 +1256,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, + }; + // 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) + } + fn pruned_schema(&self) -> SchemaRef { self.pruned_schema.clone() } @@ -1584,7 +1701,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 +1715,128 @@ 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)> { + 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)); + 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(), limit, 0)) + } else { + 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?; + 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_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?; + eprintln!( + "[sort external/spill, auron-only] repeat={repeat:>3}, auron={elapsed_auron:.3}s" + ); + } + Ok(()) + } + #[tokio::test] async fn fuzztest_in_mem_sorting() -> Result<()> { let time_start = Instant::now();