You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/04 17:32:34 UTC
(datafusion) branch gh-readonly-queue/main/pr-24034-9114e78953bad29dc0acf5ba230a8b839b1fe605 created (now 179b32c9b6)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-24034-9114e78953bad29dc0acf5ba230a8b839b1fe605
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 179b32c9b6 perf: gather Linear-mode window input more efficiently (#24034)
This branch includes the following new commits:
new 179b32c9b6 perf: gather Linear-mode window input more efficiently (#24034)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: perf: gather Linear-mode window input more efficiently (#24034)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24034-9114e78953bad29dc0acf5ba230a8b839b1fe605
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 179b32c9b60103d9c4e6a4364f10f6286c963904
Author: Neil Conway <ne...@gmail.com>
AuthorDate: Tue Aug 4 13:32:13 2026 -0400
perf: gather Linear-mode window input more efficiently (#24034)
## Which issue does this PR close?
- Part of #23982
## Rationale for this change
`LinearSearch::evaluate_partition_batches` issued one
`take_record_batch` call per partition present in the input batch. For
batches with many partitions, this is inefficient.
Instead, we can build a Vec of the batch's row indices that groups the
rows by partition, gather all rows with a single `take_record_batch`
call, and hand each partition a slice of the result. Batches that
contain a single partition skip the gather entirely.
Memory caveat: the emitted slices share the gathered batch's buffers. A
partition that never receives rows again retains its slice and therefore
pins the gathered batch's buffers (up to one input batch worth of memory
per input batch in the worst case). This could be addressed, e.g., with
a compaction pass to copy long-lived slices into owned buffers, but I
have omitted that for now.
Benchmarks: (using #24032)
- linear 100 partitions: 44.3 ms -> 44.1 ms (within noise)
- linear 10000 partitions: 199.8 ms -> 170.1 ms (-14.9%)
- linear sparse 32768 partitions: 224.6 ms -> 205.3 ms (-8.6%)
- linear rows 10000 partitions: 169.0 ms -> 142.0 ms (-15.9%)
- linear multi 10000 partitions: 295.9 ms -> 268.2 ms (-9.4%)
- sorted 10000 partitions: 34.0 ms -> 34.5 ms (+1.3%; unchanged code
path)
## What changes are included in this PR?
- Rewrite `get_per_partition_indices` and rename to
`compute_partition_permutation`
- Add focused unit test for computing partition permutations correctly
## Are these changes tested?
Yes, new test added.
## Are there any user-facing changes?
No.
---
.../src/windows/bounded_window_agg_exec.rs | 184 +++++++++++++++++----
1 file changed, 150 insertions(+), 34 deletions(-)
diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs
index 0ed751f506..03a8e9867c 100644
--- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs
+++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs
@@ -43,7 +43,7 @@ use crate::{
use arrow::compute::take_record_batch;
use arrow::{
- array::{Array, ArrayRef, RecordBatchOptions, UInt32Builder},
+ array::{Array, ArrayRef, RecordBatchOptions, UInt32Array, UInt32Builder},
compute::{concat, concat_batches, sort_to_indices, take_arrays},
datatypes::SchemaRef,
record_batch::RecordBatch,
@@ -664,17 +664,25 @@ impl PartitionSearcher for LinearSearch {
evaluate_partition_by_column_values(record_batch, window_expr)?;
// NOTE: In Linear or PartiallySorted modes, we are sure that
// `partition_bys` are not empty.
- // Calculate indices for each partition and construct a new record
- // batch from the rows at these indices for each partition:
- self.get_per_partition_indices(&partition_bys, record_batch)?
+ let (mut keys, permutation, bounds) =
+ self.compute_partition_permutation(&partition_bys, record_batch)?;
+ if keys.len() == 1 {
+ // The batch contains a single partition, so the gather below
+ // would be an identity permutation; use the batch as-is.
+ let key = keys.remove(0);
+ return Ok(vec![(key, record_batch.clone())]);
+ }
+ // Reorder the batch with a single `take` so that each partition's
+ // rows become contiguous, then hand each partition a zero-copy slice
+ // of the result. The slices share the gathered batch's buffers;
+ // `PartitionBatchState::extend` copies out of them the next time the
+ // partition receives rows.
+ let gathered = take_record_batch(record_batch, &UInt32Array::from(permutation))?;
+ Ok(keys
.into_iter()
- .map(|(row, indices)| {
- let mut new_indices = UInt32Builder::with_capacity(indices.len());
- new_indices.append_slice(&indices);
- let indices = new_indices.finish();
- Ok((row, take_record_batch(record_batch, &indices)?))
- })
- .collect()
+ .zip(bounds.windows(2))
+ .map(|(key, bound)| (key, gathered.slice(bound[0], bound[1] - bound[0])))
+ .collect())
}
fn prune(&mut self, n_out: usize) {
@@ -728,42 +736,66 @@ impl LinearSearch {
}
}
- /// Calculate indices of each partition (according to PARTITION BY expression)
- /// `columns` contain partition by expression results.
- fn get_per_partition_indices(
+ /// Splits the rows of `batch` by partition, according to the PARTITION BY
+ /// expression results in `columns`. Returns the distinct partition keys
+ /// in first-appearance order, a permutation of the row indices of
+ /// `batch` that groups each partition's rows together, and the
+ /// boundaries of each partition's run of rows within that permutation:
+ /// partition `p` occupies `permutation[bounds[p]..bounds[p + 1]]`, and
+ /// its indices are in ascending (stream) order.
+ fn compute_partition_permutation(
&mut self,
columns: &[ArrayRef],
batch: &RecordBatch,
- ) -> Result<Vec<(PartitionKey, Vec<u32>)>> {
- let mut batch_hashes = vec![0; batch.num_rows()];
+ ) -> Result<(Vec<PartitionKey>, Vec<u32>, Vec<usize>)> {
+ let num_rows = batch.num_rows();
+ let mut batch_hashes = vec![0; num_rows];
create_hashes(columns, &self.random_state, &mut batch_hashes)?;
self.input_buffer_hashes.extend(&batch_hashes);
// reset row_map for new calculation
self.row_map_batch.clear();
- // res stores PartitionKey and row indices (indices where these partition occurs in the `batch`) for each partition.
- let mut result: Vec<(PartitionKey, Vec<u32>)> = vec![];
+ let mut keys: Vec<PartitionKey> = vec![];
+ // Partition id of each row, in row order:
+ let mut row_partition_ids = Vec::with_capacity(num_rows);
+ // Number of rows in each partition:
+ let mut counts: Vec<usize> = vec![];
for (hash, row_idx) in batch_hashes.into_iter().zip(0u32..) {
let entry = self.row_map_batch.find_mut(hash, |(_, group_idx)| {
- // We can safely get the first index of the partition indices
- // since partition indices has one element during initialization.
let row = get_row_at_idx(columns, row_idx as usize).unwrap();
- // Handle hash collusions with an equality check:
- row.eq(&result[*group_idx].0)
+ // Handle hash collisions with an equality check:
+ row == keys[*group_idx]
});
- if let Some((_, group_idx)) = entry {
- result[*group_idx].1.push(row_idx)
+ let group_idx = if let Some((_, group_idx)) = entry {
+ *group_idx
} else {
- self.row_map_batch.insert_unique(
- hash,
- (hash, result.len()),
- |(hash, _)| *hash,
- );
- let row = get_row_at_idx(columns, row_idx as usize)?;
- // This is a new partition its only index is row_idx for now.
- result.push((row, vec![row_idx]));
- }
+ let group_idx = keys.len();
+ self.row_map_batch
+ .insert_unique(hash, (hash, group_idx), |(hash, _)| *hash);
+ keys.push(get_row_at_idx(columns, row_idx as usize)?);
+ counts.push(0);
+ group_idx
+ };
+ row_partition_ids.push(group_idx);
+ counts[group_idx] += 1;
+ }
+ // A prefix sum over the counts gives each partition's run boundaries
+ // in the permutation.
+ let mut bounds = Vec::with_capacity(counts.len() + 1);
+ let mut total = 0;
+ bounds.push(0);
+ for count in counts {
+ total += count;
+ bounds.push(total);
+ }
+ // Scatter each row's index into its partition's run. Visiting rows
+ // in ascending order keeps each run in ascending row order.
+ let mut cursors: Vec<usize> = bounds[..bounds.len() - 1].to_vec();
+ let mut permutation = vec![0u32; num_rows];
+ for (row_idx, group_idx) in row_partition_ids.into_iter().enumerate() {
+ permutation[cursors[group_idx]] = row_idx as u32;
+ cursors[group_idx] += 1;
}
- Ok(result)
+ Ok((keys, permutation, bounds))
}
/// Calculates partition keys and result indices for each partition.
@@ -1939,4 +1971,88 @@ mod tests {
));
Ok(())
}
+
+ /// Checks the per-partition batches that `LinearSearch` splits an input
+ /// batch into: partitions appear in first-appearance order, rows within a
+ /// partition keep their stream order, NULL keys form their own partition,
+ /// and a single-partition batch is passed through without copying.
+ #[test]
+ fn test_linear_search_evaluate_partition_batches() -> Result<()> {
+ use super::{LinearSearch, PartitionSearcher};
+ use arrow::array::{Int32Array, Int64Array};
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Int64, false),
+ ]));
+ let window_expr = create_window_expr(
+ &WindowFunctionDefinition::AggregateUDF(count_udaf()),
+ "count".to_string(),
+ &[col("b", &schema)?],
+ &[col("a", &schema)?],
+ &[],
+ Arc::new(WindowFrame::new(None)),
+ Arc::clone(&schema),
+ false,
+ false,
+ None,
+ )?;
+ let mut searcher = LinearSearch::new(vec![], Arc::clone(&schema));
+
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(vec![
+ Some(1),
+ Some(2),
+ Some(1),
+ None,
+ Some(2),
+ Some(1),
+ ])),
+ Arc::new(Int64Array::from(vec![10, 20, 11, 30, 21, 12])),
+ ],
+ )?;
+ let result =
+ searcher.evaluate_partition_batches(&batch, &[Arc::clone(&window_expr)])?;
+ assert_eq!(result.len(), 3);
+ let expected = [
+ (
+ ScalarValue::Int32(Some(1)),
+ vec![Some(1); 3],
+ vec![10i64, 11, 12],
+ ),
+ (ScalarValue::Int32(Some(2)), vec![Some(2); 2], vec![20, 21]),
+ (ScalarValue::Int32(None), vec![None], vec![30]),
+ ];
+ for ((key, partition_batch), (exp_key, exp_a, exp_b)) in
+ result.iter().zip(expected)
+ {
+ assert_eq!(key, &vec![exp_key]);
+ let exp_batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(exp_a)),
+ Arc::new(Int64Array::from(exp_b)),
+ ],
+ )?;
+ assert_eq!(partition_batch, &exp_batch);
+ }
+
+ let single = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(vec![Some(7), Some(7)])),
+ Arc::new(Int64Array::from(vec![70, 71])),
+ ],
+ )?;
+ let result = searcher.evaluate_partition_batches(&single, &[window_expr])?;
+ assert_eq!(result.len(), 1);
+ assert_eq!(result[0].0, vec![ScalarValue::Int32(Some(7))]);
+ assert_eq!(result[0].1, single);
+ // The whole batch belongs to one partition, so its columns are reused
+ // rather than gathered into a new batch.
+ assert!(Arc::ptr_eq(result[0].1.column(0), single.column(0)));
+ Ok(())
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org