diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 24711f12..3f6ab1c4 100644 --- a/src/substrait/type_inference.py +++ b/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) diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index e45b3e7a..e67cfe65 100644 --- a/tests/test_type_inference.py +++ b/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) and the required inner string element are preserved. + assert infer_rel_schema(rel) == _struct(_NULL, _NULL, _NULL)