snowpark-connect 0.20.2__py3-none-any.whl → 0.21.0__py3-none-any.whl

snowflake/snowpark_connect/relation/catalogs/snowflake_catalog.py

@@ -10,7 +10,6 @@ import pyspark.sql.connect.proto.common_pb2 as common_proto
 import pyspark.sql.connect.proto.types_pb2 as types_proto
 from snowflake.core.exceptions import NotFoundError
 
-from snowflake import snowpark
 from snowflake.snowpark import functions
 from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     quote_name_without_upper_casing,
@@ -18,8 +17,11 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
 )
 from snowflake.snowpark.functions import lit
 from snowflake.snowpark.types import BooleanType, StringType
-from snowflake.snowpark_connect import column_name_handler
-from snowflake.snowpark_connect.config import auto_uppercase_ddl, global_config
+from snowflake.snowpark_connect.config import (
+    auto_uppercase_non_column_identifiers,
+    global_config,
+)
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.catalogs.abstract_spark_catalog import (
     AbstractSparkCatalog,
     _get_current_snowflake_schema,
@@ -27,7 +29,7 @@ from snowflake.snowpark_connect.relation.catalogs.abstract_spark_catalog import
     _process_multi_layer_identifier,
 )
 from snowflake.snowpark_connect.type_mapping import proto_to_snowpark_type
-from snowflake.snowpark_connect.utils.attribute_handling import (
+from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
 )
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
@@ -40,7 +42,7 @@ from snowflake.snowpark_connect.utils.udf_cache import cached_udf
 def _normalize_identifier(identifier: str | None) -> str | None:
     if identifier is None:
         return None
-    return identifier.upper() if auto_uppercase_ddl() else identifier
+    return identifier.upper() if auto_uppercase_non_column_identifiers() else identifier
 
 
 def sf_quote(name: str | None) -> str | None:
@@ -345,7 +347,7 @@ class SnowflakeCatalog(AbstractSparkCatalog):
     def dropGlobalTempView(
         self,
         spark_view_name: str,
-    ) -> snowpark.DataFrame:
+    ) -> DataFrameContainer:
         session = get_or_create_snowpark_session()
         schema = global_config.spark_sql_globalTempDatabase
         result_df = session.sql(
@@ -358,14 +360,17 @@ class SnowflakeCatalog(AbstractSparkCatalog):
             )
         )
         columns = ["value"]
-        return column_name_handler.with_column_map(
-            result_df, columns, columns, [BooleanType()]
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=result_df,
+            spark_column_names=columns,
+            snowpark_column_names=columns,
+            snowpark_column_types=[BooleanType()],
         )
 
     def dropTempView(
         self,
         spark_view_name: str,
-    ) -> snowpark.DataFrame:
+    ) -> DataFrameContainer:
         """Drop the current temporary view."""
         session = get_or_create_snowpark_session()
         result = session.sql(
@@ -377,8 +382,11 @@ class SnowflakeCatalog(AbstractSparkCatalog):
         )
         result_df = session.createDataFrame([(view_was_dropped,)], schema=["value"])
         columns = ["value"]
-        return column_name_handler.with_column_map(
-            result_df, columns, columns, [BooleanType()]
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=result_df,
+            spark_column_names=columns,
+            snowpark_column_names=columns,
+            snowpark_column_types=[BooleanType()],
         )
 
     def createTable(
@@ -389,7 +397,7 @@ class SnowflakeCatalog(AbstractSparkCatalog):
         schema: types_proto.DataType,
         description: str,
         **options: typing.Any,
-    ) -> snowpark.DataFrame:
+    ) -> DataFrameContainer:
         """Create either an external, or a managed table.
 
         If path is supplied in which the data for this table exists. When path is specified, an external table is
@@ -422,11 +430,11 @@ class SnowflakeCatalog(AbstractSparkCatalog):
         table_name_parts = split_fully_qualified_spark_name(tableName)
         qualifiers = [table_name_parts for _ in columns]
         column_types = [f.datatype for f in sp_schema.fields]
-        return column_name_handler.with_column_map(
-            session.createDataFrame([], sp_schema),
-            columns,
-            columns,
-            column_types,
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=session.createDataFrame([], sp_schema),
+            spark_column_names=columns,
+            snowpark_column_names=columns,
+            snowpark_column_types=column_types,
             column_qualifiers=qualifiers,
         )
 

snowflake/snowpark_connect/relation/map_aggregate.py

@@ -11,8 +11,8 @@ from snowflake import snowpark
 from snowflake.snowpark.types import DataType
 from snowflake.snowpark_connect.column_name_handler import (
     make_column_names_snowpark_compatible,
-    with_column_map,
 )
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.expression.literal import get_literal_field_and_name
 from snowflake.snowpark_connect.expression.map_expression import (
     map_single_column_expression,
@@ -20,103 +20,124 @@ from snowflake.snowpark_connect.expression.map_expression import (
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.typed_column import TypedColumn
-from snowflake.snowpark_connect.utils.context import temporary_pivot_expression
+from snowflake.snowpark_connect.utils.context import (
+    set_current_grouping_columns,
+    temporary_pivot_expression,
+)
 
 
-def map_group_by_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_group_by_aggregate(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Groups the DataFrame using the specified columns.
 
     Aggregations come in as expressions, which are mapped to `snowpark.Column`
     objects.
     """
