diff --git a/src/snowflake/snowpark/_internal/type_utils.py b/src/snowflake/snowpark/_internal/type_utils.py index 5e22ace8c0..b531c8a51c 100644 --- a/src/snowflake/snowpark/_internal/type_utils.py +++ b/src/snowflake/snowpark/_internal/type_utils.py @@ -1481,6 +1481,18 @@ def format_year_month_interval_for_display( years = str(int(parts[0])) months = str(int(parts[1])) + elif (cell.startswith("+") or cell.startswith("-")) and len(cell) > 1: + # Newer connector behavior: single-field strings like "+5" / "-5" + is_negative = cell.startswith("-") + v = str(int(cell[1:])) + if ( + start_field == YearMonthIntervalType.MONTH + and end_field == YearMonthIntervalType.MONTH + ): + months = v + else: + years = v + # Format based on start/end field sign_prefix = "-" if is_negative else "" diff --git a/src/snowflake/snowpark/dataframe_writer.py b/src/snowflake/snowpark/dataframe_writer.py index c1d2c4da41..36222055bf 100644 --- a/src/snowflake/snowpark/dataframe_writer.py +++ b/src/snowflake/snowpark/dataframe_writer.py @@ -53,6 +53,7 @@ from snowflake.snowpark.functions import sql_expr from snowflake.snowpark.mock._connection import MockServerConnection from snowflake.snowpark.row import Row +import snowflake.snowpark.context as context # Python 3.8 needs to use typing.Iterable because collections.abc.Iterable is not subscriptable # Python 3.9 can use both @@ -537,43 +538,54 @@ def save_as_table( SaveMode.APPEND, SaveMode.TRUNCATE, ] or (save_mode == SaveMode.OVERWRITE and overwrite_condition is not None) - if ( - table_exists is None - and not isinstance(session._conn, MockServerConnection) - and needs_table_exists_check - ): - # whether the table already exists in the database - # determines the compiled SQL for APPEND, TRUNCATE, and OVERWRITE with overwrite_condition - # if the table does not exist, we need to create it first; - # if the table exists, we can skip the creation step and insert data directly - table_exists = session._table_exists(table_name) - - create_table_logic_plan = SnowflakeCreateTable( - table_name, - column_names, - save_mode, - self._dataframe._plan, - TableCreationSource.OTHERS, - table_type, - clustering_exprs, - comment, - enable_schema_evolution, - data_retention_time, - max_data_extension_time, - change_tracking, - copy_grants, - iceberg_config, - table_exists, - overwrite_condition_expr, - ) - snowflake_plan = session._analyzer.resolve(create_table_logic_plan) - result = session._conn.execute( - snowflake_plan, - _statement_params=statement_params, - block=block, - data_type=_AsyncResultType.NO_RESULT, - **kwargs, - ) + + def _execute_save_as_table_plan(table_exists_flag: Optional[bool]): + create_table_logic_plan = SnowflakeCreateTable( + table_name, + column_names, + save_mode, + self._dataframe._plan, + TableCreationSource.OTHERS, + table_type, + clustering_exprs, + comment, + enable_schema_evolution, + data_retention_time, + max_data_extension_time, + change_tracking, + copy_grants, + iceberg_config, + table_exists_flag, + overwrite_condition_expr, + ) + snowflake_plan = session._analyzer.resolve(create_table_logic_plan) + return session._conn.execute( + snowflake_plan, + _statement_params=statement_params, + block=block, + data_type=_AsyncResultType.NO_RESULT, + **kwargs, + ) + + if context._is_snowpark_connect_compatible_mode: + if save_mode == SaveMode.APPEND: + result = _execute_save_as_table_plan(False) + else: + result = _execute_save_as_table_plan(None) + + else: + if ( + table_exists is None + and not isinstance(session._conn, MockServerConnection) + and needs_table_exists_check + ): + # whether the table already exists in the database + # determines the compiled SQL for APPEND, TRUNCATE, and OVERWRITE with overwrite_condition + # if the table does not exist, we need to create it first; + # if the table exists, we can skip the creation step and insert data directly + table_exists = session._table_exists(table_name) + + result = _execute_save_as_table_plan(table_exists) return result if not block else None @overload diff --git a/tests/unit/test_interval_display_formatting.py b/tests/unit/test_interval_display_formatting.py new file mode 100644 index 0000000000..7110aaaf4e --- /dev/null +++ b/tests/unit/test_interval_display_formatting.py @@ -0,0 +1,43 @@ +# +# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved. +# + +import pytest + +from snowflake.snowpark._internal.type_utils import ( + format_year_month_interval_for_display, +) +from snowflake.snowpark.types import YearMonthIntervalType + + +@pytest.mark.parametrize("cell", ["+5-00", "+5", "-5-00", "-5"]) +def test_format_year_month_interval_for_display_accepts_single_field_year(cell): + # Connector behavior change: single-field intervals may come back as "+5" / "-5" + # instead of "+5-00" / "-5-00". We should accept both. + formatted = format_year_month_interval_for_display( + cell, YearMonthIntervalType.YEAR, YearMonthIntervalType.YEAR + ) + assert formatted in {"INTERVAL '5' YEAR", "INTERVAL '-5' YEAR"} + + +@pytest.mark.parametrize("cell", ["+5-00", "+5"]) +def test_format_year_month_interval_for_display_year_to_month_defaults_month_to_zero( + cell, +): + assert ( + format_year_month_interval_for_display( + cell, YearMonthIntervalType.YEAR, YearMonthIntervalType.MONTH + ) + == "INTERVAL '5-0' YEAR TO MONTH" + ) + + +@pytest.mark.parametrize("cell", ["+5", "-5"]) +def test_format_year_month_interval_for_display_month_only_single_field(cell): + expected = "INTERVAL '5' MONTH" if cell == "+5" else "INTERVAL '-5' MONTH" + assert ( + format_year_month_interval_for_display( + cell, YearMonthIntervalType.MONTH, YearMonthIntervalType.MONTH + ) + == expected + )