fix(snowflake): handle udf function naming

ibis-project · Feb 12, 2024 · 926eda6 · 926eda6
1 parent 0b9a40c
commit 926eda6
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/ibis/backends/snowflake/__init__.py b/ibis/backends/snowflake/__init__.py
@@ -305,6 +305,7 @@ def _get_udf_source(self, udf_node: ops.ScalarUDF):
         return dict(
             source=source,
             name=name,
+            func_name=udf_node.__func_name__,
             preamble="\n".join(preamble_lines).format(
                 name=name,
                 signature=signature,
@@ -351,7 +352,7 @@ def _compile_pyarrow_udf(self, udf_node: ops.ScalarUDF) -> None:
     def _compile_python_udf(self, udf_node: ops.ScalarUDF) -> str:
         return """\
 {preamble}
-HANDLER = '{name}'
+HANDLER = '{func_name}'
 AS $$
 from __future__ import annotations
 
@@ -376,7 +377,7 @@ def _compile_pandas_udf(self, udf_node: ops.ScalarUDF) -> str:
 
 @_snowflake.vectorized(input=pd.DataFrame)
 def wrapper(df):
-    return {name}(*(col for _, col in df.items()))
+    return {func_name}(*(col for _, col in df.items()))
 $$"""
         return template.format(**self._get_udf_source(udf_node))