yashaggarwal85d commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4366065885
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -732,4 +744,87 @@ mod tests {
assert_eq!(pool.reserved(), 0);
Ok(())
}
+
+ /// Runs a final `SUM(k) GROUP BY k` over two partial-state batches sorted
on
+ /// `k` that share a boundary key. Returns the output and the most memory
the
+ /// stream held while handing off output.
+ async fn run_fully_sorted(
+ pool: Arc<dyn MemoryPool>,
+ ) -> Result<(Vec<RecordBatch>, usize)> {
+ let schema = Arc::new(Schema::new(vec![Field::new("k",
DataType::Int64, false)]));
+ let groups = PhysicalGroupBy::new_single(vec![(col("k", &schema)?,
"k".into())]);
+ let expressions = vec![Arc::new(
+ AggregateExprBuilder::new(sum_udaf(), vec![col("k", &schema)?])
+ .schema(Arc::clone(&schema))
+ .alias("sum")
+ .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],
+ empty,
+ Arc::clone(&schema),
+ )?;
+ let partial_schema = partial.schema();
+ let batches = [0..=256, 256..=512]
+ .into_iter()
+ .map(|keys| {
+ let k = Int64Array::from_iter_values(keys);
+ RecordBatch::try_new(
+ Arc::clone(&partial_schema),
+ vec![Arc::new(k.clone()), Arc::new(k)],
+ )
+ })
+ .collect::<Result<Vec<_>, _>>()?;
+ let ordering =
+ LexOrdering::new([PhysicalSortExpr::new_default(col("k",
&partial_schema)?)])
+ .unwrap();
+ let input =
+ TestMemoryExec::try_new(&[batches], Arc::clone(&partial_schema),
None)?
+ .try_with_sort_information(vec![ordering])?;
+ let aggregate = AggregateExec::try_new(
+ AggregateMode::Final,
+ groups.as_final(),
+ expressions,
+ vec![None],
+ Arc::new(TestMemoryExec::update_cache(&Arc::new(input))),
+ schema,
+ )?;
+ assert!(matches!(
+ aggregate.group_clustering_mode(),
+ GroupClusteringMode::Full
Review Comment:
Added `partially_clustered_memory_pressure_without_spilling` for both
streams: a two-column key `(a, b)` sorted on `a`, with the disk manager
disabled, asserting `GroupClusteringMode::Partial([0])`. While writing it I
found that under partial clustering the group values keep their capacity after
`EmitTo::First`. `memory_size()` stays the same after the emit (31192 bytes
before and after for 512 → 256 groups in this test), so emitting can't bring
the table under a limit it already exceeded. The test therefore pins that the
branch still returns `ResourcesExhausted` cleanly one byte below the unlimited
footprint, rather than claiming success. I checked that it actually reaches the
new branch.
--
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]
