avantgardnerio commented on code in PR #24035:
URL: https://github.com/apache/datafusion/pull/24035#discussion_r3741808151
##########
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:
Thanks, sorry I left that in.
--
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]
