Fix inspect issue of pandas udf in SP by using kwargs (#274)

snowflakedb · Mar 23, 2022 · c8d5e2a · c8d5e2a
1 parent 95ab331
commit c8d5e2a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 8 deletions.
diff --git a/src/snowflake/snowpark/functions.py b/src/snowflake/snowpark/functions.py
@@ -2148,6 +2148,7 @@ def add_df_one(df: pd.DataFrame) -> pd.Series:
             replace=replace,
             parallel=parallel,
             max_batch_size=max_batch_size,
+            _from_pandas_udf_function=True,
         )
     else:
         return session.udf.register(
@@ -2162,6 +2163,7 @@ def add_df_one(df: pd.DataFrame) -> pd.Series:
             replace=replace,
             parallel=parallel,
             max_batch_size=max_batch_size,
+            _from_pandas_udf_function=True,
         )
 
 

diff --git a/src/snowflake/snowpark/udf.py b/src/snowflake/snowpark/udf.py
@@ -4,7 +4,6 @@
 # Copyright (c) 2012-2022 Snowflake Computing Inc. All rights reserved.
 #
 """User-defined functions (UDFs) in Snowpark."""
-import inspect
 from types import ModuleType
 from typing import Callable, Iterable, List, Optional, Tuple, Union
 
@@ -429,12 +428,6 @@ def register(
             TempObjectType.FUNCTION, name, is_permanent, stage_location, parallel
         )
 
-        # whether called from pandas_udf
-        caller_frame = inspect.getouterframes(inspect.currentframe())
-        from_pandas_udf_function = (
-            len(caller_frame) > 1 and caller_frame[1].function == "_pandas_udf"
-        )
-
         # register udf
         return self.__do_register_udf(
             func,
@@ -447,7 +440,7 @@ def register(
             replace,
             parallel,
             kwargs.get("max_batch_size"),
-            from_pandas_udf_function,
+            kwargs.get("_from_pandas_udf_function", False),
         )
 
     def register_from_file(