feat: allow passing multiple functions to function builders

Author: tokoko

src/substrait/builders/extended_expression.py

@@ -204,12 +204,12 @@ def resolve(
 
 
 def scalar_function(
-    uri: str,
-    function: str,
+    function: Union[Iterable[str], str],
     expressions: Iterable[ExtendedExpressionOrUnbound],
     alias: Union[Iterable[str], str] = None,
 ):
     """Builds a resolver for ExtendedExpression containing a ScalarFunction expression"""
+    functions = [function] if isinstance(function, str) else function
 
     def resolve(
         base_schema: stp.NamedStruct, registry: ExtensionRegistry
@@ -224,23 +224,30 @@ def resolve(
 
         signature = [typ for es in expression_schemas for typ in es.types]
 
-        func = registry.lookup_function(uri, function, signature)
+        for f in functions:
+            uri, name = f.split(":")
+            func = registry.lookup_function(uri, name, signature)
+            if func:
+                break
 
         if not func:
             raise Exception(f"Unknown function {function} for {signature}")
 
+        resolved_func, return_type = func
+
         func_extension_uris = [
             ste.SimpleExtensionURI(
-                extension_uri_anchor=registry.lookup_uri(uri), uri=uri
+                extension_uri_anchor=registry.lookup_uri(resolved_func.uri),
+                uri=resolved_func.uri,
             )
         ]
 
         func_extensions = [
             ste.SimpleExtensionDeclaration(
                 extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction(
-                    extension_uri_reference=registry.lookup_uri(uri),
-                    function_anchor=func[0].anchor,
-                    name=str(func[0]),
+                    extension_uri_reference=registry.lookup_uri(resolved_func.uri),
+                    function_anchor=resolved_func.anchor,
+                    name=str(resolved_func),
                 )
             )
         ]
@@ -258,14 +265,14 @@ def resolve(
                 stee.ExpressionReference(
                     expression=stalg.Expression(
                         scalar_function=stalg.Expression.ScalarFunction(
-                            function_reference=func[0].anchor,
+                            function_reference=resolved_func.anchor,
                             arguments=[
                                 stalg.FunctionArgument(
                                     value=e.referred_expr[0].expression
                                 )
                                 for e in bound_expressions
                             ],
-                            output_type=func[1],
+                            output_type=return_type,
                         )
                     ),
                     output_names=_alias_or_inferred(
@@ -284,12 +291,12 @@ def resolve(
 
 
 def aggregate_function(
-    uri: str,
-    function: str,
+    function: Union[Iterable[str], str],
     expressions: Iterable[ExtendedExpressionOrUnbound],
     alias: Union[Iterable[str], str] = None,
 ):
     """Builds a resolver for ExtendedExpression containing a AggregateFunction measure"""
+    functions = [function] if isinstance(function, str) else function
 
     def resolve(
         base_schema: stp.NamedStruct, registry: ExtensionRegistry
@@ -304,23 +311,30 @@ def resolve(
 
         signature = [typ for es in expression_schemas for typ in es.types]
 
-        func = registry.lookup_function(uri, function, signature)
+        for f in functions:
+            uri, name = f.split(":")
+            func = registry.lookup_function(uri, name, signature)
+            if func:
+                break
 
         if not func:
             raise Exception(f"Unknown function {function} for {signature}")
 
+        resolved_func, return_type = func
+
         func_extension_uris = [
             ste.SimpleExtensionURI(
-                extension_uri_anchor=registry.lookup_uri(uri), uri=uri
+                extension_uri_anchor=registry.lookup_uri(resolved_func.uri),
+                uri=resolved_func.uri,
             )
         ]
 
         func_extensions = [
             ste.SimpleExtensionDeclaration(
                 extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction(
-                    extension_uri_reference=registry.lookup_uri(uri),
-                    function_anchor=func[0].anchor,
-                    name=str(func[0]),
+                    extension_uri_reference=registry.lookup_uri(resolved_func.uri),
+                    function_anchor=resolved_func.anchor,
+                    name=str(resolved_func),
                 )
             )
         ]
@@ -342,7 +356,7 @@ def resolve(
                             stalg.FunctionArgument(value=e.referred_expr[0].expression)
                             for e in bound_expressions
                         ],
-                        output_type=func[1],
+                        output_type=return_type,
                     ),
                     output_names=_alias_or_inferred(
                         alias,
@@ -361,13 +375,13 @@ def resolve(
 
 # TODO bounds, sorts
 def window_function(
-    uri: str,
-    function: str,
+    function: Union[Iterable[str], str],
     expressions: Iterable[ExtendedExpressionOrUnbound],
     partitions: Iterable[ExtendedExpressionOrUnbound] = [],
     alias: Union[Iterable[str], str] = None,
 ):
     """Builds a resolver for ExtendedExpression containing a WindowFunction expression"""
+    functions = [function] if isinstance(function, str) else function
 
     def resolve(
         base_schema: stp.NamedStruct, registry: ExtensionRegistry
@@ -386,23 +400,30 @@ def resolve(
 
         signature = [typ for es in expression_schemas for typ in es.types]
 
-        func = registry.lookup_function(uri, function, signature)
+        for f in functions:
+            uri, name = f.split(":")
+            func = registry.lookup_function(uri, name, signature)
+            if func:
+                break
 
         if not func:
             raise Exception(f"Unknown function {function} for {signature}")
 
+        resolved_func, return_type = func
+
         func_extension_uris = [
             ste.SimpleExtensionURI(
-                extension_uri_anchor=registry.lookup_uri(uri), uri=uri
+                extension_uri_anchor=registry.lookup_uri(resolved_func.uri),
+                uri=resolved_func.uri,
             )
         ]
 
         func_extensions = [
             ste.SimpleExtensionDeclaration(
                 extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction(
-                    extension_uri_reference=registry.lookup_uri(uri),
-                    function_anchor=func[0].anchor,
-                    name=str(func[0]),
+                    extension_uri_reference=registry.lookup_uri(resolved_func.uri),
+                    function_anchor=resolved_func.anchor,
+                    name=str(resolved_func),
                 )
             )
         ]
@@ -431,7 +452,7 @@ def resolve(
                                 )
                                 for e in bound_expressions
                             ],
-                            output_type=func[1],
+                            output_type=return_type,
                             partitions=[
                                 e.referred_expr[0].expression for e in bound_partitions
                             ],

src/substrait/sql/sql_to_substrait.py

@@ -28,18 +28,21 @@
 SchemaResolver = Callable[[str], stt.NamedStruct]
 
 function_mapping = {
-    "Plus": ("functions_arithmetic.yaml", "add"),
-    "Minus": ("functions_arithmetic.yaml", "subtract"),
-    "Gt": ("functions_comparison.yaml", "gt"),
-    "GtEq": ("functions_comparison.yaml", "gte"),
-    "Lt": ("functions_comparison.yaml", "lt"),
-    "Eq": ("functions_comparison.yaml", "equal"),
+    "Plus": ["functions_arithmetic.yaml:add", "functions_arithmetic_decimal.yaml:add"],
+    "Minus": [
+        "functions_arithmetic.yaml:subtract",
+        "functions_arithmetic_decimal.yaml:subtract",
+    ],
+    "Gt": ["functions_comparison.yaml:gt"],
+    "GtEq": ["functions_comparison.yaml:gte"],
+    "Lt": ["functions_comparison.yaml:lt"],
+    "Eq": ["functions_comparison.yaml:equal"],
 }
 
-aggregate_function_mapping = {"SUM": ("functions_arithmetic.yaml", "sum")}
+aggregate_function_mapping = {"SUM": ["functions_arithmetic.yaml:sum"]}
 
 window_function_mapping = {
-    "row_number": ("functions_arithmetic.yaml", "row_number"),
+    "row_number": ["functions_arithmetic.yaml:row_number"],
 }
 
 
@@ -105,7 +108,7 @@ def translate_expression(
             ),
         ]
         func = function_mapping[ast["op"]]
-        return scalar_function(func[0], func[1], expressions=expressions, alias=alias)
+        return scalar_function(func, expressions=expressions, alias=alias)
     elif op == "Value":
         return literal(
             int(ast["value"]["Number"][0]), stt.Type(i64=stt.Type.I64()), alias=alias
@@ -138,7 +141,7 @@ def translate_expression(
             random_name = "".join(
                 random.choices(string.ascii_uppercase + string.digits, k=5)
             )  # TODO make this deterministic
-            aggr = aggregate_function(func[0], func[1], expressions, alias=random_name)
+            aggr = aggregate_function(func, expressions, alias=random_name)
             measures.append((aggr, ast, random_name))
             return column(random_name, alias=alias)
         elif name in window_function_mapping:
@@ -156,7 +159,7 @@ def translate_expression(
             ]
 
             return window_function(
-                func[0], func[1], expressions, partitions=partitions, alias=alias
+                func, expressions, partitions=partitions, alias=alias
             )
 
         else:

tests/builders/extended_expression/test_aggregate_function.py

@@ -41,8 +41,7 @@
 
 def test_aggregate_count():
     e = aggregate_function(
-        "test_uri",
-        "count",
+        "test_uri:count",
         expressions=[
             literal(
                 10,

tests/builders/extended_expression/test_scalar_function.py

@@ -45,8 +45,7 @@
 
 def test_sclar_add():
     e = scalar_function(
-        "test_uri",
-        "test_func",
+        "test_uri:test_func",
         expressions=[
             literal(
                 10,
@@ -68,7 +67,7 @@ def test_sclar_add():
         extensions=[
             ste.SimpleExtensionDeclaration(
                 extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction(
-                    extension_uri_reference=1, function_anchor=1, name="test_func:i8"   
+                    extension_uri_reference=1, function_anchor=1, name="test_func:i8"
                 )
             )
         ],
@@ -98,23 +97,24 @@ def test_sclar_add():
                         ),
                     )
                 ),
-                output_names=["test_func(Literal(10),Literal(20))"],
+                output_names=["test_uri:test_func(Literal(10),Literal(20))"],
             )
         ],
         base_schema=named_struct,
     )
 
+    print(e)
+    print(expected)
+
     assert e == expected
 
 
 def test_nested_scalar_calls():
     e = scalar_function(
-        "test_uri",
-        "is_positive",
+        "test_uri:is_positive",
         expressions=[
             scalar_function(
-                "test_uri",
-                "test_func",
+                "test_uri:test_func",
                 expressions=[
                     literal(
                         10,

tests/builders/extended_expression/test_window_function.py

@@ -46,7 +46,7 @@
 
 
 def test_row_number():
-    e = window_function("test_uri", "row_number", expressions=[], alias="rn")(
+    e = window_function("test_uri:row_number", expressions=[], alias="rn")(
         named_struct, registry
     )
 

tests/builders/plan/test_aggregate.py

@@ -38,7 +38,7 @@ def test_aggregate():
 
     group_expr = column("id")
     measure_expr = aggregate_function(
-        "test_uri", "count", expressions=[column("is_applicable")], alias=["count"]
+        "test_uri:count", expressions=[column("is_applicable")], alias=["count"]
     )
 
     actual = aggregate(