From 168fb4c8dd9e13860e2b53aa4e327c01c8cbc297 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Mon, 5 Oct 2026 14:48:06 +0300 Subject: [PATCH 1/4] fix(type-inference): combine Expand switching field nullability --- src/substrait/type_inference.py | 63 ++++++++-- tests/builders/plan/test_expand.py | 59 +++++++++ tests/test_type_inference.py | 184 +++++++++++++++++++++++++++++ 3 files changed, 293 insertions(+), 13 deletions(-) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 444057ef..10a31ea7 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -146,6 +146,12 @@ def schema_of(self, anchor, registry) -> stt.Type.Struct: "outer_schemas", default=() ) +# Lambda parameter structs, outermost first. Root references still resolve +# against the incoming row supplied to infer_expression_type. +_lambda_schemas: contextvars.ContextVar = contextvars.ContextVar( + "lambda_schemas", default=() +) + # The anchor scope (an ``_AnchorScope``) for the plan / correlated sub-tree # currently being inferred, or None outside inference. Set for the duration of a # whole-plan inference (``infer_plan_schema``) so a ``rel_reference`` anywhere in @@ -606,8 +612,21 @@ def infer_expression_type( "outer reference outside an enclosing (correlated) query" ) schema = stack[len(stack) - steps].struct + elif root_type == "lambda_parameter_reference": + stack = _lambda_schemas.get() + if stack: + steps = expression.selection.lambda_parameter_reference.steps_out + if steps >= len(stack): + raise Exception( + "lambda parameter reference outside an enclosing lambda scope" + ) + schema = stack[len(stack) - steps - 1] + else: + # Builders also infer unwrapped parameter expressions with their + # parameter struct supplied as parent_schema. + schema = parent_schema else: - assert root_type in ("root_reference", "lambda_parameter_reference") + assert root_type == "root_reference" schema = parent_schema reference_type = expression.selection.WhichOneof("reference_type") @@ -654,13 +673,19 @@ def infer_expression_type( return infer_nested_type( expression.nested, parent_schema, registry=registry, subtrees=subtrees ) - elif rex_type == "lambda": - # A lambda's type is func body_type>; the body's parameter - # references resolve against the lambda's own parameter struct. - lam = getattr(expression, "lambda") - body_type = infer_expression_type( - lam.body, lam.parameters, registry=registry, subtrees=subtrees - ) + elif rex_type in ("lambda", "lambda_invocation"): + source = expression if rex_type == "lambda" else expression.lambda_invocation + lam = getattr(source, "lambda") + stack = _lambda_schemas.get() + token = _lambda_schemas.set((*stack, lam.parameters)) + try: + body_type = infer_expression_type( + lam.body, parent_schema, registry=registry, subtrees=subtrees + ) + finally: + _lambda_schemas.reset(token) + if rex_type == "lambda_invocation": + return body_type return stt.Type( func=stt.Type.Func( parameter_types=list(lam.parameters.types), @@ -1163,16 +1188,28 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None, subtrees=()) -> stt.Type. "expand switching field has no duplicate expressions; its " "output type cannot be inferred" ) - # All duplicates of a switching field share one type; the first - # determines the output column type. - field_types.append( + # Duplicates share a type class but may differ in nullability. + # The output is nullable if any duplicate is nullable. + duplicate_types = [ infer_expression_type( - duplicates[0], + duplicate, parent_schema, registry=registry, subtrees=subtrees, ) - ) + for duplicate in duplicates + ] + output_type = duplicate_types[0] + # Unbound placeholders carry no nullability. + if any( + t.WhichOneof("kind") != "unbound" + and _field_nullability(t) == stt.Type.NULLABILITY_NULLABLE + for t in duplicate_types + ): + output_type = _with_field_nullability( + output_type, stt.Type.NULLABILITY_NULLABLE + ) + field_types.append(output_type) # Expand appends an i32 column with the index of the duplicate the row # is derived from. field_types.append( diff --git a/tests/builders/plan/test_expand.py b/tests/builders/plan/test_expand.py index 2fb95a92..5369e979 100644 --- a/tests/builders/plan/test_expand.py +++ b/tests/builders/plan/test_expand.py @@ -1,3 +1,4 @@ +import pytest import substrait.algebra_pb2 as stalg import substrait.plan_pb2 as stp import substrait.type_pb2 as stt @@ -86,6 +87,64 @@ def test_expand_schema_inference(): assert kinds == ["string", "fp64", "i32"] +@pytest.mark.parametrize( + "nullabilities", + [ + pytest.param((False, True), id="nullable-last"), + pytest.param((True, False), id="nullable-first"), + pytest.param((False, False), id="all-required"), + pytest.param((True, True), id="all-nullable"), + pytest.param((False, False, True), id="nullable-third"), + pytest.param((False,), id="single-required"), + pytest.param((True,), id="single-nullable"), + ], +) +def test_expand_switching_field_nullability(nullabilities): + # algebra.proto: a switching field is nullable if any duplicate is nullable. + names = [f"value_{i}" for i in range(len(nullabilities))] + input_schema = stt.NamedStruct( + names=["region", *names], + struct=stt.Type.Struct( + types=[string(nullable=False)] + [fp64(nullable=n) for n in nullabilities], + nullability=stt.Type.NULLABILITY_REQUIRED, + ), + ) + plan = expand( + read_named_table("sales", input_schema), + fields=[ + ("consistent", column("region")), + ("switching", [column(name) for name in names]), + ], + names=["region", "value", "idx"], + )(None) + original = stp.Plan() + original.CopyFrom(plan) + + schema = infer_plan_schema(plan) + + assert schema.struct.types[0] == string(nullable=False) + assert schema.struct.types[1] == fp64(nullable=any(nullabilities)) + assert plan == original + + +def test_expand_switching_field_preserves_unbound_type(): + # A partially bound plan may carry a placeholder with no nullability. + unbound = stt.Type(unbound=stt.Type.Unbound()) + input_schema = stt.NamedStruct( + names=["u"], + struct=stt.Type.Struct( + types=[unbound], nullability=stt.Type.NULLABILITY_REQUIRED + ), + ) + plan = expand( + read_named_table("partial", input_schema), + fields=[("switching", [column("u"), column("u")])], + names=["value", "idx"], + )(None) + + assert infer_plan_schema(plan).struct.types[0] == unbound + + def test_expand_empty_switching_field_raises_clear_error(): # An empty switching field has no expression to derive a type from; schema # inference must raise a clear error rather than an opaque IndexError. diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index c435d998..59d9b807 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -702,6 +702,190 @@ def test_infer_expression_type_literal(): assert result == expected +def _lambda_ref(field=0, *, steps_out=0): + return stalg.Expression( + selection=stalg.Expression.FieldReference( + lambda_parameter_reference=stalg.Expression.FieldReference.LambdaParameterReference( + steps_out=steps_out + ), + direct_reference=stalg.Expression.ReferenceSegment( + struct_field=stalg.Expression.ReferenceSegment.StructField(field=field) + ), + ) + ) + + +def _root_ref(field=0): + 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) + ), + ) + ) + + +def _lambda_expression(parameters, body, *, invoke=False, arguments=()): + lam = stalg.Expression.Lambda( + parameters=stt.Type.Struct(types=parameters, nullability=_REQ), body=body + ) + if invoke: + return stalg.Expression( + lambda_invocation=stalg.Expression.LambdaInvocation( + **{ + "lambda": lam, + "arguments": stalg.Expression.Nested.Struct(fields=arguments), + } + ) + ) + return stalg.Expression(**{"lambda": lam}) + + +@pytest.mark.parametrize("nullable", [False, True]) +def test_infer_lambda_invocation_literal_body(nullable): + body = stalg.Expression(literal=stalg.Expression.Literal(i32=7, nullable=nullable)) + expr = _lambda_expression([], body, invoke=True) + + assert infer_expression_type(expr, struct) == stt.Type( + i32=stt.Type.I32(nullability=_NULL if nullable else _REQ) + ) + + +@pytest.mark.parametrize("invoke", [False, True]) +@pytest.mark.parametrize("reference", ["root", "parameter"]) +def test_infer_lambda_keeps_row_and_parameter_scopes_distinct(invoke, reference): + # Same field index, different types: a root reference captures the input row. + row_type = stt.Type(string=stt.Type.String(nullability=_NULL)) + parameter = stt.Type(i64=stt.Type.I64(nullability=_REQ)) + body = _root_ref() if reference == "root" else _lambda_ref() + expr = _lambda_expression( + [parameter], + body, + invoke=invoke, + arguments=[stalg.Expression(literal=stalg.Expression.Literal(i64=42))], + ) + original = expr.SerializeToString() + result = infer_expression_type( + expr, stt.Type.Struct(types=[row_type], nullability=_REQ) + ) + expected = row_type if reference == "root" else parameter + + if invoke: + assert result == expected + else: + assert result == stt.Type( + func=stt.Type.Func( + parameter_types=[parameter], return_type=expected, nullability=_REQ + ) + ) + assert expr.SerializeToString() == original + + +@pytest.mark.parametrize("steps_out", [0, 1]) +@pytest.mark.parametrize("invoke", [False, True]) +def test_infer_nested_lambda_parameter_scopes(steps_out, invoke): + outer = stt.Type(string=stt.Type.String(nullability=_NULL)) + inner = stt.Type(i64=stt.Type.I64(nullability=_REQ)) + inner_expr = _lambda_expression( + [inner], + _lambda_ref(steps_out=steps_out), + invoke=invoke, + arguments=[stalg.Expression(literal=stalg.Expression.Literal(i64=42))], + ) + expr = _lambda_expression( + [outer], + inner_expr, + invoke=invoke, + arguments=[ + stalg.Expression( + literal=stalg.Expression.Literal(string="x", nullable=True) + ) + ], + ) + result = infer_expression_type(expr, struct) + expected = inner if steps_out == 0 else outer + + if invoke: + assert result == expected + else: + assert result.func.return_type.func.return_type == expected + + +def test_lambda_parameter_reference_cannot_reach_a_missing_scope(): + parameter = stt.Type(i32=stt.Type.I32(nullability=_REQ)) + expr = _lambda_expression([parameter], _lambda_ref(steps_out=1)) + + with pytest.raises(Exception, match="outside an enclosing lambda scope"): + infer_expression_type(expr, struct) + + +@pytest.mark.parametrize("fail", [False, True]) +def test_lambda_parameter_scope_does_not_leak(fail): + parameter = stt.Type(i32=stt.Type.I32(nullability=_REQ)) + # The higher-order builders also infer parameter expressions directly, + # supplying their parameter struct before wrapping the body in a lambda. + supplied = stt.Type.Struct( + types=[stt.Type(string=stt.Type.String(nullability=_NULL))], nullability=_REQ + ) + body = _root_ref(99) if fail else _lambda_ref() + expr = _lambda_expression([parameter], body) + if fail: + with pytest.raises(IndexError): + infer_expression_type(expr, supplied) + else: + assert infer_expression_type(expr, supplied).func.return_type == parameter + + assert infer_expression_type(_lambda_ref(), supplied) == supplied.types[0] + + +def _expand_switching(types, duplicates): + inp = stalg.Rel( + read=stalg.ReadRel( + named_table=stalg.ReadRel.NamedTable(names=["t"]), + base_schema=stt.NamedStruct( + names=[f"c{i}" for i in range(len(types))], + struct=stt.Type.Struct(types=types, nullability=_REQ), + ), + ) + ) + fields = [ + stalg.ExpandRel.ExpandField( + switching_field=stalg.ExpandRel.SwitchingField(duplicates=duplicates) + ) + ] + [ + stalg.ExpandRel.ExpandField(consistent_field=_root_ref(i)) + for i in range(1, len(types)) + ] + return stalg.Rel(expand=stalg.ExpandRel(input=inp, fields=fields)) + + +@pytest.mark.parametrize("nullable", [False, True]) +def test_expand_switching_field_with_lambda_invocation(nullable): + integer = stt.Type(i32=stt.Type.I32(nullability=_REQ)) + lit = stalg.Expression(literal=stalg.Expression.Literal(i32=7)) + body = stalg.Expression(literal=stalg.Expression.Literal(i32=7, nullable=nullable)) + rel = _expand_switching([integer], [lit, _lambda_expression([], body, invoke=True)]) + original = rel.SerializeToString() + + assert infer_rel_schema(rel).types[0] == stt.Type( + i32=stt.Type.I32(nullability=_NULL if nullable else _REQ) + ) + assert rel.SerializeToString() == original + + +def test_expand_switching_lambda_captures_input_row(): + integer = stt.Type(i32=stt.Type.I32(nullability=_REQ)) + function = stt.Type(func=stt.Type.Func(return_type=integer, nullability=_REQ)) + rel = _expand_switching( + [function, integer], [_root_ref(), _lambda_expression([], _root_ref(1))] + ) + original = rel.SerializeToString() + + assert infer_rel_schema(rel).types[0] == function + assert rel.SerializeToString() == original + + def test_infer_expression_type_selection(): """Test infer_expression_type with a field selection expression.""" expr = stalg.Expression( From 902eb61d17129b570eeca32360be5025e2951273 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Mon, 5 Oct 2026 17:48:41 +0300 Subject: [PATCH 2/4] fix(type-inference): reject standalone outer lambda references --- src/substrait/type_inference.py | 6 +++++- tests/test_type_inference.py | 5 +++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 10a31ea7..71ab4233 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -614,8 +614,8 @@ def infer_expression_type( schema = stack[len(stack) - steps].struct elif root_type == "lambda_parameter_reference": stack = _lambda_schemas.get() + steps = expression.selection.lambda_parameter_reference.steps_out if stack: - steps = expression.selection.lambda_parameter_reference.steps_out if steps >= len(stack): raise Exception( "lambda parameter reference outside an enclosing lambda scope" @@ -624,6 +624,10 @@ def infer_expression_type( else: # Builders also infer unwrapped parameter expressions with their # parameter struct supplied as parent_schema. + if steps != 0: + raise Exception( + "lambda parameter reference outside an enclosing lambda scope" + ) schema = parent_schema else: assert root_type == "root_reference" diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index 59d9b807..8a907cdd 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -820,6 +820,11 @@ def test_lambda_parameter_reference_cannot_reach_a_missing_scope(): infer_expression_type(expr, struct) +def test_standalone_lambda_parameter_reference_cannot_reach_a_missing_scope(): + with pytest.raises(Exception, match="outside an enclosing lambda scope"): + infer_expression_type(_lambda_ref(steps_out=1), struct) + + @pytest.mark.parametrize("fail", [False, True]) def test_lambda_parameter_scope_does_not_leak(fail): parameter = stt.Type(i32=stt.Type.I32(nullability=_REQ)) From fe931894fe51f2fb84364807afd917941d3bc3dd Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Tue, 6 Oct 2026 10:59:39 +0300 Subject: [PATCH 3/4] refactor: separate lambda inference from the Expand fix --- src/substrait/type_inference.py | 45 ++------ tests/test_type_inference.py | 189 -------------------------------- 2 files changed, 8 insertions(+), 226 deletions(-) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 71ab4233..a7b5bdbc 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -146,12 +146,6 @@ def schema_of(self, anchor, registry) -> stt.Type.Struct: "outer_schemas", default=() ) -# Lambda parameter structs, outermost first. Root references still resolve -# against the incoming row supplied to infer_expression_type. -_lambda_schemas: contextvars.ContextVar = contextvars.ContextVar( - "lambda_schemas", default=() -) - # The anchor scope (an ``_AnchorScope``) for the plan / correlated sub-tree # currently being inferred, or None outside inference. Set for the duration of a # whole-plan inference (``infer_plan_schema``) so a ``rel_reference`` anywhere in @@ -612,25 +606,8 @@ def infer_expression_type( "outer reference outside an enclosing (correlated) query" ) schema = stack[len(stack) - steps].struct - elif root_type == "lambda_parameter_reference": - stack = _lambda_schemas.get() - steps = expression.selection.lambda_parameter_reference.steps_out - if stack: - if steps >= len(stack): - raise Exception( - "lambda parameter reference outside an enclosing lambda scope" - ) - schema = stack[len(stack) - steps - 1] - else: - # Builders also infer unwrapped parameter expressions with their - # parameter struct supplied as parent_schema. - if steps != 0: - raise Exception( - "lambda parameter reference outside an enclosing lambda scope" - ) - schema = parent_schema else: - assert root_type == "root_reference" + assert root_type in ("root_reference", "lambda_parameter_reference") schema = parent_schema reference_type = expression.selection.WhichOneof("reference_type") @@ -677,19 +654,13 @@ def infer_expression_type( return infer_nested_type( expression.nested, parent_schema, registry=registry, subtrees=subtrees ) - elif rex_type in ("lambda", "lambda_invocation"): - source = expression if rex_type == "lambda" else expression.lambda_invocation - lam = getattr(source, "lambda") - stack = _lambda_schemas.get() - token = _lambda_schemas.set((*stack, lam.parameters)) - try: - body_type = infer_expression_type( - lam.body, parent_schema, registry=registry, subtrees=subtrees - ) - finally: - _lambda_schemas.reset(token) - if rex_type == "lambda_invocation": - return body_type + elif rex_type == "lambda": + # A lambda's type is func body_type>; the body's parameter + # references resolve against the lambda's own parameter struct. + lam = getattr(expression, "lambda") + body_type = infer_expression_type( + lam.body, lam.parameters, registry=registry, subtrees=subtrees + ) return stt.Type( func=stt.Type.Func( parameter_types=list(lam.parameters.types), diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index 8a907cdd..c435d998 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -702,195 +702,6 @@ def test_infer_expression_type_literal(): assert result == expected -def _lambda_ref(field=0, *, steps_out=0): - return stalg.Expression( - selection=stalg.Expression.FieldReference( - lambda_parameter_reference=stalg.Expression.FieldReference.LambdaParameterReference( - steps_out=steps_out - ), - direct_reference=stalg.Expression.ReferenceSegment( - struct_field=stalg.Expression.ReferenceSegment.StructField(field=field) - ), - ) - ) - - -def _root_ref(field=0): - 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) - ), - ) - ) - - -def _lambda_expression(parameters, body, *, invoke=False, arguments=()): - lam = stalg.Expression.Lambda( - parameters=stt.Type.Struct(types=parameters, nullability=_REQ), body=body - ) - if invoke: - return stalg.Expression( - lambda_invocation=stalg.Expression.LambdaInvocation( - **{ - "lambda": lam, - "arguments": stalg.Expression.Nested.Struct(fields=arguments), - } - ) - ) - return stalg.Expression(**{"lambda": lam}) - - -@pytest.mark.parametrize("nullable", [False, True]) -def test_infer_lambda_invocation_literal_body(nullable): - body = stalg.Expression(literal=stalg.Expression.Literal(i32=7, nullable=nullable)) - expr = _lambda_expression([], body, invoke=True) - - assert infer_expression_type(expr, struct) == stt.Type( - i32=stt.Type.I32(nullability=_NULL if nullable else _REQ) - ) - - -@pytest.mark.parametrize("invoke", [False, True]) -@pytest.mark.parametrize("reference", ["root", "parameter"]) -def test_infer_lambda_keeps_row_and_parameter_scopes_distinct(invoke, reference): - # Same field index, different types: a root reference captures the input row. - row_type = stt.Type(string=stt.Type.String(nullability=_NULL)) - parameter = stt.Type(i64=stt.Type.I64(nullability=_REQ)) - body = _root_ref() if reference == "root" else _lambda_ref() - expr = _lambda_expression( - [parameter], - body, - invoke=invoke, - arguments=[stalg.Expression(literal=stalg.Expression.Literal(i64=42))], - ) - original = expr.SerializeToString() - result = infer_expression_type( - expr, stt.Type.Struct(types=[row_type], nullability=_REQ) - ) - expected = row_type if reference == "root" else parameter - - if invoke: - assert result == expected - else: - assert result == stt.Type( - func=stt.Type.Func( - parameter_types=[parameter], return_type=expected, nullability=_REQ - ) - ) - assert expr.SerializeToString() == original - - -@pytest.mark.parametrize("steps_out", [0, 1]) -@pytest.mark.parametrize("invoke", [False, True]) -def test_infer_nested_lambda_parameter_scopes(steps_out, invoke): - outer = stt.Type(string=stt.Type.String(nullability=_NULL)) - inner = stt.Type(i64=stt.Type.I64(nullability=_REQ)) - inner_expr = _lambda_expression( - [inner], - _lambda_ref(steps_out=steps_out), - invoke=invoke, - arguments=[stalg.Expression(literal=stalg.Expression.Literal(i64=42))], - ) - expr = _lambda_expression( - [outer], - inner_expr, - invoke=invoke, - arguments=[ - stalg.Expression( - literal=stalg.Expression.Literal(string="x", nullable=True) - ) - ], - ) - result = infer_expression_type(expr, struct) - expected = inner if steps_out == 0 else outer - - if invoke: - assert result == expected - else: - assert result.func.return_type.func.return_type == expected - - -def test_lambda_parameter_reference_cannot_reach_a_missing_scope(): - parameter = stt.Type(i32=stt.Type.I32(nullability=_REQ)) - expr = _lambda_expression([parameter], _lambda_ref(steps_out=1)) - - with pytest.raises(Exception, match="outside an enclosing lambda scope"): - infer_expression_type(expr, struct) - - -def test_standalone_lambda_parameter_reference_cannot_reach_a_missing_scope(): - with pytest.raises(Exception, match="outside an enclosing lambda scope"): - infer_expression_type(_lambda_ref(steps_out=1), struct) - - -@pytest.mark.parametrize("fail", [False, True]) -def test_lambda_parameter_scope_does_not_leak(fail): - parameter = stt.Type(i32=stt.Type.I32(nullability=_REQ)) - # The higher-order builders also infer parameter expressions directly, - # supplying their parameter struct before wrapping the body in a lambda. - supplied = stt.Type.Struct( - types=[stt.Type(string=stt.Type.String(nullability=_NULL))], nullability=_REQ - ) - body = _root_ref(99) if fail else _lambda_ref() - expr = _lambda_expression([parameter], body) - if fail: - with pytest.raises(IndexError): - infer_expression_type(expr, supplied) - else: - assert infer_expression_type(expr, supplied).func.return_type == parameter - - assert infer_expression_type(_lambda_ref(), supplied) == supplied.types[0] - - -def _expand_switching(types, duplicates): - inp = stalg.Rel( - read=stalg.ReadRel( - named_table=stalg.ReadRel.NamedTable(names=["t"]), - base_schema=stt.NamedStruct( - names=[f"c{i}" for i in range(len(types))], - struct=stt.Type.Struct(types=types, nullability=_REQ), - ), - ) - ) - fields = [ - stalg.ExpandRel.ExpandField( - switching_field=stalg.ExpandRel.SwitchingField(duplicates=duplicates) - ) - ] + [ - stalg.ExpandRel.ExpandField(consistent_field=_root_ref(i)) - for i in range(1, len(types)) - ] - return stalg.Rel(expand=stalg.ExpandRel(input=inp, fields=fields)) - - -@pytest.mark.parametrize("nullable", [False, True]) -def test_expand_switching_field_with_lambda_invocation(nullable): - integer = stt.Type(i32=stt.Type.I32(nullability=_REQ)) - lit = stalg.Expression(literal=stalg.Expression.Literal(i32=7)) - body = stalg.Expression(literal=stalg.Expression.Literal(i32=7, nullable=nullable)) - rel = _expand_switching([integer], [lit, _lambda_expression([], body, invoke=True)]) - original = rel.SerializeToString() - - assert infer_rel_schema(rel).types[0] == stt.Type( - i32=stt.Type.I32(nullability=_NULL if nullable else _REQ) - ) - assert rel.SerializeToString() == original - - -def test_expand_switching_lambda_captures_input_row(): - integer = stt.Type(i32=stt.Type.I32(nullability=_REQ)) - function = stt.Type(func=stt.Type.Func(return_type=integer, nullability=_REQ)) - rel = _expand_switching( - [function, integer], [_root_ref(), _lambda_expression([], _root_ref(1))] - ) - original = rel.SerializeToString() - - assert infer_rel_schema(rel).types[0] == function - assert rel.SerializeToString() == original - - def test_infer_expression_type_selection(): """Test infer_expression_type with a field selection expression.""" expr = stalg.Expression( From 96cc4da94a1f91abe593dfbc0a201bd3f9fddc46 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Tue, 6 Oct 2026 11:26:32 +0300 Subject: [PATCH 4/4] fix(type-inference): validate Expand duplicate types and preserve unknown nullability --- src/substrait/type_inference.py | 25 ++++++++++--- tests/builders/plan/test_expand.py | 60 ++++++++++++++++++++++++++++++ tests/dataframe/test_frame.py | 8 ++++ 3 files changed, 87 insertions(+), 6 deletions(-) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index a7b5bdbc..741fdcd0 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -1174,16 +1174,29 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None, subtrees=()) -> stt.Type. ) for duplicate in duplicates ] - output_type = duplicate_types[0] - # Unbound placeholders carry no nullability. + bound = [ + t for t in duplicate_types if t.WhichOneof("kind") != "unbound" + ] + unbound = [ + t for t in duplicate_types if t.WhichOneof("kind") == "unbound" + ] + if len({t.WhichOneof("kind") for t in bound}) > 1: + raise ValueError( + "expand switching field duplicates must share one type class" + ) if any( - t.WhichOneof("kind") != "unbound" - and _field_nullability(t) == stt.Type.NULLABILITY_NULLABLE - for t in duplicate_types + _field_nullability(t) == stt.Type.NULLABILITY_NULLABLE + for t in bound ): + # Nullable whatever an unbound placeholder later binds to. output_type = _with_field_nullability( - output_type, stt.Type.NULLABILITY_NULLABLE + bound[0], stt.Type.NULLABILITY_NULLABLE ) + elif unbound: + # Otherwise the nullability depends on the placeholder. + output_type = unbound[0] + else: + output_type = bound[0] field_types.append(output_type) # Expand appends an i32 column with the index of the duplicate the row # is derived from. diff --git a/tests/builders/plan/test_expand.py b/tests/builders/plan/test_expand.py index 5369e979..affc221d 100644 --- a/tests/builders/plan/test_expand.py +++ b/tests/builders/plan/test_expand.py @@ -145,6 +145,66 @@ def test_expand_switching_field_preserves_unbound_type(): assert infer_plan_schema(plan).struct.types[0] == unbound +def _switching_plan(types): + names = [f"value_{i}" for i in range(len(types))] + schema = stt.NamedStruct( + names=names, + struct=stt.Type.Struct(types=types, nullability=stt.Type.NULLABILITY_REQUIRED), + ) + return expand( + read_named_table("partial", schema), + fields=[("switching", [column(name) for name in names])], + names=["value", "idx"], + )(None) + + +@pytest.mark.parametrize( + "bound_nullabilities, unbound_position", + [ + pytest.param((False,), 0, id="unbound-required"), + pytest.param((False,), 1, id="required-unbound"), + pytest.param((True,), 0, id="unbound-nullable"), + pytest.param((True,), 1, id="nullable-unbound"), + pytest.param((False, True), 0, id="unbound-required-nullable"), + pytest.param((False, True), 1, id="required-unbound-nullable"), + pytest.param((False, True), 2, id="required-nullable-unbound"), + ], +) +def test_expand_switching_field_combines_unbound_duplicates( + bound_nullabilities, unbound_position +): + unbound = stt.Type(unbound=stt.Type.Unbound()) + types = [fp64(nullable=n) for n in bound_nullabilities] + types.insert(unbound_position, unbound) + plan = _switching_plan(types) + original = plan.SerializeToString() + + result = infer_plan_schema(plan).struct.types[0] + + # A known nullable duplicate fixes nullability regardless of the placeholder. + # Otherwise the placeholder may still bind to a nullable type. + expected = fp64(nullable=True) if any(bound_nullabilities) else unbound + assert result == expected + assert plan.SerializeToString() == original + + +@pytest.mark.parametrize("reverse", [False, True]) +@pytest.mark.parametrize("unbound_position", [None, 0, 1, 2]) +def test_expand_switching_field_rejects_mixed_type_classes(reverse, unbound_position): + types = [fp64(nullable=False), string(nullable=True)] + if reverse: + types.reverse() + if unbound_position is not None: + types.insert(unbound_position, stt.Type(unbound=stt.Type.Unbound())) + plan = _switching_plan(types) + original = plan.SerializeToString() + + with pytest.raises(ValueError, match="duplicates must share one type class"): + infer_plan_schema(plan) + + assert plan.SerializeToString() == original + + def test_expand_empty_switching_field_raises_clear_error(): # An empty switching field has no expression to derive a type from; schema # inference must raise a clear error rather than an opaque IndexError. diff --git a/tests/dataframe/test_frame.py b/tests/dataframe/test_frame.py index 2e1760b8..38aa3843 100644 --- a/tests/dataframe/test_frame.py +++ b/tests/dataframe/test_frame.py @@ -828,6 +828,14 @@ def test_unpivot_schema_inference_allows_chaining(): assert plan.relations[-1].root.input.HasField("filter") +@pytest.mark.parametrize("columns", [["a", "b"], ["b", "a"]]) +def test_unpivot_rejects_mixed_type_classes(columns): + frame = sub.read_named_table("mixed", {"a": sub.i32, "b": sub.string}) + + with pytest.raises(ValueError, match="duplicates must share one type class"): + frame.unpivot(columns).to_plan() + + def test_unpivot_requires_on(): with pytest.raises(ValueError, match="at least one column"): _wide_df().unpivot([], index="region")