From 9470831b51dcf2d453a342841f44358809604559 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Tue, 6 Oct 2026 10:59:29 +0300 Subject: [PATCH 1/4] fix(type-inference)!: resolve lambda scopes and invocation types --- src/substrait/type_inference.py | 45 ++++++-- tests/test_type_inference.py | 186 ++++++++++++++++++++++++++++++++ 2 files changed, 223 insertions(+), 8 deletions(-) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 444057e..73f6033 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,25 @@ 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 in ("root_reference", "lambda_parameter_reference") + assert root_type == "root_reference" schema = parent_schema reference_type = expression.selection.WhichOneof("reference_type") @@ -654,13 +677,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), diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index c435d99..c73a684 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -702,6 +702,192 @@ 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)) + body = stalg.Expression(literal=stalg.Expression.Literal(i32=7, nullable=nullable)) + rel = _expand_switching([integer], [_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], [_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 de67ff8615c25d8b0f843a467eac5d81e23c76c2 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Tue, 6 Oct 2026 11:57:23 +0300 Subject: [PATCH 2/4] fix(type-inference): bind builder callbacks within lambda scopes --- src/substrait/dataframe/expr.py | 8 ++-- src/substrait/type_inference.py | 36 +++++++-------- tests/dataframe/test_frame.py | 81 +++++++++++++++++++++++++++++++++ tests/test_type_inference.py | 21 +++++++-- 4 files changed, 118 insertions(+), 28 deletions(-) diff --git a/src/substrait/dataframe/expr.py b/src/substrait/dataframe/expr.py index 761f2a5..43b58e3 100644 --- a/src/substrait/dataframe/expr.py +++ b/src/substrait/dataframe/expr.py @@ -66,7 +66,7 @@ from substrait.builders.extended_expression import ( set_predicate as _set_predicate, ) -from substrait.type_inference import infer_extended_expression_schema +from substrait.type_inference import infer_extended_expression_schema, lambda_scope # Standard Substrait function-extension URNs used by the operators below. FUNCTIONS_COMPARISON = "extension:io.substrait:functions_comparison" @@ -643,7 +643,6 @@ def resolve(base_schema, registry): param_struct = stp.Type.Struct( types=[element_type], nullability=stp.Type.NULLABILITY_REQUIRED ) - param_ns = stp.NamedStruct(names=["element"], struct=param_struct) param_ref = stalg.Expression( selection=stalg.Expression.FieldReference( lambda_parameter_reference=( @@ -665,10 +664,11 @@ def resolve(base_schema, registry): expression=param_ref, output_names=["element"] ) ], - base_schema=param_ns, + base_schema=base_schema, ) ) - body = Expr._coerce(callback(element))._unbound(param_ns, registry) + with lambda_scope(param_struct): + body = Expr._coerce(callback(element))._unbound(base_schema, registry) lambda_expr = stalg.Expression( **{ "lambda": stalg.Expression.Lambda( diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 73f6033..bfa6b3e 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -152,6 +152,17 @@ def schema_of(self, anchor, registry) -> stt.Type.Struct: "lambda_schemas", default=() ) + +@contextlib.contextmanager +def lambda_scope(parameters: stt.Type.Struct): + """Bind a lambda's parameter struct while building or inferring its body.""" + token = _lambda_schemas.set((*_lambda_schemas.get(), parameters)) + try: + yield + finally: + _lambda_schemas.reset(token) + + # 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 @@ -615,20 +626,11 @@ def infer_expression_type( 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 + if steps >= len(stack): + raise Exception( + "lambda parameter reference outside an enclosing lambda scope" + ) + schema = stack[len(stack) - steps - 1] else: assert root_type == "root_reference" schema = parent_schema @@ -680,14 +682,10 @@ def infer_expression_type( 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: + with lambda_scope(lam.parameters): 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( diff --git a/tests/dataframe/test_frame.py b/tests/dataframe/test_frame.py index 2e1760b..deb54db 100644 --- a/tests/dataframe/test_frame.py +++ b/tests/dataframe/test_frame.py @@ -997,6 +997,87 @@ def test_list_transform_schema_inference(): assert t.list.type.WhichOneof("kind") == "i64" +def test_list_transform_captures_input_row(): + from substrait.type_inference import infer_plan_schema + + plan = ( + sub.read_named_table( + "t", {"arr": sub.list_(sub.i64.non_null), "label": sub.string} + ) + .select(sub.col("arr").list_transform(lambda _: sub.col("label"))) + .to_plan() + ) + assert infer_plan_schema(plan).struct.types[0].list.type == string() + + +def test_nested_list_transform_resolves_enclosing_parameter(): + import substrait.extended_expression_pb2 as stee + + from substrait.type_inference import infer_plan_schema + + # An explicit steps_out=1 capture is inferred while the builders bind the + # inner callback. Outer and inner elements have different types. + capture = stalg.Expression( + selection=stalg.Expression.FieldReference( + lambda_parameter_reference=stalg.Expression.FieldReference.LambdaParameterReference( + steps_out=1 + ), + direct_reference=stalg.Expression.ReferenceSegment( + struct_field=stalg.Expression.ReferenceSegment.StructField(field=0) + ), + ) + ) + outer_parameter = sub.Expr( + lambda schema, _: stee.ExtendedExpression( + base_schema=schema, + referred_expr=[ + stee.ExpressionReference(expression=capture, output_names=["outer"]) + ], + ) + ) + plan = ( + sub.read_named_table( + "t", + {"arr": sub.list_(sub.i64.non_null), "inner": sub.list_(fp64())}, + ) + .select( + sub.col("arr").list_transform( + lambda _: sub.col("inner").list_transform(lambda _: outer_parameter + 1) + ) + ) + .to_plan() + ) + assert infer_plan_schema(plan).struct.types[0].list.type.list.type == i64() + + +@pytest.mark.parametrize("fail", [False, True]) +def test_higher_order_lambda_scope_does_not_leak(fail): + from substrait.type_inference import infer_expression_type + + def callback(element): + if fail: + raise ValueError("callback failed") + return element + 1 + + df = _list_df().select(sub.col("arr").list_transform(callback)) + if fail: + with pytest.raises(ValueError, match="callback failed"): + df.to_plan() + else: + df.to_plan() + + reference = stalg.Expression( + selection=stalg.Expression.FieldReference( + lambda_parameter_reference=stalg.Expression.FieldReference.LambdaParameterReference(), + direct_reference=stalg.Expression.ReferenceSegment( + struct_field=stalg.Expression.ReferenceSegment.StructField(field=0) + ), + ) + ) + with pytest.raises(Exception, match="outside an enclosing lambda scope"): + infer_expression_type(reference, struct([i64()]).struct) + + # -- Niche close-out: equi-joins, params, UDT, options, extension --------- diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index c73a684..4651b24 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -820,16 +820,26 @@ 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(): +@pytest.mark.parametrize("steps_out", [0, 1]) +def test_standalone_lambda_parameter_reference_cannot_reach_a_missing_scope(steps_out): + with pytest.raises(Exception, match="outside an enclosing lambda scope"): + infer_expression_type(_lambda_ref(steps_out=steps_out), struct) + + +def test_project_lambda_parameter_reference_requires_a_lambda_scope(): + rel = stalg.Rel( + project=stalg.ProjectRel( + input=stalg.Rel(read=stalg.ReadRel(base_schema=named_struct)), + expressions=[_lambda_ref()], + ) + ) with pytest.raises(Exception, match="outside an enclosing lambda scope"): - infer_expression_type(_lambda_ref(steps_out=1), struct) + infer_rel_schema(rel) @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 ) @@ -841,7 +851,8 @@ def test_lambda_parameter_scope_does_not_leak(fail): else: assert infer_expression_type(expr, supplied).func.return_type == parameter - assert infer_expression_type(_lambda_ref(), supplied) == supplied.types[0] + with pytest.raises(Exception, match="outside an enclosing lambda scope"): + infer_expression_type(_lambda_ref(), supplied) def _expand_switching(types, duplicates): From 33a5a7c01c462984ed75f3c721e911faa99ff0b5 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Tue, 6 Oct 2026 12:27:52 +0300 Subject: [PATCH 3/4] fix(type-inference): preserve nested lambda captures and check invocation arguments --- src/substrait/dataframe/expr.py | 43 ++++++++++++++---------- src/substrait/type_inference.py | 14 ++++++++ tests/dataframe/test_frame.py | 55 ++++++++++++++++++++++++++++++ tests/test_type_inference.py | 59 ++++++++++++++++++++++++--------- 4 files changed, 138 insertions(+), 33 deletions(-) diff --git a/src/substrait/dataframe/expr.py b/src/substrait/dataframe/expr.py index 43b58e3..f39a55b 100644 --- a/src/substrait/dataframe/expr.py +++ b/src/substrait/dataframe/expr.py @@ -66,7 +66,11 @@ from substrait.builders.extended_expression import ( set_predicate as _set_predicate, ) -from substrait.type_inference import infer_extended_expression_schema, lambda_scope +from substrait.type_inference import ( + infer_extended_expression_schema, + lambda_depth, + lambda_scope, +) # Standard Substrait function-extension URNs used by the operators below. FUNCTIONS_COMPARISON = "extension:io.substrait:functions_comparison" @@ -643,30 +647,33 @@ def resolve(base_schema, registry): param_struct = stp.Type.Struct( types=[element_type], nullability=stp.Type.NULLABILITY_REQUIRED ) - param_ref = stalg.Expression( - selection=stalg.Expression.FieldReference( - lambda_parameter_reference=( - stalg.Expression.FieldReference.LambdaParameterReference( - steps_out=0 - ) - ), - direct_reference=stalg.Expression.ReferenceSegment( - struct_field=stalg.Expression.ReferenceSegment.StructField( - field=0 - ) - ), + depth = lambda_depth() + 1 + + def resolve_element(schema, _registry): + param_ref = stalg.Expression( + selection=stalg.Expression.FieldReference( + lambda_parameter_reference=( + stalg.Expression.FieldReference.LambdaParameterReference( + steps_out=lambda_depth() - depth + ) + ), + direct_reference=stalg.Expression.ReferenceSegment( + struct_field=stalg.Expression.ReferenceSegment.StructField( + field=0 + ) + ), + ) ) - ) - element = Expr( - lambda _bs, _r: stee.ExtendedExpression( + return stee.ExtendedExpression( referred_expr=[ stee.ExpressionReference( expression=param_ref, output_names=["element"] ) ], - base_schema=base_schema, + base_schema=schema, ) - ) + + element = Expr(resolve_element) with lambda_scope(param_struct): body = Expr._coerce(callback(element))._unbound(base_schema, registry) lambda_expr = stalg.Expression( diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index bfa6b3e..9e805d3 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -153,6 +153,11 @@ def schema_of(self, anchor, registry) -> stt.Type.Struct: ) +def lambda_depth() -> int: + """Number of lambda scopes enclosing the expression currently being bound.""" + return len(_lambda_schemas.get()) + + @contextlib.contextmanager def lambda_scope(parameters: stt.Type.Struct): """Bind a lambda's parameter struct while building or inferring its body.""" @@ -687,6 +692,15 @@ def infer_expression_type( lam.body, parent_schema, registry=registry, subtrees=subtrees ) if rex_type == "lambda_invocation": + arguments = expression.lambda_invocation.arguments.fields + if len(arguments) != len(lam.parameters.types): + raise Exception( + "lambda invocation needs one argument per lambda parameter" + ) + for argument in arguments: + infer_expression_type( + argument, parent_schema, registry=registry, subtrees=subtrees + ) return body_type return stt.Type( func=stt.Type.Func( diff --git a/tests/dataframe/test_frame.py b/tests/dataframe/test_frame.py index deb54db..a562497 100644 --- a/tests/dataframe/test_frame.py +++ b/tests/dataframe/test_frame.py @@ -1050,6 +1050,61 @@ def test_nested_list_transform_resolves_enclosing_parameter(): assert infer_plan_schema(plan).struct.types[0].list.type.list.type == i64() +@pytest.mark.parametrize("capture_outer", [False, True]) +def test_nested_list_transform_captures_callback_element(capture_outer): + from substrait.type_inference import infer_plan_schema + + plan = ( + sub.read_named_table( + "t", + {"arr": sub.list_(i64(nullable=False)), "inner": sub.list_(fp64())}, + ) + .select( + sub.col("arr").list_transform( + lambda x: sub.col("inner").list_transform( + lambda y: x if capture_outer else y + ) + ) + ) + .to_plan() + ) + outer = getattr( + plan.relations[-1] + .root.input.project.expressions[0] + .scalar_function.arguments[1] + .value, + "lambda", + ) + inner = getattr(outer.body.scalar_function.arguments[1].value, "lambda") + assert inner.body.selection.lambda_parameter_reference.steps_out == int( + capture_outer + ) + expected = i64(nullable=False) if capture_outer else fp64() + assert infer_plan_schema(plan).struct.types[0].list.type.list.type == expected + + +def test_list_transform_captures_parameter_two_scopes_out(): + from substrait.type_inference import infer_plan_schema + + plan = ( + sub.read_named_table( + "t", + {"arr": sub.list_(i64(nullable=False)), "inner": sub.list_(fp64())}, + ) + .select( + sub.col("arr").list_transform( + lambda x: sub.col("inner").list_transform( + lambda _: sub.col("inner").list_transform(lambda _: x) + ) + ) + ) + .to_plan() + ) + assert infer_plan_schema(plan).struct.types[0].list.type.list.type.list.type == i64( + nullable=False + ) + + @pytest.mark.parametrize("fail", [False, True]) def test_higher_order_lambda_scope_does_not_leak(fail): from substrait.type_inference import infer_expression_type diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index 4651b24..9ce43c1 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -715,17 +715,6 @@ def _lambda_ref(field=0, *, steps_out=0): ) -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 @@ -752,13 +741,51 @@ def test_infer_lambda_invocation_literal_body(nullable): ) +@pytest.mark.parametrize("argument_count", [0, 2]) +def test_lambda_invocation_requires_one_argument_per_parameter(argument_count): + expr = _lambda_expression( + [struct.types[0]], + _lambda_ref(), + invoke=True, + arguments=[_field_reference(0)] * argument_count, + ) + with pytest.raises(Exception, match="one argument per lambda parameter"): + infer_expression_type(expr, struct) + + +@pytest.mark.parametrize("steps_out", [0, 1]) +@pytest.mark.parametrize("argument_index", [0, 1]) +def test_lambda_invocation_arguments_cannot_use_the_invoked_lambda_scope( + steps_out, argument_index +): + arguments = [_field_reference(0), _field_reference(0)] + arguments[argument_index] = _lambda_ref(steps_out=steps_out) + expr = _lambda_expression( + [struct.types[0], struct.types[0]], + _lambda_ref(), + invoke=True, + arguments=arguments, + ) + with pytest.raises(Exception, match="outside an enclosing lambda scope"): + infer_expression_type(expr, struct) + + +def test_nested_lambda_invocation_argument_uses_the_callers_scope(): + outer = struct.types[1] + inner = _lambda_expression( + [outer], _lambda_ref(), invoke=True, arguments=[_lambda_ref()] + ) + expr = _lambda_expression([outer], inner) + assert infer_expression_type(expr, struct).func.return_type == outer + + @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() + body = _field_reference(0) if reference == "root" else _lambda_ref() expr = _lambda_expression( [parameter], body, @@ -843,7 +870,7 @@ def test_lambda_parameter_scope_does_not_leak(fail): supplied = stt.Type.Struct( types=[stt.Type(string=stt.Type.String(nullability=_NULL))], nullability=_REQ ) - body = _root_ref(99) if fail else _lambda_ref() + body = _field_reference(99) if fail else _lambda_ref() expr = _lambda_expression([parameter], body) if fail: with pytest.raises(IndexError): @@ -870,7 +897,7 @@ def _expand_switching(types, duplicates): switching_field=stalg.ExpandRel.SwitchingField(duplicates=duplicates) ) ] + [ - stalg.ExpandRel.ExpandField(consistent_field=_root_ref(i)) + stalg.ExpandRel.ExpandField(consistent_field=_field_reference(i)) for i in range(1, len(types)) ] return stalg.Rel(expand=stalg.ExpandRel(input=inp, fields=fields)) @@ -892,7 +919,9 @@ def test_expand_switching_field_with_lambda_invocation(nullable): 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], [_lambda_expression([], _root_ref(1))]) + rel = _expand_switching( + [function, integer], [_lambda_expression([], _field_reference(1))] + ) original = rel.SerializeToString() assert infer_rel_schema(rel).types[0] == function From ac83190d3ab45dc0ac6df5013d8c9336d869bc99 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Tue, 6 Oct 2026 13:07:31 +0300 Subject: [PATCH 4/4] fix(type-inference): require matching lambda invocation argument types --- src/substrait/type_inference.py | 8 +++++-- tests/test_type_inference.py | 37 +++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 9e805d3..8d11cfe 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -697,10 +697,14 @@ def infer_expression_type( raise Exception( "lambda invocation needs one argument per lambda parameter" ) - for argument in arguments: - infer_expression_type( + for parameter_type, argument in zip(lam.parameters.types, arguments): + argument_type = infer_expression_type( argument, parent_schema, registry=registry, subtrees=subtrees ) + if argument_type != parameter_type: + raise Exception( + "lambda invocation argument type must match parameter type" + ) return body_type return stt.Type( func=stt.Type.Func( diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index 9ce43c1..696e988 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -753,6 +753,43 @@ def test_lambda_invocation_requires_one_argument_per_parameter(argument_count): infer_expression_type(expr, struct) +@pytest.mark.parametrize( + "parameter_type, argument_type", + [ + (struct.types[0], struct.types[1]), + (struct.types[0], stt.Type(i64=stt.Type.I64(nullability=_NULL))), + ( + stt.Type(decimal=stt.Type.Decimal(precision=10, scale=2, nullability=_REQ)), + stt.Type(decimal=stt.Type.Decimal(precision=10, scale=3, nullability=_REQ)), + ), + ( + struct.types[0], + stt.Type(i64=stt.Type.I64(nullability=_REQ, type_variation_reference=1)), + ), + ], +) +@pytest.mark.parametrize("argument_index", [0, 1]) +@pytest.mark.parametrize("matches", [False, True]) +def test_lambda_invocation_argument_types_must_match_parameters( + parameter_type, argument_type, argument_index, matches +): + argument_types = [parameter_type, parameter_type] + if not matches: + argument_types[argument_index] = argument_type + expr = _lambda_expression( + [parameter_type, parameter_type], + _lambda_ref(), + invoke=True, + arguments=[_field_reference(0), _field_reference(1)], + ) + row = stt.Type.Struct(types=argument_types, nullability=_REQ) + if matches: + assert infer_expression_type(expr, row) == parameter_type + else: + with pytest.raises(Exception, match="argument type must match parameter type"): + infer_expression_type(expr, row) + + @pytest.mark.parametrize("steps_out", [0, 1]) @pytest.mark.parametrize("argument_index", [0, 1]) def test_lambda_invocation_arguments_cannot_use_the_invoked_lambda_scope(