[xavlee]

Which issue does this PR relate to?

• Part of Supporting analytics over large amounts of pre‑partitioned data #24438.
• Builds on test: cover unsorted contiguous groups in one partition #24737 and refactor: use group clustering terminology in aggregate execution #24697, both merged.

Rationale for this change

Aggregation can recognize completed groups when all rows with the same grouping-key tuple occur in one contiguous run within a partition. Sorted input provides this guarantee, and the runs can also appear in an arbitrary order:

AAABBBCCCC -> ordered and grouped
CCCAAABBB  -> grouped; each group can be completed when its run ends
AAABBBAACC -> A appears in separate runs

This PR represents that guarantee in EquivalenceProperties, allowing execution-plan producers to declare grouping tuples and consumers to check whether an input satisfies a grouping requirement.

What changes are included in this PR?

Semantics and propagation

A grouping assertion describes a complete expression tuple within every output partition. [a, b] and [b, a] identify the same groups. The guarantee applies to the complete pair; a guarantee for a subset such as [a] requires its own proof.

Grouping assertions are correctness guarantees supplied by producers. DataFusion trusts the asserted row layout. Producers can declare derived tuples such as [bhandle, date_bin(timestamp)] directly.

• Store grouping tuples alongside ordering, equality, and constant information in EquivalenceProperties.
• Normalize tuples using expression equivalences and constants, removing duplicate expressions.
• Check grouping requirements against declared tuples and grouping implied by lexicographical ordering prefixes.
• Project a tuple when every expression can be represented in the output, including explicitly declared derived expressions.
• Remap grouping expressions during schema replacement.
• Preserve a declared grouping tuple through UnionExec when every input satisfies it, explicitly or through ordering.
• Clear explicit grouping assertions at aggregation, sorting, repartitioning, and partition-merging boundaries where execution can invalidate them.

Public API additions

• GroupingEquivalenceClass represents complete grouping tuples and supports construction, insertion, lookup, clearing, schema rewriting, iteration, and display.
• EquivalenceProperties::{geq_class, normalized_geq_class, add_grouping, add_groupings, clear_groupings, grouping_satisfy} expose grouping metadata and its analysis.
• TestMemoryExec::try_with_grouping_information attaches grouping assertions in physical-plan tests and examples.

Execution-plan producers attach these assertions through the EquivalenceProperties held by PlanProperties. The follow-up #24497 consumes them to derive GroupClusteringMode::Full for eligible aggregates.

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 ← this PR
4. #24497 — feat: stream aggregates over grouped input

What is the testing strategy for this PR?

Tests cover complete-tuple semantics, equivalent permutations, normalization through equivalences and constants, ordering-derived grouping, complete and incomplete projection mappings, explicitly declared derived expressions, operator boundaries, and UNION propagation with selective tuple retention, mixed grouping/ordering proofs, schema rewrites, and input permutations.

aggregate_clears_explicit_grouping covers partial and single aggregation with reordered and aliased grouping keys. Both cases fail without the aggregate boundary clear and pass with it.

Local validation passed the extended workspace suite, including 2,997 physical-plan tests and 546 SQL logic files, and all-target/all-feature Clippy with warnings denied.

Are there any user-facing changes?

Yes. The additive Rust APIs above let plan producers describe contiguous grouping tuples and let consumers evaluate grouping requirements. The streaming aggregate optimization follows in #24497.

Review this layer

View this PR's diff

[codecov-commenter]

Codecov Report

❌ Patch coverage is 79.05028% with 150 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.44%. Comparing base (ebdb657) to head (c2fe167).

Files with missing lines	Patch %	Lines
datafusion/physical-expr/src/equivalence/mod.rs	58.97%	3 Missing and 61 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 ⚠️
datafusion/physical-plan/src/aggregates/mod.rs	82.19%	5 Missing and 8 partials ⚠️
...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 ⚠️
...n/physical-plan/src/sorts/sort_preserving_merge.rs	85.00%	0 Missing and 3 partials ⚠️
datafusion/functions/src/datetime/date_bin.rs	84.61%	0 Missing and 2 partials ⚠️
... and 6 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #24698      +/-   ##
==========================================
- Coverage   82.45%   82.44%   -0.01%     
==========================================
  Files        1140     1141       +1     
  Lines      436434   437085     +651     
  Branches   436434   437085     +651     
==========================================
+ Hits       359840   360355     +515     
- Misses      54840    54879      +39     
- Partials    21754    21851      +97     

☔ 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.

[xavlee]

Hey @jayzhan211 , this is ready for another look when you get the chance 🙇

[jayzhan211]

@xavlee thanks for the updates, both earlier points look good now. One remaining gap: an operator that reorders rows and calls only clear_orderings() still forwards its input's groupings. That's the same thing AggregateExec hit in the previous revision, and it's the pattern custom operators downstream already follow. Could clear_orderings() also clear groupings? Then every operator that already drops orderings stays correct without having to know about groupings:

     pub fn clear_orderings(&mut self) {
         self.oeq_class.clear();
         self.oeq_cache.clear();
+        self.clear_groupings();
     }

         eq_properties.clear_orderings();
