fix: combine set-operation nullability across all inputs

src/substrait/type_inference.py

@@ -386,6 +386,67 @@ def _join_output_struct(type_name: str, left_rel, right_rel) -> stt.Type.Struct:
     return stt.Type.Struct(types=types, nullability=required)
 
 
+def _field_nullability(t: stt.Type):
+    """The nullability of a (concrete) field type, or UNSPECIFIED if it has none."""
+    kind = t.WhichOneof("kind")
+    return getattr(t, kind).nullability if kind else stt.Type.NULLABILITY_UNSPECIFIED
+
+
+def _with_field_nullability(t: stt.Type, nullability) -> stt.Type:
+    """A copy of ``t`` with its nullability replaced (unchanged if it has no kind)."""
+    out = stt.Type()
+    out.CopyFrom(t)
+    kind = out.WhichOneof("kind")
+    if kind:
+        getattr(out, kind).nullability = nullability
+    return out
+
+
+def _combine_set_nullability(op_name: str, nullabilities: list):
+    """Combine one field's nullability across a set operation's inputs.
+
+    Set inputs share field *types* but may differ in nullability; the output
+    nullability is combined across all inputs per the operation (matching the
+    Substrait spec's set-operation output-type derivation). ``nullabilities`` is
+    the field's nullability in each input, primary first.
+    """
+    nullable = stt.Type.NULLABILITY_NULLABLE
+    required = stt.Type.NULLABILITY_REQUIRED
+    primary, secondaries = nullabilities[0], nullabilities[1:]
+    if op_name in ("SET_OP_UNION_DISTINCT", "SET_OP_UNION_ALL"):
+        # Nullable if nullable in any input.
+        return nullable if nullable in nullabilities else required
+    if op_name == "SET_OP_INTERSECTION_PRIMARY":
+        # Nullable only if nullable in the primary and in some secondary input.
+        return nullable if primary == nullable and nullable in secondaries else required
+    if op_name in ("SET_OP_INTERSECTION_MULTISET", "SET_OP_INTERSECTION_MULTISET_ALL"):
+        # Required if required in any input.
+        return required if required in nullabilities else nullable
+    # MINUS_* (and unspecified): the same as the primary input.
+    return primary
+
+
+def _set_output_struct(op_name: str, inputs: list) -> stt.Type.Struct:
+    """The output struct of a set operation over already-inferred ``inputs``.
+
+    Field types are taken from the primary input (the spec requires identical
+    field types across inputs); each field's nullability is combined across all
+    inputs according to ``op_name`` via :func:`_combine_set_nullability`.
+    """
+    primary = inputs[0]
+    types = [
+        _with_field_nullability(
+            field,
+            _combine_set_nullability(
+                op_name,
+                [_field_nullability(s.types[i]) for s in inputs if i < len(s.types)],
+            ),
+        )
+        for i, field in enumerate(primary.types)
+    ]
+    return stt.Type.Struct(types=types, nullability=primary.nullability)
+
+
 def infer_rel_schema(rel: stalg.Rel) -> stt.Type.Struct:
     rel_type = rel.WhichOneof("rel_type")
 
@@ -429,7 +490,11 @@ def infer_rel_schema(rel: stalg.Rel) -> stt.Type.Struct:
 
         (common, struct) = (rel.project.common, raw_schema)
     elif rel_type == "set":
-        (common, struct) = (rel.set.common, infer_rel_schema(rel.set.inputs[0]))
+        input_structs = [infer_rel_schema(i) for i in rel.set.inputs]
+        (common, struct) = (
+            rel.set.common,
+            _set_output_struct(stalg.SetRel.SetOp.Name(rel.set.op), input_structs),
+        )
     elif rel_type == "cross":
         left_schema = infer_rel_schema(rel.cross.left)
         right_schema = infer_rel_schema(rel.cross.right)

tests/test_type_inference.py

@@ -1,3 +1,4 @@
+import pytest
 import substrait.algebra_pb2 as stalg
 import substrait.type_pb2 as stt
 
@@ -7,6 +8,9 @@
     infer_rel_schema,
 )
 
+_REQ = stt.Type.NULLABILITY_REQUIRED
+_NULL = stt.Type.NULLABILITY_NULLABLE
+
 struct = stt.Type.Struct(
     types=[
         stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_REQUIRED)),
@@ -480,3 +484,119 @@ def test_infer_nested_type_map():
         )
     )
     assert result == expected
