[xavlee]

Which issue does this PR close?

• Closes Supporting analytics over large amounts of pre‑partitioned data #24438.
• Builds on test: cover unsorted contiguous groups in one partition #24737, refactor: use group clustering terminology in aggregate execution #24697, and feat: add grouping equivalence properties #24698 (all merged).

Rationale for this change

An aggregate can emit a completed group once it knows that the same complete grouping tuple cannot appear again in its input partition. EquivalenceProperties::grouping_satisfy now provides that guarantee for both lexicographically ordered input and explicitly grouped input.

For example, these complete (key, time_bin) groups are contiguous even though their values are not globally sorted:

(A, 20), (A, 20) | (B, 20), (B, 20) | (A, 0), (A, 0)

The aggregate can therefore release each group as its run ends. Its clustering mode describes that guarantee, while equivalence properties describe any actual output ordering.

What changes are included in this PR?

• Query the input equivalence properties for the complete GROUP BY tuple during AggregateExec construction.
• Select GroupClusteringMode::Full when complete grouping tuple is satisfied, at least one grouping expression is non-constant, and grouping uses one non-nulling mask.
• Select clustered execution for grouped input, allowing completed groups to be emitted incrementally.
• Use input group metadata for clustered execution in Single, SinglePartitioned, and Partial aggregations. In a two-stage plan, the partial aggregate clears explicit grouping metadata from its output, so this property alone does not enabled clustered final aggregation. The final stage can still use ordering provided by its own input.

input grouping:  [[key, time_bin]]
GROUP BY:        [key, time_bin]
output ordering: none

grouping_satisfy(GROUP BY)
  -> GroupClusteringMode::Full
  -> EmissionType::Incremental
  -> ClusteredSingleAggregateStream

Stack

1. #24737 — test: cover unsorted contiguous groups in one partition (merged)
2. #24697 — refactor: use group clustering terminology in aggregate execution (merged)
3. #24698 — feat: add grouping equivalence properties (merged)
4. #24497 — feat: stream aggregates over grouped input ← this PR

Are these changes tested?

The characterization test from #24737 declares its complete grouping tuple and verifies GroupClusteringMode::Full, EmissionType::Incremental, and ClusteredSingleAggregateStream. It uses BarrierExec to hold back the second input batch until the first completed group is emitted, then checks the full result, including a group spanning both batches.

Regression tests cover PartialReduce with ordered and explicitly grouped input, and constant-only grouping with both a literal and a column made constant by a filter. A projection test verifies that a semantically equivalent date_bin grouping tuple reaches AggregateExec and enables incremental emission.

A repeated-mask regression compares grouped and unannotated inputs and verifies that both produce the same merged groups with GroupClusteringMode::None.

Are there any user-facing changes?

Yes. Aggregates whose complete grouping tuple is satisfied by the input equivalence properties can execute with full group clustering and incremental emission on unsorted input.

Review this layer

View only this PR layer

[codecov-commenter]

Codecov Report

❌ Patch coverage is 79.40854% with 188 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.45%. Comparing base (ebdb657) to head (9c62dea).
⚠️ Report is 21 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-expr/src/equivalence/mod.rs	58.97%	3 Missing and 61 partials ⚠️
datafusion/physical-plan/src/aggregates/mod.rs	79.75%	12 Missing and 38 partials ⚠️
...tafusion/physical-expr/src/equivalence/grouping.rs	82.17%	13 Missing and 5 partials ⚠️
datafusion/physical-plan/src/test.rs	60.46%	16 Missing and 1 partial ⚠️
...on/physical-expr/src/equivalence/properties/mod.rs	90.00%	7 Missing and 5 partials ⚠️
datafusion/physical-plan/src/repartition/mod.rs	78.57%	0 Missing and 6 partials ⚠️
...atafusion/physical-plan/src/coalesce_partitions.rs	84.61%	0 Missing and 4 partials ⚠️
datafusion/physical-expr/src/expressions/cast.rs	81.25%	0 Missing and 3 partials ⚠️
...tafusion/physical-plan/src/aggregates/order/mod.rs	78.57%	3 Missing ⚠️
...n/physical-plan/src/sorts/sort_preserving_merge.rs	85.00%	0 Missing and 3 partials ⚠️
... and 6 more