-    input_df, columns = map_aggregate_helper(rel)
+    input_df_container, columns = map_aggregate_helper(rel)
+    input_df_actual = input_df_container.dataframe
+
     if len(columns.grouping_expressions()) == 0:
-        result = input_df.agg(*columns.aggregation_expressions())
+        result = input_df_actual.agg(*columns.aggregation_expressions())
     else:
-        result = input_df.group_by(*columns.grouping_expressions()).agg(
+        result = input_df_actual.group_by(*columns.grouping_expressions()).agg(
             *columns.aggregation_expressions()
         )
-    return with_column_map(
-        result,
-        columns.spark_names(),
-        columns.snowpark_names(),
-        columns.data_types(),
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=columns.spark_names(),
+        snowpark_column_names=columns.snowpark_names(),
+        snowpark_column_types=columns.data_types(),
         column_qualifiers=columns.get_qualifiers(),
-        parent_column_name_map=input_df._column_map,
+        parent_column_name_map=input_df_container.column_map,
     )
 
 
-def map_rollup_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_rollup_aggregate(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Create a multidimensional rollup for the current DataFrame using the specified columns.
 
     Aggregations come in as expressions, which are mapped to `snowpark.Column`
     objects.
     """
-    input_df, columns = map_aggregate_helper(rel)
+    input_container, columns = map_aggregate_helper(rel)
+    input_df_actual = input_container.dataframe
+
     if len(columns.grouping_expressions()) == 0:
-        result = input_df.agg(*columns.aggregation_expressions())
+        result = input_df_actual.agg(*columns.aggregation_expressions())
     else:
-        result = input_df.rollup(*columns.grouping_expressions()).agg(
+        result = input_df_actual.rollup(*columns.grouping_expressions()).agg(
             *columns.aggregation_expressions()
         )
-    return with_column_map(
-        result,
-        columns.spark_names(),
-        columns.snowpark_names(),
-        columns.data_types(),
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=columns.spark_names(),
+        snowpark_column_names=columns.snowpark_names(),
+        snowpark_column_types=columns.data_types(),
         column_qualifiers=columns.get_qualifiers(),
-        parent_column_name_map=input_df._column_map,
+        parent_column_name_map=input_container.column_map,
     )
 
 
-def map_cube_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_cube_aggregate(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Create a multidimensional cube for the current DataFrame using the specified columns.
 
     Aggregations come in as expressions, which are mapped to `snowpark.Column`
     objects.
     """
-    input_df, columns = map_aggregate_helper(rel)
+    input_container, columns = map_aggregate_helper(rel)
+    input_df_actual = input_container.dataframe
+
     if len(columns.grouping_expressions()) == 0:
-        result = input_df.agg(*columns.aggregation_expressions())
+        result = input_df_actual.agg(*columns.aggregation_expressions())
     else:
-        result = input_df.cube(*columns.grouping_expressions()).agg(
+        result = input_df_actual.cube(*columns.grouping_expressions()).agg(
             *columns.aggregation_expressions()
         )
-    return with_column_map(
-        result,
-        columns.spark_names(),
-        columns.snowpark_names(),
-        columns.data_types(),
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=columns.spark_names(),
+        snowpark_column_names=columns.snowpark_names(),
+        snowpark_column_types=columns.data_types(),
         column_qualifiers=columns.get_qualifiers(),
-        parent_column_name_map=input_df._column_map,
+        parent_column_name_map=input_container.column_map,
     )
 
 
-def map_pivot_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_pivot_aggregate(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Pivots a column of the current DataFrame and performs the specified aggregation.
 
     There are 2 versions of the pivot function: one that requires the caller to specify the list of the distinct values
     to pivot on and one that does not.
     """
-    input_df, columns = map_aggregate_helper(rel, pivot=True, skip_alias=True)
+    input_container, columns = map_aggregate_helper(rel, pivot=True, skip_alias=True)
+    input_df_actual = input_container.dataframe
+
     pivot_column = map_single_column_expression(
-        rel.aggregate.pivot.col, input_df._column_map, ExpressionTyper(input_df)
+        rel.aggregate.pivot.col,
+        input_container.column_map,
+        ExpressionTyper(input_df_actual),
     )
     pivot_values = [
         get_literal_field_and_name(lit)[0] for lit in rel.aggregate.pivot.values
     ]
 
     if len(columns.grouping_expressions()) == 0:
-        result = input_df.pivot(
+        result = input_df_actual.pivot(
             pivot_column[1].col, pivot_values if pivot_values else None
         ).agg(*columns.aggregation_expressions())
     else:
         result = (
-            input_df.group_by(*columns.grouping_expressions())
+            input_df_actual.group_by(*columns.grouping_expressions())
             .pivot(pivot_column[1].col, pivot_values if pivot_values else None)
             .agg(*columns.aggregation_expressions())
         )
@@ -124,7 +145,7 @@ def map_pivot_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
     spark_columns = []
     for col in [string_parser(s) for s in result.columns]:
         spark_col = (
-            input_df._column_map.get_spark_column_name_from_snowpark_column_name(
+            input_container.column_map.get_spark_column_name_from_snowpark_column_name(
                 col, allow_non_exists=True
             )
         )
@@ -135,15 +156,15 @@ def map_pivot_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
             spark_columns.append(col)
 
     agg_name_list = [c.spark_name for c in columns.grouping_columns]
-    return with_column_map(
-        result,
-        agg_name_list + spark_columns[len(agg_name_list) :],
-        result.columns,
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=agg_name_list + spark_columns[len(agg_name_list) :],
+        snowpark_column_names=result.columns,
         column_qualifiers=(
             columns.get_qualifiers()[: len(agg_name_list)]
             + [[]] * (len(spark_columns) - len(agg_name_list))
         ),
-        parent_column_name_map=input_df._column_map,
+        parent_column_name_map=input_container.column_map,
     )
 
 
@@ -246,7 +267,8 @@ class _Columns:
 def map_aggregate_helper(
     rel: relation_proto.Relation, pivot: bool = False, skip_alias: bool = False
 ):
-    input_df = map_relation(rel.aggregate.input)
+    input_container = map_relation(rel.aggregate.input)
+    input_df = input_container.dataframe
     grouping_expressions = rel.aggregate.grouping_expressions
     expressions = rel.aggregate.aggregate_expressions
     groupings: list[_ColumnMetadata] = []
@@ -258,7 +280,7 @@ def map_aggregate_helper(
     with temporary_pivot_expression(pivot):
         for exp in grouping_expressions:
             new_name, snowpark_column = map_single_column_expression(
-                exp, input_df._column_map, typer
+                exp, input_container.column_map, typer
             )
             alias = make_column_names_snowpark_compatible(
                 [new_name], rel.common.plan_id, len(groupings)
@@ -275,9 +297,12 @@ def map_aggregate_helper(
                 )
             )
 
+        grouping_cols = [g.spark_name for g in groupings]
+        set_current_grouping_columns(grouping_cols)
+
         for exp in expressions:
             new_name, snowpark_column = map_single_column_expression(
-                exp, input_df._column_map, typer
+                exp, input_container.column_map, typer
             )
             alias = make_column_names_snowpark_compatible(
                 [new_name], rel.common.plan_id, len(groupings) + len(aggregations)
@@ -313,7 +338,7 @@ def map_aggregate_helper(
             )
 
         return (
-            input_df,
+            input_container,
             _Columns(
                 grouping_columns=groupings,
                 aggregation_columns=aggregations,

snowflake/snowpark_connect/relation/map_catalog.py

@@ -7,7 +7,7 @@ import re
 import pandas
 import pyspark.sql.connect.proto.catalog_pb2 as catalog_proto
 
-from snowflake import snowpark
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.catalogs import CATALOGS
 from snowflake.snowpark_connect.relation.catalogs.utils import (
     CURRENT_CATALOG_NAME,
@@ -22,7 +22,7 @@ from snowflake.snowpark_connect.utils.telemetry import (
 
 def map_catalog(
     rel: catalog_proto.Catalog,
-) -> pandas.DataFrame | snowpark.DataFrame:
+) -> DataFrameContainer | pandas.DataFrame:
     match rel.WhichOneof("cat_type"):
         # Database related APIs
         case "current_database":