alamb commented on code in PR #24035:
URL: https://github.com/apache/datafusion/pull/24035#discussion_r3738845164
##########
datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs:
##########
@@ -1955,6 +2063,411 @@ mod tests {
Ok(())
}
+ type Observation = (usize, PartitionKey, Vec<Option<Vec<ScalarValue>>>);
+
+ /// Test [`WindowStateObserver`] that records every callback into a shared
+ /// `Vec` for later assertion.
+ struct RecordingObserver {
+ sink: Arc<std::sync::Mutex<Vec<Observation>>>,
+ }
+
+ impl WindowStateObserver for RecordingObserver {
+ fn finalized(
+ &self,
+ partition_idx: usize,
+ partition_key: &PartitionKey,
+ states: &[Option<Vec<ScalarValue>>],
+ ) -> Result<()> {
+ self.sink.lock().unwrap().push((
+ partition_idx,
+ partition_key.clone(),
+ states.to_vec(),
+ ));
+ Ok(())
+ }
+ }
+
+ #[tokio::test]
+ async fn test_finalized_state_observer_fires_at_partition_close() ->
Result<()> {
+ use std::sync::Mutex;
+
+ let task_ctx = Arc::new(TaskContext::default());
+ let schema = test_schema();
+
+ // Two PARTITION BY groups: hash=1 [sn=1,2,3] then hash=2 [sn=4,5,6].
+ // Input is sorted by (hash, sn) so we can run in Sorted mode; in that
+ // mode `mark_partition_end` closes the leading group mid-stream and
+ // EOS closes the tail — both should fire the observer.
+ let mut sn_b = UInt64Builder::with_capacity(6);
Review Comment:
these tests appear to have a lot of repeated setup, which makes it hard to
see what is different between them. Could we extract it to a helper perhaps?
##########
datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs:
##########
@@ -1042,9 +1103,45 @@ pub struct BoundedWindowAggStream {
/// partitions, so finished partitions are pruned eagerly instead and no
/// such bound is needed.
most_recent_row: Option<RecordBatch>,
+ /// Output partition index this stream serves; passed as the first
+ /// argument to [`WindowStateObserver::finalized`].
+ partition_idx: usize,
+ /// If set, invoked from [`Self::publish_finalized_states`] with the
+ /// finalized per-window-expression state for every partition key that is
+ /// about to be dropped.
+ state_observer: Option<Arc<dyn WindowStateObserver>>,
}
impl BoundedWindowAggStream {
+ /// Fire the [`WindowStateObserver`] for every partition key whose
+ /// `WindowAggState::is_end` is true.
+ fn publish_finalized_states(&mut self) -> Result<()> {
Review Comment:
SOmething doesn't seem right to me that this method takes a `&mut` and
modifies the internal state, but only when there is an observer 🤔 that seems
complicated to reason about
##########
datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs:
##########
@@ -140,9 +187,17 @@ impl BoundedWindowAggExec {
ordered_partition_by_indices,
cache: Arc::new(cache),
can_repartition,
+ state_observer: None,
})
}
+ /// Install a [`WindowStateObserver`] that receives each PARTITION BY
+ /// group's finalized window state at partition close.
+ pub fn with_state_observer(mut self, observer: Arc<dyn
WindowStateObserver>) -> Self {
Review Comment:
I think a more "idomatic" API here would be to take the same type as the
field (aka `Arc<dyn WindowStateObserver>`) -- that would allow both:
1. clearing the field
2. Skipping the checks below that are looking for Some
```rust
pub fn with_state_observer(mut self, observer: Option<Arc<dyn
WindowStateObserver>>) -> Self {
```
Which I think would let
```rust
let mut new = BoundedWindowAggExec::try_new(
self.window_expr.clone(),
Arc::clone(&children[0]),
self.input_order_mode.clone(),
self.can_repartition,
)?;
if let Some(observer) = &self.state_observer {
new = new.with_state_observer(Arc::clone(observer));
}
```
Become
```rust
let mut new = BoundedWindowAggExec::try_new(
self.window_expr.clone(),
Arc::clone(&children[0]),
self.input_order_mode.clone(),
self.can_repartition,
)?
.with_state_observer(observer.cloned())
```
##########
datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs:
##########
@@ -1042,9 +1103,45 @@ pub struct BoundedWindowAggStream {
/// partitions, so finished partitions are pruned eagerly instead and no
/// such bound is needed.
most_recent_row: Option<RecordBatch>,
+ /// Output partition index this stream serves; passed as the first
+ /// argument to [`WindowStateObserver::finalized`].
+ partition_idx: usize,
+ /// If set, invoked from [`Self::publish_finalized_states`] with the
+ /// finalized per-window-expression state for every partition key that is
+ /// about to be dropped.
+ state_observer: Option<Arc<dyn WindowStateObserver>>,
}
impl BoundedWindowAggStream {
+ /// Fire the [`WindowStateObserver`] for every partition key whose
+ /// `WindowAggState::is_end` is true.
+ fn publish_finalized_states(&mut self) -> Result<()> {
+ let Some(observer) = &self.state_observer else {
+ return Ok(());
+ };
+ let Some((first, rest)) = self.window_agg_states.split_first_mut()
else {
+ return Ok(());
+ };
+ for (key, ws) in first.iter_mut() {
+ if !ws.state.is_end {
+ continue;
+ }
+ let mut states: Vec<Option<Vec<ScalarValue>>> =
Review Comment:
it is unfortunate that we have to make a new allocation for each partition
Would it work to call `observer.finalized(self.partition_idx, key,
ws.aggregate_state()?)?
(aka call the observer once per window function?)
##########
datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs:
##########
@@ -76,8 +76,29 @@ use hashbrown::hash_table::HashTable;
use indexmap::IndexMap;
use log::debug;
+/// Callback receiver for per-partition window state.
+pub trait WindowStateObserver: Send + Sync {
+ /// Invoked once per (output-partition-index, PARTITION BY tuple) as each
+ /// PARTITION BY group closes.
+ ///
+ /// # Arguments
+ ///
+ /// * `partition_idx` - Output partition index of the
[`BoundedWindowAggExec`]
+ /// stream firing this callback.
+ /// * `partition_key` - The PARTITION BY tuple that just closed.
+ /// * `states` - One entry per window expression on the exec, in the same
+ /// order as [`BoundedWindowAggExec::window_expr`]; `None` for
+ /// non-aggregate window functions.
+ fn finalized(
+ &self,
Review Comment:
How will you know what window function this method call belongs to if there
are multiple window functions ? I see the note about the order on
[`BoundedWindowAggExec::window_expr`] but then you would have to know about
this before hand
Is it possible to call this with the relevant
[`BoundedWindowAggExec::window_expr`] ?
--
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]
