feat: add Plan.ExecutionBehavior support

src/substrait/builders/plan.py

@@ -45,12 +45,51 @@ def _create_default_version():
 _create_default_version()
 
 
-def _merge_extensions(*objs):
-    """Merge extension URNs and declarations from multiple plan/expression objects."""
-    return {
+def _merge_plan_metadata(*objs):
+    """Collect the plan-level metadata a builder carries over from its inputs.
+
+    ``objs`` is a mix of input Plans and bound ExtendedExpressions. Extension
+    URNs and declarations are merged from all of them; the plan-level execution
+    behavior is carried over from the first input Plan that declares one
+    (expressions have no such field). Because every relational builder routes
+    its inputs through here, an execution behavior set anywhere upstream is
+    preserved on the freshly-constructed output Plan -- so it is order
+    independent across a pipeline rather than only surviving as the last step.
+    """
+    metadata = {
         "extension_urns": merge_extension_urns(*[b.extension_urns for b in objs if b]),
         "extensions": merge_extension_declarations(*[b.extensions for b in objs if b]),
     }
+    for b in objs:
+        if isinstance(b, stp.Plan) and b.HasField("execution_behavior"):
+            metadata["execution_behavior"] = b.execution_behavior
+            break
+    return metadata
+
+
+def with_execution_behavior(
+    plan: PlanOrUnbound,
+    variable_eval_mode: "stp.ExecutionBehavior.VariableEvaluationMode.ValueType",
+) -> UnboundPlan:
+    """Return a copy of ``plan`` with its plan-level execution behavior set.
+
+    ``variable_eval_mode`` controls how often execution context variables (such
+    as those from ``extended_expression.execution_context_variable``) are
+    evaluated: once per plan (``VARIABLE_EVALUATION_MODE_PER_PLAN``) or once per
+    record (``VARIABLE_EVALUATION_MODE_PER_RECORD``). The setting is carried over
+    by subsequent builders (see ``_merge_plan_metadata``), so it may be applied
+    at any point in a pipeline rather than only as the final step.
+    """
+
+    def resolve(registry: ExtensionRegistry) -> stp.Plan:
+        bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry)
+
+        result = stp.Plan()
+        result.CopyFrom(bound_plan)
+        result.execution_behavior.variable_eval_mode = variable_eval_mode
+        return result
+
+    return resolve
 
 
 def read_named_table(
@@ -225,7 +264,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_extensions(_plan, *bound_expressions),
+            **_merge_plan_metadata(_plan, *bound_expressions),
         )
 
     return resolve
@@ -281,7 +320,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_extensions(_plan, *bound_expressions),
+            **_merge_plan_metadata(_plan, *bound_expressions),
         )
 
     return resolve
@@ -312,7 +351,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_extensions(bound_plan, bound_expression),
+            **_merge_plan_metadata(bound_plan, bound_expression),
         )
 
     return resolve
@@ -359,7 +398,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=ns.names))],
-            **_merge_extensions(bound_plan, *[e[0] for e in bound_expressions]),
+            **_merge_plan_metadata(bound_plan, *[e[0] for e in bound_expressions]),
         )
 
     return resolve
@@ -383,7 +422,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
                     )
                 )
             ],
-            **_merge_extensions(*bound_inputs),
+            **_merge_plan_metadata(*bound_inputs),
         )
 
     return resolve
@@ -427,7 +466,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
                     )
                 )
             ],
-            **_merge_extensions(bound_plan, bound_offset, bound_count),
+            **_merge_plan_metadata(bound_plan, bound_offset, bound_count),
         )
 
     return resolve
@@ -486,7 +525,9 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=out_names))],
-            **_merge_extensions(bound_left, bound_right, bound_expression, bound_post),
+            **_merge_plan_metadata(
+                bound_left, bound_right, bound_expression, bound_post
+            ),
         )
 
     return resolve
@@ -522,7 +563,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=ns.names))],
-            **_merge_extensions(bound_left, bound_right),
+            **_merge_plan_metadata(bound_left, bound_right),
         )
 
     return resolve
