gene-bordegaray commented on code in PR #24497:
URL: https://github.com/apache/datafusion/pull/24497#discussion_r4084207456
##########
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:
collect will make this go till EOF. Could we check the first completed group
is returned before releasing the rest of the input then assert results
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -1096,6 +1103,15 @@ impl AggregateExec {
input_order_mode = InputOrderMode::Linear;
}
+ let group_completion_mode = if !group_by.has_grouping_set()
Review Comment:
I think this can override what is done in the lines above since
`grouping_satisfy` can make a `PartialReduce` become `Full` which shoundt be
allowed
we could just check this extra conditiion here
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -1096,6 +1103,15 @@ impl AggregateExec {
input_order_mode = InputOrderMode::Linear;
}
+ let group_completion_mode = if !group_by.has_grouping_set()
Review Comment:
I think we also need to make sure we have some non constant group by exprs
here or this will use intremental but shoudl use Final
--
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]
