From d07a4417af48da7935da4764df72e7168d143911 Mon Sep 17 00:00:00 2001 From: MBWhite Date: Thu, 16 Oct 2025 15:20:44 +0100 Subject: [PATCH] feat: advanced_extensions Add support to the builders for suppling an (optional) advanced extension. Signed-off-by: MBWhite --- src/substrait/builders/plan.py | 32 ++++++++++++++++++++++++++----- tests/builders/plan/test_read.py | 33 ++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 5 deletions(-) diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index ed487c7a..9f852eb3 100644 --- a/src/substrait/builders/plan.py +++ b/src/substrait/builders/plan.py @@ -5,9 +5,10 @@ See `examples/builder_example.py` for usage. """ -from typing import Iterable, Union, Callable +from typing import Iterable, Optional, Union, Callable import substrait.gen.proto.algebra_pb2 as stalg +from substrait.gen.proto.extensions.extensions_pb2 import AdvancedExtension import substrait.gen.proto.plan_pb2 as stp import substrait.gen.proto.type_pb2 as stt import substrait.gen.proto.extended_expression_pb2 as stee @@ -32,7 +33,9 @@ def _merge_extensions(*objs): def read_named_table( - names: Union[str, Iterable[str]], named_struct: stt.NamedStruct + names: Union[str, Iterable[str]], + named_struct: stt.NamedStruct, + extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: if named_struct.struct.nullability is stt.Type.NULLABILITY_NULLABLE: raise Exception("NamedStruct must not contain a nullable struct") @@ -47,6 +50,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: common=stalg.RelCommon(direct=stalg.RelCommon.Direct()), base_schema=named_struct, named_table=stalg.ReadRel.NamedTable(names=_names), + advanced_extension=extension, ) ) @@ -60,7 +64,9 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: def project( - plan: PlanOrUnbound, expressions: Iterable[ExtendedExpressionOrUnbound] + plan: PlanOrUnbound, + expressions: Iterable[ExtendedExpressionOrUnbound], + extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: _plan = plan if isinstance(plan, stp.Plan) else plan(registry) @@ -86,6 +92,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: expressions=[ e.expression for ee in bound_expressions for e in ee.referred_expr ], + advanced_extension=extension, ) ) @@ -97,7 +104,11 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: return resolve -def filter(plan: PlanOrUnbound, expression: ExtendedExpressionOrUnbound) -> UnboundPlan: +def filter( + plan: PlanOrUnbound, + expression: ExtendedExpressionOrUnbound, + extension: Optional[AdvancedExtension] = None, +) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry) ns = infer_plan_schema(bound_plan) @@ -109,6 +120,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: filter=stalg.FilterRel( input=bound_plan.relations[-1].root.input, condition=bound_expression.referred_expr[0].expression, + advanced_extension=extension, ) ) @@ -130,6 +142,7 @@ def sort( tuple[ExtendedExpressionOrUnbound, stalg.SortField.SortDirection.ValueType], ] ], + extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry) @@ -155,7 +168,8 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: ) for e in bound_expressions ], - ) + advanced_extension=extension, + ), ) return stp.Plan( @@ -193,6 +207,7 @@ def fetch( plan: PlanOrUnbound, offset: ExtendedExpressionOrUnbound, count: ExtendedExpressionOrUnbound, + extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry) @@ -208,6 +223,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: if bound_offset else None, count_expr=bound_count.referred_expr[0].expression, + advanced_extension=extension, ) ) @@ -230,6 +246,7 @@ def join( right: PlanOrUnbound, expression: ExtendedExpressionOrUnbound, type: stalg.JoinRel.JoinType, + extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_left = left if isinstance(left, stp.Plan) else left(registry) @@ -254,6 +271,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: right=bound_right.relations[-1].root.input, expression=bound_expression.referred_expr[0].expression, type=type, + advanced_extension=extension, ) ) @@ -268,6 +286,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: def cross( left: PlanOrUnbound, right: PlanOrUnbound, + extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_left = left if isinstance(left, stp.Plan) else left(registry) @@ -287,6 +306,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: cross=stalg.CrossRel( left=bound_left.relations[-1].root.input, right=bound_right.relations[-1].root.input, + advanced_extension=extension, ) ) @@ -303,6 +323,7 @@ def aggregate( input: PlanOrUnbound, grouping_expressions: Iterable[ExtendedExpressionOrUnbound], measures: Iterable[ExtendedExpressionOrUnbound], + extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_input = input if isinstance(input, stp.Plan) else input(registry) @@ -332,6 +353,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: stalg.AggregateRel.Measure(measure=m.referred_expr[0].measure) for m in bound_measures ], + advanced_extension=extension, ) ) diff --git a/tests/builders/plan/test_read.py b/tests/builders/plan/test_read.py index 909216ad..20cc0c49 100644 --- a/tests/builders/plan/test_read.py +++ b/tests/builders/plan/test_read.py @@ -4,6 +4,9 @@ from substrait.builders.type import boolean, i64 from substrait.builders.plan import read_named_table import pytest +from substrait.gen.proto.extensions.extensions_pb2 import AdvancedExtension +from google.protobuf import any +from google.protobuf.wrappers_pb2 import StringValue struct = stt.Type.Struct( types=[i64(nullable=False), boolean()], @@ -74,3 +77,33 @@ def test_read_rel_schema_nullable(): Exception, match=r"NamedStruct must not contain a nullable struct" ): read_named_table("example_table", named_struct)(None) + + +def test_read_rel_ae(): + extension = AdvancedExtension(optimization=[any.pack(StringValue(value="Opt1"))]) + + actual = read_named_table(["example_db", "example_table"], named_struct, extension)( + None + ) + + expected = stp.Plan( + relations=[ + stp.PlanRel( + root=stalg.RelRoot( + input=stalg.Rel( + read=stalg.ReadRel( + common=stalg.RelCommon(direct=stalg.RelCommon.Direct()), + base_schema=named_struct, + named_table=stalg.ReadRel.NamedTable( + names=["example_db", "example_table"] + ), + advanced_extension=extension, + ) + ), + names=["id", "is_applicable"], + ) + ) + ] + ) + + assert actual == expected