-        if input_partitions > 1 {
-            eq_properties.clear_groupings();
-        }
         eq_properties.clear_per_partition_constants();

I ran this locally. The only test that changes is the single-partition half of grouping_is_cleared_only_when_partitions_are_merged, which now expects the grouping to be cleared, the same as orderings. With this change aggregate_clears_explicit_grouping no longer exercises the explicit clear in AggregateExec, so I'd add a clustered case too. It passes with the explicit clear and fails without it:

#[test]
fn clustered_aggregate_clears_explicit_grouping() -> Result<()> {
    let schema = Arc::new(Schema::new(vec![
        Field::new("a", DataType::Int32, true),
        Field::new("c", DataType::Int32, true),
        Field::new("value", DataType::Int64, true),
    ]));
    let a = col("a", &schema)?;
    let c = col("c", &schema)?;
    let input = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)?
        .try_with_sort_information(vec![
            [PhysicalSortExpr::new_default(Arc::clone(&a))].into(),
        ])?
        .try_with_grouping_information(vec![vec![Arc::clone(&c)]])?;
    let aggregate = AggregateExec::try_new(
        AggregateMode::Single,
        PhysicalGroupBy::new_single(vec![
            (a, "group_a".to_string()),
            (c, "group_c".to_string()),
        ]),
        vec![Arc::new(
            AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?])
                .schema(Arc::clone(&schema))
                .alias("SUM(value)")
                .build()?,
        )],
        vec![None],
        Arc::new(input),
        schema,
    )?;

    assert!(matches!(
        aggregate.group_clustering_mode(),
        GroupClusteringMode::Partial(_)
    ));
    let properties = aggregate.properties().equivalence_properties();
    assert!(!properties.oeq_class().is_empty());
    assert!(properties.geq_class().is_empty());
    assert!(!properties.grouping_satisfy([col("group_c", &aggregate.schema())?])?);
    Ok(())
}

If you'd rather keep clear_orderings() as it is, could you note in its doc comment and in add_groupings that operators which reorder rows must also call clear_groupings()? A follow-up PR is fine for either option.

[xavlee]

Thanks @jayzhan211. Assuming we have no existing operators that would change ordering but somehow maintain groupings, the conservative route to clear groupings in clear orderings makes sense to me.
I added this as a follow-up PR here: #26852

[xavlee]

Hi @jayzhan211, I came across another finding- we may also want to propagate grouping properties through UnionExec.

Currently, we construct a new EquivalenceProperties and only add back constants and orderings. However, since UnionExec maintains isolated partitions, grouping properties should also be maintained. On the other hand, InterleaveExec should not propagate groupings since it merges partitions, potentially breaking contiguity.

Please let me know if this makes sense, and if I should add this as part of this PR or in another followup 🙇

[alamb]

@mhilton is this a PR you think would help us at Influx?

[mhilton]

It might be that the ordering maintains the grouping, if every column in the grouping is included in the ordering, and therefore the groupings could filtered rather than just cleared. I suspect however that this is unnecessary as ordering is the stronger signal and will be checked first.

[xavlee]

Thanks for the review @mhilton. That's the intent here 🙂 grouping implied by the resulting ordering remains discoverable when determining a ClusteringMode, so we don't need to retain the grouping assertions here.

[jayzhan211]

@xavlee Yes, that makes sense. UnionExec turns each child partition into its own output partition, so a tuple holds on the output whenever it holds on every input, and InterleaveExec is right to drop it. It's a small change, so could you add it to this PR? This rule passes the existing union tests locally. It keeps a tuple from either side when both sides satisfy it, so an ordering on one side counts too:

+    // Each union output partition is one input partition, so a grouping holds
+    // when every input satisfies it, explicitly or through an ordering.
+    let mut groupings = vec![];
+    for grouping in lhs.geq_class().iter().chain(rhs.geq_class().iter()) {
+        if lhs.grouping_satisfy(grouping.iter().cloned())?
+            && rhs.grouping_satisfy(grouping.iter().cloned())?
+        {
+            groupings.push(grouping.clone());
+        }
+    }
+
     let mut eq_properties = EquivalenceProperties::new(lhs.schema);
     eq_properties.add_constants(constants)?;
     eq_properties.add_orderings(orderings);
+    eq_properties.add_groupings(groupings);

Please add tests for these cases: grouped on [a, b] with grouped keeps it, grouped with an unannotated input drops it, and sorted on [a, b] with grouped keeps it.

[xavlee]

@jayzhan211 I updated the PR with the union change/tests. Please re-review when you get the chance 🙇

[jayzhan211]

Thanks @xavlee @gene-bordegaray @mhilton 🚀 !

[alamb]

epic