diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index a564cd37..24711f12 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -429,7 +429,7 @@ 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])) + (common, struct) = (rel.set.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) diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index af61903b..e45b3e7a 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -66,6 +66,27 @@ def test_inference_project_emit(): assert infer_rel_schema(rel) == expected +def test_inference_set_emit(): + # A SetRel's own RelCommon.Emit must drive the output schema; regression for + # the branch reading rel.fetch.common instead of rel.set.common (issue #217). + rel = stalg.Rel( + set=stalg.SetRel( + inputs=[read_rel, read_rel], + op=stalg.SetRel.SET_OP_UNION_ALL, + common=stalg.RelCommon(emit=stalg.RelCommon.Emit(output_mapping=[2, 0])), + ) + ) + + expected = stt.Type.Struct( + types=[ + stt.Type(fp32=stt.Type.FP32(nullability=stt.Type.NULLABILITY_NULLABLE)), + stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_REQUIRED)), + ] + ) + + assert infer_rel_schema(rel) == expected + + def test_inference_project_literal(): rel = stalg.Rel( project=stalg.ProjectRel(