SNOW-3384967: reduce describe query generated by alias (#4183)

Author: sfc-gh-yuwang

CHANGELOG.md

@@ -18,6 +18,10 @@
 
 - Added `artifact_repository` support to `udtf_configs` in `session.read.dbapi()`, enabling users to specify a custom artifact repository (e.g. PyPI) for packages used by the internal UDTF during distributed ingestion.
 
+#### Improvements
+
+- When `Session.reduce_describe_query_enabled` is enabled, fewer DESCRIBE queries are issued when the outer query only projects or renames columns from an inner subquery whose column types are already known.
+
 #### Bug Fixes
 
 - Fixed a bug where `TRY_CAST` reader option is ignored when calling `DataFrameReader.schema().csv()`.

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

@@ -85,8 +85,9 @@
     has_invalid_projection_merge_functions,
 )
 from snowflake.snowpark._internal.utils import (
-    is_sql_select_statement,
     ExprAliasUpdateDict,
+    is_sql_select_statement,
+    quote_name,
 )
 import snowflake.snowpark.context as context
 
@@ -1592,6 +1593,72 @@ def select(self, cols: List[Expression]) -> "SelectStatement":
                 )
             )
 
+        # When describe reduction is on and the inner select already has resolved
+        # attributes, infer new.attributes for this outer select by reusing datatype and
+        # nullable from the subquery: (0) skip if parent column names collide, (1) index
+        # attributes by quote_name (Snowflake identifier rules; invalid delimited forms
+        # raise), (2) walk new.projection, (3) only handle plain columns or Alias(column),
+        # (4) resolve source via the same quote_name key lookup, (5) assign only if every
+        # output column was inferred (length matches projection).
+        if self._session.reduce_describe_query_enabled and self.attributes is not None:
+            parent_attributes = self.attributes
+            projection = new.projection
+            inferred_attributes: Optional[List[Attribute]] = None
+            # Skip: no projection to walk (do not assert; leave new.attributes unchanged).
+            if projection is not None:
+                # Skip: duplicate output names on the parent — dict/lookup would be ambiguous.
+                attributes_by_normalized: Dict[str, Attribute] = {}
+                collision = False
+                for attr in parent_attributes:
+                    key = quote_name(attr.name)
+                    existing = attributes_by_normalized.get(key)
+                    # Skip: two parent columns map to the same quote_name key.
+                    if existing is not None and existing is not attr:
+                        collision = True
+                        break
+                    attributes_by_normalized[key] = attr
+                if not collision:
+                    inferred_attributes = []
+                    for expr in projection:
+                        source_column_name: Optional[str] = None
+                        projected_column_name: Optional[str] = None
+                        if isinstance(expr, (Attribute, UnresolvedAttribute)):
+                            source_column_name = expr.name
+                            projected_column_name = expr.name
+                        elif isinstance(expr, Alias) and isinstance(
+                            expr.child, (Attribute, UnresolvedAttribute)
+                        ):
+                            source_column_name = expr.child.name
+                            projected_column_name = expr.name
+                        else:
+                            # Skip: not a plain column or Alias(Attribute|UnresolvedAttribute).
+                            inferred_attributes = []
+                            break
+
+                        if source_column_name is None or projected_column_name is None:
+                            # Skip: missing projected output name.
+                            inferred_attributes = []
+                            break
+                        source_attr = attributes_by_normalized.get(
+                            quote_name(source_column_name)
+                        )
+                        # Skip: no parent column for this source name.
+                        if source_attr is None:
+                            inferred_attributes = []
+                            break
+                        inferred_attributes.append(
+                            Attribute(
+                                projected_column_name,
+                                source_attr.datatype,
+                                source_attr.nullable,
+                            )
+                        )
+                    if len(inferred_attributes) != len(projection):
+                        # Skip: incomplete inference (includes defensive mismatch).
+                        inferred_attributes = None
+            if inferred_attributes is not None:
+                new.attributes = inferred_attributes
+
         new.flatten_disabled = disable_next_level_flatten
         assert new.projection is not None
         new._column_states = derive_column_states_from_subquery(

tests/integ/test_cte.py

@@ -259,7 +259,7 @@ def test_binary(session, type, action):
 
 def test_join_with_alias_dataframe(session):
     expected_describe_count = (
-        3
+        2
         if (session.reduce_describe_query_enabled and session.sql_simplifier_enabled)
         else 4
     )