feat: add type signature to scalar function name definitions

As per a recent change in the spec upstream (and also to enable better compatibility with other substrait libraries), extension function names now have the expected input types included in the name of the function. This can seem very weird, but helps to make it much easier for consumers to parse a plan. Basically, if you call a scalar function `add` and pass in two integer 32s, the "name" of the function will be `add:i32_i32`. We do not validate that the particular signature exists in the extension.yaml files, but we do perform the requisite upcasting to make sure that only integers and floats of the same precision are passed to functions (also in line with the Substrait specification).
ibis-project · Aug 30, 2023 · fe19d3c · fe19d3c
1 parent d832b4d
commit fe19d3c
Show file tree

Hide file tree

Showing 20 changed files with 164 additions and 143 deletions.
diff --git a/ibis_substrait/compiler/core.py b/ibis_substrait/compiler/core.py
@@ -97,6 +97,7 @@ def extension_lookup(
             )
 
         anykey = ("any",) * len([arg for arg in op.args if arg is not None])
+        sigkey = anykey
 
         # First check if `any` is an option
         # This function will take arguments of any type
@@ -136,7 +137,12 @@ def extension_lookup(
 
         extension_uri = self.register_extension_uri(function_extension.uri)
 
-        extension_function = self.create_extension_function(extension_uri, op_name)
+        # format signature key for extension name
+        extension_type_signature = "_".join(sigkey)
+
+        extension_function = self.create_extension_function(
+            extension_uri, f"{op_name}:{extension_type_signature}"
+        )
 
         return extension_function
 

diff --git a/ibis_substrait/compiler/mapping.py b/ibis_substrait/compiler/mapping.py
@@ -111,28 +111,28 @@
     "Int64": "i64",
     "Float32": "fp32",
     "Float64": "fp64",
-    "String": "varchar",
-    "string": "string",
-    "Boolean": "boolean",
+    "String": "str",
+    "string": "str",
+    "Boolean": "bool",
     "Date": "date",
-    "Decimal": "decimal",
+    "Decimal": "dec",
 }
 
 _normalized_key_names = {
     # decimal precision and scale aren't part of the
     # extension signature they're passed in separately
-    "decimal<p, s>": "decimal",
-    "decimal<p,s>": "decimal",
-    "decimal<p1,s1>": "decimal",
-    "decimal<p2,s2>": "decimal",
+    "decimal<p, s>": "dec",
+    "decimal<p,s>": "dec",
+    "decimal<p1,s1>": "dec",
+    "decimal<p2,s2>": "dec",
     # we don't care about string length
-    "fixedchar<l1>": "fixedchar",
-    "fixedchar<l2>": "fixedchar",
-    "varchar<l1>": "varchar",
-    "varchar<l2>": "varchar",
-    "varchar<l3>": "varchar",
+    "fixedchar<l1>": "str",
+    "fixedchar<l2>": "str",
+    "varchar<l1>": "str",
+    "varchar<l2>": "str",
+    "varchar<l3>": "str",
     # for now ignore nullability marker
-    "boolean?": "boolean",
+    "boolean?": "bool",
     # why is there a 1?
     "any1": "any",
     "Date": "date",

diff --git a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h01/tpc_h01.json b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h01/tpc_h01.json
@@ -18,49 +18,49 @@
       "extensionFunction": {
         "extensionUriReference": 1,
         "functionAnchor": 1,
-        "name": "lte"
+        "name": "lte:any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 2,
-        "name": "sum"
+        "name": "sum:dec"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 3,
-        "name": "multiply"
+        "name": "multiply:dec_dec"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 4,
-        "name": "subtract"
+        "name": "subtract:dec_dec"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 5,
-        "name": "add"
+        "name": "add:dec_dec"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 6,
-        "name": "avg"
+        "name": "avg:dec"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 3,
         "functionAnchor": 7,
-        "name": "count"
+        "name": "count:any"
       }
     }
   ],

