yashaggarwal85d commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4366068022
##########
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
+ ));
+ let context = Arc::new(
+ TaskContext::default().with_runtime(
+ RuntimeEnvBuilder::new()
+ .with_memory_pool(Arc::clone(&pool))
+ .build_arc()?,
+ ),
+ );
+ let mut stream =
+ ClusteredFinalAggregateStream::new(&aggregate, &context,
0)?.into_stream();
+ // Each output batch is a single slice, so the stream releases it
before
+ // handing it off and holds only the table.
+ let mut held = pool.reserved();
+ let mut output = vec![];
+ while let Some(batch) = stream.next().await.transpose()? {
+ held = held.max(pool.reserved());
+ output.push(batch);
+ }
+ Ok((output, held))
+ }
+
+ /// The limit fits the table with only the active group, but not the table
+ /// before the completed groups are emitted. The result must not change.
+ #[tokio::test]
+ async fn completed_groups_under_memory_pressure() -> Result<()> {
Review Comment:
Added `memory_pressure_without_completed_groups` for both streams. A single
key never completes a group before EOF, and the limit fits only the initial
table, so the refused reservation must come back as `ResourcesExhausted`.
--
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]
