avantgardnerio commented on code in PR #24148:
URL: https://github.com/apache/datafusion/pull/24148#discussion_r3757463181
##########
datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs:
##########
@@ -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;
Review Comment:
nit: `split_first_mut()` would be slightly more idiomatic
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
