From 37f7cee19b9fafa5a5c74046f7a67277c9d0c057 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Fri, 25 Sep 2026 09:33:02 +0300 Subject: [PATCH 1/2] fix!: make a grouping key nullable when a grouping set leaves it out The aggregate schema copied each grouping expression's type as is. A key missing from some grouping set is null in that set's records, and the spec makes its column nullable; a key in every set keeps its own nullability. BREAKING CHANGE: inferred aggregate schemas now mark a grouping key nullable when some grouping set leaves it out. --- src/substrait/type_inference.py | 12 ++++++++++ tests/test_type_inference.py | 41 +++++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index c17b9df..fdd43eb 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -990,6 +990,18 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None, subtrees=()) -> stt.Type. ) for g in rel.aggregate.grouping_expressions ] + # A grouping expression missing from any grouping set is null in the + # records of that set, so its column is nullable. + if rel.aggregate.groupings: + in_every_set = set.intersection( + *(set(g.expression_references) for g in rel.aggregate.groupings) + ) + grouping_types = [ + t + if i in in_every_set + 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 = ( diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index b40f2d6..8f53d50 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -236,6 +236,47 @@ 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]], [_REQ, _NULL]), + ([[0], [1]], [_NULL, _NULL, _REQ]), + ([[0, 1], [0]], [_REQ, _NULL, _REQ]), + ([[0, 1], [1]], [_NULL, _NULL, _REQ]), + ([[0, 1], []], [_NULL, _NULL, _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. + rel = stalg.Rel( + aggregate=stalg.AggregateRel( + input=read_rel, + grouping_expressions=[_field_reference(0), _field_reference(2)], + groupings=[ + stalg.AggregateRel.Grouping(expression_references=s) for s in sets + ], + ) + ) + + schema = infer_rel_schema(rel) + + assert [getattr(t, t.WhichOneof("kind")).nullability for t in schema.types] == ( + expected + ) + + def test_inference_cross(): rel = stalg.Rel(cross=stalg.CrossRel(left=read_rel, right=right_read_rel)) From be8217f65a4d9a61c42ba61830391946288496d9 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Mon, 28 Sep 2026 12:05:53 +0300 Subject: [PATCH 2/2] fix: check grouping set references and pin both key positions in the test --- src/substrait/type_inference.py | 23 ++++++++++++---------- tests/test_type_inference.py | 35 +++++++++++++++++++++++++-------- 2 files changed, 40 insertions(+), 18 deletions(-) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index fdd43eb..80354f5 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -990,18 +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. - if rel.aggregate.groupings: - in_every_set = set.intersection( - *(set(g.expression_references) for g in rel.aggregate.groupings) - ) - grouping_types = [ - t - if i in in_every_set - else _with_field_nullability(t, stt.Type.NULLABILITY_NULLABLE) - for i, t in enumerate(grouping_types) - ] + 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 = ( diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index 8f53d50..c0387aa 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -250,20 +250,22 @@ def _field_reference(field: int) -> stalg.Expression: @pytest.mark.parametrize( "sets, expected", [ - ([[0, 1]], [_REQ, _NULL]), - ([[0], [1]], [_NULL, _NULL, _REQ]), - ([[0, 1], [0]], [_REQ, _NULL, _REQ]), - ([[0, 1], [1]], [_NULL, _NULL, _REQ]), - ([[0, 1], []], [_NULL, _NULL, _REQ]), + ([[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(2)], + grouping_expressions=[_field_reference(0), _field_reference(0)], groupings=[ stalg.AggregateRel.Grouping(expression_references=s) for s in sets ], @@ -272,10 +274,27 @@ def test_inference_aggregate_grouping_key_absent_from_a_set(sets, expected): schema = infer_rel_schema(rel) - assert [getattr(t, t.WhichOneof("kind")).nullability for t in schema.types] == ( - expected + 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))