@@ -598,7 +639,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_extensions(
+            **_merge_plan_metadata(
                 bound_input,
                 *bound_grouping_expressions,
                 *bound_measures,
@@ -639,7 +680,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             relations=[
                 stp.PlanRel(root=stalg.RelRoot(input=write_rel, names=ns.names))
             ],
-            **_merge_extensions(bound_input),
+            **_merge_plan_metadata(bound_input),
         )
 
     return resolve
@@ -689,7 +730,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=ddl_rel, names=out_names))],
-            **_merge_extensions(*merge_sources),
+            **_merge_plan_metadata(*merge_sources),
         )
 
     return resolve
@@ -744,7 +785,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
                     )
                 )
             ],
-            **_merge_extensions(*merge_sources),
+            **_merge_plan_metadata(*merge_sources),
         )
 
     return resolve
@@ -833,7 +874,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_extensions(
+            **_merge_plan_metadata(
                 bound_plan,
                 *bound_partitions,
                 *[e[0] for e in bound_sorts],
@@ -892,7 +933,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=list(names)))],
-            **_merge_extensions(*merge_sources),
+            **_merge_plan_metadata(*merge_sources),
         )
 
     return resolve
@@ -939,7 +980,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=out_names))],
-            **_merge_extensions(bound_left, bound_right, bound_expression),
+            **_merge_plan_metadata(bound_left, bound_right, bound_expression),
         )
 
     return resolve
@@ -1009,7 +1050,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             return stp.Plan(
                 version=default_version,
                 relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-                **_merge_extensions(bound_left, bound_right),
+                **_merge_plan_metadata(bound_left, bound_right),
             )
 
         return resolve
@@ -1073,7 +1114,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_extensions(bound_plan),
+            **_merge_plan_metadata(bound_plan),
         )
 
     return resolve
@@ -1097,7 +1138,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_extensions(*bound_inputs),
+            **_merge_plan_metadata(*bound_inputs),
         )
 
     return resolve
@@ -1137,7 +1178,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
                     )
                 )
             ],
-            **_merge_extensions(bound_plan),
+            **_merge_plan_metadata(bound_plan),
         )
 
     return resolve
@@ -1198,7 +1239,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
                     )
                 )
             ],
