kosiew commented on code in PR #25424:
URL: https://github.com/apache/datafusion/pull/25424#discussion_r4046043832
##########
datafusion/physical-plan/src/aggregates/hash_stream.rs:
##########
@@ -1607,4 +1623,247 @@ mod tests {
Ok(())
}
+
+ #[derive(Clone, Copy)]
+ enum Finish {
+ Collect,
+ DropDuringReplay,
+ InputError,
+ }
+
+ #[tokio::test]
+ async fn final_hash_spill_replay_with_other_partitions_holding_state() ->
Result<()> {
+ for input_batches in [40, 55, 70] {
+ run_shared_pool_case(input_batches, 1024 * 1024,
Finish::Collect).await?;
+ }
+ // The same input also produces the reference results without spilling.
+ run_shared_pool_case(70, 10 * 1024 * 1024, Finish::Collect).await
+ }
+
+ #[tokio::test]
+ async fn final_hash_spill_replay_releases_memory_on_drop() -> Result<()> {
+ run_shared_pool_case(55, 1024 * 1024, Finish::DropDuringReplay).await
+ }
+
+ #[tokio::test]
+ async fn final_hash_spill_releases_memory_on_input_error() -> Result<()> {
+ run_shared_pool_case(55, 1024 * 1024, Finish::InputError).await
+ }
+
+ /// Partition 0 spills and replays while partitions 1..3 keep their
+ /// aggregate state in the same greedy pool.
+ async fn run_shared_pool_case(
+ input_batches: i64,
+ limit: usize,
+ finish: Finish,
+ ) -> Result<()> {
+ const PARTITIONS: usize = 4;
+ const HELD_BATCHES: i64 = 20;
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ Field::new("v", DataType::Int64, false),
+ Field::new("s", DataType::Utf8View, false),
+ ]));
+ let groups = PhysicalGroupBy::new_single(vec![
+ (col("a", &schema)?, "a".into()),
+ (col("b", &schema)?, "b".into()),
+ ]);
+ let expressions = vec![
+ Arc::new(
+ AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?])
+ .schema(Arc::clone(&schema))
+ .alias("sum")
+ .build()?,
+ ),
+ Arc::new(
+ AggregateExprBuilder::new(min_udaf(), vec![col("s", &schema)?])
+ .schema(Arc::clone(&schema))
+ .alias("min")
+ .build()?,
+ ),
+ ];
+ let empty = TestMemoryExec::try_new_exec(&[vec![]],
Arc::clone(&schema), None)?;
+ let partial = AggregateExec::try_new(
+ AggregateMode::Partial,
+ groups.clone(),
+ expressions.clone(),
+ vec![None; 2],
+ empty,
+ Arc::clone(&schema),
+ )?;
+ let partial_schema = partial.schema();
+ let input = TestMemoryExec::try_new_exec(
+ &vec![vec![]; PARTITIONS],
+ Arc::clone(&partial_schema),
+ None,
+ )?;
+ let aggregate = AggregateExec::try_new(
+ AggregateMode::FinalPartitioned,
+ groups.as_final(),
+ expressions,
+ vec![None; 2],
+ input,
+ schema,
+ )?;
+ assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear);
+
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(limit));
+ let context = Arc::new(
+ TaskContext::default()
+ .with_session_config(SessionConfig::new().with_batch_size(128))
+ .with_runtime(
+ RuntimeEnvBuilder::new()
+ .with_max_spill_merge_fan_in(2)
Review Comment:
Optional clarity suggestion: removing `.with_max_spill_merge_fan_in(2)`
here, which otherwise defaults to `0` for unbounded fan-in, still lets all
three tests pass. These tests seem focused on shared-pool replay lifecycle and
correctness, while Case G in `aggregate_memory_spill.slt` already owns the
SQL-level regression coverage for #25423.
Could we either mention that distinction in `run_shared_pool_case`'s doc
comment, or remove the fan-in setting if it does not have another intended
role? I would avoid turning this unit test into a duplicate stress reproduction
since the SLT seems like the better boundary for that behavior.
--
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]
