SNOW-2370108: snowpark fails when group by agg called twice on the same df (#3819)

Author: sfc-gh-aling

.github/workflows/daily_precommit.yml

@@ -655,7 +655,7 @@ jobs:
             .tox/coverage.xml
 
   test-enable-fix-join-alias:
-    name: Test Fixing Join Alias py-${{ matrix.os }}-${{ matrix.python-version }}
+    name: Test Fixing Join Alias py-${{ matrix.os }}-${{ matrix.python-version }}-${{ matrix.cloud-provider }}
     needs: build
     runs-on: ${{ matrix.os }}
     strategy:
@@ -707,7 +707,7 @@ jobs:
       - name: Install tox
         run: uv pip install tox --system
       # we only run doctest on macos
-      - if: ${{ matrix.os == 'macos-latest' && matrix.python-version != '3.12'}}
+      - if: ${{ matrix.os == 'macos-latest'}}
         name: Run doctests
         run: python -m tox -e "py${PYTHON_VERSION}-doctest-notudf-ci"
         env:
@@ -719,7 +719,7 @@ jobs:
           # For example, see https://github.com/snowflakedb/snowpark-python/pull/681
         shell: bash
       # do not run other tests for macos
-      - if: ${{ matrix.os != 'macos-latest' && matrix.python-version != '3.12' }}
+      - if: ${{ matrix.os != 'macos-latest'}}
         name: Run tests (excluding doctests)
         run: python -m tox -e "py${PYTHON_VERSION/\./}-notdoctest-ci"
         env:
@@ -730,18 +730,6 @@ jobs:
           SNOWPARK_PYTHON_API_TEST_BUCKET_PATH: ${{ secrets.SNOWPARK_PYTHON_API_TEST_BUCKET_PATH }}
           SNOWPARK_PYTHON_API_S3_STORAGE_INTEGRATION: ${{ vars.SNOWPARK_PYTHON_API_S3_STORAGE_INTEGRATION }}
         shell: bash
-      - if: ${{ matrix.python-version == '3.12' }}
-        name: Run tests (excluding doctests and udf tests)
-        run: python -m tox -e "py${PYTHON_VERSION/\./}-notudfdoctest-ci"
-        env:
-          PYTHON_VERSION: ${{ matrix.python-version }}
-          cloud_provider: ${{ matrix.cloud-provider }}
-          PYTEST_ADDOPTS: --color=yes --tb=short --join_alias_fix
-          TOX_PARALLEL_NO_SPINNER: 1
-          SNOWFLAKE_IS_PYTHON_RUNTIME_TEST: 1
-          SNOWPARK_PYTHON_API_TEST_BUCKET_PATH: ${{ secrets.SNOWPARK_PYTHON_API_TEST_BUCKET_PATH }}
-          SNOWPARK_PYTHON_API_S3_STORAGE_INTEGRATION: ${{ vars.SNOWPARK_PYTHON_API_S3_STORAGE_INTEGRATION }}
-        shell: bash
       - name: Combine coverages
         run: python -m tox -e coverage --skip-missing-interpreters false
         shell: bash

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

@@ -808,6 +808,7 @@ def __deepcopy__(self, memodict={}) -> "SnowflakePlan":  # noqa: B006
 
     def add_aliases(self, to_add: Dict) -> None:
         if self.session._join_alias_fix:
+            self.expr_to_alias = self.expr_to_alias.copy()
             self.expr_to_alias.update(to_add)
         else:
             self.expr_to_alias = {**self.expr_to_alias, **to_add}

tests/integ/scala/test_dataframe_join_suite.py

@@ -17,7 +17,14 @@
     SnowparkSQLException,
     SnowparkSQLInvalidIdException,
 )
-from snowflake.snowpark.functions import coalesce, col, count, is_null, lit
+from snowflake.snowpark.functions import (
+    coalesce,
+    col,
+    count,
+    is_null,
+    lit,
+    sum as sp_sum,
+)
 from snowflake.snowpark.types import (
     IntegerType,
     StringType,
@@ -1626,3 +1633,37 @@ def test_dataframe_join_and_select_same_column_name_from_one_df(session):
     assert df1.join(df2,).select(
         df2.col("a")
     ).collect() == [Row(2)]
+
+
+@pytest.mark.skipif(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="SNOW-1373887: The join alias fix is not supported in Local Testing",
+)
+def test_dataframe_alias_map_unmodified(session):
+    origin = session._join_alias_fix
+    try:
+        session._join_alias_fix = True
+        df = session.create_dataframe([None], ["__DUMMY"])
+
+        cols = [lit("James"), lit(3000)]
+        df = (
+            df.with_columns(["name", "salary"], cols)
+            .select(*cols)
+            .toDF(*["name", "salary"])
+        )
+
+        def aggregate(input):
+            source_expr_to_alias = input._plan.expr_to_alias
+            ret_df = input.group_by(input.col("name").alias("new_name")).agg(
+                sp_sum(input.col("salary"))
+            )
+            assert (
+                source_expr_to_alias == input._plan.expr_to_alias
+            )  # ensure the original df alias map is not changed
+            return ret_df
+
+        Utils.check_answer(aggregate(df), [Row("James", 3000)])
+        # execute twice to make sure no side effect
+        Utils.check_answer(aggregate(df), [Row("James", 3000)])
+    finally:
+        session._join_alias_fix = origin

tests/integ/test_cte.py

@@ -809,7 +809,6 @@ def test_sql_simplifier(session):
         describe_count=0,
         union_count=0,
         join_count=2,
-        describe_count_for_optimized=1 if session._join_alias_fix else None,
     )
     with SqlCounter(query_count=0, describe_count=0):
         # When adding a lsuffix, expr alias map will be updated, so df2 and df3 are considered