From f643c84c298f1e20577b9957d37a243857dc879e Mon Sep 17 00:00:00 2001 From: sfc-gh-mvashishtha Date: Mon, 29 Apr 2024 16:46:22 -0700 Subject: [PATCH] Skip pivot warnings for snowpark pandas users, and improve messages Signed-off-by: sfc-gh-mvashishtha --- src/snowflake/snowpark/_internal/utils.py | 33 ++++++++------ .../snowpark/modin/plugin/__init__.py | 13 ++++++ .../plugin/_internal/ordered_dataframe.py | 45 ++++--------------- .../scala/test_dataframe_aggregate_suite.py | 11 ++++- 4 files changed, 50 insertions(+), 52 deletions(-) diff --git a/src/snowflake/snowpark/_internal/utils.py b/src/snowflake/snowpark/_internal/utils.py index e20e7bb9225..ff6689ab75e 100644 --- a/src/snowflake/snowpark/_internal/utils.py +++ b/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") diff --git a/src/snowflake/snowpark/modin/plugin/__init__.py b/src/snowflake/snowpark/modin/plugin/__init__.py index bd7cab8b4d9..35042684921 100644 --- a/src/snowflake/snowpark/modin/plugin/__init__.py +++ b/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 +) diff --git a/src/snowflake/snowpark/modin/plugin/_internal/ordered_dataframe.py b/src/snowflake/snowpark/modin/plugin/_internal/ordered_dataframe.py index 70fb958287c..08fed491a3b 100644 --- a/src/snowflake/snowpark/modin/plugin/_internal/ordered_dataframe.py +++ b/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( diff --git a/tests/integ/scala/test_dataframe_aggregate_suite.py b/tests/integ/scala/test_dataframe_aggregate_suite.py index a5f6acc3b7e..5bdf12a7846 100644 --- a/tests/integ/scala/test_dataframe_aggregate_suite.py +++ b/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