Skip to content

perf: prune window state only for partitions that made progress #24148

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
neilconway wants to merge 1 commit into
base: main
Choose a base branch
from
+66 −54
Open

perf: prune window state only for partitions that made progress #24148

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion datafusion/expr/src/window_state.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,7 @@ pub struct PartitionBatchState {
pub record_batch: RecordBatch,
/// Flag indicating whether we have received all data for this partition
pub is_end: bool,
/// Number of rows emitted for each partition
/// Number of rows emitted for this partition since the last pruning pass
pub n_out_row: usize,
}

Expand Down
118 changes: 65 additions & 53 deletions datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1005,14 +1005,9 @@ pub struct BoundedWindowAggStream {
/// The record batch executor receives as input (i.e. the columns needed
/// while calculating aggregation results).
input_buffer: RecordBatch,
/// We separate `input_buffer` based on partitions (as
/// determined by PARTITION BY columns) and store them per partition
/// in `partition_batches`. We use this variable when calculating results
/// for each window expression. This enables us to use the same batch for
/// different window expressions without copying.
// Note that we could keep record batches for each window expression in
// `PartitionWindowAggStates`. However, this would use more memory (as
// many times as the number of window expressions).
/// Each partition's rows, accumulated across input batches. All window
/// expressions calculate their results against these shared rows without
/// copying.
partition_buffers: PartitionBatches,
/// An executor can run multiple window expressions if the PARTITION BY
/// and ORDER BY sections are same. We keep state of the each window
Expand Down Expand Up @@ -1049,13 +1044,13 @@ impl BoundedWindowAggStream {
/// results (as determined by window frame boundaries and number of results generated).
// For instance, if first `n` (not necessarily same with `n_out`) elements are no longer needed to
// calculate window expression result (outside the window frame boundary) we retract first `n` elements
// from `self.partition_batches` in corresponding partition.
// from the corresponding partition's batch in `self.partition_buffers`.
// For instance, if `n_out` number of rows are calculated, we can remove
// first `n_out` rows from `self.input_buffer`.
fn prune_state(&mut self, n_out: usize) -> Result<()> {
// Prune `self.window_agg_states`:
self.prune_out_columns();
// Prune `self.partition_batches`:
// Prune `self.partition_buffers`:
self.prune_partition_batches();
// Prune `self.input_buffer`:
self.prune_input_batch(n_out)?;
Expand Down Expand Up @@ -1191,51 +1186,62 @@ impl BoundedWindowAggStream {
}
}

/// Prunes the sections of the record batch (for each partition)
/// that we no longer need to calculate the window function result.
/// Removes partitions that have ended. For the remaining partitions,
/// drops buffered rows that no window expression will need again.
fn prune_partition_batches(&mut self) {
// Remove partitions which we know already ended (is_end flag is true).
// Since the retain method preserves insertion order, we still have
// ordering in between partitions after removal.
self.partition_buffers
.retain(|_, partition_batch_state| !partition_batch_state.is_end);

// The data in `self.partition_batches` is used by all window expressions.
// Therefore, when removing from `self.partition_batches`, we need to remove
// from the earliest range boundary among all window expressions. Variable
// `n_prune_each_partition` fill the earliest range boundary information for
// each partition. This way, we can delete the no-longer-needed sections from
// `self.partition_batches`.
// For instance, if window frame one uses [10, 20] and window frame two uses
// [5, 15]; we only prune the first 5 elements from the corresponding record
// batch in `self.partition_batches`.

// Calculate how many elements to prune for each partition batch
// Calculate how many rows to prune from each partition's batch. For a
// single window expression, rows before min(window_frame_range.start,
// last_calculated_index) are prunable: their results are already
// calculated, and frame boundaries never move backwards, so no future
// frame can include them. All window expressions share the partition
// batch, so a row can only be pruned once every expression is done with
// it: the count to prune is the minimum across expressions. A partition
// missing from the map has nothing to prune.
let mut n_prune_each_partition = HashMap::new();
let mut first = true;
for window_agg_state in self.window_agg_states.iter_mut() {
window_agg_state.retain(|_, WindowState { state, .. }| !state.is_end);
for (partition_row, WindowState { state: value, .. }) in window_agg_state {
let n_prune =
min(value.window_frame_range.start, value.last_calculated_index);
if let Some(current) = n_prune_each_partition.get_mut(partition_row) {
if n_prune < *current {
*current = n_prune;
if first {
// First window expression seeds the prune-count map
first = false;
for (partition_row, WindowState { state, .. }) in window_agg_state.iter()
{
let n_prune =
min(state.window_frame_range.start, state.last_calculated_index);
if n_prune > 0 {
n_prune_each_partition.insert(partition_row.clone(), n_prune);
}
} else {
n_prune_each_partition.insert(partition_row.clone(), n_prune);
}
} else {
// Take the per-partition min of the prune-count for each
// additional window expression
n_prune_each_partition.retain(|partition_row, current| {
let Some(WindowState { state, .. }) =
window_agg_state.get(partition_row)
else {
return false;
};
let n_prune =
min(state.window_frame_range.start, state.last_calculated_index);
*current = min(*current, n_prune);
*current > 0
});
}
}

// Retract no longer needed parts during window calculations from partition batch:
// Drop the prunable prefix of each partition's buffered batch:
for (partition_row, n_prune) in n_prune_each_partition.iter() {
debug_assert!(
*n_prune > 0,
"prune-count map must only contain positive entries"
);
let pb_state = &mut self.partition_buffers[partition_row];
pb_state.n_out_row = 0;

// If there is nothing to prune, leave the batch as-is
if *n_prune == 0 {
continue;
}

let batch = &pb_state.record_batch;
pb_state.record_batch = batch.slice(*n_prune, batch.num_rows() - n_prune);
Expand Down Expand Up @@ -1272,23 +1278,29 @@ impl BoundedWindowAggStream {
// field of `WindowAggState`. Given how many rows are emitted, we remove
// these sections from state.
for partition_window_agg_states in self.window_agg_states.iter_mut() {
// Remove `n_out` entries from the `out_col` field of `WindowAggState`.
// `n_out` is stored in `self.partition_buffers` for each partition.
// If `is_end` is set, directly remove them; this shrinks the hash map.
// If `is_end` is set, directly remove the entry; this shrinks the
// hash map.
partition_window_agg_states
.retain(|_, partition_batch_state| !partition_batch_state.state.is_end);
for (
partition_key,
WindowState {
state: WindowAggState { out_col, .. },
..
},
) in partition_window_agg_states
{
let partition_batch = &mut self.partition_buffers[partition_key];
let n_to_del = partition_batch.n_out_row;
let n_to_keep = out_col.len() - n_to_del;
*out_col = out_col.slice(n_to_del, n_to_keep);
}
// Only partitions that emitted rows since the previous pruning pass
// have output columns to shrink. Their emitted-row counts are
// consumed and reset here, so partitions that emitted nothing keep
// a count of zero and are passed over without any hash lookups.
for (partition_key, partition_batch) in self.partition_buffers.iter_mut() {
let n_emitted = partition_batch.n_out_row;
if n_emitted == 0 {
continue;
}
partition_batch.n_out_row = 0;
for partition_window_agg_states in self.window_agg_states.iter_mut() {
if let Some(WindowState { state, .. }) =
partition_window_agg_states.get_mut(partition_key)
{
let out_col = &mut state.out_col;
let n_to_keep = out_col.len() - n_emitted;
*out_col = out_col.slice(n_emitted, n_to_keep);
}
}
}
}
Expand Down
Loading