feat: add inferene for set, cross and aggregate rels

Author: tokoko

src/substrait/type_inference.py

@@ -152,6 +152,17 @@ def infer_rel_schema(rel: stalg.Rel) -> stt.Type.Struct:
         (common, struct) = (rel.filter.common, infer_rel_schema(rel.filter.input))
     elif rel_type == "fetch":
         (common, struct) = (rel.fetch.common, infer_rel_schema(rel.fetch.input))
+    elif rel_type == "aggregate":
+        parent_schema = infer_rel_schema(rel.aggregate.input)
+        grouping_types = [infer_expression_type(g, parent_schema) for g in rel.aggregate.grouping_expressions]
+        measure_types = [m.measure.output_type for m in rel.aggregate.measures]
+
+        raw_schema = stt.Type.Struct(
+            types=grouping_types + measure_types,
+            nullability=parent_schema.nullability,
+        )
+
+        (common, struct) = (rel.aggregate.common, raw_schema)
     elif rel_type == "sort":
         (common, struct) = (rel.sort.common, infer_rel_schema(rel.sort.input))
     elif rel_type == "project":
@@ -165,6 +176,18 @@ def infer_rel_schema(rel: stalg.Rel) -> stt.Type.Struct:
         )
 
         (common, struct) = (rel.project.common, raw_schema)
+    elif rel_type == "set":
+        (common, struct) = (rel.fetch.common, infer_rel_schema(rel.set.inputs[0]))
+    elif rel_type == "cross":
+        left_schema = infer_rel_schema(rel.cross.left)
+        right_schema = infer_rel_schema(rel.cross.right)
+
+        raw_schema = stt.Type.Struct(
+            types=left_schema.types + right_schema.types,
+            nullability=stt.Type.Nullability.NULLABILITY_REQUIRED,
+        )
+
+        (common, struct) = (rel.cross.common, raw_schema)
     else:
         raise Exception(f"Unhandled rel_type {rel_type}")
 

tests/test_type_inference.py

@@ -97,3 +97,48 @@ def test_inference_project_scalar_function():
     )
 
     assert infer_rel_schema(rel) == expected
+
+def test_inference_aggregate():
+    rel = stalg.Rel(
+        aggregate=stalg.AggregateRel(
+            input=read_rel,
+            grouping_expressions=[
+                stalg.Expression(
+                    selection=stalg.Expression.FieldReference(
+                        root_reference=stalg.Expression.FieldReference.RootReference(),
+                        direct_reference=stalg.Expression.ReferenceSegment(
+                            struct_field=stalg.Expression.ReferenceSegment.StructField(
+                                field=1,
+                            ),
+                        ),
+                    )
+                )
+            ],
+            groupings=[
+                stalg.AggregateRel.Grouping(
+                    expression_references=[0]
+                )
+            ],
+            measures=[
+                stalg.AggregateRel.Measure(
+                    measure=stalg.AggregateFunction(
+                        function_reference=0,
+                        output_type=stt.Type(
+                            bool=stt.Type.Boolean(
+                                nullability=stt.Type.NULLABILITY_REQUIRED
+                            )
+                        ),
+                    )
+                )
+            ],
+        )
+    )
+
+    expected = stt.Type.Struct(
+        types=[
+            stt.Type(string=stt.Type.String(nullability=stt.Type.NULLABILITY_NULLABLE)),
+            stt.Type(bool=stt.Type.Boolean(nullability=stt.Type.NULLABILITY_REQUIRED)),
+        ]
+    )
+
+    assert infer_rel_schema(rel) == expected