gene-bordegaray commented on code in PR #24737:
URL: https://github.com/apache/datafusion/pull/24737#discussion_r3896780329
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -4539,6 +4539,60 @@ mod tests {
Ok(())
}
+ #[test]
+ fn unsorted_contiguous_groups_use_final_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 concatenated into one DataFusion
partition.
+ // Every distinct grouping tuple occupies one contiguous range, but
tuple
+ // order resets between runs, so (key, time_bin) is not globally
sorted.
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(vec![1, 1, 2, 2, 1, 1, 2, 2])),
+ Arc::new(Int64Array::from(vec![20, 20, 20, 20, 0, 0, 0, 0])),
+ Arc::new(Int64Array::from(vec![10, 20, 30, 40, 50, 60, 70,
80])),
+ ],
+ )?;
+ let group_by = PhysicalGroupBy::new_single(vec![
+ (col("key", &schema)?, "key".to_string()),
+ (col("time_bin", &schema)?, "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: Arc<dyn ExecutionPlan> =
+ TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema),
None)?;
+ 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);
+ // This captures the behavior before #24438. When the source can
declare
+ // `(key, time_bin)` group-contiguous, the corresponding case can use
+ // `EmissionType::Incremental`.
+ assert_eq!(aggregate.cache().emission_type, EmissionType::Final);
+
+ let task_ctx = new_migrated_hash_ctx(1024);
+ let stream = aggregate.execute_typed(0, &task_ctx)?;
Review Comment:
non blockiung: can we make this a tokiop test and colelct the stream that
actaulla ssert the results
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -4539,6 +4539,60 @@ mod tests {
Ok(())
}
+ #[test]
+ fn unsorted_contiguous_groups_use_final_emission() -> Result<()> {
Review Comment:
done at a batch boundary
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -4539,6 +4539,60 @@ mod tests {
Ok(())
}
+ #[test]
+ fn unsorted_contiguous_groups_use_final_emission() -> Result<()> {
Review Comment:
for this use case i think it would also be good if there were two batches to
show that this works over multiple polls 👍
--
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]