Additional details and impacted files

@@           Coverage Diff            @@
##             main   #24497    +/-   ##
========================================
  Coverage   82.45%   82.45%            
========================================
  Files        1140     1141     +1     
  Lines      436434   437417   +983     
  Branches   436434   437417   +983     
========================================
+ Hits       359840   360655   +815     
- Misses      54840    54883    +43     
- Partials    21754    21879   +125     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[alamb]

This PR models group contiguity as its own correctness property. It is intentionally neither an ordering guarantee nor an output-distribution guarantee.

Why does it need a new property? I think this notion is designed to be covered y the existing ordering / monotonic analyses

[alamb]

Oh, I see, somehow the external system knows the data is not sorted but is non overlapping

I think this is going to be really hard to manage / ensure through the plan -- we will need to ensure that every operator properly reports if it will propagate this property or not

I am not sure this is something we want to complicate datafusion with

[alamb]

@NGA-TRAN can you help evaluate this PR for its impact and if we will be able to keep this property in tact?

[xavlee]

Hi Andrew, thanks for taking an initial look. Apologies that this is still verymuch a draft.

I agree that propagating another physical property through the plan would be a little complicated. I was hoping we would narrow the group_contiguous_exprs assertion s.t. it:

• is declared explicitly by the data source
• defaults to absent on every execution operator
• may pass only through ProjectionExec when every expression maps
• is consumed by the only the first AggregateExec (and not present on the aggregate output)

[NGA-TRAN]

@alamb : I added a comment in the ticket

@NGA-TRAN can you help evaluate this PR for its impact

The impact of this PR is huge for telemetry use cases of AI frontiers as I described in the comment above

and if we will be able to keep this property in tact?

This property, like some properties, will be no longer available after certain operators so I think it would wok the same. I agree the propagation is a bit more complicated than usual but we work together to split this PR into smaller ones and will look into design carefully to avoid a lot of side effect. I think we would be able to make the design simpler

[NGA-TRAN]

@jayzhan211 : This the forth PR in the series you are reviewing. Hopefully, you can help us review this, too

[github-actions]

Marked as draft after 14 days without activity: merge conflicts.

Please address these items and mark Ready for review when ready.

[xavlee]

@jayzhan211 bumping the final PR in the stack for your review 🙇

[jayzhan211]

@xavlee thanks for the stack, and sorry for the slow turnaround. I ran a differential check of grouped vs. ungrouped input across Single and Partial→Final, both aggregate implementations, null keys, groups spanning batches, and superset/subset GROUP BY keys. Everything matched, and every operator I checked that reorders rows clears the grouping. Two small asks below: an upgrade-guide note and a partial-stage test.

In docs/source/library-user-guide/upgrading/56.0.0.md around line 81, the guide maps Full to the old Sorted, and that text ships in 56. With this PR, group_clustering_mode=Full can appear on unsorted input. Could you add a 57 entry (new 57.0.0.md plus an index.rst link)? Something like:

### `GroupClusteringMode::Full` no longer implies sorted input

`AggregateExec` now selects `GroupClusteringMode::Full` when the input's
equivalence properties declare the complete `GROUP BY` tuple as grouped
(contiguous within each partition), even if the input is not sorted.
`Full` describes group contiguity only. Code that used
`group_clustering_mode() == GroupClusteringMode::Full` to infer sorted input
or output should check `properties().output_ordering()` instead.

[jayzhan211]

Could you also cover the partial stage? Multi-partition plans run ClusteredPartialAggregateStream followed by a final aggregate, but only Single is exercised here. Something like this, using the same grouped input:

let partial = Arc::new(AggregateExec::try_new(
    AggregateMode::Partial,
    group_by.clone(),
    vec![Arc::clone(&aggr_expr)],
    vec![None],
    Arc::clone(&input) as Arc<dyn ExecutionPlan>,
    Arc::clone(&schema),
)?);
assert_eq!(partial.group_clustering_mode(), &GroupClusteringMode::Full);
assert!(matches!(
    partial.execute_typed(0, &task_ctx)?,
    StreamType::ClusteredPartialAggregate(_)
));
let final_agg = AggregateExec::try_new(
    AggregateMode::Final,
    group_by.as_final(),
    vec![aggr_expr],
    vec![None],
    Arc::new(CoalescePartitionsExec::new(partial)),
    schema,
)?;
let output = collect(final_agg.execute(0, task_ctx)?).await?;
// same snapshot as the single-stage result

