SNOW-3590670: fix that agg stats are not copied in df.copy (#4253)

Author: sfc-gh-aling

CHANGELOG.md

@@ -10,6 +10,7 @@
 
 #### Bug Fixes
 
+- Fixed a bug where copying a `DataFrame` via `copy.copy()` lost post-aggregate state, causing subsequent `.limit()` or `.sort()` to generate incorrect SQL.
 - Fixed a bug where calling `DataFrame.alias()` twice on the same DataFrame (e.g. for a self-join) caused both aliases to share the same internal column-mapping dictionary. This made `col("R", "col")` resolve to the same column as `col("L", "col")`, producing incorrect join conditions and filter expressions.
 - Fixed a bug where `cloudpickle` could not be resolved when registering a Python stored procedure or UDF with `runtime_version='3.13'`.
 

src/snowflake/snowpark/dataframe.py

@@ -1047,9 +1047,21 @@ def _copy_plan(self) -> LogicalPlan:
         else:
             return copy.copy(self._plan)
 
+    def _copy_agg_state(self, target: "DataFrame") -> None:
+        """Copy post-aggregate state to target so operations like .limit() and
+        .sort() on the copy go through _build_post_agg_df and generate correct
+        SQL (ORDER BY inside the aggregate subquery, not on the outer query)."""
+        target._ops_after_agg = self._ops_after_agg
+        target._agg_base_plan = self._agg_base_plan
+        target._agg_base_select_statement = self._agg_base_select_statement
+        target._pending_havings = self._pending_havings
+        target._pending_order_bys = self._pending_order_bys
+
     def _copy_without_ast(self) -> "DataFrame":
         """Returns a shallow copy of the DataFrame without AST generation."""
-        return DataFrame(self._session, self._copy_plan(), _emit_ast=False)
+        result = DataFrame(self._session, self._copy_plan(), _emit_ast=False)
+        self._copy_agg_state(result)
+        return result
 
     def __copy__(self) -> "DataFrame":
         """Implements shallow copy protocol for copy.copy(...)."""
@@ -1058,12 +1070,14 @@ def __copy__(self) -> "DataFrame":
             stmt = self._session._ast_batch.bind()
             with_src_position(stmt.expr.dataframe_ref, stmt)
             self._set_ast_ref(stmt.expr)
-        return DataFrame(
+        result = DataFrame(
             self._session,
             self._copy_plan(),
             _ast_stmt=stmt,
             _emit_ast=self._session.ast_enabled,
         )
+        self._copy_agg_state(result)
+        return result
 
     if installed_pandas:
         import pandas  # pragma: no cover

tests/integ/test_df_aggregate.py

@@ -2,6 +2,7 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
+import copy
 import decimal
 import math
 from unittest import mock
@@ -32,6 +33,7 @@
 )
 from snowflake.snowpark.mock._snowflake_data_type import ColumnEmulator, ColumnType
 from snowflake.snowpark.types import DoubleType, IntegerType, StructType, StructField
+from snowflake.snowpark._internal.utils import is_ast_enabled
 from tests.utils import Utils
 
 
@@ -1370,3 +1372,68 @@ def test_group_by_exclude_grouping_columns(session):
     )
     assert len(result_builtin_exclude[0]) == 1  # only sum
     Utils.check_answer(result_builtin_exclude, [Row(6), Row(15)])
+
+
+@pytest.mark.skipif(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="ORDER BY append and limit append are not supported in local testing mode",
+)
+def test_copy_preserves_agg_state(session):
+    """copy.copy() and _copy_without_ast() must preserve post-aggregate state so
+    that .limit() and .sort() on the copy go through _build_post_agg_df and
+    generate correct SQL (ORDER BY inside the aggregate subquery, not lost on
+    the outer wrapper)."""
+    if is_ast_enabled():
+        pytest.skip(
+            "_copy_without_ast() leaves _ast_id=None; calling limit() on the copy "
+            "crashes in AST mode because publicapi injects _emit_ast=True via the "
+            "global is_ast_enabled() which bypasses the Session.ast_enabled mock."
+        )
+    # Disable AST: copy.copy(df).limit() triggers debug_check_missing_ast because
+    # the copy carries the source's API usage with no corresponding AST entries.
+    with mock.patch(
+        "snowflake.snowpark.context._is_snowpark_connect_compatible_mode", True
+    ):
+        df = session.create_dataframe(
+            [
+                ("a", 3),
+                ("b", 1),
+                ("a", 1),
+                ("b", 2),
+                ("c", 10),
+            ],
+            ["k", "v"],
+        )
+        agg_sorted = (
+            df.group_by("k")
+            .agg(sum_("v").alias("total"))
+            .filter(col("total") > 1)
+            .sort(col("total").desc())
+        )
+
+        for copied in (copy.copy(agg_sorted), agg_sorted._copy_without_ast()):
+            # Internal state must be carried over so _build_post_agg_df fires correctly
+            assert (
+                copied._ops_after_agg
+                and copied._ops_after_agg == agg_sorted._ops_after_agg
+            )
+            assert (
+                copied._agg_base_plan
+                and copied._agg_base_plan == agg_sorted._agg_base_plan
+            )
+            assert (
+                copied._agg_base_select_statement
+                and copied._agg_base_select_statement
+                is agg_sorted._agg_base_select_statement
+            )
+            assert (
+                copied._pending_order_bys
+                and copied._pending_order_bys == agg_sorted._pending_order_bys
+            )
+            assert (
+                copied._pending_havings
+                and copied._pending_havings == agg_sorted._pending_havings
+            )
+
+            # Observable result: ORDER BY must be respected under LIMIT
+            Utils.check_answer(copied.limit(2), [Row("c", 10), Row("a", 4)])