alamb commented on code in PR #24035:
URL: https://github.com/apache/datafusion/pull/24035#discussion_r3723357151
##########
datafusion/physical-expr/src/window/window_expr.rs:
##########
@@ -647,6 +647,19 @@ pub struct WindowState {
pub state: WindowAggState,
pub window_fn: WindowFn,
}
+
+impl WindowState {
+ /// `Accumulator::state()` if this window function is an aggregate, `None`
+ /// otherwise (built-in functions like `row_number`, `rank`, `lead`/`lag`
+ /// have no serializable accumulator state).
+ pub fn aggregate_state(&mut self) -> Result<Option<Vec<ScalarValue>>> {
+ match &mut self.window_fn {
+ WindowFn::Aggregate(accumulator) => accumulator.state().map(Some),
+ WindowFn::Builtin(_) => Ok(None),
Review Comment:
I was confused about what `BuiltIn` was -- it turns out I think it is just a
normal window function (the name probably dates to the time when we had built
in and user defined window functions)
##########
datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs:
##########
@@ -140,9 +176,20 @@ impl BoundedWindowAggExec {
ordered_partition_by_indices,
cache: Arc::new(cache),
can_repartition,
+ finalized_state_observer: None,
})
}
+ /// Install a callback that receives each PARTITION BY group's finalized
+ /// window state at partition close.
+ pub fn with_finalized_state_observer(
Review Comment:
As above, I recommend making this more general (`with_state_observer`) to
make it flexible going forward
##########
datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs:
##########
@@ -76,8 +76,18 @@ use hashbrown::hash_table::HashTable;
use indexmap::IndexMap;
use log::debug;
+/// Called by [`BoundedWindowAggExec`] when a PARTITION BY group closes,
+/// once per (output-partition-index, PARTITION BY tuple). The third argument
+/// is one entry per window expression on the exec, in the same order as
Review Comment:
For example, it may make sense to have a a function for each kind of window
function that also passes along information about the window function /
aggregate
Then I could imagine we could add functions to the various aggregates and
window functions to manipulate this state
If you don't pass along the function, I think you would have to special case
the states as well which seems much harder to integrate
So for example, something like
```rust
/// called at the end of each partition for each window aggegate
fn finalize_window_aggregate(&self, aggregate_func, state) {}
```
Which then leave space for a function like
```rust
/// called at the end of each partition for each window function
fn finalize_window_function(&self, window_func, window_state) {}
```
In the future and makes it easier to explain what the third argument does
##########
datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs:
##########
@@ -1053,6 +1113,33 @@ impl BoundedWindowAggStream {
// 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<()> {
+ // `WindowAggState::is_end` is copied from
`PartitionBatchState::is_end`
+ // during `evaluate_stateful`, so every window_expr's state map agrees
+ // on which partition keys have just closed. Publish those keys before
+ // the retains in `prune_out_columns` / `prune_partition_batches` drop
+ // them.
+ if let Some(observer) = self.finalized_state_observer.clone()
Review Comment:
I don't think you need to clone
```suggestion
if let Some(observer) = &self.finalized_state_observer
```
##########
datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs:
##########
@@ -76,8 +76,18 @@ use hashbrown::hash_table::HashTable;
use indexmap::IndexMap;
use log::debug;
+/// Called by [`BoundedWindowAggExec`] when a PARTITION BY group closes,
+/// once per (output-partition-index, PARTITION BY tuple). The third argument
+/// is one entry per window expression on the exec, in the same order as
Review Comment:
I suggest that this is a real trait rather than a typedef:
```rust
pub trait WindowStateObserver {
fn finalized(...)
}
```
That will:
1. Let us potentially extend it in the future with other functions (e.g.
maybe window start?)
2. Let us name the parameters (though this PR does explain them in the
comments)
3. It would mean the types are `Option<Arc<dyn
FinalizedWindowStateObserver>>` which I think would make it clearer that
cloning it just an Arc::clone
##########
datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs:
##########
@@ -1053,6 +1113,33 @@ impl BoundedWindowAggStream {
// 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<()> {
+ // `WindowAggState::is_end` is copied from
`PartitionBatchState::is_end`
+ // during `evaluate_stateful`, so every window_expr's state map agrees
+ // on which partition keys have just closed. Publish those keys before
+ // the retains in `prune_out_columns` / `prune_partition_batches` drop
+ // them.
+ if let Some(observer) = self.finalized_state_observer.clone()
+ && !self.window_agg_states.is_empty()
+ {
+ let closed_keys: Vec<PartitionKey> = self.window_agg_states[0]
Review Comment:
It looks to me like this is copying *all* the intermediate state, which
could be quite large and will be called once per partition. I realize in your
case there is one partition so maybe this overhead isn't that bad, but I think
it would be better to pass in the references (e.g `&[PartitionKey]` or
something) and the observer can copy them if it wants, rather than forcing a
copy
Same thing for the agg states
--
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]
