xavlee commented on code in PR #24497:
URL: https://github.com/apache/datafusion/pull/24497#discussion_r4086380978
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -5530,6 +5571,151 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn unsorted_contiguous_groups_use_incremental_emission() ->
Result<()> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("key", DataType::Int32, false),
+ Field::new("time_bin", DataType::Int64, false),
+ Field::new("value", DataType::Int64, false),
+ ]));
+ // Two sorted logical runs are emitted as batches in one DataFusion
+ // partition. Every distinct grouping tuple occupies one contiguous
range,
+ // but tuple order resets at the batch boundary, so (key, time_bin) is
not
+ // globally sorted.
+ let input_batches = vec![
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(vec![1, 1, 2, 2])),
+ Arc::new(Int64Array::from(vec![20, 20, 20, 20])),
+ Arc::new(Int64Array::from(vec![10, 20, 30, 40])),
+ ],
+ )?,
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(vec![1, 1, 2, 2])),
+ Arc::new(Int64Array::from(vec![0, 0, 0, 0])),
+ Arc::new(Int64Array::from(vec![50, 60, 70, 80])),
+ ],
+ )?,
+ ];
+ let key = col("key", &schema)?;
+ let time_bin = col("time_bin", &schema)?;
+ let group_by = PhysicalGroupBy::new_single(vec![
+ (Arc::clone(&key), "key".to_string()),
+ (Arc::clone(&time_bin), "time_bin".to_string()),
+ ]);
+ let aggr_expr = Arc::new(
+ AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?])
+ .schema(Arc::clone(&schema))
+ .alias("SUM(value)")
+ .build()?,
+ );
+ let input = TestMemoryExec::try_new(&[input_batches],
Arc::clone(&schema), None)?
+ .try_with_grouping_information(vec![vec![key, time_bin]])?;
+ let input: Arc<dyn ExecutionPlan> = Arc::new(input);
+ assert_eq!(input.output_partitioning().partition_count(), 1);
+
+ let aggregate = AggregateExec::try_new(
+ AggregateMode::Single,
+ group_by,
+ vec![aggr_expr],
+ vec![None],
+ input,
+ schema,
+ )?;
+
+ assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear);
+ assert_eq!(aggregate.group_completion_mode, GroupCompletionMode::Full);
+ assert_eq!(aggregate.cache().emission_type, EmissionType::Incremental);
+ assert!(
+ aggregate
+ .cache()
+ .equivalence_properties()
+ .geq_class()
+ .is_empty()
+ );
+ assert!(aggregate.cache().output_ordering().is_none());
+
+ let task_ctx = new_migrated_hash_ctx(1024);
+ let stream = aggregate.execute_typed(0, &task_ctx)?;
+ assert!(matches!(stream, StreamType::OrderedSingleAggregate(_)));
+ let stream: SendableRecordBatchStream = stream.into();
+ let output = collect(stream).await?;
Review Comment:
Added, thx.
--
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]
