add tests

Author: sfc-gh-yuwang

tests/integ/test_reduce_describe_query.py

@@ -2,15 +2,15 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
-from typing import List
+from typing import Dict, List
 import copy
 import pytest
 
 import snowflake.snowpark._internal.analyzer.snowflake_plan as snowflake_plan
 
 from unittest.mock import patch
 
-from snowflake.snowpark import DataFrame
+from snowflake.snowpark import DataFrame, Row
 from snowflake.snowpark._internal.analyzer.expression import Attribute, Star
 from snowflake.snowpark._internal.analyzer.unary_expression import UnresolvedAlias
 from snowflake.snowpark._internal.analyzer.unary_plan_node import Project
@@ -34,9 +34,9 @@
     _PYTHON_SNOWPARK_REDUCE_DESCRIBE_QUERY_ENABLED,
     Session,
 )
-from snowflake.snowpark.types import LongType, StructField, StructType
+from snowflake.snowpark.types import LongType, StringType, StructField, StructType
 from tests.integ.utils.sql_counter import SqlCounter
-from tests.utils import IS_IN_STORED_PROC, TestData
+from tests.utils import IS_IN_STORED_PROC, TestData, Utils
 
 pytestmark = [
     pytest.mark.skipif(
@@ -228,6 +228,10 @@ def check_attributes_equality(attrs1: List[Attribute], attrs2: List[Attribute])
         assert attr1.nullable == attr2.nullable
 
 
+def _attrs_by_name(parent_attributes: List[Attribute]) -> Dict[str, Attribute]:
+    return {attr.name: attr for attr in parent_attributes}
+
+
 def has_star_in_projection(df: DataFrame) -> bool:
     plan = df._plan.source_plan
     return isinstance(plan, Project) and any(
@@ -441,6 +445,142 @@ def test_project_alias_infers_attributes_from_parent_metadata(session):
         check_attributes_equality(df2._plan.attributes, expected_attributes)
 
 
+def test_swap_column_aliases_infers_types_from_source_names(session):
+    """n-1: columns a, b; n: b AS a, a AS b — metadata follows source column types."""
+    df = session.create_dataframe([[1, 2]], schema=["a", "b"])
+    _ = df.schema
+
+    parent_attributes = df._plan._metadata.attributes
+    assert parent_attributes is not None
+    by_name = _attrs_by_name(parent_attributes)
+    expected_attributes = [
+        by_name['"B"'].with_name("a"),
+        by_name['"A"'].with_name("b"),
+    ]
+
+    df2 = df.select(col("b").alias("a"), col("a").alias("b"))
+    Utils.check_answer(df2, [Row(A=2, B=1)], sort=False)
+
+    if session.reduce_describe_query_enabled:
+        check_attributes_equality(df2._plan._metadata.attributes, expected_attributes)
+        expected_describe_count = 0
+    else:
+        assert df2._plan._metadata.attributes is None
+        expected_describe_count = 1
+
+    with SqlCounter(query_count=0, describe_count=expected_describe_count):
+        check_attributes_equality(df2._plan.attributes, expected_attributes)
+
+
+def test_swap_mixed_column_types_inference_follows_source(session):
+    """Swapped output names must take datatypes from the referenced source column."""
+    schema = StructType(
+        [
+            StructField("a", LongType(), nullable=True),
+            StructField("b", StringType(), nullable=True),
+        ]
+    )
+    df = session.create_dataframe([(1, "x")], schema=schema)
+    _ = df.schema
+
+    parent_attributes = df._plan._metadata.attributes
+    assert parent_attributes is not None
+    by_name = _attrs_by_name(parent_attributes)
+    expected_attributes = [
+        by_name['"B"'].with_name("a"),
+        by_name['"A"'].with_name("b"),
+    ]
+
+    df2 = df.select(col("b").alias("a"), col("a").alias("b"))
+    Utils.check_answer(df2, [Row(A="x", B=1)], sort=False)
+
+    if session.reduce_describe_query_enabled:
+        check_attributes_equality(df2._plan._metadata.attributes, expected_attributes)
+        expected_describe_count = 0
+    else:
+        assert df2._plan._metadata.attributes is None
+        expected_describe_count = 1
+
+    with SqlCounter(query_count=0, describe_count=expected_describe_count):
+        check_attributes_equality(df2._plan.attributes, expected_attributes)
+
+
+def test_column_permutation_inference_name_keyed_lookup(session):
+    """Non-swap rename: c->a, a->x, b->y — each output type matches its source column."""
+    df = session.create_dataframe([[1, 2, 3]], schema=["a", "b", "c"])
+    _ = df.schema
+
+    parent_attributes = df._plan._metadata.attributes
+    assert parent_attributes is not None
+    by_name = _attrs_by_name(parent_attributes)
+    expected_attributes = [
+        by_name['"C"'].with_name("a"),
+        by_name['"A"'].with_name("x"),
+        by_name['"B"'].with_name("y"),
+    ]
+
+    df2 = df.select(col("c").alias("a"), col("a").alias("x"), col("b").alias("y"))
+    Utils.check_answer(df2, [Row(A=3, X=1, Y=2)], sort=False)
+
+    if session.reduce_describe_query_enabled:
+        check_attributes_equality(df2._plan._metadata.attributes, expected_attributes)
+        expected_describe_count = 0
+    else:
+        assert df2._plan._metadata.attributes is None
+        expected_describe_count = 1
+
+    with SqlCounter(query_count=0, describe_count=expected_describe_count):
+        check_attributes_equality(df2._plan.attributes, expected_attributes)
+
+
+def test_chained_simple_renames_infer_from_previous_metadata(session):
+    """Second select's parent already has inferred attributes from the first rename."""
+    df = session.create_dataframe([[10, 20]], schema=["a", "b"])
+    _ = df.schema
+
+    df1 = df.select(col("a").alias("p"), col("b").alias("q"))
+    if session.reduce_describe_query_enabled:
+        assert df1._plan._metadata.attributes is not None
+    mid_attrs = df1._plan._metadata.attributes
+    assert mid_attrs is not None or not session.reduce_describe_query_enabled
+
+    df2 = df1.select(col("p").alias("x"), col("q").alias("y"))
+    _ = df1.schema
+
+    if session.reduce_describe_query_enabled:
+        assert df2._plan._metadata.attributes is not None
+        by_mid = _attrs_by_name(df1._plan._metadata.attributes or [])
+        expected_attributes = [
+            by_mid['"P"'].with_name("x"),
+            by_mid['"Q"'].with_name("y"),
+        ]
+        check_attributes_equality(df2._plan._metadata.attributes, expected_attributes)
+        with SqlCounter(query_count=0, describe_count=0):
+            check_attributes_equality(df2._plan.attributes, expected_attributes)
+    else:
+        assert df2._plan._metadata.attributes is None
+        with SqlCounter(query_count=0, describe_count=1):
+            _ = df2._plan.attributes
+
+
+def test_non_simple_projection_skips_metadata_inference(session):
+    """Expressions other than plain column or simple alias(column) do not infer attributes."""
+    df = session.create_dataframe([[1, 2]], schema=["a", "b"])
+    _ = df.schema
+
+    df2 = df.select((col("a") + lit(1)).alias("ap1"), "b")
+
+    assert df2._plan._metadata.attributes is None
+
+    with SqlCounter(query_count=0, describe_count=1):
+        _ = df2._plan.attributes
+
+    df3 = df.select(col("a"), (col("b") + lit(1)).alias("b"))
+    assert df3._plan._metadata.attributes is None
+    with SqlCounter(query_count=0, describe_count=1):
+        _ = df3._plan.attributes
+
+
 @pytest.mark.skipif(IS_IN_STORED_PROC, reason="Can't create a session in SP")
 def test_reduce_describe_query_enabled_on_session(db_parameters):
     with Session.builder.configs(db_parameters).create() as new_session: