NGA-TRAN commented on code in PR #24697:
URL: https://github.com/apache/datafusion/pull/24697#discussion_r3928461756
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -4555,6 +4436,10 @@ mod tests {
aggregate.input_order_mode(),
InputOrderMode::PartiallySorted(_)
));
+ assert_eq!(
+ aggregate.group_completion_mode,
+ GroupCompletionMode::Partial(vec![0])
+ );
Review Comment:
Nice
##########
datafusion/physical-plan/src/aggregates/order/mod.rs:
##########
@@ -28,6 +28,50 @@ use crate::InputOrderMode;
pub use full::GroupOrderingFull;
pub use partial::GroupOrderingPartial;
+/// Describes how an aggregate can determine that groups are complete.
+///
+/// This is distinct from [`InputOrderMode`], which describes the ordering of
+/// the input relative to the grouping expressions. Input ordering is one way
+/// to establish a group-completion mode, but the execution machinery only
+/// needs to know when it can safely emit completed groups.
+///
+/// For example, when grouping by `key`, both inputs have fully contiguous
+/// groups within the input partition:
+///
+/// ```text
+/// sorted: A A B B C C
+/// not sorted: C C A A B B
+/// ```
+///
+/// In both cases, once the key changes, the previous key will not appear
again,
+/// so its group is complete and can be emitted.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub(crate) enum GroupCompletionMode {
+ /// No group can be known complete before the input ends.
+ None,
+ /// Rows with the same values at these grouping-expression indices form one
+ /// contiguous range. When those values change, every group in the previous
+ /// range is complete and can be emitted.
+ ///
+ /// For example, with `GROUP BY (a, b)`, `Partial(vec![0])` means all rows
+ /// for each value of `a` are contiguous, while an `(a, b)` tuple may recur
+ /// within that range.
Review Comment:
Does this mean I have 2 keys `(a, b)` and data is sorted on (a) only?
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -4657,162 +4533,104 @@ mod tests {
Ok(())
}
- /// Partial-reduce hash aggregation emits its accumulated partial states
early
- /// under memory pressure instead of failing, and the early-emitted states
- /// still merge into the correct result.
+ /// Spilling behavior is not implemented for partial-reduce stream yet, so
fall
+ /// back to the existing `GroupedHashAggregateStream`
#[tokio::test]
- async fn partial_reduce_aggregate_with_memory_limit_emits_early() ->
Result<()> {
- let num_input_batches = 3;
- let partial_reduce =
- partial_reduce_test_aggregate_with_batches(num_input_batches)?;
+ async fn partial_reduce_aggregate_with_memory_limit_planning() ->
Result<()> {
+ let partial_reduce = partial_reduce_test_aggregate()?;
let runtime = RuntimeEnvBuilder::new()
.with_memory_limit(1, 1.0)
.build_arc()?;
- // A batch size smaller than the number of flushed groups also covers
- // splitting one flush across several output batches.
- let batch_size = 2;
- let task_ctx = Arc::new(
- TaskContext::default()
- .with_session_config(migrated_hash_session_config(batch_size))
- .with_runtime(runtime),
- );
+ let task_ctx =
+ Arc::new(
+ TaskContext::default()
+ .with_session_config(SessionConfig::new().set_bool(
+ "datafusion.execution.enable_migration_aggregate",
+ true,
+ ))
+ .with_runtime(runtime),
+ );
let stream = partial_reduce.execute_typed(0, &task_ctx)?;
- assert!(matches!(stream, StreamType::PartialReduceHash(_)));
- let stream: SendableRecordBatchStream = stream.into();
- let output = collect(stream).await?;
-
- // The table is flushed after every input batch, so each of the three
- // groups is emitted once per input batch instead of being merged into
a
- // single row. Each flush is sliced into batches of 2 and 1 rows.
- assert_eq!(output.len(), 2 * num_input_batches);
- assert_snapshot!(batches_to_string(&output), @r"
- +---+-------------+
- | a | SUM(b)[sum] |
- +---+-------------+
- | 1 | 50.0 |
- | 2 | 20.0 |
- | 3 | 30.0 |
- | 1 | 50.0 |
- | 2 | 20.0 |
- | 3 | 30.0 |
- | 1 | 50.0 |
- | 2 | 20.0 |
- | 3 | 30.0 |
- +---+-------------+
- ");
+ assert!(matches!(stream, StreamType::GroupedHash(_)));
Ok(())
}
- /// Same shape as [`partial_reduce_test_aggregate_with_batches`], but with
multiple
- /// group keys.
- fn partial_reduce_test_aggregate_rows_multi_group_keys(
- num_input_batches: usize,
- ) -> Result<AggregateExec> {
+ #[tokio::test]
+ async fn unsorted_contiguous_groups_use_final_emission() -> Result<()> {
let schema = Arc::new(Schema::new(vec![
- Field::new("a", DataType::UInt32, false),
- Field::new("n", DataType::Null, true),
- Field::new("b", DataType::Float64, false),
+ 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 group_by = PhysicalGroupBy::new_single(vec![
- (col("a", &schema)?, "a".to_string()),
- (col("n", &schema)?, "n".to_string()),
+ (col("key", &schema)?, "key".to_string()),
+ (col("time_bin", &schema)?, "time_bin".to_string()),
]);
- let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![Arc::new(
- AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?])
+ let aggr_expr = Arc::new(
+ AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?])
.schema(Arc::clone(&schema))
- .alias("SUM(b)")
+ .alias("SUM(value)")
.build()?,
- )];
-
- let empty_input =
- TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema),
None)?;
- let partial = AggregateExec::try_new(
- AggregateMode::Partial,
- group_by.clone(),
- aggregates.clone(),
- vec![None],
- empty_input,
- Arc::clone(&schema),
- )?;
- let partial_schema = partial.schema();
- let partial_state_batch = RecordBatch::try_new(
- Arc::clone(&partial_schema),
- vec![
- Arc::new(UInt32Array::from(vec![1, 2, 1, 3])),
- Arc::new(NullArray::new(4)),
- Arc::new(Float64Array::from(vec![10.0, 20.0, 40.0, 30.0])),
- ],
- )?;
- let partial_reduce_input = TestMemoryExec::try_new_exec(
- &[vec![partial_state_batch; num_input_batches]],
- Arc::clone(&partial_schema),
- None,
- )?;
+ );
+ let input: Arc<dyn ExecutionPlan> =
+ TestMemoryExec::try_new_exec(&[input_batches],
Arc::clone(&schema), None)?;
+ assert_eq!(input.output_partitioning().partition_count(), 1);
- AggregateExec::try_new(
- AggregateMode::PartialReduce,
+ let aggregate = AggregateExec::try_new(
+ AggregateMode::Single,
group_by,
- aggregates,
+ vec![aggr_expr],
vec![None],
- partial_reduce_input,
- partial_schema,
- )
- }
-
- #[tokio::test]
- async fn
partial_reduce_aggregate_with_memory_limit_emits_early_multi_group_keys()
- -> Result<()> {
- let num_input_batches = 3;
- let partial_reduce =
-
partial_reduce_test_aggregate_rows_multi_group_keys(num_input_batches)?;
-
- // Pin the representation: this is exactly the condition
- // `new_group_values` uses to pick `GroupValuesRows` over
- // `GroupValuesColumn`. If a `Null` `GroupColumn` is ever added, this
- // assertion fires and the test stops covering the row-encoded path.
- let group_schema = partial_reduce
- .group_by
- .group_schema(&partial_reduce.schema())?;
- assert!(
- !group_values::multi_group_by::supported_schema(&group_schema),
- "expected the Null group column to force the GroupValuesRows
fallback"
- );
+ input,
+ schema,
+ )?;
- let runtime = RuntimeEnvBuilder::new()
- .with_memory_limit(1, 1.0)
- .build_arc()?;
- let batch_size = 2;
- let task_ctx = Arc::new(
- TaskContext::default()
- .with_session_config(migrated_hash_session_config(batch_size))
- .with_runtime(runtime),
- );
+ assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear);
+ assert_eq!(aggregate.group_completion_mode, GroupCompletionMode::None);
+ // 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);
Review Comment:
👍
--
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]
