fix!: make a grouping key nullable when a grouping set leaves it out

src/substrait/type_inference.py

@@ -990,6 +990,21 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None, subtrees=()) -> stt.Type.
             )
             for g in rel.aggregate.grouping_expressions
         ]
+        for grouping in rel.aggregate.groupings:
+            for ref in grouping.expression_references:
+                if ref >= len(grouping_types):
+                    raise ValueError(
+                        f"Grouping expression_references index {ref} is out of "
+                        f"range for {len(grouping_types)} grouping expression(s)"
+                    )
+        # A grouping expression missing from any grouping set is null in the
+        # records of that set, so its column is nullable.
+        grouping_types = [
+            t
+            if all(i in g.expression_references for g in rel.aggregate.groupings)
+            else _with_field_nullability(t, stt.Type.NULLABILITY_NULLABLE)
+            for i, t in enumerate(grouping_types)
+        ]
         measure_types = [m.measure.output_type for m in rel.aggregate.measures]
 
         grouping_identifier_types = (

tests/test_type_inference.py

@@ -236,6 +236,66 @@ def test_inference_aggregate_multiple_groupings():
     assert infer_rel_schema(rel) == expected
 
 
+def _field_reference(field: int) -> stalg.Expression:
+    return stalg.Expression(
+        selection=stalg.Expression.FieldReference(
+            root_reference=stalg.Expression.FieldReference.RootReference(),
+            direct_reference=stalg.Expression.ReferenceSegment(
+                struct_field=stalg.Expression.ReferenceSegment.StructField(field=field)
+            ),
+        )
+    )
+
+
+@pytest.mark.parametrize(
+    "sets, expected",
+    [
+        ([[0, 1]], [("i64", _REQ), ("i64", _REQ)]),
+        ([[0], [1]], [("i64", _NULL), ("i64", _NULL), ("i32", _REQ)]),
+        ([[0, 1], [0]], [("i64", _REQ), ("i64", _NULL), ("i32", _REQ)]),
+        ([[0, 1], [1]], [("i64", _NULL), ("i64", _REQ), ("i32", _REQ)]),
+        ([[0, 1], []], [("i64", _NULL), ("i64", _NULL), ("i32", _REQ)]),
+    ],
+)
+def test_inference_aggregate_grouping_key_absent_from_a_set(sets, expected):
+    # A grouping key missing from any grouping set is null in that set's records,
+    # so its column is nullable; a key in every set keeps its own nullability.
+    # Both keys reference the REQUIRED field 0 so that a failure to widen is
+    # observable at either position.
+    rel = stalg.Rel(
+        aggregate=stalg.AggregateRel(
+            input=read_rel,
+            grouping_expressions=[_field_reference(0), _field_reference(0)],
+            groupings=[
+                stalg.AggregateRel.Grouping(expression_references=s) for s in sets
+            ],
+        )
+    )
+
+    schema = infer_rel_schema(rel)
+
+    assert [
+        (t.WhichOneof("kind"), getattr(t, t.WhichOneof("kind")).nullability)
+        for t in schema.types
+    ] == expected
+
+
+def test_inference_aggregate_grouping_reference_out_of_range():
+    rel = stalg.Rel(
+        aggregate=stalg.AggregateRel(
+            input=read_rel,
+            grouping_expressions=[_field_reference(0)],
+            groupings=[
+                stalg.AggregateRel.Grouping(expression_references=[0]),
+                stalg.AggregateRel.Grouping(expression_references=[1]),
+            ],
+        )
+    )
+
+    with pytest.raises(ValueError, match="index 1 is out of range"):
+        infer_rel_schema(rel)
+
+
 def test_inference_cross():
     rel = stalg.Rel(cross=stalg.CrossRel(left=read_rel, right=right_read_rel))