The input needs at least two partitions so the final stage actually merges, and a TestMemoryExec input with grouping information instead of BarrierExec keeps it simple.

[jayzhan211]

This test is still needed. I restricted the new Full selection to AggregateMode::Single locally and all 2959 datafusion-physical-plan lib tests still passed, so nothing covers Partial or SinglePartitioned over grouped input. SinglePartitioned is the plan a key-partitioned source gets, which is the shape #24438 is after. Could you add this test in this PR? It passes on this head.

Test covering Partial → Final and SinglePartitioned over grouped input

#[tokio::test]
async fn grouped_input_uses_clustered_partial_and_single_partitioned() -> Result<()> {
    let schema = Arc::new(Schema::new(vec![
        Field::new("key", DataType::Int32, false),
        Field::new("value", DataType::Int64, false),
    ]));
    let batch = |keys: Vec<i32>, values: Vec<i64>| {
        RecordBatch::try_new(
            Arc::clone(&schema),
            vec![
                Arc::new(Int32Array::from(keys)),
                Arc::new(Int64Array::from(values)),
            ],
        )
    };
    // Groups are contiguous within each partition but not sorted. Key 1 spans
    // a batch boundary, and keys 2 and 3 occur in both partitions.
    let partitions = vec![
        vec![
            batch(vec![2, 1], vec![1, 2])?,
            batch(vec![1, 3], vec![3, 4])?,
        ],
        vec![batch(vec![3, 3, 2], vec![5, 6, 7])?],
    ];
    let key = col("key", &schema)?;
    let input: Arc<dyn ExecutionPlan> = Arc::new(
        TestMemoryExec::try_new(&partitions, Arc::clone(&schema), None)?
            .try_with_grouping_information(vec![vec![Arc::clone(&key)]])?,
    );
    let group_by = PhysicalGroupBy::new_single(vec![(key, "key".to_string())]);
    let aggr_expr = Arc::new(
        AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?])
            .schema(Arc::clone(&schema))
            .alias("SUM(value)")
            .build()?,
    );
    let task_ctx = new_migrated_hash_ctx(1024);

    let single_partitioned = AggregateExec::try_new(
        AggregateMode::SinglePartitioned,
        group_by.clone(),
        vec![Arc::clone(&aggr_expr)],
        vec![None],
        Arc::clone(&input),
        Arc::clone(&schema),
    )?;
    assert_eq!(
        single_partitioned.group_clustering_mode(),
        &GroupClusteringMode::Full
    );
    assert!(matches!(
        single_partitioned.execute_typed(0, &task_ctx)?,
        StreamType::ClusteredSingleAggregate(_)
    ));

    let partial = Arc::new(AggregateExec::try_new(
        AggregateMode::Partial,
        group_by.clone(),
        vec![Arc::clone(&aggr_expr)],
        vec![None],
        input,
        Arc::clone(&schema),
    )?);
    assert_eq!(partial.group_clustering_mode(), &GroupClusteringMode::Full);
    assert!(matches!(
        partial.execute_typed(0, &task_ctx)?,
        StreamType::ClusteredPartialAggregate(_)
    ));

    let final_aggregate = AggregateExec::try_new(
        AggregateMode::Final,
        group_by.as_final(),
        vec![aggr_expr],
        vec![None],
        Arc::new(CoalescePartitionsExec::new(partial)),
        schema,
    )?;
    let output = collect(final_aggregate.execute(0, task_ctx)?).await?;
    assert_snapshot!(batches_to_sort_string(&output), @r"
    +-----+------------+
    | key | SUM(value) |
    +-----+------------+
    | 1   | 5          |
    | 2   | 8          |
    | 3   | 15         |
    +-----+------------+
    ");
    Ok(())
}

[xavlee]

Added this test 🙇

[jayzhan211]

