SNOW-1569005: Deduplicate queries and post actions with building queries from binary plan

CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Snowpark Python API Updates
 
 #### New Features
+
 - Added support for `snowflake.snowpark.testing.assert_dataframe_equal` that is a util function to check the equality of two Snowpark DataFrames.
 - Added support for `Resampler.fillna` and `Resampler.bfill`.
 - Added limited support for the `Timedelta` type, including creating `Timedelta` columns and `to_pandas`.
@@ -26,6 +27,7 @@
 - Fixed a bug in `DataFrame.lineage.trace` to split the quoted feature view's name and version correctly.
 - Fixed a bug in `Column.isin` that caused invalid sql generation when passed an empty list.
 - Fixed a bug that fails to raise NotImplementedError while setting cell with list like item.
+- Fixed a bug in query generation from set operations that allowed generation of duplicate queries when children have common subqueries.
 
 ### Snowpark Local Testing Updates
 

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

@@ -1224,11 +1224,15 @@ def __init__(self, *set_operands: SetOperand, analyzer: "Analyzer") -> None:
             if operand.selectable.pre_actions:
                 if not self.pre_actions:
                     self.pre_actions = []
-                self.pre_actions.extend(operand.selectable.pre_actions)
+                for action in operand.selectable.pre_actions:
+                    if action not in self.pre_actions:
+                        self.pre_actions.append(copy(action))
             if operand.selectable.post_actions:
                 if not self.post_actions:
                     self.post_actions = []
-                self.post_actions.extend(operand.selectable.post_actions)
+                for action in operand.selectable.post_actions:
+                    if action not in self.post_actions:
+                        self.post_actions.append(copy(action))
             self._nodes.append(operand.selectable)
 
     def __deepcopy__(self, memodict={}) -> "SetStatement":  # noqa: B006

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

@@ -637,8 +637,15 @@ def build_binary(
                 if post_action not in post_actions:
                     post_actions.append(copy.copy(post_action))
         else:
-            merged_queries = select_left.queries[:-1] + select_right.queries[:-1]
-            post_actions = select_left.post_actions + select_right.post_actions
+            merged_queries = select_left.queries[:-1]
+            for query in select_right.queries[:-1]:
+                if query not in merged_queries:
+                    merged_queries.append(copy.copy(query))
+
+            post_actions = select_left.post_actions
+            for action in select_right.post_actions:
+                if action not in post_actions:
+                    post_actions.append(copy.copy(action))
 
         queries = merged_queries + [
             Query(

src/snowflake/snowpark/_internal/compiler/query_generator.py

@@ -89,8 +89,12 @@ def generate_queries(
             plan_queries = get_snowflake_plan_queries(
                 snowflake_plan, self.resolved_with_query_block
             )
-            queries.extend(plan_queries[PlanQueryType.QUERIES])
-            post_actions.extend(plan_queries[PlanQueryType.POST_ACTIONS])
+            for query in plan_queries[PlanQueryType.QUERIES]:
+                if query not in queries:
+                    queries.append(query)
+            for action in plan_queries[PlanQueryType.POST_ACTIONS]:
+                if action not in post_actions:
+                    post_actions.append(action)
 
         return {
             PlanQueryType.QUERIES: queries,

tests/integ/scala/test_snowflake_plan_suite.py

@@ -128,7 +128,7 @@ def test_multiple_queries(session):
         Utils.drop_table(session, table_name2)
 
 
-def test_execution_queries_and_queries(session):
+def test_execution_queries_and_post_actions(session):
     df = session.create_dataframe([[1, 2], [3, 4]], schema=["a", "b"])
     df1 = df.select("a", "b")
     # create a df where cte optimization can be applied

tests/integ/test_cte.py

@@ -116,7 +116,7 @@ def test_unary(session, action):
         lambda x, y: x.join(y.select("a"), how="left", rsuffix="_y"),
     ],
 )
-def test_binary(session, action, sql_simplifier_enabled):
+def test_binary(session, action):
     df = session.create_dataframe([[1, 2], [3, 4]], schema=["a", "b"])
     check_result(session, action(df, df), expect_cte_optimized=True)
 
@@ -134,19 +134,8 @@ def test_binary(session, action, sql_simplifier_enabled):
     check_result(session, df3, expect_cte_optimized=True)
     plan_queries = df3.queries
     # check the number of queries
-    binary_build_used = (not sql_simplifier_enabled) or (
-        "JOIN" in plan_queries["queries"][-1]
-    )
-    # TODO (SNOW-1569005): Deduplicate queries during plan build for binary operators. This currently
-    #   is only fixed when _query_compilation_stage_enabled for build_binary.
-    if session._query_compilation_stage_enabled and binary_build_used:
-        num_queries = 3
-        num_post_actions = 1
-    else:
-        num_queries = 5
-        num_post_actions = 2
-    assert len(plan_queries["queries"]) == num_queries
-    assert len(plan_queries["post_actions"]) == num_post_actions
+    assert len(plan_queries["queries"]) == 3
+    assert len(plan_queries["post_actions"]) == 1
 
 
 @pytest.mark.parametrize(