Fix limit() after sort() in aggregation query (#3596)

Author: sfc-gh-jdu

src/snowflake/snowpark/_internal/analyzer/analyzer.py

@@ -1166,9 +1166,10 @@ def do_resolve_with_resolved_children(
             )
 
         if isinstance(logical_plan, Limit):
-            on_top_of_order_by = isinstance(
-                logical_plan.child, SnowflakePlan
-            ) and isinstance(logical_plan.child.source_plan, Sort)
+            on_top_of_order_by = logical_plan.is_limit_append or (
+                isinstance(logical_plan.child, SnowflakePlan)
+                and isinstance(logical_plan.child.source_plan, Sort)
+            )
             return self.plan_builder.limit(
                 self.to_sql_try_avoid_cast(
                     logical_plan.limit_expr, df_aliased_col_name_to_real_col_name

src/snowflake/snowpark/_internal/analyzer/snowflake_plan_node.py

@@ -288,13 +288,18 @@ def individual_node_complexity(self) -> Dict[PlanNodeCategory, int]:
 
 class Limit(LogicalPlan):
     def __init__(
-        self, limit_expr: Expression, offset_expr: Expression, child: LogicalPlan
+        self,
+        limit_expr: Expression,
+        offset_expr: Expression,
+        child: LogicalPlan,
+        is_limit_append: bool = False,
     ) -> None:
         super().__init__()
         self.limit_expr = limit_expr
         self.offset_expr = offset_expr
         self.child = child
         self.children.append(child)
+        self.is_limit_append = is_limit_append
 
     @property
     def individual_node_complexity(self) -> Dict[PlanNodeCategory, int]:

src/snowflake/snowpark/dataframe.py

@@ -617,7 +617,7 @@ def __init__(
 
         self._statement_params = None
         self.is_cached: bool = is_cached  #: Whether the dataframe is cached.
-        self._is_grouped_by_and_aggregated = False
+        self._ops_after_agg = None
 
         # Whether all columns are VARIANT data type,
         # which support querying nested fields via dot notations
@@ -1970,7 +1970,8 @@ def filter(
         # the filtering for dataframe after aggregation without nesting using HAVING
         if (
             context._is_snowpark_connect_compatible_mode
-            and self._is_grouped_by_and_aggregated
+            and self._ops_after_agg is not None
+            and "filter" not in self._ops_after_agg
         ):
             having_plan = Filter(filter_col_expr, self._plan, is_having=True)
             if self._select_statement:
@@ -1985,7 +1986,8 @@ def filter(
                 )
             else:
                 df = self._with_plan(having_plan, _ast_stmt=stmt)
-            df._is_grouped_by_and_aggregated = True
+            df._ops_after_agg = self._ops_after_agg.copy()
+            df._ops_after_agg.add("filter")
             return df
         else:
             if self._select_statement:
@@ -2134,7 +2136,8 @@ def sort(
         # the sorting for dataframe after aggregation without nesting
         if (
             context._is_snowpark_connect_compatible_mode
-            and self._is_grouped_by_and_aggregated
+            and self._ops_after_agg is not None
+            and "sort" not in self._ops_after_agg
         ):
             sort_plan = Sort(sort_exprs, self._plan, is_order_by_append=True)
             if self._select_statement:
@@ -2149,7 +2152,8 @@ def sort(
                 )
             else:
                 df = self._with_plan(sort_plan, _ast_stmt=stmt)
-            df._is_grouped_by_and_aggregated = True
+            df._ops_after_agg = self._ops_after_agg.copy()
+            df._ops_after_agg.add("sort")
             return df
         else:
             df = (
@@ -2855,13 +2859,39 @@ def limit(
         else:
             stmt = None
 
-        if self._select_statement:
+        # In snowpark_connect_compatible mode, we need to handle
+        # the limit for dataframe after aggregation without nesting
+        if (
+            context._is_snowpark_connect_compatible_mode
+            and self._ops_after_agg is not None
+            and "limit" not in self._ops_after_agg
+        ):
+            limit_plan = Limit(
+                Literal(n), Literal(offset), self._plan, is_limit_append=True
+            )
+            if self._select_statement:
+                df = self._with_plan(
+                    self._session._analyzer.create_select_statement(
+                        from_=self._session._analyzer.create_select_snowflake_plan(
+                            limit_plan, analyzer=self._session._analyzer
+                        ),
+                        analyzer=self._session._analyzer,
+                    ),
+                    _ast_stmt=stmt,
+                )
+            else:
+                df = self._with_plan(limit_plan, _ast_stmt=stmt)
+            df._ops_after_agg = self._ops_after_agg.copy()
+            df._ops_after_agg.add("limit")
+            return df
+        else:
+            if self._select_statement:
+                return self._with_plan(
+                    self._select_statement.limit(n, offset=offset), _ast_stmt=stmt
+                )
             return self._with_plan(
-                self._select_statement.limit(n, offset=offset), _ast_stmt=stmt
+                Limit(Literal(n), Literal(offset), self._plan), _ast_stmt=stmt
             )
-        return self._with_plan(
-            Limit(Literal(n), Literal(offset), self._plan), _ast_stmt=stmt
-        )
 
     @df_api_usage
     @publicapi

src/snowflake/snowpark/relational_grouped_dataframe.py

@@ -324,7 +324,7 @@ def agg(
                     agg_exprs.append(_str_to_expr(e[1], _emit_ast)(col_expr))
 
         df = self._to_df(agg_exprs, _emit_ast=False)
-        df._is_grouped_by_and_aggregated = True
+        df._ops_after_agg = set()
 
         if _emit_ast:
             df._ast_id = stmt.uid
@@ -515,7 +515,7 @@ def end_partition(
             ),
             _emit_ast=False,
         )
-        df._is_grouped_by_and_aggregated = True
+        df._ops_after_agg = set()
 
         if _emit_ast:
             stmt = working_dataframe._session._ast_batch.bind()
@@ -694,7 +694,7 @@ def count(self, _emit_ast: bool = True) -> DataFrame:
             ],
             _emit_ast=False,
         )
-        df._is_grouped_by_and_aggregated = True
+        df._ops_after_agg = set()
 
         # TODO: count seems similar to mean, min, .... Can we unify implementation here?
         if _emit_ast:
@@ -730,7 +730,7 @@ def _function(
             )._expression
             agg_exprs.append(expr)
         df = self._to_df(agg_exprs)
-        df._is_grouped_by_and_aggregated = True
+        df._ops_after_agg = set()
 
         if _emit_ast:
             stmt = self._dataframe._session._ast_batch.bind()

tests/integ/test_df_aggregate.py

@@ -841,3 +841,90 @@ def test_agg_filter_and_sort_with_grouping_snowpark_connect_compatible(session):
         assert results6[0][2] == 1  # gc=1 for NULL course
     finally:
         context._is_snowpark_connect_compatible_mode = original_value
+
+
+@pytest.mark.skipif(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="HAVING, ORDER BY append, and limit append are not supported in local testing mode",
+)
+def test_filter_sort_limit_snowpark_connect_compatible(session):
+    original_value = context._is_snowpark_connect_compatible_mode
+
+    try:
+        context._is_snowpark_connect_compatible_mode = True
+        df = session.create_dataframe(
+            [(1, 2, 3), (3, 2, 1), (3, 2, 1)], ["a", "b", "c"]
+        )
+
+        # Basic aggregation with filter, sort, limit - should be in same level
+        agg_df = df.group_by("a").agg(
+            sum_("b").alias("sum_b"), count("c").alias("count_c")
+        )
+        result_df1 = agg_df.filter(col("sum_b") > 1).sort("a").limit(10)
+
+        # Check the result
+        Utils.check_answer(result_df1, [Row(1, 2, 1), Row(3, 4, 2)])
+
+        # Check that filter, sort, and limit are in the same query level (single SELECT)
+        query1 = result_df1.queries["queries"][-1]
+        # Count SELECT statements - should be 3 for operations in same level
+        assert query1.upper().count("SELECT") == 3
+        assert "ORDER BY" in query1.upper()
+        assert "LIMIT" in query1.upper()
+        assert "HAVING" in query1.upper()
+
+        # Duplicate sort operations - second sort should be in next level
+        result_df2 = agg_df.sort("a").sort("sum_b")
+
+        # Check the result
+        Utils.check_answer(result_df2, [Row(1, 2, 1), Row(3, 4, 2)])
+
+        # Check that the second sort creates a new query level
+        query2 = result_df2.queries["queries"][-1]
+        # Should have 4 SELECT statements for nested query
+        assert query2.upper().count("SELECT") == 4
+
+        # filter.sort().limit().sort() - last sort should be in next level
+        result_df3 = (
+            agg_df.filter(col("count_c") >= 1)
+            .sort("a")
+            .limit(10)
+            .sort("sum_b", ascending=False)
+        )
+
+        # Check the result
+        Utils.check_answer(result_df3, [Row(3, 4, 2), Row(1, 2, 1)])
+
+        # Check query structure - should have nested SELECT due to sort after limit
+        query3 = result_df3.queries["queries"][-1]
+        assert query3.upper().count("SELECT") == 4
+
+        # limit().limit() - second limit should create new level
+        result_df5 = agg_df.limit(10).limit(1)
+
+        # Check the result (should return only first row)
+        assert result_df5.count() == 1
+
+        # Check query structure - nested due to second limit
+        query5 = result_df5.queries["queries"][-1]
+        assert query5.upper().count("SELECT") == 4
+
+        # Complex chain - filter().sort().limit().filter().sort()
+        result_df6 = (
+            agg_df.filter(col("sum_b") >= 2)
+            .sort("a")
+            .limit(10)
+            .filter(col("count_c") > 1)
+            .sort("sum_b", ascending=False)
+        )
+
+        # Check the result
+        Utils.check_answer(result_df6, [Row(3, 4, 2)])
+
+        # Check query structure - should have multiple levels due to operations after limit
+        query6 = result_df6.queries["queries"][-1]
+        # Should have 4 SELECT statements
+        assert query6.upper().count("SELECT") == 4
+
+    finally:
+        context._is_snowpark_connect_compatible_mode = original_value