From bbb1521e4b2f09784a50750bd2869e15d6a20dda Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Fri, 17 Jul 2026 10:55:05 +0200 Subject: [PATCH 1/2] feat: add Plan.ExecutionBehavior support Exposes the plan-level execution behavior introduced in Substrait 0.87.0 (Plan.execution_behavior.variable_eval_mode), which governs how often execution context variables (current_timestamp / current_date / current_timezone) are evaluated -- once per plan or once per record. - builders.plan.with_execution_behavior(plan, variable_eval_mode) sets it - execution behavior is now carried over by every relational builder (folded into the renamed _merge_plan_metadata helper), so it survives subsequent operations and is order independent across a pipeline - DataFrame.with_execution_behavior("per_plan" | "per_record") wires it into the native fluent frame Resolves #185. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/substrait/builders/plan.py | 87 +++++++++++++----- src/substrait/dataframe/frame.py | 25 +++++ .../plan/test_with_execution_behavior.py | 92 +++++++++++++++++++ tests/dataframe/test_frame.py | 77 ++++++++++++++++ 4 files changed, 258 insertions(+), 23 deletions(-) create mode 100644 tests/builders/plan/test_with_execution_behavior.py diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index 1e510546..8885ca72 100644 --- a/src/substrait/builders/plan.py +++ b/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, diff --git a/src/substrait/dataframe/frame.py b/src/substrait/dataframe/frame.py index 8b0bc244..e69fde16 100644 --- a/src/substrait/dataframe/frame.py +++ b/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``. diff --git a/tests/builders/plan/test_with_execution_behavior.py b/tests/builders/plan/test_with_execution_behavior.py new file mode 100644 index 00000000..97b254ad --- /dev/null +++ b/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") diff --git a/tests/dataframe/test_frame.py b/tests/dataframe/test_frame.py index 86381d64..ee77a994 100644 --- a/tests/dataframe/test_frame.py +++ b/tests/dataframe/test_frame.py @@ -6,6 +6,7 @@ import pytest import substrait.algebra_pb2 as stalg +import substrait.plan_pb2 as stp import substrait.dataframe as sub from substrait.builders.extended_expression import ( @@ -1345,3 +1346,79 @@ def test_drop_unknown_column_raises(): def test_drop_all_columns_raises(): with pytest.raises(ValueError, match="every column"): people_df().drop("id", "age", "name").to_plan() + + +# -- Execution behavior (Plan.execution_behavior) ------------------------- + + +@pytest.mark.parametrize( + "mode, expected", + [ + ("per_plan", stp.ExecutionBehavior.VARIABLE_EVALUATION_MODE_PER_PLAN), + ("per_record", stp.ExecutionBehavior.VARIABLE_EVALUATION_MODE_PER_RECORD), + ], +) +def test_with_execution_behavior(mode, expected): + plan = ( + sub.read_named_table("t", {"id": sub.i64}) + .with_execution_behavior(mode) + .to_plan() + ) + assert plan.execution_behavior.variable_eval_mode == expected + + +def test_execution_behavior_preserved_through_later_ops(): + # Set it first, then keep building -- the setting must survive downstream ops. + plan = ( + sub.read_named_table("t", {"id": sub.i64, "age": sub.i64}) + .with_execution_behavior("per_record") + .filter(sub.col("age") > 25) + .select("id") + .to_plan() + ) + assert ( + plan.execution_behavior.variable_eval_mode + == stp.ExecutionBehavior.VARIABLE_EVALUATION_MODE_PER_RECORD + ) + + +def test_execution_behavior_set_mid_chain(): + # Applying it in the middle of a chain is equivalent -- it is order independent. + plan = ( + sub.read_named_table("t", {"id": sub.i64, "age": sub.i64}) + .filter(sub.col("age") > 25) + .with_execution_behavior("per_plan") + .select("id") + .to_plan() + ) + assert ( + plan.execution_behavior.variable_eval_mode + == stp.ExecutionBehavior.VARIABLE_EVALUATION_MODE_PER_PLAN + ) + + +def test_execution_behavior_unset_by_default(): + plan = sub.read_named_table("t", {"id": sub.i64}).to_plan() + assert not plan.HasField("execution_behavior") + + +def test_with_execution_behavior_invalid_mode(): + with pytest.raises(ValueError, match="unknown execution behavior mode 'nonsense'"): + sub.read_named_table("t", {"id": sub.i64}).with_execution_behavior("nonsense") + + +def test_execution_behavior_with_context_variable_end_to_end(): + # The common pairing: a per-record execution behavior alongside a + # current_timestamp execution context variable. + plan = ( + sub.read_named_table("t", {"id": sub.i64}) + .with_columns(now=sub.current_timestamp()) + .with_execution_behavior("per_record") + .to_plan() + ) + assert ( + plan.execution_behavior.variable_eval_mode + == stp.ExecutionBehavior.VARIABLE_EVALUATION_MODE_PER_RECORD + ) + projected = plan.relations[-1].root.input.project.expressions[0] + assert projected.WhichOneof("rex_type") == "execution_context_variable" From ae055ca2c2a9f1cd1d9e2495b085b07f6ac96f7f Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Mon, 20 Jul 2026 08:24:18 +0200 Subject: [PATCH 2/2] refactor: rename with_execution_behavior param to variable_eval_mode Addresses review feedback: 'mode' was too generic for the DataFrame method. Renaming to variable_eval_mode also matches the builders.plan parameter name. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/substrait/dataframe/frame.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/substrait/dataframe/frame.py b/src/substrait/dataframe/frame.py index e69fde16..2b7cd6ff 100644 --- a/src/substrait/dataframe/frame.py +++ b/src/substrait/dataframe/frame.py @@ -523,20 +523,20 @@ def resolve(registry: ExtensionRegistry): return self._next(resolve) - def with_execution_behavior(self, mode: str) -> "DataFrame": + def with_execution_behavior(self, variable_eval_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. + ``variable_eval_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] + eval_mode = _VARIABLE_EVAL_MODES[variable_eval_mode] except KeyError: raise ValueError( - f"unknown execution behavior mode {mode!r}; " + f"unknown execution behavior mode {variable_eval_mode!r}; " f"expected one of {sorted(_VARIABLE_EVAL_MODES)}" ) from None return self._next(_plan.with_execution_behavior(self._plan, eval_mode))