Skip pivot warnings for snowpark pandas users, and improve messages

Signed-off-by: sfc-gh-mvashishtha <[email protected]>

src/snowflake/snowpark/_internal/utils.py

@@ -182,8 +182,12 @@
 
 SUPPORTED_TABLE_TYPES = ["temp", "temporary", "transient"]
 
-PIVOT_VALUES_NONE_OR_DATAFRAME_WARNING = "Parameter values is Optional or DataFrame is in private preview since v1.15.0.  Do not use it in production."
-PIVOT_DEFAULT_ON_NULL_WARNING = "Parameter default_on_null is not None is in private preview since v1.15.0.  Do not use it in production."
+PIVOT_VALUES_NONE_OR_DATAFRAME_WARNING = (
+    "Calling pivot() with the `value` parameter set to None or to a Snowpark "
+    + "DataFrame is in private preview since v1.15.0. Do not use this feature "
+    + "in production."
+)
+PIVOT_DEFAULT_ON_NULL_WARNING = "Calling pivot() with the default_on_null is not None is in private preview since v1.15.0.  Do not use it in production."
 
 
 class TempObjectType(Enum):
@@ -879,6 +883,9 @@ def escape_quotes(unescaped: str) -> str:
     return unescaped.replace(DOUBLE_QUOTE, DOUBLE_QUOTE + DOUBLE_QUOTE)
 
 
+should_warn_dynamic_pivot_is_in_private_preview = True
+
+
 def prepare_pivot_arguments(
     df: "snowflake.snowpark.DataFrame",
     df_name: str,
@@ -899,17 +906,17 @@ def prepare_pivot_arguments(
     """
     from snowflake.snowpark.dataframe import DataFrame
 
-    if values is None or isinstance(values, DataFrame):
-        warning(
-            df_name,
-            PIVOT_VALUES_NONE_OR_DATAFRAME_WARNING,
-        )
-
-    if default_on_null is not None:
-        warning(
-            df_name,
-            PIVOT_DEFAULT_ON_NULL_WARNING,
-        )
+    if should_warn_dynamic_pivot_is_in_private_preview:
+        if values is None or isinstance(values, DataFrame):
+            warning(
+                df_name,
+                PIVOT_VALUES_NONE_OR_DATAFRAME_WARNING,
+            )
+        if default_on_null is not None:
+            warning(
+                df_name,
+                PIVOT_DEFAULT_ON_NULL_WARNING,
+            )
 
     if values is not None and not values:
         raise ValueError("values cannot be empty")

src/snowflake/snowpark/modin/plugin/__init__.py

@@ -7,6 +7,8 @@
 
 from packaging import version
 
+import snowflake.snowpark._internal.utils
+
 if sys.version_info.major == 3 and sys.version_info.minor == 8:
     raise RuntimeError(
         "Snowpark pandas does not support Python 3.8. Please update to Python 3.9 or later."
@@ -60,3 +62,14 @@
 from snowflake.snowpark.modin.plugin import docstrings  # isort: skip  # noqa: E402
 
 DocModule.put(docstrings.__name__)
+
+
+# Don't warn the user about our internal usage of private preview pivot
+# features. The user should have already been warned that Snowpark pandas
+# is in public or private preview. They likely don't know or care that we are
+# using Snowpark DataFrame pivot() internally, let alone that we are using
+# private preview features of Snowpark Python.
+
+snowflake.snowpark._internal.utils.should_warn_dynamic_pivot_is_in_private_preview = (
+    False
+)

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

@@ -15,11 +15,7 @@
     ColumnOrSqlExpr,
     LiteralType,
 )
-from snowflake.snowpark._internal.utils import (
-    PIVOT_DEFAULT_ON_NULL_WARNING,
-    PIVOT_VALUES_NONE_OR_DATAFRAME_WARNING,
-    parse_positional_args_to_list,
-)
+from snowflake.snowpark._internal.utils import parse_positional_args_to_list
 from snowflake.snowpark.column import Column
 from snowflake.snowpark.dataframe import DataFrame as SnowparkDataFrame
 from snowflake.snowpark.dataframe_writer import DataFrameWriter
@@ -822,41 +818,16 @@ def pivot(
         See detailed docstring in Snowpark DataFrame's pivot.
         """
         snowpark_dataframe = self.to_projected_snowpark_dataframe()
-
-        # Don't warn the user about our internal usage of private preview
-        # pivot features. The user should have already been warned that
-        # Snowpark pandas is in public or private preview. They likely don't
-        # know or care that we are using Snowpark DataFrame pivot() internally,
-        # let alone that we are using private preview features of Snowpark
-        # Python.
-
-        class NoPivotWarningFilter(logging.Filter):
-            def filter(self, record: logging.LogRecord) -> bool:
-                message = record.getMessage()
-                return (
-                    PIVOT_DEFAULT_ON_NULL_WARNING not in message
-                    and PIVOT_VALUES_NONE_OR_DATAFRAME_WARNING not in message
-                )
-
-        filter = NoPivotWarningFilter()
-
-        class NoPivotWarningContext:
-            def __enter__(self):
-                logging.getLogger("snowflake.snowpark").addFilter(filter)
-
-            def __exit__(self, exc_type, exc_val, exc_tb):
-                logging.getLogger("snowflake.snowpark").removeFilter(filter)
-
-        with NoPivotWarningContext():
-            pivoted_snowpark_dataframe = snowpark_dataframe.pivot(
-                pivot_col=pivot_col,
-                values=values,
-                default_on_null=default_on_null,
-            )
         return OrderedDataFrame(
             # the pivot result columns for dynamic pivot are data dependent, a schema call is required
             # to know all the quoted identifiers for the pivot result.
-            DataFrameReference(pivoted_snowpark_dataframe.agg(*agg_exprs))
+            DataFrameReference(
+                snowpark_dataframe.pivot(
+                    pivot_col=pivot_col,
+                    values=values,
+                    default_on_null=default_on_null,
+                ).agg(*agg_exprs)
+            )
         )
 
     def unpivot(

tests/integ/scala/test_dataframe_aggregate_suite.py

@@ -3,6 +3,7 @@
 # Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
 #
 
+import sys
 from decimal import Decimal
 from math import sqrt
 from typing import NamedTuple
@@ -115,7 +116,10 @@ def test_group_by_pivot_dynamic_any(session, caplog):
         sort=False,
     )
 
-    assert PIVOT_VALUES_NONE_OR_DATAFRAME_WARNING in caplog.text
+    if "snowflake.snowpark.modin.plugin" not in sys.modules:
+        # Snowpark pandas users don't get warnings about dynamic pivot
+        # features. See SNOW-1344848.
+        assert PIVOT_VALUES_NONE_OR_DATAFRAME_WARNING in caplog.text
 
     Utils.check_answer(
         TestData.monthly_sales_with_team(session)
@@ -332,7 +336,10 @@ class MonthlySales(NamedTuple):
             sort=False,
         )
 
-    assert PIVOT_DEFAULT_ON_NULL_WARNING in caplog.text
+    if "snowflake.snowpark.modin.plugin" not in sys.modules:
+        # Snowpark pandas users don't get warnings about dynamic pivot
+        # features. See SNOW-1344848.
+        assert PIVOT_DEFAULT_ON_NULL_WARNING in caplog.text
 
 
 @pytest.mark.localtest