Optional cleanup: grouping_satisfy already covers the ordering-based Full case, since it drops constants and then checks find_longest_permutation. The mode can therefore be picked in one place. This also puts back the groups.len() == 1 check that the old Full branch had. Without it, a hand-built PhysicalGroupBy with two groups and has_grouping_set = false over grouped input currently gets Full and emits duplicate groups. I ran this locally and all aggregates tests pass:

-        let mut group_clustering_mode = if indices.len() == num_non_constant_groupby_exprs
-            && !indices.is_empty()
-            && group_by.groups.len() == 1
-        {
-            GroupClusteringMode::Full
-        } else if !indices.is_empty() {
-            GroupClusteringMode::Partial(indices)
-        } else {
-            GroupClusteringMode::None
-        };
-
-        // Grouping sets can change group keys. PartialReduce combines intermediate
-        // states without using group boundaries to recognize completed groups.
-        if group_by.has_grouping_set() || mode == AggregateMode::PartialReduce {
-            group_clustering_mode = GroupClusteringMode::None;
-        } else if num_non_constant_groupby_exprs > 0
-            && input_eq_properties.grouping_satisfy(groupby_exprs.iter().cloned())?
-        {
-            group_clustering_mode = GroupClusteringMode::Full;
-        }
+        // Grouping sets can change group keys. PartialReduce combines intermediate
+        // states without using group boundaries to recognize completed groups.
+        let group_clustering_mode = if group_by.has_grouping_set()
+            || group_by.groups.len() != 1
+            || mode == AggregateMode::PartialReduce
+        {
+            GroupClusteringMode::None
+        } else if num_non_constant_groupby_exprs > 0
+            && input_eq_properties.grouping_satisfy(groupby_exprs.iter().cloned())?
+        {
+            GroupClusteringMode::Full
+        } else if !indices.is_empty() {
+            GroupClusteringMode::Partial(indices)
+        } else {
+            GroupClusteringMode::None
+        };

[xavlee]

Added 🙇

[jayzhan211]

@xavlee thanks for folding the mode selection into one place and adding the mask guard. The upgrade note from my last review is still needed: docs/source/library-user-guide/upgrading/56.0.0.md around line 69 maps Full to the old Sorted, and with this PR Full can appear on unsorted input. Could you add a 57 entry in this PR (new 57.0.0.md plus an index.rst link)? Something like:

### `GroupClusteringMode::Full` no longer implies sorted input

`AggregateExec` now selects `GroupClusteringMode::Full` when the input's
equivalence properties declare the complete `GROUP BY` tuple as grouped
(contiguous within each partition), even if the input is not sorted.
`Full` describes group contiguity only. Code that used
`group_clustering_mode() == GroupClusteringMode::Full` to infer sorted input
or output should check `properties().output_ordering()` instead.

Optional cleanup: in datafusion/physical-plan/src/aggregates/mod.rs around line 40, the module docs still describe clustering as coming only from ordered input.

-//! When the input is ordered by group keys, rows are clustered by those keys.
+//! When the input is ordered by group keys, or its equivalence properties
+//! declare them grouped, rows are clustered by those keys.

Could you also adjust the PR description, which says the grouped path applies to final aggregation? A final stage only gets Full if its own input declares grouping, because the partial stage clears it. In the plans I printed, only Single and SinglePartitioned report incremental emission at the top; a multi-partition source without key partitioning runs Partial[Full] into a hash final stage. Saying so tells users which plans actually stream.

[jayzhan211]

Optional cleanup: this guard still lets through a single mask that nulls a key when has_grouping_set is false. For example, vec![vec![true, false]] over grouped (a, b) rows (1, 1), (2, 2), (3, 1) in separate batches gets Full and returns (NULL, 1) twice, where hash mode merges them. Only a hand-built PhysicalGroupBy can reach this, since the physical planner and proto both carry the flag. Comparing against the all-false mask closes it, and the aggregates tests still pass. To cover it, the mask above should give GroupClusteringMode::None.

Suggested change

	|| group_by.groups.len() != 1
	|| group_by.groups != [vec![false; group_by.expr.len()]]

[xavlee]

I updated the guard here, thanks Jay

[xavlee]

Hi @jayzhan211, I added the update note along with the new 57.0.0.md file, and implemented the various documentation/description cleanups. Let me know if there are any other changes needed 🙇

[jayzhan211]

Thanks @xavlee and @gene-bordegaray!