fix(type-inference)!: resolve lambda scopes and invocation types

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
+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,32 +647,35 @@ 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=(
-                        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=param_ns,
+                    base_schema=schema,
                 )
-            )
-            body = Expr._coerce(callback(element))._unbound(param_ns, registry)
+
+            element = Expr(resolve_element)
+            with lambda_scope(param_struct):
+                body = Expr._coerce(callback(element))._unbound(base_schema, registry)
             lambda_expr = stalg.Expression(
                 **{
                     "lambda": stalg.Expression.Lambda(

src/substrait/type_inference.py

@@ -146,6 +146,28 @@ 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=()
+)
+
+
+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."""
+    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
@@ -606,8 +628,16 @@ 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 steps >= len(stack):
+                raise Exception(
+                    "lambda parameter reference outside an enclosing lambda scope"
+                )
+            schema = stack[len(stack) - steps - 1]
         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 +684,28 @@ 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<param_types -> 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")
+        with lambda_scope(lam.parameters):
+            body_type = 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 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(
                 parameter_types=list(lam.parameters.types),

tests/dataframe/test_frame.py

@@ -997,6 +997,142 @@ 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("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
+
+    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 ---------