-            **_merge_extensions(
+            **_merge_plan_metadata(
                 bound_plan,
                 *[s for s, _ in bound_sorts],
                 bound_count,

src/substrait/dataframe/frame.py

@@ -32,6 +32,7 @@
 from typing import Any, Iterable, Optional, Union
 
 import substrait.algebra_pb2 as stalg
+import substrait.plan_pb2 as stplan
 import substrait.type_pb2 as stp
 
 from substrait.builders import plan as _plan
@@ -82,6 +83,12 @@
     (True, True): stalg.SortField.SORT_DIRECTION_DESC_NULLS_LAST,
 }
 
+# How often execution context variables (current_timestamp, ...) are evaluated.
+_VARIABLE_EVAL_MODES = {
+    "per_plan": stplan.ExecutionBehavior.VARIABLE_EVALUATION_MODE_PER_PLAN,
+    "per_record": stplan.ExecutionBehavior.VARIABLE_EVALUATION_MODE_PER_RECORD,
+}
+
 
 def _per_column(value: Any, n: int, name: str) -> "list[bool]":
     """Broadcast a bool, or validate a per-column list of bools, to length ``n``."""
@@ -516,6 +523,24 @@ def resolve(registry: ExtensionRegistry):
 
         return self._next(resolve)
 
+    def with_execution_behavior(self, mode: str) -> "DataFrame":
+        """Set how often execution context variables are evaluated (plan-level).
+
+        ``mode`` is ``"per_plan"`` (evaluate once for the whole plan) or
+        ``"per_record"`` (evaluate once per record), controlling variables such
+        as :func:`substrait.dataframe.current_timestamp`. The setting is carried
+        across subsequent operations, so it may be applied at any point in the
+        chain.
+        """
+        try:
+            eval_mode = _VARIABLE_EVAL_MODES[mode]
+        except KeyError:
+            raise ValueError(
+                f"unknown execution behavior mode {mode!r}; "
+                f"expected one of {sorted(_VARIABLE_EVAL_MODES)}"
+            ) from None
+        return self._next(_plan.with_execution_behavior(self._plan, eval_mode))
+
     def union(self, *others: "DataFrame", distinct: bool = False) -> "DataFrame":
         """Concatenate rows of this DataFrame with ``others``.
 

tests/builders/plan/test_with_execution_behavior.py

@@ -0,0 +1,92 @@
+import substrait.algebra_pb2 as stalg
+import substrait.plan_pb2 as stp
+import substrait.type_pb2 as stt
+
+from substrait.builders.plan import (
+    default_version,
+    exchange,
+    read_named_table,
+    set,
+    with_execution_behavior,
+)
+from substrait.builders.type import boolean, i64
+
+struct = stt.Type.Struct(
+    types=[i64(nullable=False), boolean()],
+    nullability=stt.Type.Nullability.NULLABILITY_REQUIRED,
+)
+
+named_struct = stt.NamedStruct(names=["id", "is_applicable"], struct=struct)
+
+PER_PLAN = stp.ExecutionBehavior.VARIABLE_EVALUATION_MODE_PER_PLAN
+PER_RECORD = stp.ExecutionBehavior.VARIABLE_EVALUATION_MODE_PER_RECORD
+
+
+def test_with_execution_behavior_per_record():
+    base = read_named_table("example_table", named_struct)
+
+    actual = with_execution_behavior(base, PER_RECORD)(None)
+
+    expected = read_named_table("example_table", named_struct)(None)
+    expected.execution_behavior.CopyFrom(
+        stp.ExecutionBehavior(variable_eval_mode=PER_RECORD)
+    )
+
+    assert actual == expected
+    assert actual.version == default_version
+
+
+def test_with_execution_behavior_per_plan_preserves_relations():
+    base = read_named_table("example_table", named_struct)
+
+    actual = with_execution_behavior(base, PER_PLAN)(None)
+
+    assert actual.execution_behavior.variable_eval_mode == PER_PLAN
+    # the underlying relation is left untouched
+    assert actual.relations == base(None).relations
+
+
+def test_with_execution_behavior_accepts_bound_plan_without_mutating_it():
+    bound = read_named_table("example_table", named_struct)(None)
+
+    actual = with_execution_behavior(bound, PER_PLAN)(None)
+
+    assert actual.execution_behavior.variable_eval_mode == PER_PLAN
+    # the caller's input plan is not mutated in place
+    assert not bound.HasField("execution_behavior")
+
+
+def test_execution_behavior_survives_subsequent_builder():
+    # Applying a builder after the setting must not drop it: the fresh output
+    # Plan carries the execution behavior over from its input.
+    base = with_execution_behavior(read_named_table("t", named_struct), PER_RECORD)
+
+    actual = exchange(base, partition_count=2)(None)
+
+    assert actual.execution_behavior.variable_eval_mode == PER_RECORD
+
+
+def test_with_execution_behavior_overwrites_existing():
+    # Applying it again replaces the previous mode (last wins), rather than
+    # merging or keeping the earlier value.
+    base = with_execution_behavior(read_named_table("t", named_struct), PER_PLAN)
+
+    actual = with_execution_behavior(base, PER_RECORD)(None)
+
+    assert actual.execution_behavior.variable_eval_mode == PER_RECORD
+
+
+def test_execution_behavior_taken_from_first_input():
+    # For multi-input builders the first input that declares one wins.
+    left = with_execution_behavior(read_named_table("l", named_struct), PER_RECORD)
+    right = with_execution_behavior(read_named_table("r", named_struct), PER_PLAN)
+
+    actual = set([left, right], stalg.SetRel.SET_OP_UNION_ALL)(None)
+
+    assert actual.execution_behavior.variable_eval_mode == PER_RECORD
+
+
+def test_plans_have_no_execution_behavior_by_default():
+    plan = read_named_table("example_table", named_struct)(None)
+
+    assert not plan.HasField("execution_behavior")