SNOW-1856438: Directly support Snowflake Cortex functions Summarize and Sentiment with apply

CHANGELOG.md

@@ -13,7 +13,6 @@
   - `map_contains_key`
   - `map_keys`
   - `nullifzero`
-  - `snowflake_cortex_sentiment`
   - `acosh`
   - `asinh`
   - `atanh`
@@ -89,14 +88,21 @@
 - Fixed a bug in `replace` and `lit` which raised type hint assertion error when passing `Column` expression objects.
 - Fixed a bug in `pandas_udf` and `pandas_udtf` where `session` parameter was erroneously ignored.
 
+#### Removed Features
+
+- Removed support for Snowpark Python function `snowflake_cortex_summarize`. Users can use the Summarize function from Snowflake.Cortex instead.
+
 ### Snowpark pandas API Updates
 
+#### Dependency Updates
+
+- Added a test dependency for `snowflake-ml-python`.
+
 #### New Features
 
 - Added support for `Series.str.ljust` and `Series.str.rjust`.
 - Added support for `Series.str.center`.
 - Added support for `Series.str.pad`.
-- Added support for applying Snowpark Python function `snowflake_cortex_sentiment`.
 - Added support for `DataFrame.map`.
 - Added support for `DataFrame.from_dict` and `DataFrame.from_records`.
 - Added support for mixed case field names in struct type columns.
@@ -121,6 +127,7 @@
 - Added support for aggregations `"count"`, `"median"`, `np.median`,
   `"skew"`, `"std"`, `np.std` `"var"`, and `np.var` in
   `pd.pivot_table()`, `DataFrame.pivot_table()`, and `pd.crosstab()`.
+- Added support for applying Snowflake Cortex functions `Summarize`, `Sentiment`, and `ClassifyText`.
 
 #### Bug Fixes
 

docs/source/snowpark/functions.rst

@@ -308,8 +308,6 @@ Functions
     sinh
     size
     skew
-    snowflake_cortex_sentiment
-    snowflake_cortex_summarize
     sort_array
     soundex
     split

setup.py

@@ -202,6 +202,7 @@ def run(self):
             "scikit-learn",  # Snowpark pandas 3rd party library testing
             # plotly version restricted due to foreseen change in query counts in version 6.0.0+
             "plotly<6.0.0",  # Snowpark pandas 3rd party library testing
