From 847189a5f0871ec62355af98068a75852c7dcda8 Mon Sep 17 00:00:00 2001 From: linfeng Date: Thu, 3 Sep 2026 23:54:18 +0800 Subject: [PATCH 1/3] perf: optimize prefix processing in PartialSortExec --- datafusion/physical-plan/Cargo.toml | 5 + .../physical-plan/benches/partial_sort.rs | 214 +++++++++++ .../physical-plan/src/sorts/partial_sort.rs | 349 ++++++++++++++---- .../sqllogictest/test_files/group_by.slt | 73 ++++ 4 files changed, 576 insertions(+), 65 deletions(-) create mode 100644 datafusion/physical-plan/benches/partial_sort.rs diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 1adac74ecb3c5..9493ef0a1c037 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -109,6 +109,11 @@ tokio = { workspace = true, features = [ harness = false name = "partial_ordering" +[[bench]] +harness = false +name = "partial_sort" +required-features = ["test_utils"] + [[bench]] harness = false name = "union_schema" diff --git a/datafusion/physical-plan/benches/partial_sort.rs b/datafusion/physical-plan/benches/partial_sort.rs new file mode 100644 index 0000000000000..7dda8303a187e --- /dev/null +++ b/datafusion/physical-plan/benches/partial_sort.rs @@ -0,0 +1,214 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, UInt64Array}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col}; +use datafusion_physical_plan::sorts::partial_sort::PartialSortExec; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::{ExecutionPlan, collect}; + +const NUM_BATCHES: usize = 32; +const BATCH_SIZE: usize = 8192; +const NUM_ROWS: usize = NUM_BATCHES * BATCH_SIZE; + +#[derive(Clone, Copy)] +enum PrefixLayout { + BatchesPerPrefix(usize), + PrefixesPerBatch(usize), + RowsPerPrefix(usize), +} + +impl PrefixLayout { + fn prefix_and_suffix(self, batch_idx: usize, row_idx: usize) -> (u64, u64) { + match self { + Self::BatchesPerPrefix(batches_per_prefix) => { + let rows_per_prefix = batches_per_prefix * BATCH_SIZE; + let offset = (batch_idx % batches_per_prefix) * BATCH_SIZE + row_idx; + ( + (batch_idx / batches_per_prefix) as u64, + (rows_per_prefix - offset - 1) as u64, + ) + } + Self::PrefixesPerBatch(prefixes_per_batch) => { + let rows_per_prefix = BATCH_SIZE / prefixes_per_batch; + let prefix_in_batch = row_idx / rows_per_prefix; + ( + (batch_idx * prefixes_per_batch + prefix_in_batch) as u64, + (rows_per_prefix - row_idx % rows_per_prefix - 1) as u64, + ) + } + Self::RowsPerPrefix(rows_per_prefix) => { + let global_idx = batch_idx * BATCH_SIZE + row_idx; + ( + (global_idx / rows_per_prefix) as u64, + (rows_per_prefix - global_idx % rows_per_prefix - 1) as u64, + ) + } + } + } +} + +fn schema() -> SchemaRef { + Arc::new(Schema::new( + [ + "prefix", + "suffix", + "payload_0", + "payload_1", + "payload_2", + "payload_3", + ] + .into_iter() + .map(|name| Field::new(name, DataType::UInt64, false)) + .collect::>(), + )) +} + +fn make_batches(layout: PrefixLayout) -> Vec { + let schema = schema(); + (0..NUM_BATCHES) + .map(|batch_idx| { + let rows = 0..BATCH_SIZE; + let prefix = UInt64Array::from_iter_values( + rows.clone() + .map(|row_idx| layout.prefix_and_suffix(batch_idx, row_idx).0), + ); + let suffix = UInt64Array::from_iter_values( + rows.clone() + .map(|row_idx| layout.prefix_and_suffix(batch_idx, row_idx).1), + ); + let payload = |row_idx: usize| (batch_idx * BATCH_SIZE + row_idx) as u64; + let payload_0 = UInt64Array::from_iter_values(rows.clone().map(payload)); + let payload_1 = UInt64Array::from_iter_values( + rows.clone() + .map(|row_idx| payload(row_idx).wrapping_mul(31)), + ); + let payload_2 = UInt64Array::from_iter_values( + rows.clone().map(|row_idx| payload(row_idx).rotate_left(13)), + ); + let payload_3 = UInt64Array::from_iter_values( + rows.map(|row_idx| payload(row_idx).wrapping_neg()), + ); + + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(prefix) as ArrayRef, + Arc::new(suffix) as ArrayRef, + Arc::new(payload_0) as ArrayRef, + Arc::new(payload_1) as ArrayRef, + Arc::new(payload_2) as ArrayRef, + Arc::new(payload_3) as ArrayRef, + ], + ) + .unwrap() + }) + .collect() +} + +fn make_plan(batches: &[RecordBatch]) -> Arc { + let schema = batches[0].schema(); + let input = + TestMemoryExec::try_new_exec(&[batches.to_vec()], Arc::clone(&schema), None) + .unwrap(); + let ordering = LexOrdering::new([ + PhysicalSortExpr::new(col("prefix", &schema).unwrap(), SortOptions::default()), + PhysicalSortExpr::new(col("suffix", &schema).unwrap(), SortOptions::default()), + ]) + .unwrap(); + Arc::new(PartialSortExec::new(ordering, input, 1)) +} + +fn partial_sort_benchmark(c: &mut Criterion) { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let task_ctx = Arc::new(TaskContext::default()); + let mut group = c.benchmark_group("partial_sort"); + group.sample_size(10); + + let cases = [ + ( + BenchmarkId::new("prefix_across_batches", "batches_per_prefix=1"), + PrefixLayout::BatchesPerPrefix(1), + ), + ( + BenchmarkId::new("prefix_across_batches", "batches_per_prefix=2"), + PrefixLayout::BatchesPerPrefix(2), + ), + ( + BenchmarkId::new("prefix_across_batches", "batches_per_prefix=4"), + PrefixLayout::BatchesPerPrefix(4), + ), + ( + BenchmarkId::new("prefixes_within_batch", "prefixes_per_batch=2"), + PrefixLayout::PrefixesPerBatch(2), + ), + ( + BenchmarkId::new("prefixes_within_batch", "prefixes_per_batch=8"), + PrefixLayout::PrefixesPerBatch(8), + ), + ( + BenchmarkId::new("prefixes_within_batch", "prefixes_per_batch=32"), + PrefixLayout::PrefixesPerBatch(32), + ), + ( + BenchmarkId::new("prefixes_within_batch", "prefixes_per_batch=128"), + PrefixLayout::PrefixesPerBatch(128), + ), + ( + BenchmarkId::new("unaligned_mixed", "rows_per_prefix=1000"), + PrefixLayout::RowsPerPrefix(1000), + ), + ]; + + for (id, layout) in cases { + let batches = make_batches(layout); + let validation = runtime + .block_on(collect(make_plan(&batches), Arc::clone(&task_ctx))) + .unwrap(); + assert_eq!( + validation.iter().map(RecordBatch::num_rows).sum::(), + NUM_ROWS + ); + drop(validation); + + group.bench_function(id, |b| { + b.iter_batched( + || make_plan(&batches), + |plan| { + let output = runtime + .block_on(collect(plan, Arc::clone(&task_ctx))) + .unwrap(); + black_box(output); + }, + BatchSize::LargeInput, + ); + }); + } + + group.finish(); +} + +criterion_group!(benches, partial_sort_benchmark); +criterion_main!(benches); diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index 80520d58dfa9b..20e8be6c22485 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -48,10 +48,12 @@ //! +---+---+---+ //! ``` //! -//! The plan concats incoming data with such last rows of previous input -//! and continues partial sorting of the segments. +//! The plan buffers the trailing prefix group across incoming batches and +//! continues partial sorting once that group is complete. +use std::cmp::Ordering; use std::fmt::Debug; +use std::ops::Range; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -66,12 +68,14 @@ use crate::{ SendableRecordBatchStream, Statistics, validate_child_count, }; -use arrow::compute::concat_batches; +use arrow::compute::{ + SortColumn, concat, concat_batches, interleave_record_batch, lexsort_to_indices, +}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::Result; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::utils::evaluate_partition_ranges; +use datafusion_common::utils::{compare_rows, evaluate_partition_ranges, get_row_at_idx}; use datafusion_execution::{RecordBatchStream, TaskContext}; use datafusion_physical_expr::{LexOrdering, PhysicalExpr}; @@ -149,9 +153,9 @@ use log::trace; /// +---+---+---+ /// ``` /// -/// Once known complete, the buffered rows are sorted by the full `(a, b, c)` -/// ordering and emitted as a [`RecordBatch`]; Any rows from the most recently -/// seen prefix remain buffered (as more rows with the same prefix may arrive in +/// Once known complete, each prefix group is sorted by the remaining suffix +/// ordering and emitted as a [`RecordBatch`]. Any rows from the most recently +/// seen prefix remain buffered, as more rows with the same prefix may arrive in /// future batches. /// /// ```text @@ -162,9 +166,9 @@ use log::trace; /// | 0 | 0 | 1 | <-- completed group /// | 0 | 0 | 2 | /// | 0 | 0 | 3 | +/// | 0 | 1 | 1 | <-- completed group /// | 0 | 2 | 0 | <-- completed group /// | 0 | 2 | 4 | -/// | 0 | 1 | 1 | <-- completed group /// +---+---+---+ /// /// Buffer @@ -455,7 +459,10 @@ impl ExecutionPlan for PartialSortExec { input, expr: self.expr.clone(), common_prefix_length: self.common_prefix_length, - in_mem_batch: RecordBatch::new_empty(Arc::clone(&self.schema())), + suffix_expr: LexOrdering::new( + self.expr.iter().skip(self.common_prefix_length).cloned(), + ), + in_mem_batches: vec![], fetch: self.fetch, is_closed: false, baseline_metrics: BaselineMetrics::new(&self.metrics_set, partition), @@ -487,8 +494,10 @@ struct PartialSortStream { /// Length of prefix common to input ordering and required ordering of plan /// should be more than 0 otherwise PartialSort is not applicable common_prefix_length: usize, - /// Used as a buffer for part of the input not ready for sort - in_mem_batch: RecordBatch, + /// Sort expressions not already satisfied by the input ordering + suffix_expr: Option, + /// Fragments of the trailing prefix group that is not ready for sort + in_mem_batches: Vec, /// Fetch top N results fetch: Option, /// Whether the stream has finished returning all of its data or not @@ -497,6 +506,23 @@ struct PartialSortStream { baseline_metrics: BaselineMetrics, } +#[derive(Debug)] +struct PrefixBatchRange { + batch_idx: usize, + range: Range, +} + +fn whole_batches_prefix(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .enumerate() + .map(|(batch_idx, batch)| PrefixBatchRange { + batch_idx, + range: 0..batch.num_rows(), + }) + .collect() +} + impl Stream for PartialSortStream { type Item = Result; @@ -540,26 +566,72 @@ impl PartialSortStream { match ready!(self.input.poll_next_unpin(cx)) { Some(Ok(batch)) => { - // Merge new batch into in_mem_batch - self.in_mem_batch = concat_batches( - &self.schema(), - &[self.in_mem_batch.clone(), batch], - )?; + if batch.num_rows() == 0 { + continue; + } - // Check if we have a slice point, otherwise keep accumulating in `self.in_mem_batch`. - if let Some(slice_point) = self - .get_slice_point(self.common_prefix_length, &self.in_mem_batch)? - { - let sorted = self.in_mem_batch.slice(0, slice_point); - self.in_mem_batch = self.in_mem_batch.slice( - slice_point, - self.in_mem_batch.num_rows() - slice_point, - ); - let sorted_batch = sort_batch(&sorted, &self.expr, self.fetch)?; - if let Some(fetch) = self.fetch.as_mut() { - *fetch -= sorted_batch.num_rows(); + let prefix_ranges = self.get_prefix_ranges(&batch)?; + let boundary_changed = self + .in_mem_batches + .last() + .map(|previous| { + self.prefix_changed_at_batch_boundary(previous, &batch) + }) + .transpose()? + .unwrap_or(false); + + let completed = if prefix_ranges.len() >= 2 { + let trailing_range = prefix_ranges.last().unwrap(); + let trailing_batch = + batch.slice(trailing_range.start, trailing_range.len()); + + let mut completed_batches = + std::mem::take(&mut self.in_mem_batches); + let current_batch_idx = completed_batches.len(); + let mut completed_prefixes = if completed_batches.is_empty() { + vec![] + } else { + vec![whole_batches_prefix(&completed_batches)] + }; + + for (idx, range) in prefix_ranges[..prefix_ranges.len() - 1] + .iter() + .cloned() + .enumerate() + { + let current_prefix = PrefixBatchRange { + batch_idx: current_batch_idx, + range, + }; + if idx == 0 + && !boundary_changed + && let Some(previous_prefix) = + completed_prefixes.last_mut() + { + previous_prefix.push(current_prefix); + continue; + } + completed_prefixes.push(vec![current_prefix]); } + completed_batches.push(batch); + self.in_mem_batches.push(trailing_batch); + Some((completed_batches, completed_prefixes)) + } else if boundary_changed { + let completed_batches = std::mem::take(&mut self.in_mem_batches); + let completed_prefix = whole_batches_prefix(&completed_batches); + self.in_mem_batches.push(batch); + Some((completed_batches, vec![completed_prefix])) + } else { + self.in_mem_batches.push(batch); + None + }; + + if let Some((completed_batches, completed_prefixes)) = completed { + let sorted_batch = self.sort_completed_prefixes( + completed_batches, + completed_prefixes, + )?; if sorted_batch.num_rows() > 0 { return Poll::Ready(Some(Ok(sorted_batch))); } @@ -571,8 +643,16 @@ impl PartialSortStream { // Release the input pipeline's resources before sorting. let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - // Once input is consumed, sort the rest of the inserted batches - let remaining_batch = self.sort_in_mem_batch()?; + // Once input is consumed, the trailing prefix is complete. + let completed_batches = std::mem::take(&mut self.in_mem_batches); + if completed_batches.is_empty() { + return Poll::Ready(None); + } + let completed_prefix = whole_batches_prefix(&completed_batches); + let remaining_batch = self.sort_completed_prefixes( + completed_batches, + vec![completed_prefix], + )?; return if remaining_batch.num_rows() > 0 { Poll::Ready(Some(Ok(remaining_batch))) } else { @@ -583,50 +663,189 @@ impl PartialSortStream { } } - /// Returns a sorted RecordBatch from in_mem_batches and clears in_mem_batches - /// - /// If fetch is specified for PartialSortStream `sort_in_mem_batch` will limit - /// the last RecordBatch returned and will mark the stream as closed - fn sort_in_mem_batch(self: &mut Pin<&mut Self>) -> Result { - let input_batch = self.in_mem_batch.clone(); - self.in_mem_batch = RecordBatch::new_empty(self.schema()); - let result = sort_batch(&input_batch, &self.expr, self.fetch)?; + fn sort_completed_prefixes( + self: &mut Pin<&mut Self>, + completed_batches: Vec, + completed_prefixes: Vec>, + ) -> Result { + debug_assert!(!completed_prefixes.is_empty()); + let result = if completed_prefixes.len() == 1 { + // A single prefix may span batches, so concatenate it at most once. + let prefix = completed_prefixes.into_iter().next().unwrap(); + let mut completed_batches = + completed_batches.into_iter().map(Some).collect::>(); + let mut prefix_batches = prefix + .into_iter() + .map(|prefix_range| { + let batch = completed_batches[prefix_range.batch_idx] + .take() + .expect("each source batch occurs once in a single prefix"); + if prefix_range.range == (0..batch.num_rows()) { + batch + } else { + batch.slice(prefix_range.range.start, prefix_range.range.len()) + } + }) + .collect::>(); + let batch = if prefix_batches.len() == 1 { + prefix_batches.pop().unwrap() + } else { + concat_batches(&self.schema(), &prefix_batches)? + }; + + if let Some(suffix_expr) = &self.suffix_expr { + sort_batch(&batch, suffix_expr, self.fetch)? + } else { + let row_count = self + .fetch + .unwrap_or_else(|| batch.num_rows()) + .min(batch.num_rows()); + batch.slice(0, row_count) + } + } else { + // Evaluate suffix expressions once per source batch, then sort each + // prefix independently and materialize the result with one interleave. + let sort_columns_by_batch = completed_batches + .iter() + .map(|batch| { + self.suffix_expr + .iter() + .flat_map(|exprs| exprs.iter()) + .map(|expr| expr.evaluate_to_sort_column(batch)) + .collect::>>() + }) + .collect::>>()?; + let mut remaining_fetch = self.fetch.unwrap_or(usize::MAX); + let mut interleave_indices = vec![]; + + for prefix in completed_prefixes { + if remaining_fetch == 0 { + break; + } + + let row_count = prefix + .iter() + .map(|prefix_range| prefix_range.range.len()) + .sum::(); + let prefix_fetch = remaining_fetch.min(row_count); + let offsets = prefix + .iter() + .scan(0, |offset, prefix_range| { + let current = *offset; + *offset += prefix_range.range.len(); + Some(current) + }) + .collect::>(); + + let sorted_indices = if let Some(suffix_expr) = &self.suffix_expr { + let sort_columns = suffix_expr + .iter() + .enumerate() + .map(|(sort_idx, expr)| { + let values = prefix + .iter() + .map(|prefix_range| { + let values = &sort_columns_by_batch + [prefix_range.batch_idx][sort_idx] + .values; + if prefix_range.range + == (0..completed_batches[prefix_range.batch_idx] + .num_rows()) + { + Arc::clone(values) + } else { + values.slice( + prefix_range.range.start, + prefix_range.range.len(), + ) + } + }) + .collect::>(); + let values = if values.len() == 1 { + Arc::clone(&values[0]) + } else { + let values = values + .iter() + .map(|values| values.as_ref()) + .collect::>(); + concat(&values)? + }; + Ok(SortColumn { + values, + options: Some(expr.options), + }) + }) + .collect::>>()?; + lexsort_to_indices(&sort_columns, Some(prefix_fetch))? + .values() + .iter() + .map(|idx| *idx as usize) + .collect::>() + } else { + (0..prefix_fetch).collect() + }; + + interleave_indices.extend(sorted_indices.into_iter().map(|idx| { + let range_idx = offsets.partition_point(|offset| *offset <= idx) - 1; + let prefix_range = &prefix[range_idx]; + ( + prefix_range.batch_idx, + prefix_range.range.start + idx - offsets[range_idx], + ) + })); + remaining_fetch -= prefix_fetch; + } + + if interleave_indices.is_empty() { + RecordBatch::new_empty(self.schema()) + } else { + let completed_batches = completed_batches.iter().collect::>(); + interleave_record_batch(&completed_batches, &interleave_indices)? + } + }; + if let Some(remaining_fetch) = self.fetch { - // remaining_fetch - result.num_rows() is always be >= 0 - // because result length of sort_batch with limit cannot be - // more than the requested limit self.fetch = Some(remaining_fetch - result.num_rows()); - if remaining_fetch == result.num_rows() { - self.is_closed = true; - } } Ok(result) } - /// Return the end index of the second last partition if the batch - /// can be partitioned based on its already sorted columns - /// - /// Return None if the batch cannot be partitioned, which means the - /// batch does not have the information for a safe sort - fn get_slice_point( + fn get_prefix_ranges(&self, batch: &RecordBatch) -> Result>> { + let common_prefix_sort_keys = (0..self.common_prefix_length) + .map(|idx| self.expr[idx].evaluate_to_sort_column(batch)) + .collect::>>()?; + evaluate_partition_ranges(batch.num_rows(), &common_prefix_sort_keys) + } + + fn prefix_changed_at_batch_boundary( &self, - common_prefix_len: usize, + previous_batch: &RecordBatch, batch: &RecordBatch, - ) -> Result> { - let common_prefix_sort_keys = (0..common_prefix_len) - .map(|idx| self.expr[idx].evaluate_to_sort_column(batch)) + ) -> Result { + let previous = previous_batch.slice(previous_batch.num_rows() - 1, 1); + let next = batch.slice(0, 1); + let previous_columns = (0..self.common_prefix_length) + .map(|idx| self.expr[idx].evaluate_to_sort_column(&previous)) .collect::>>()?; - let partition_points = - evaluate_partition_ranges(batch.num_rows(), &common_prefix_sort_keys)?; - // If partition points are [0..100], [100..200], [200..300] - // we should return 200, which is the safest and furthest partition boundary - // Please note that we shouldn't return 300 (which is number of rows in the batch), - // because this boundary may change with new data. - if partition_points.len() >= 2 { - Ok(Some(partition_points[partition_points.len() - 2].end)) - } else { - Ok(None) - } + let next_columns = (0..self.common_prefix_length) + .map(|idx| self.expr[idx].evaluate_to_sort_column(&next)) + .collect::>>()?; + let previous_values = previous_columns + .iter() + .map(|column| Arc::clone(&column.values)) + .collect::>(); + let next_values = next_columns + .iter() + .map(|column| Arc::clone(&column.values)) + .collect::>(); + let previous_row = get_row_at_idx(&previous_values, 0)?; + let next_row = get_row_at_idx(&next_values, 0)?; + let sort_options = self.expr[..self.common_prefix_length] + .iter() + .map(|expr| expr.options) + .collect::>(); + + Ok(compare_rows(&previous_row, &next_row, &sort_options)? != Ordering::Equal) } } diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 38d1b7821451d..56eb73128f4a0 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -2410,6 +2410,79 @@ ORDER BY a, b, d; 1 3 4 1 3 4 +# Multiple prefix groups complete in one batch. PartialSortExec sorts each +# suffix independently and materializes the result through the multi-group path. +statement ok +set datafusion.execution.batch_size = 128; + +query IIII +SELECT a, b, c, d +FROM annotated_data_infinite2 +ORDER BY a, b, d, c +LIMIT 8; +---- +0 0 0 0 +0 0 2 0 +0 0 3 0 +0 0 6 0 +0 0 20 0 +0 0 22 0 +0 0 23 0 +0 0 4 1 + +# One prefix group spans three input batches before the next prefix completes it. +statement ok +set datafusion.execution.batch_size = 10; + +query IIII +SELECT a, b, c, d +FROM annotated_data_infinite2 +ORDER BY a, b, d, c +LIMIT 6; +---- +0 0 0 0 +0 0 2 0 +0 0 3 0 +0 0 6 0 +0 0 20 0 +0 0 22 0 + +# The first prefix continues in the second batch, which also completes another +# prefix. The limit spans prefix groups and exercises the multi-source path. +statement ok +set datafusion.execution.batch_size = 40; + +query IIII +SELECT a, b, c, d +FROM annotated_data_infinite2 +WHERE c BETWEEN 20 AND 79 +ORDER BY a, b, d, c +LIMIT 20; +---- +0 0 20 0 +0 0 22 0 +0 0 23 0 +0 0 21 3 +0 0 24 4 +0 1 25 0 +0 1 27 0 +0 1 35 0 +0 1 38 0 +0 1 45 0 +0 1 46 0 +0 1 28 1 +0 1 29 1 +0 1 36 1 +0 1 39 1 +0 1 40 1 +0 1 47 1 +0 1 48 1 +0 1 26 2 +0 1 32 2 + +statement ok +reset datafusion.execution.batch_size; + statement ok drop table annotated_data_infinite2; From 3b56538a8d5c4a71956abe2e71647d5e39b4b568 Mon Sep 17 00:00:00 2001 From: linfeng <33561138+lyne7-sc@users.noreply.github.com> Date: Fri, 4 Sep 2026 15:50:01 +0800 Subject: [PATCH 2/3] refactor: simplify completed prefix tracking in PartialSortExec --- .../physical-plan/src/sorts/partial_sort.rs | 524 ++++++++++-------- 1 file changed, 306 insertions(+), 218 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index 20e8be6c22485..0ad11cd183d85 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -51,7 +51,6 @@ //! The plan buffers the trailing prefix group across incoming batches and //! continues partial sorting once that group is complete. -use std::cmp::Ordering; use std::fmt::Debug; use std::ops::Range; use std::pin::Pin; @@ -72,10 +71,11 @@ use arrow::compute::{ SortColumn, concat, concat_batches, interleave_record_batch, lexsort_to_indices, }; use arrow::datatypes::SchemaRef; -use arrow::record_batch::RecordBatch; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; +use arrow_ord::ord::make_comparator; use datafusion_common::Result; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::utils::{compare_rows, evaluate_partition_ranges, get_row_at_idx}; +use datafusion_common::utils::evaluate_partition_ranges; use datafusion_execution::{RecordBatchStream, TaskContext}; use datafusion_physical_expr::{LexOrdering, PhysicalExpr}; @@ -459,10 +459,10 @@ impl ExecutionPlan for PartialSortExec { input, expr: self.expr.clone(), common_prefix_length: self.common_prefix_length, - suffix_expr: LexOrdering::new( + suffix_ordering: LexOrdering::new( self.expr.iter().skip(self.common_prefix_length).cloned(), ), - in_mem_batches: vec![], + pending_prefix_batches: vec![], fetch: self.fetch, is_closed: false, baseline_metrics: BaselineMetrics::new(&self.metrics_set, partition), @@ -495,9 +495,9 @@ struct PartialSortStream { /// should be more than 0 otherwise PartialSort is not applicable common_prefix_length: usize, /// Sort expressions not already satisfied by the input ordering - suffix_expr: Option, + suffix_ordering: Option, /// Fragments of the trailing prefix group that is not ready for sort - in_mem_batches: Vec, + pending_prefix_batches: Vec, /// Fetch top N results fetch: Option, /// Whether the stream has finished returning all of its data or not @@ -506,21 +506,168 @@ struct PartialSortStream { baseline_metrics: BaselineMetrics, } +/// Describes completed prefix groups without concatenating their backing batches. +/// +/// The batches are treated as one logical batch. For example, batches with 3 and +/// 5 rows have `batch_ends = [3, 8]`. With `prefix_ends = [5, 8]`, the first +/// prefix consists of rows `0..3` from the first batch and rows `0..2` from the +/// second batch, while the second prefix consists of rows `2..5` from the second +/// batch. +/// +/// This mapping lets the multi-prefix path concatenate only suffix sort columns +/// for prefixes spanning batches. Payload columns remain in their original +/// batches until the final interleave materializes the output. #[derive(Debug)] -struct PrefixBatchRange { - batch_idx: usize, - range: Range, +struct CompletedPrefixes { + /// Backing batches containing only completed prefix groups. + batches: Vec, + /// Exclusive logical row offset for each backing batch. + batch_ends: Vec, + /// Exclusive logical row offset for each completed prefix group. + prefix_ends: Vec, } -fn whole_batches_prefix(batches: &[RecordBatch]) -> Vec { - batches - .iter() - .enumerate() - .map(|(batch_idx, batch)| PrefixBatchRange { - batch_idx, - range: 0..batch.num_rows(), +impl CompletedPrefixes { + fn new(batches: Vec, prefix_ends: Vec) -> Self { + debug_assert!(!batches.is_empty()); + debug_assert!(batches.iter().all(|batch| batch.num_rows() > 0)); + debug_assert!(prefix_ends.windows(2).all(|ends| ends[0] < ends[1])); + + let mut total_rows = 0; + let batch_ends = batches + .iter() + .map(|batch| { + total_rows += batch.num_rows(); + total_rows + }) + .collect::>(); + debug_assert_eq!(prefix_ends.last().copied(), Some(total_rows)); + + Self { + batches, + batch_ends, + prefix_ends, + } + } + + fn single(batches: Vec) -> Self { + let total_rows = batches.iter().map(RecordBatch::num_rows).sum(); + Self::new(batches, vec![total_rows]) + } + + fn num_prefixes(&self) -> usize { + self.prefix_ends.len() + } + + fn into_batches(self) -> Vec { + self.batches + } + + /// Returns each prefix range in the logical concatenation of `batches`. + fn prefix_ranges(&self) -> impl Iterator> + '_ { + self.prefix_ends.iter().copied().scan(0, |start, end| { + let range = *start..end; + *start = end; + Some(range) + }) + } + + /// Splits a logical row range into its physical batch-local ranges. + fn fragments( + &self, + logical_range: Range, + ) -> impl Iterator)> + '_ { + debug_assert!(logical_range.start < logical_range.end); + debug_assert!(logical_range.end <= self.batch_ends.last().copied().unwrap()); + let first_batch = self + .batch_ends + .partition_point(|batch_end| *batch_end <= logical_range.start); + let end_batch = self + .batch_ends + .partition_point(|batch_end| *batch_end < logical_range.end) + + 1; + (first_batch..end_batch).map(move |batch_index| { + let batch_end = self.batch_ends[batch_index]; + let batch_start = if batch_index == 0 { + 0 + } else { + self.batch_ends[batch_index - 1] + }; + let fragment_start = logical_range.start.max(batch_start); + let fragment_end = logical_range.end.min(batch_end); + ( + batch_index, + (fragment_start - batch_start)..(fragment_end - batch_start), + ) }) - .collect() + } + + /// Maps a logical row offset to the `(batch index, row index)` required by + /// the final interleave. + fn source_row(&self, logical_row: usize) -> (usize, usize) { + debug_assert!(logical_row < self.batch_ends.last().copied().unwrap()); + let last_batch_index = self.batches.len() - 1; + let last_batch_start = if last_batch_index == 0 { + 0 + } else { + self.batch_ends[last_batch_index - 1] + }; + if logical_row >= last_batch_start { + return (last_batch_index, logical_row - last_batch_start); + } + + let batch_index = self + .batch_ends + .partition_point(|batch_end| *batch_end <= logical_row); + let batch_start = if batch_index == 0 { + 0 + } else { + self.batch_ends[batch_index - 1] + }; + (batch_index, logical_row - batch_start) + } + + /// Builds the suffix sort columns for one prefix, concatenating only when + /// that prefix spans more than one backing batch. + fn gather_sort_columns( + &self, + prefix_range: Range, + sort_columns_by_batch: &[Vec], + suffix_ordering: &LexOrdering, + ) -> Result> { + let fragments = self.fragments(prefix_range).collect::>(); + suffix_ordering + .iter() + .enumerate() + .map(|(sort_index, expr)| { + let values = fragments + .iter() + .map(|(batch_index, row_range)| { + let values = + &sort_columns_by_batch[*batch_index][sort_index].values; + if row_range == &(0..self.batches[*batch_index].num_rows()) { + Arc::clone(values) + } else { + values.slice(row_range.start, row_range.len()) + } + }) + .collect::>(); + let values = if values.len() == 1 { + Arc::clone(&values[0]) + } else { + let values = values + .iter() + .map(|values| values.as_ref()) + .collect::>(); + concat(&values)? + }; + Ok(SortColumn { + values, + options: Some(expr.options), + }) + }) + .collect() + } } impl Stream for PartialSortStream { @@ -572,7 +719,7 @@ impl PartialSortStream { let prefix_ranges = self.get_prefix_ranges(&batch)?; let boundary_changed = self - .in_mem_batches + .pending_prefix_batches .last() .map(|previous| { self.prefix_changed_at_batch_boundary(previous, &batch) @@ -584,54 +731,44 @@ impl PartialSortStream { let trailing_range = prefix_ranges.last().unwrap(); let trailing_batch = batch.slice(trailing_range.start, trailing_range.len()); + let completed_batch = batch.slice(0, trailing_range.start); let mut completed_batches = - std::mem::take(&mut self.in_mem_batches); - let current_batch_idx = completed_batches.len(); - let mut completed_prefixes = if completed_batches.is_empty() { - vec![] - } else { - vec![whole_batches_prefix(&completed_batches)] - }; - - for (idx, range) in prefix_ranges[..prefix_ranges.len() - 1] + std::mem::take(&mut self.pending_prefix_batches); + let pending_rows = completed_batches .iter() - .cloned() - .enumerate() - { - let current_prefix = PrefixBatchRange { - batch_idx: current_batch_idx, - range, - }; - if idx == 0 - && !boundary_changed - && let Some(previous_prefix) = - completed_prefixes.last_mut() - { - previous_prefix.push(current_prefix); - continue; - } - completed_prefixes.push(vec![current_prefix]); + .map(RecordBatch::num_rows) + .sum::(); + let mut completed_prefix_ends = Vec::with_capacity( + prefix_ranges.len() + usize::from(boundary_changed), + ); + if pending_rows > 0 && boundary_changed { + completed_prefix_ends.push(pending_rows); } + completed_prefix_ends.extend( + prefix_ranges[..prefix_ranges.len() - 1] + .iter() + .map(|range| pending_rows + range.end), + ); - completed_batches.push(batch); - self.in_mem_batches.push(trailing_batch); - Some((completed_batches, completed_prefixes)) + completed_batches.push(completed_batch); + self.pending_prefix_batches.push(trailing_batch); + Some(CompletedPrefixes::new( + completed_batches, + completed_prefix_ends, + )) } else if boundary_changed { - let completed_batches = std::mem::take(&mut self.in_mem_batches); - let completed_prefix = whole_batches_prefix(&completed_batches); - self.in_mem_batches.push(batch); - Some((completed_batches, vec![completed_prefix])) + let completed_batches = + std::mem::take(&mut self.pending_prefix_batches); + self.pending_prefix_batches.push(batch); + Some(CompletedPrefixes::single(completed_batches)) } else { - self.in_mem_batches.push(batch); + self.pending_prefix_batches.push(batch); None }; - if let Some((completed_batches, completed_prefixes)) = completed { - let sorted_batch = self.sort_completed_prefixes( - completed_batches, - completed_prefixes, - )?; + if let Some(completed) = completed { + let sorted_batch = self.sort_completed_prefixes(completed)?; if sorted_batch.num_rows() > 0 { return Poll::Ready(Some(Ok(sorted_batch))); } @@ -644,15 +781,13 @@ impl PartialSortStream { let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); // Once input is consumed, the trailing prefix is complete. - let completed_batches = std::mem::take(&mut self.in_mem_batches); + let completed_batches = + std::mem::take(&mut self.pending_prefix_batches); if completed_batches.is_empty() { return Poll::Ready(None); } - let completed_prefix = whole_batches_prefix(&completed_batches); - let remaining_batch = self.sort_completed_prefixes( - completed_batches, - vec![completed_prefix], - )?; + let completed = CompletedPrefixes::single(completed_batches); + let remaining_batch = self.sort_completed_prefixes(completed)?; return if remaining_batch.num_rows() > 0 { Poll::Ready(Some(Ok(remaining_batch))) } else { @@ -665,149 +800,108 @@ impl PartialSortStream { fn sort_completed_prefixes( self: &mut Pin<&mut Self>, - completed_batches: Vec, - completed_prefixes: Vec>, + completed: CompletedPrefixes, ) -> Result { - debug_assert!(!completed_prefixes.is_empty()); - let result = if completed_prefixes.len() == 1 { - // A single prefix may span batches, so concatenate it at most once. - let prefix = completed_prefixes.into_iter().next().unwrap(); - let mut completed_batches = - completed_batches.into_iter().map(Some).collect::>(); - let mut prefix_batches = prefix - .into_iter() - .map(|prefix_range| { - let batch = completed_batches[prefix_range.batch_idx] - .take() - .expect("each source batch occurs once in a single prefix"); - if prefix_range.range == (0..batch.num_rows()) { - batch - } else { - batch.slice(prefix_range.range.start, prefix_range.range.len()) - } - }) - .collect::>(); - let batch = if prefix_batches.len() == 1 { - prefix_batches.pop().unwrap() - } else { - concat_batches(&self.schema(), &prefix_batches)? - }; + let result = if completed.num_prefixes() == 1 { + self.sort_single_prefix(completed.into_batches())? + } else { + self.sort_multiple_prefixes(&completed)? + }; - if let Some(suffix_expr) = &self.suffix_expr { - sort_batch(&batch, suffix_expr, self.fetch)? - } else { - let row_count = self - .fetch - .unwrap_or_else(|| batch.num_rows()) - .min(batch.num_rows()); - batch.slice(0, row_count) - } + if let Some(remaining_fetch) = self.fetch { + self.fetch = Some(remaining_fetch - result.num_rows()); + } + Ok(result) + } + + fn sort_single_prefix( + &self, + completed_batches: Vec, + ) -> Result { + // A single prefix may span batches, so concatenate it at most once. + let batch = if completed_batches.len() == 1 { + completed_batches.into_iter().next().unwrap() } else { - // Evaluate suffix expressions once per source batch, then sort each - // prefix independently and materialize the result with one interleave. - let sort_columns_by_batch = completed_batches - .iter() - .map(|batch| { - self.suffix_expr - .iter() - .flat_map(|exprs| exprs.iter()) - .map(|expr| expr.evaluate_to_sort_column(batch)) - .collect::>>() - }) - .collect::>>()?; - let mut remaining_fetch = self.fetch.unwrap_or(usize::MAX); - let mut interleave_indices = vec![]; + concat_batches(&self.schema(), &completed_batches)? + }; - for prefix in completed_prefixes { - if remaining_fetch == 0 { - break; - } + if let Some(suffix_ordering) = &self.suffix_ordering { + sort_batch(&batch, suffix_ordering, self.fetch) + } else { + let row_count = self + .fetch + .unwrap_or_else(|| batch.num_rows()) + .min(batch.num_rows()); + Ok(batch.slice(0, row_count)) + } + } - let row_count = prefix - .iter() - .map(|prefix_range| prefix_range.range.len()) - .sum::(); - let prefix_fetch = remaining_fetch.min(row_count); - let offsets = prefix + fn sort_multiple_prefixes( + &self, + completed: &CompletedPrefixes, + ) -> Result { + // Evaluate suffix expressions once per source batch, then sort each + // prefix independently and materialize the result with one interleave. + let sort_columns_by_batch = completed + .batches + .iter() + .map(|batch| { + self.suffix_ordering .iter() - .scan(0, |offset, prefix_range| { - let current = *offset; - *offset += prefix_range.range.len(); - Some(current) - }) - .collect::>(); - - let sorted_indices = if let Some(suffix_expr) = &self.suffix_expr { - let sort_columns = suffix_expr - .iter() - .enumerate() - .map(|(sort_idx, expr)| { - let values = prefix - .iter() - .map(|prefix_range| { - let values = &sort_columns_by_batch - [prefix_range.batch_idx][sort_idx] - .values; - if prefix_range.range - == (0..completed_batches[prefix_range.batch_idx] - .num_rows()) - { - Arc::clone(values) - } else { - values.slice( - prefix_range.range.start, - prefix_range.range.len(), - ) - } - }) - .collect::>(); - let values = if values.len() == 1 { - Arc::clone(&values[0]) - } else { - let values = values - .iter() - .map(|values| values.as_ref()) - .collect::>(); - concat(&values)? - }; - Ok(SortColumn { - values, - options: Some(expr.options), - }) - }) - .collect::>>()?; - lexsort_to_indices(&sort_columns, Some(prefix_fetch))? - .values() - .iter() - .map(|idx| *idx as usize) - .collect::>() - } else { - (0..prefix_fetch).collect() - }; + .flat_map(|exprs| exprs.iter()) + .map(|expr| expr.evaluate_to_sort_column(batch)) + .collect::>>() + }) + .collect::>>()?; + let mut remaining_fetch = self.fetch.unwrap_or(usize::MAX); + let mut interleave_indices = vec![]; - interleave_indices.extend(sorted_indices.into_iter().map(|idx| { - let range_idx = offsets.partition_point(|offset| *offset <= idx) - 1; - let prefix_range = &prefix[range_idx]; - ( - prefix_range.batch_idx, - prefix_range.range.start + idx - offsets[range_idx], - ) - })); - remaining_fetch -= prefix_fetch; + for prefix_range in completed.prefix_ranges() { + if remaining_fetch == 0 { + break; } - if interleave_indices.is_empty() { - RecordBatch::new_empty(self.schema()) + let row_count = prefix_range.len(); + let prefix_fetch = remaining_fetch.min(row_count); + let sorted_indices = if let Some(suffix_ordering) = &self.suffix_ordering { + let sort_columns = completed.gather_sort_columns( + prefix_range.clone(), + &sort_columns_by_batch, + suffix_ordering, + )?; + let fetch = (prefix_fetch < row_count).then_some(prefix_fetch); + lexsort_to_indices(&sort_columns, fetch)? + .values() + .iter() + .map(|idx| *idx as usize) + .collect::>() } else { - let completed_batches = completed_batches.iter().collect::>(); - interleave_record_batch(&completed_batches, &interleave_indices)? - } - }; + (0..prefix_fetch).collect() + }; - if let Some(remaining_fetch) = self.fetch { - self.fetch = Some(remaining_fetch - result.num_rows()); + interleave_indices.extend( + sorted_indices + .into_iter() + .map(|index| completed.source_row(prefix_range.start + index)), + ); + remaining_fetch -= prefix_fetch; + } + + if completed.batches[0].num_columns() == 0 { + let options = + RecordBatchOptions::new().with_row_count(Some(interleave_indices.len())); + Ok(RecordBatch::try_new_with_options( + self.schema(), + vec![], + &options, + )?) + } else { + let completed_batches = completed.batches.iter().collect::>(); + Ok(interleave_record_batch( + &completed_batches, + &interleave_indices, + )?) } - Ok(result) } fn get_prefix_ranges(&self, batch: &RecordBatch) -> Result>> { @@ -822,30 +916,24 @@ impl PartialSortStream { previous_batch: &RecordBatch, batch: &RecordBatch, ) -> Result { - let previous = previous_batch.slice(previous_batch.num_rows() - 1, 1); - let next = batch.slice(0, 1); - let previous_columns = (0..self.common_prefix_length) - .map(|idx| self.expr[idx].evaluate_to_sort_column(&previous)) - .collect::>>()?; - let next_columns = (0..self.common_prefix_length) - .map(|idx| self.expr[idx].evaluate_to_sort_column(&next)) - .collect::>>()?; - let previous_values = previous_columns - .iter() - .map(|column| Arc::clone(&column.values)) - .collect::>(); - let next_values = next_columns - .iter() - .map(|column| Arc::clone(&column.values)) - .collect::>(); - let previous_row = get_row_at_idx(&previous_values, 0)?; - let next_row = get_row_at_idx(&next_values, 0)?; - let sort_options = self.expr[..self.common_prefix_length] - .iter() - .map(|expr| expr.options) - .collect::>(); - - Ok(compare_rows(&previous_row, &next_row, &sort_options)? != Ordering::Equal) + debug_assert!(previous_batch.num_rows() > 0); + debug_assert!(batch.num_rows() > 0); + let previous_row = previous_batch.slice(previous_batch.num_rows() - 1, 1); + let current_row = batch.slice(0, 1); + + for sort_expr in self.expr.iter().take(self.common_prefix_length) { + let previous_key = sort_expr.evaluate_to_sort_column(&previous_row)?; + let current_key = sort_expr.evaluate_to_sort_column(¤t_row)?; + let comparator = make_comparator( + previous_key.values.as_ref(), + current_key.values.as_ref(), + sort_expr.options, + )?; + if !comparator(0, 0).is_eq() { + return Ok(true); + } + } + Ok(false) } } From ddd03a04b9a4e1fd450bcd19df5f991278ce8763 Mon Sep 17 00:00:00 2001 From: linfeng <33561138+lyne7-sc@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:09:34 +0800 Subject: [PATCH 3/3] bench: simplify partial sort benchmark --- .../physical-plan/benches/partial_sort.rs | 105 ++---------------- 1 file changed, 12 insertions(+), 93 deletions(-) diff --git a/datafusion/physical-plan/benches/partial_sort.rs b/datafusion/physical-plan/benches/partial_sort.rs index 7dda8303a187e..c418b219bbd09 100644 --- a/datafusion/physical-plan/benches/partial_sort.rs +++ b/datafusion/physical-plan/benches/partial_sort.rs @@ -31,44 +31,7 @@ use datafusion_physical_plan::{ExecutionPlan, collect}; const NUM_BATCHES: usize = 32; const BATCH_SIZE: usize = 8192; -const NUM_ROWS: usize = NUM_BATCHES * BATCH_SIZE; - -#[derive(Clone, Copy)] -enum PrefixLayout { - BatchesPerPrefix(usize), - PrefixesPerBatch(usize), - RowsPerPrefix(usize), -} - -impl PrefixLayout { - fn prefix_and_suffix(self, batch_idx: usize, row_idx: usize) -> (u64, u64) { - match self { - Self::BatchesPerPrefix(batches_per_prefix) => { - let rows_per_prefix = batches_per_prefix * BATCH_SIZE; - let offset = (batch_idx % batches_per_prefix) * BATCH_SIZE + row_idx; - ( - (batch_idx / batches_per_prefix) as u64, - (rows_per_prefix - offset - 1) as u64, - ) - } - Self::PrefixesPerBatch(prefixes_per_batch) => { - let rows_per_prefix = BATCH_SIZE / prefixes_per_batch; - let prefix_in_batch = row_idx / rows_per_prefix; - ( - (batch_idx * prefixes_per_batch + prefix_in_batch) as u64, - (rows_per_prefix - row_idx % rows_per_prefix - 1) as u64, - ) - } - Self::RowsPerPrefix(rows_per_prefix) => { - let global_idx = batch_idx * BATCH_SIZE + row_idx; - ( - (global_idx / rows_per_prefix) as u64, - (rows_per_prefix - global_idx % rows_per_prefix - 1) as u64, - ) - } - } - } -} +const ROWS_PER_PREFIX: &[usize] = &[100, 1_000, 5_000, 8_192, 10_000, 20_000]; fn schema() -> SchemaRef { Arc::new(Schema::new( @@ -86,20 +49,20 @@ fn schema() -> SchemaRef { )) } -fn make_batches(layout: PrefixLayout) -> Vec { +fn make_batches(rows_per_prefix: usize) -> Vec { let schema = schema(); (0..NUM_BATCHES) .map(|batch_idx| { - let rows = 0..BATCH_SIZE; + let rows = batch_idx * BATCH_SIZE..(batch_idx + 1) * BATCH_SIZE; let prefix = UInt64Array::from_iter_values( rows.clone() - .map(|row_idx| layout.prefix_and_suffix(batch_idx, row_idx).0), + .map(|row_idx| (row_idx / rows_per_prefix) as u64), ); - let suffix = UInt64Array::from_iter_values( - rows.clone() - .map(|row_idx| layout.prefix_and_suffix(batch_idx, row_idx).1), - ); - let payload = |row_idx: usize| (batch_idx * BATCH_SIZE + row_idx) as u64; + let suffix = + UInt64Array::from_iter_values(rows.clone().map(|row_idx| { + (rows_per_prefix - row_idx % rows_per_prefix - 1) as u64 + })); + let payload = |row_idx: usize| row_idx as u64; let payload_0 = UInt64Array::from_iter_values(rows.clone().map(payload)); let payload_1 = UInt64Array::from_iter_values( rows.clone() @@ -147,53 +110,9 @@ fn partial_sort_benchmark(c: &mut Criterion) { let mut group = c.benchmark_group("partial_sort"); group.sample_size(10); - let cases = [ - ( - BenchmarkId::new("prefix_across_batches", "batches_per_prefix=1"), - PrefixLayout::BatchesPerPrefix(1), - ), - ( - BenchmarkId::new("prefix_across_batches", "batches_per_prefix=2"), - PrefixLayout::BatchesPerPrefix(2), - ), - ( - BenchmarkId::new("prefix_across_batches", "batches_per_prefix=4"), - PrefixLayout::BatchesPerPrefix(4), - ), - ( - BenchmarkId::new("prefixes_within_batch", "prefixes_per_batch=2"), - PrefixLayout::PrefixesPerBatch(2), - ), - ( - BenchmarkId::new("prefixes_within_batch", "prefixes_per_batch=8"), - PrefixLayout::PrefixesPerBatch(8), - ), - ( - BenchmarkId::new("prefixes_within_batch", "prefixes_per_batch=32"), - PrefixLayout::PrefixesPerBatch(32), - ), - ( - BenchmarkId::new("prefixes_within_batch", "prefixes_per_batch=128"), - PrefixLayout::PrefixesPerBatch(128), - ), - ( - BenchmarkId::new("unaligned_mixed", "rows_per_prefix=1000"), - PrefixLayout::RowsPerPrefix(1000), - ), - ]; - - for (id, layout) in cases { - let batches = make_batches(layout); - let validation = runtime - .block_on(collect(make_plan(&batches), Arc::clone(&task_ctx))) - .unwrap(); - assert_eq!( - validation.iter().map(RecordBatch::num_rows).sum::(), - NUM_ROWS - ); - drop(validation); - - group.bench_function(id, |b| { + for &rows_per_prefix in ROWS_PER_PREFIX { + let batches = make_batches(rows_per_prefix); + group.bench_function(BenchmarkId::new("rows_per_prefix", rows_per_prefix), |b| { b.iter_batched( || make_plan(&batches), |plan| {