+
+
+# Three set inputs, one i64 column per (required/nullable) combination across
+# them, matching the worked example in the Substrait spec's set-operation
+# "Output Type Derivation" table (issue #219). Columns, per (primary, s1, s2):
+#   RRR  RRN  RNR  RNN  NRR  NRN  NNR  NNN
+_SET_INPUT_NULLABILITIES = [
+    [_REQ, _REQ, _REQ, _REQ, _NULL, _NULL, _NULL, _NULL],  # primary
+    [_REQ, _REQ, _NULL, _NULL, _REQ, _REQ, _NULL, _NULL],  # secondary
+    [_REQ, _NULL, _REQ, _NULL, _REQ, _NULL, _REQ, _NULL],  # secondary
+]
+
+
+@pytest.mark.parametrize(
+    ("op", "expected"),
+    [
+        # MINUS variants inherit the primary input's nullability.
+        (
+            stalg.SetRel.SET_OP_MINUS_PRIMARY,
+            [_REQ, _REQ, _REQ, _REQ, _NULL, _NULL, _NULL, _NULL],
+        ),
+        (
+            stalg.SetRel.SET_OP_MINUS_PRIMARY_ALL,
+            [_REQ, _REQ, _REQ, _REQ, _NULL, _NULL, _NULL, _NULL],
+        ),
+        (
+            stalg.SetRel.SET_OP_MINUS_MULTISET,
+            [_REQ, _REQ, _REQ, _REQ, _NULL, _NULL, _NULL, _NULL],
+        ),
+        # Nullable only if nullable in the primary and some secondary.
+        (
+            stalg.SetRel.SET_OP_INTERSECTION_PRIMARY,
+            [_REQ, _REQ, _REQ, _REQ, _REQ, _NULL, _NULL, _NULL],
+        ),
+        # Required if required in any input.
+        (
+            stalg.SetRel.SET_OP_INTERSECTION_MULTISET,
+            [_REQ, _REQ, _REQ, _REQ, _REQ, _REQ, _REQ, _NULL],
+        ),
+        (
+            stalg.SetRel.SET_OP_INTERSECTION_MULTISET_ALL,
+            [_REQ, _REQ, _REQ, _REQ, _REQ, _REQ, _REQ, _NULL],
+        ),
+        # Nullable if nullable in any input.
+        (
+            stalg.SetRel.SET_OP_UNION_DISTINCT,
+            [_REQ, _NULL, _NULL, _NULL, _NULL, _NULL, _NULL, _NULL],
+        ),
+        (
+            stalg.SetRel.SET_OP_UNION_ALL,
+            [_REQ, _NULL, _NULL, _NULL, _NULL, _NULL, _NULL, _NULL],
+        ),
+    ],
+)
+def test_inference_set_nullability(op, expected):
+    inputs = [
+        stalg.Rel(
+            read=stalg.ReadRel(
+                base_schema=stt.NamedStruct(
+                    names=[f"c{i}" for i in range(len(nullabilities))],
+                    struct=stt.Type.Struct(
+                        types=[
+                            stt.Type(i64=stt.Type.I64(nullability=n))
+                            for n in nullabilities
+                        ]
+                    ),
+                )
+            )
+        )
+        for nullabilities in _SET_INPUT_NULLABILITIES
+    ]
+
+    rel = stalg.Rel(set=stalg.SetRel(inputs=inputs, op=op))
+
+    result = [t.i64.nullability for t in infer_rel_schema(rel).types]
+    assert result == expected
+
+
+def test_inference_set_nullability_preserves_field_types():
+    # Combining nullability must keep each field's full type (parameters and
+    # nested element types) intact -- only the top-level nullability changes.
+    def _struct(dec_null, vc_null, list_null):
+        return stt.Type.Struct(
+            types=[
+                stt.Type(
+                    decimal=stt.Type.Decimal(
+                        precision=10, scale=2, nullability=dec_null
+                    )
+                ),
+                stt.Type(varchar=stt.Type.VarChar(length=5, nullability=vc_null)),
+                stt.Type(
+                    list=stt.Type.List(
+                        type=stt.Type(string=stt.Type.String(nullability=_REQ)),
+                        nullability=list_null,
+                    )
+                ),
+            ],
+            nullability=_REQ,
+        )
+
+    def _read(struct):
+        return stalg.Rel(
+            read=stalg.ReadRel(
+                base_schema=stt.NamedStruct(names=["d", "v", "l"], struct=struct)
+            )
+        )
+
+    primary = _read(_struct(_NULL, _REQ, _REQ))
+    secondary = _read(_struct(_REQ, _NULL, _NULL))
+    rel = stalg.Rel(
+        set=stalg.SetRel(inputs=[primary, secondary], op=stalg.SetRel.SET_OP_UNION_ALL)
+    )
+
+    # UNION -> nullable if nullable in any input; types (decimal 10/2, varchar 5,
+    # list<string>) and the required inner string element are preserved.
+    assert infer_rel_schema(rel) == _struct(_NULL, _NULL, _NULL)