readd extension_uris to extended expression + add test

Author: benbellick

src/substrait/builders/extended_expression.py

@@ -524,6 +524,12 @@ def resolve(
 
         bound_else = resolve_expression(_else, base_schema, registry)
 
+        extension_uris = merge_extension_uris(
+            *[b[0].extension_uris for b in bound_ifs],
+            *[b[1].extension_uris for b in bound_ifs],
+            bound_else.extension_uris,
+        )
+
         extension_urns = merge_extension_urns(
             *[b[0].extension_urns for b in bound_ifs],
             *[b[1].extension_urns for b in bound_ifs],
@@ -575,6 +581,7 @@ def resolve(
                 )
             ],
             base_schema=base_schema,
+            extension_uris=extension_uris,
             extension_urns=extension_urns,
             extensions=extensions,
         )

tests/builders/extended_expression/test_if_then.py

@@ -1,7 +1,9 @@
 import substrait.gen.proto.algebra_pb2 as stalg
 import substrait.gen.proto.type_pb2 as stt
 import substrait.gen.proto.extended_expression_pb2 as stee
-from substrait.builders.extended_expression import if_then, literal
+import substrait.gen.proto.extensions.extensions_pb2 as ste
+from substrait.builders.extended_expression import if_then, literal, scalar_function, column
+from substrait.extension_registry import ExtensionRegistry
 
 
 struct = stt.Type.Struct(
@@ -75,3 +77,142 @@ def test_if_else():
     )
 
     assert actual == expected
+
+def test_if_then_with_extension():
+    """Test if_then with scalar function to verify both URI and URN are present."""
+    import yaml
+
+    # Use a minimal registry with only the comparison extension
+    registry = ExtensionRegistry(load_default_extensions=False)
+    content = """%YAML 1.2
+---
+urn: extension:io.substrait:functions_comparison
+scalar_functions:
+  - name: "gt"
+    description: ""
+    impls:
+      - args:
+          - value: fp32
+          - value: fp32
+        return: boolean
+"""
+    registry.register_extension_dict(
+        yaml.safe_load(content),
+        uri="https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml"
+    )
+
+    # Create if_then: if order_total > 100 then "expensive" else "cheap"
+    actual = if_then(
+        ifs=[
+            (
+                scalar_function(
+                    "extension:io.substrait:functions_comparison",
+                    "gt",
+                    expressions=[
+                        column("order_total"),
+                        literal(100.0, type=stt.Type(
+                            fp32=stt.Type.FP32(nullability=stt.Type.NULLABILITY_REQUIRED)
+                        )),
+                    ],
+                ),
+                literal(
+                    "expensive",
+                    type=stt.Type(
+                        string=stt.Type.String(nullability=stt.Type.NULLABILITY_REQUIRED)
+                    ),
+                ),
+            )
+        ],
+        _else=literal(
+            "cheap",
+            type=stt.Type(
+                string=stt.Type.String(nullability=stt.Type.NULLABILITY_REQUIRED)
+            ),
+        ),
+    )(named_struct, registry)
+
+    expected = stee.ExtendedExpression(
+        extension_uris=[
+            ste.SimpleExtensionURI(
+                extension_uri_anchor=1,
+                uri="https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml"
+            )
+        ],
+        extension_urns=[
+            ste.SimpleExtensionURN(
+                extension_urn_anchor=1,
+                urn="extension:io.substrait:functions_comparison"
+            )
+        ],
+        extensions=[
+            ste.SimpleExtensionDeclaration(
+                extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction(
+                    extension_uri_reference=1,
+                    extension_urn_reference=1,
+                    function_anchor=1,
+                    name="gt:fp32_fp32"
+                )
+            )
+        ],
+        referred_expr=[
+            stee.ExpressionReference(
+                expression=stalg.Expression(
+                    if_then=stalg.Expression.IfThen(
+                        **{
+                            "ifs": [
+                                stalg.Expression.IfThen.IfClause(
+                                    **{
+                                        "if": stalg.Expression(
+                                            scalar_function=stalg.Expression.ScalarFunction(
+                                                function_reference=1,
+                                                output_type=stt.Type(
+                                                    bool=stt.Type.Boolean(
+                                                        nullability=stt.Type.NULLABILITY_NULLABLE
+                                                    )
+                                                ),
+                                                arguments=[
+                                                    stalg.FunctionArgument(
+                                                        value=stalg.Expression(
+                                                            selection=stalg.Expression.FieldReference(
+                                                                direct_reference=stalg.Expression.ReferenceSegment(
+                                                                    struct_field=stalg.Expression.ReferenceSegment.StructField(
+                                                                        field=2
+                                                                    )
+                                                                ),
+                                                                root_reference=stalg.Expression.FieldReference.RootReference()
+                                                            )
+                                                        )
+                                                    ),
+                                                    stalg.FunctionArgument(
+                                                        value=stalg.Expression(
+                                                            literal=stalg.Expression.Literal(
+                                                                fp32=100.0
+                                                            )
+                                                        )
+                                                    ),
+                                                ]
+                                            )
+                                        ),
+                                        "then": stalg.Expression(
+                                            literal=stalg.Expression.Literal(
+                                                string="expensive"
+                                            )
+                                        ),
+                                    }
+                                )
+                            ],
+                            "else": stalg.Expression(
+                                literal=stalg.Expression.Literal(
+                                    string="cheap"
+                                )
+                            ),
+                        }
+                    )
+                ),
+                output_names=["IfThen(gt(order_total,Literal(100.0)),Literal(expensive),Literal(cheap))"],
+            )
+        ],
+        base_schema=named_struct,
+    )
+
+    assert actual == expected