diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index c17b9df..80354f5 100644 --- a/src/substrait/type_inference.py +++ b/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 = ( diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index b40f2d6..c0387aa 100644 --- a/tests/test_type_inference.py +++ b/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))