+            "snowflake-ml-python",
         ],
         "localtest": [
             "pandas",

src/snowflake/snowpark/functions.py

@@ -10741,36 +10741,6 @@ def make_interval(
     return res
 
 
-def snowflake_cortex_summarize(text: ColumnOrLiteralStr):
-    """
-    Summarizes the given English-language input text.
-
-    Args:
-        text: A string containing the English text from which a summary should be generated.
-
-    Returns:
-        A string containing a summary of the original text.
-    """
-    sql_func_name = "snowflake.cortex.summarize"
-    text_col = _to_col_if_lit(text, sql_func_name)
-    return builtin(sql_func_name)(text_col)
-
-
-def snowflake_cortex_sentiment(text: ColumnOrLiteralStr):
-    """
-    A string containing the text for which a sentiment score should be calculated.
-
-    Args:
-        text: A string containing the English text from which a summary should be generated.
-    Returns:
-        A floating-point number from -1 to 1 (inclusive) indicating the level of negative or positive sentiment in the
-        text. Values around 0 indicate neutral sentiment.
-    """
-    sql_func_name = "snowflake.cortex.sentiment"
-    text_col = _to_col_if_lit(text, sql_func_name)
-    return builtin(sql_func_name)(text_col)
-
-
 @publicapi
 def acosh(e: ColumnOrName, _emit_ast: bool = True) -> Column:
     """

src/snowflake/snowpark/modin/plugin/_internal/apply_utils.py

@@ -14,6 +14,7 @@
 import cloudpickle
 import numpy as np
 import pandas as native_pd
+import snowflake.cortex
 from pandas._typing import AggFuncType
 from pandas.api.types import is_scalar
 
@@ -97,8 +98,12 @@
     sp_func.floor,
     sp_func.trunc,
     sp_func.sqrt,
-    sp_func.snowflake_cortex_summarize,
-    sp_func.snowflake_cortex_sentiment,
+}
+
+SUPPORTED_SNOWFLAKE_CORTEX_FUNCTIONS_IN_APPLY = {
+    snowflake.cortex.ClassifyText,
+    snowflake.cortex.Summarize,
+    snowflake.cortex.Sentiment,
 }
 
 

src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py

@@ -202,6 +202,7 @@
     is_supported_snowpark_python_function,
     sort_apply_udtf_result_columns_by_pandas_positions,
     make_series_map_snowpark_function,
+    SUPPORTED_SNOWFLAKE_CORTEX_FUNCTIONS_IN_APPLY,
 )
 from collections import defaultdict
 from snowflake.snowpark.modin.plugin._internal.binary_op_utils import (
@@ -455,6 +456,32 @@
 
 T = TypeVar("T", bound=Callable[..., Any])
 
+_CORTEX_FUNC_NAMES = [
+    "ClassifyText",
+    "Complete",
+    "CompleteOptions",
+    "EmbedText1024",
+    "EmbedText768",
+    "ExtractAnswer",
+    "Finetune",
+    "FinetuneJob",
+    "FinetuneStatus",
+    "Sentiment",
+    "Summarize",
+    "Translate",
+    "_classify_text",
+    "_complete",
+    "_embed_text_1024",
+    "_embed_text_768",
+    "_extract_answer",
+    "_finetune",
+    "_sentiment",
+    "_sse_client",
+    "_summarize",
+    "_translate",
+    "_util",
+]
+
 
 def _propagate_attrs_on_methods(cls):  # type: ignore
     """
@@ -8471,6 +8498,22 @@ def apply(
                 )
             return self._apply_snowpark_python_function_to_columns(func, kwargs)
 
+        if func in SUPPORTED_SNOWFLAKE_CORTEX_FUNCTIONS_IN_APPLY:
+            if axis != 0:
+                ErrorMessage.not_implemented(
+                    f"Snowpark pandas apply API doesn't yet support Snowflake Cortex function `{func.__name__}` with with axis = {axis}.'"
+                )
+            if raw is not False:
+                ErrorMessage.not_implemented(
+                    f"Snowpark pandas apply API doesn't yet support Snowflake Cortex function `{func.__name__}`with raw = {raw}."
+                )
+            return self._apply_snowflake_cortex_function_to_columns(func, args, kwargs)
+        else:
+            if func.__name__ in _CORTEX_FUNC_NAMES:
+                ErrorMessage.not_implemented(
+                    f"Snowpark pandas apply API doesn't yet support Snowflake Cortex function `{func.__name__}`"
+                )
+
         sf_func = NUMPY_UNIVERSAL_FUNCTION_TO_SNOWFLAKE_FUNCTION.get(func)
         if sf_func is not None:
             return self._apply_snowpark_python_function_to_columns(sf_func, kwargs)
@@ -8786,6 +8829,29 @@ def sf_function(col: SnowparkColumn) -> SnowparkColumn:
             self._modin_frame.apply_snowpark_function_to_columns(sf_function)
         )
 
+    def _apply_snowflake_cortex_function_to_columns(
+        self,
+        snowflake_function: Callable,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],  # possible named arguments which need to be added
+    ) -> "SnowflakeQueryCompiler":
+        """Apply Snowflake Cortex function to columns."""
+
+        def sf_function(col: SnowparkColumn) -> SnowparkColumn:
+            resolved_positional = []
+            resolved_positional.append(col)
+            if args:
+                resolved_positional.append(args)
+            elif kwargs:
+                for arg in kwargs:
+                    resolved_positional.append(kwargs[arg])
+
+            return snowflake_function(*resolved_positional)
+
+        return SnowflakeQueryCompiler(
+            self._modin_frame.apply_snowpark_function_to_columns(sf_function)
+        )
+
     def applymap(
         self,
         func: AggFuncType,
@@ -8819,6 +8885,18 @@ def applymap(
                 )
             return self._apply_snowpark_python_function_to_columns(func, kwargs)
 
+        if func in SUPPORTED_SNOWFLAKE_CORTEX_FUNCTIONS_IN_APPLY:
+            if na_action:
+                ErrorMessage.not_implemented(
+                    f"Snowpark pandas applymap API doesn't yet support Snowflake Cortex function `{func.__name__}` with na_action == '{na_action}'"
+                )
+            return self._apply_snowflake_cortex_function_to_columns(func, args, kwargs)
+        else:
+            if func.__name__ in _CORTEX_FUNC_NAMES:
+                ErrorMessage.not_implemented(
+                    f"Snowpark pandas apply API doesn't yet support Snowflake Cortex function `{func.__name__}`"
+                )
+
         # Check if the function is a known numpy function that can be translated
         # to Snowflake function.
         sf_func = NUMPY_UNIVERSAL_FUNCTION_TO_SNOWFLAKE_FUNCTION.get(func)