diff --git a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h03/tpc_h03.json b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h03/tpc_h03.json
@@ -18,49 +18,49 @@
       "extensionFunction": {
         "extensionUriReference": 1,
         "functionAnchor": 1,
-        "name": "equal"
+        "name": "equal:any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 2,
-        "name": "and"
+        "name": "and:bool_bool"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 1,
         "functionAnchor": 3,
-        "name": "lt"
+        "name": "lt:any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 1,
         "functionAnchor": 4,
-        "name": "gt"
+        "name": "gt:any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 3,
         "functionAnchor": 5,
-        "name": "sum"
+        "name": "sum:dec"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 3,
         "functionAnchor": 6,
-        "name": "multiply"
+        "name": "multiply:dec_dec"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 3,
         "functionAnchor": 7,
-        "name": "subtract"
+        "name": "subtract:dec_dec"
       }
     }
   ],

diff --git a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h04/tpc_h04.json b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h04/tpc_h04.json
@@ -18,35 +18,35 @@
       "extensionFunction": {
         "extensionUriReference": 1,
         "functionAnchor": 1,
-        "name": "and"
+        "name": "and:bool_bool"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 2,
-        "name": "equal"
+        "name": "equal:any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 3,
-        "name": "lt"
+        "name": "lt:any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 4,
-        "name": "gte"
+        "name": "gte:any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 3,
         "functionAnchor": 5,
-        "name": "count"
+        "name": "count:any"
       }
     }
   ],

diff --git a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h05/tpc_h05.json b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h05/tpc_h05.json
@@ -18,49 +18,49 @@
       "extensionFunction": {
         "extensionUriReference": 1,
         "functionAnchor": 1,
-        "name": "equal"
+        "name": "equal:any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 2,
-        "name": "and"
+        "name": "and:bool_bool"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 1,
         "functionAnchor": 3,
-        "name": "gte"
+        "name": "gte:any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 1,
         "functionAnchor": 4,
-        "name": "lt"
+        "name": "lt:any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 3,
         "functionAnchor": 5,
-        "name": "sum"
+        "name": "sum:dec"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 3,
         "functionAnchor": 6,
-        "name": "multiply"
+        "name": "multiply:dec_dec"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 3,
         "functionAnchor": 7,
-        "name": "subtract"
+        "name": "subtract:dec_dec"
       }
     }
   ],

diff --git a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h06/tpc_h06.json b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h06/tpc_h06.json
@@ -18,42 +18,42 @@
       "extensionFunction": {
         "extensionUriReference": 1,
         "functionAnchor": 1,
-        "name": "and"
+        "name": "and:bool_bool"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 2,
-        "name": "gte"
+        "name": "gte:any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 3,
-        "name": "lt"
+        "name": "lt:any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 4,
-        "name": "between"
+        "name": "between:any_any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 3,
         "functionAnchor": 5,
-        "name": "sum"
+        "name": "sum:dec"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 3,
         "functionAnchor": 6,
-        "name": "multiply"
+        "name": "multiply:dec_dec"
       }
     }
   ],

diff --git a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h07/tpc_h07.json b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h07/tpc_h07.json
@@ -22,56 +22,56 @@
       "extensionFunction": {
         "extensionUriReference": 1,
         "functionAnchor": 1,
-        "name": "equal"
+        "name": "equal:any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 2,
         "functionAnchor": 2,
-        "name": "extract"
+        "name": "extract:date"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 3,
         "functionAnchor": 3,
-        "name": "multiply"
+        "name": "multiply:dec_dec"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 3,
         "functionAnchor": 4,
-        "name": "subtract"
+        "name": "subtract:dec_dec"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 4,
         "functionAnchor": 5,
-        "name": "and"
+        "name": "and:bool_bool"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 4,
         "functionAnchor": 6,
-        "name": "or"
+        "name": "or:bool_bool"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 1,
         "functionAnchor": 7,
-        "name": "between"
+        "name": "between:any_any_any"
       }
     },
     {
       "extensionFunction": {
         "extensionUriReference": 3,
         "functionAnchor": 8,
-        "name": "sum"
+        "name": "sum:dec"
       }
     }
   ],