place holder

Author: sfc-gh-yuwang

src/snowflake/snowpark/__init__.py

@@ -40,6 +40,7 @@
     "QueryListener",
     "AsyncJob",
     "StoredProcedureProfiler",
+    "UDFProfiler",
 ]
 
 
@@ -54,6 +55,7 @@
 from snowflake.snowpark.async_job import AsyncJob
 from snowflake.snowpark.column import CaseExpr, Column
 from snowflake.snowpark.stored_procedure_profiler import StoredProcedureProfiler
+from snowflake.snowpark.udf_profiler import UDFProfiler
 from snowflake.snowpark.dataframe import DataFrame
 from snowflake.snowpark.dataframe_ai_functions import DataFrameAIFunctions
 from snowflake.snowpark.dataframe_analytics_functions import DataFrameAnalyticsFunctions

src/snowflake/snowpark/_internal/snowpark_profiler.py

@@ -7,17 +7,13 @@
 from typing import List, Literal, Optional
 
 import snowflake.snowpark
-from snowflake.snowpark._internal.utils import (
-    SNOWFLAKE_ANONYMOUS_CALL_WITH_PATTERN,
-)
 
 logger = logging.getLogger(__name__)
 
 
 class SnowparkProfiler:
     """
-    Set up profiler to receive profiles of stored procedures. This feature cannot be used in owner's right stored
-    procedure because owner's right stored procedure will not be able to set session-level parameters.
+    Base class for stored procedure profiler and UDF profiler
     """
 
     def __init__(
@@ -31,20 +27,22 @@ def __init__(
         self._has_target_stage = False
         self._is_enabled = False
 
-        self._active_profiler_name = ""
+        self._active_profiler_name = "ACTIVE_PYTHON_PROFILER"
         self._output_sql = ""
         self._profiler_module_name = ""
 
-    def register_modules(self, stored_procedures: Optional[List[str]] = None) -> None:
+    def register_modules(self, modules: Optional[List[str]] = None) -> None:
         """
-        Register stored procedures to generate profiles for them.
+        Register modules to generate profiles for them.
 
         Args:
             stored_procedures: List of names of stored procedures. Registered modules will be overwritten by this input.
             Input None or an empty list will remove registered modules.
         """
-        sp_string = ",".join(stored_procedures) if stored_procedures is not None else ""
-        sql_statement = f"alter session set {self._profiler_module_name}='{sp_string}'"
+        module_string = ",".join(modules) if modules is not None else ""
+        sql_statement = (
+            f"alter session set {self._profiler_module_name}='{module_string}'"
+        )
         self._session.sql(sql_statement)._internal_collect_with_tag_no_telemetry()
 
     def set_active_profiler(
@@ -91,25 +89,24 @@ def disable(self) -> None:
         self._session.sql(sql_statement)._internal_collect_with_tag_no_telemetry()
 
     @staticmethod
-    def _is_sp_call(query: str) -> bool:
-        query = query.upper().strip(" ")
-        return SNOWFLAKE_ANONYMOUS_CALL_WITH_PATTERN.match(
-            query
-        ) is not None or query.startswith("CALL")
+    def _is_procedure_or_function_call(query: str) -> bool:
+        pass
 
     def _get_last_query_id(self) -> Optional[str]:
         current_thread = threading.get_ident()
         for query in self._query_history.queries[::-1]:  # type: ignore
             query_thread = getattr(query, "thread_id", None)
-            if query_thread == current_thread and self._is_sp_call(query.sql_text):
+            if query_thread == current_thread and self._is_procedure_or_function_call(
+                query.sql_text
+            ):
                 return query.query_id
         return None
 
     def get_output(self) -> str:
         """
-        Return the profiles of last executed stored procedure in current thread. If there is no previous
-        stored procedure call, an error will be raised.
-        Please call this function right after the stored procedure you want to profile to avoid any error.
+        Return the profiles of last executed stored procedure or UDF in current thread. If there is no previous
+        stored procedure or UDF call, an error will be raised.
+        Please call this function right after the stored procedure or UDF you want to profile to avoid any error.
 
         """
         # return empty string when profiler is not enabled to not interrupt user's code
@@ -121,7 +118,7 @@ def get_output(self) -> str:
         query_id = self._get_last_query_id()
         if query_id is None:
             logger.warning(
-                "You are seeing this warning because last executed stored procedure does not exist. Please run the store procedure before get profiler output."
+                "You are seeing this warning because last executed stored procedure or UDF does not exist. Please run the store procedure or UDF before get profiler output."
             )
             return ""
         sql = self._output_sql.format(query_id=query_id)

src/snowflake/snowpark/session.py

@@ -44,6 +44,8 @@
 from snowflake.connector import ProgrammingError, SnowflakeConnection
 from snowflake.connector.options import installed_pandas, pandas, pyarrow
 from snowflake.connector.pandas_tools import write_pandas
+
+from snowflake.snowpark import UDFProfiler
 from snowflake.snowpark._internal.analyzer import analyzer_utils
 from snowflake.snowpark._internal.analyzer.analyzer import Analyzer
 from snowflake.snowpark._internal.analyzer.analyzer_utils import (
@@ -807,6 +809,7 @@ def __init__(
         self._runtime_version_from_requirement: str = None
         self._temp_table_auto_cleaner: TempTableAutoCleaner = TempTableAutoCleaner(self)
         self._sp_profiler = StoredProcedureProfiler(session=self)
+        self._udf_profiler = UDFProfiler(session=self)
         self._dataframe_profiler = DataframeProfiler(session=self)
         self._catalog = None
 
@@ -4314,6 +4317,14 @@ def stored_procedure_profiler(self) -> StoredProcedureProfiler:
         """
         return self._sp_profiler
 
+    @property
+    def udf_profiler(self) -> UDFProfiler:
+        """
+        Returns a :class:`udf_profiler.UDFProfiler` object that you can use to profile UDFs.
+        See details of how to use this object in :class:`udf_profiler.UDFProfiler`.
+        """
+        return self._udf_profiler
+
     @property
     def dataframe_profiler(self) -> DataframeProfiler:
         """

src/snowflake/snowpark/stored_procedure_profiler.py

@@ -9,6 +9,7 @@
 from snowflake.snowpark._internal.utils import (
     parse_table_name,
     strip_double_quotes_in_like_statement_in_table_name,
+    SNOWFLAKE_ANONYMOUS_CALL_WITH_PATTERN,
 )
 
 logger = logging.getLogger(__name__)
@@ -25,7 +26,6 @@ def __init__(
         session: "snowflake.snowpark.Session",
     ) -> None:
         super().__init__(session)
-        self._active_profiler_name = "ACTIVE_PYTHON_PROFILER"
         self._output_sql = (
             "select snowflake.core.get_python_profiler_output('{query_id}')"
         )
@@ -72,3 +72,10 @@ def set_active_profiler(
                 "Target stage for profiler not found, using default stage of current session."
             )
         super().set_active_profiler(active_profiler_type)
+
+    @staticmethod
+    def _is_procedure_or_function_call(query: str) -> bool:
+        query = query.upper().strip(" ")
+        return SNOWFLAKE_ANONYMOUS_CALL_WITH_PATTERN.match(
+            query
+        ) is not None or query.startswith("CALL")

src/snowflake/snowpark/udf_profiler.py

@@ -0,0 +1,38 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import logging
+import re
+
+import snowflake.snowpark
+from snowflake.snowpark._internal.snowpark_profiler import SnowparkProfiler
+
+
+logger = logging.getLogger(__name__)
+
+SNOWFLAKE_ANONYMOUS_FUNCTION_PATTERN = re.compile(
+    r"^\s*WITH\s+\w+\s+AS\s+FUNCTION", re.IGNORECASE
+)
+
+
+class UDFProfiler(SnowparkProfiler):
+    """
+    Set up profiler to receive profiles of UDFs.
+    """
+
+    def __init__(
+        self,
+        session: "snowflake.snowpark.Session",
+    ) -> None:
+        super().__init__(session)
+
+        self._output_sql = "select * from table(SNOWFLAKE.LOCAL.GET_PYTHON_UDF_PROFILER_OUTPUT('{query_id}'));"
+        self._profiler_module_name = "PYTHON_UDF_PROFILER_MODULES"
+
+    @staticmethod
+    def _is_procedure_or_function_call(query: str) -> bool:
+        query = query.upper().strip(" ")
+        return SNOWFLAKE_ANONYMOUS_FUNCTION_PATTERN.match(
+            query
+        ) is not None or query.startswith("SELECT")

tests/integ/test_stored_procedure_profiler.py

@@ -188,7 +188,7 @@ def test_set_incorrect_active_profiler(
 def test_sp_call_match(profiler_session, sp_call_sql):
     pro = profiler_session.stored_procedure_profiler
 
-    assert pro._is_sp_call(sp_call_sql)
+    assert pro._is_procedure_or_function_call(sp_call_sql)
 
 
 @pytest.mark.skipif(