From 48e4ccca367790efe91b073966aa77fae1d967e9 Mon Sep 17 00:00:00 2001 From: Jonathan Shi Date: Thu, 11 Jun 2026 00:39:24 +0000 Subject: [PATCH 1/2] SNOW-3489991: block flattening of dropped NEW columns in connect mode In can_select_statement_be_flattened, the existing guard only refused to flatten when a dropped NEW column was also is_referenced_by_same_level_column. This missed the case where a future filter/sort (not yet added at select() time) references the dropped column, generating invalid SQL. In connect mode, extend the guard to block flattening unconditionally when a NEW column is dropped. This preserves the defining subquery and keeps the column visible to subsequent WHERE/ORDER BY clauses. Non-connect behavior is unchanged. Co-Authored-By: Claude Sonnet 4.6 --- .../snowpark/_internal/analyzer/select_statement.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/snowflake/snowpark/_internal/analyzer/select_statement.py b/src/snowflake/snowpark/_internal/analyzer/select_statement.py index 15db45a247..8e9dcf70da 100644 --- a/src/snowflake/snowpark/_internal/analyzer/select_statement.py +++ b/src/snowflake/snowpark/_internal/analyzer/select_statement.py @@ -2219,7 +2219,12 @@ def can_select_statement_be_flattened( state.change_state == ColumnChangeState.DROPPED and (subquery_state := subquery_column_states.get(col)) and subquery_state.change_state == ColumnChangeState.NEW - and subquery_state.is_referenced_by_same_level_column + and ( + subquery_state.is_referenced_by_same_level_column + # In connect mode, preserve the subquery so that future + # filter/sort clauses can still reference the dropped NEW column. + or context._is_snowpark_connect_compatible_mode + ) ): return False return True From b1a9a747b497919734570ab500a897be0eb2dcdd Mon Sep 17 00:00:00 2001 From: Jonathan Shi Date: Thu, 11 Jun 2026 20:47:59 +0000 Subject: [PATCH 2/2] SNOW-3489991: replace global connect-mode guard with targeted per-statement flag Instead of checking `context._is_snowpark_connect_compatible_mode` broadly in `can_select_statement_be_flattened` (which blocked flattening for *all* dropped-NEW column cases in connect mode), introduce a `protect_dropped_new_columns` bool on `SelectStatement`. SCOS sets this flag specifically on the select statement produced by the withColumn path; `can_select_statement_be_flattened` honours it only for that statement. Non-withColumn paths (e.g. `select().order_by().select()`) are unaffected and continue to flatten normally, keeping generated SQL concise. Co-Authored-By: Claude Sonnet 4.6 --- .../_internal/analyzer/select_statement.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/src/snowflake/snowpark/_internal/analyzer/select_statement.py b/src/snowflake/snowpark/_internal/analyzer/select_statement.py index 8e9dcf70da..0fb111f538 100644 --- a/src/snowflake/snowpark/_internal/analyzer/select_statement.py +++ b/src/snowflake/snowpark/_internal/analyzer/select_statement.py @@ -269,6 +269,7 @@ def __init__( self.pre_actions: Optional[List["Query"]] = None self.post_actions: Optional[List["Query"]] = None self.flatten_disabled: bool = False + self.protect_dropped_new_columns: bool = False self._column_states: Optional[ColumnStateDict] = None self._snowflake_plan: Optional[SnowflakePlan] = None self.expr_to_alias = ( @@ -1499,7 +1500,7 @@ def select(self, cols: List[Expression]) -> "SelectStatement": can_be_flattened = False else: can_be_flattened = can_select_statement_be_flattened( - self.column_states, new_column_states + self.column_states, new_column_states, self.protect_dropped_new_columns ) if can_be_flattened: @@ -2196,7 +2197,9 @@ def parse_column_name( def can_select_statement_be_flattened( - subquery_column_states: ColumnStateDict, new_column_states: ColumnStateDict + subquery_column_states: ColumnStateDict, + new_column_states: ColumnStateDict, + protect_dropped_new_columns: bool = False, ) -> bool: for col, state in new_column_states.items(): dependent_columns = state.dependent_columns @@ -2221,9 +2224,10 @@ def can_select_statement_be_flattened( and subquery_state.change_state == ColumnChangeState.NEW and ( subquery_state.is_referenced_by_same_level_column - # In connect mode, preserve the subquery so that future - # filter/sort clauses can still reference the dropped NEW column. - or context._is_snowpark_connect_compatible_mode + # If the subquery was explicitly marked (e.g. by the SCOS withColumn + # path), preserve it so that future filter/sort clauses can still + # reference the dropped NEW column. + or protect_dropped_new_columns ) ): return False