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

snowflake/snowpark_connect/relation/map_aggregate.py

@@ -4,10 +4,14 @@
 
 import re
 from dataclasses import dataclass
+from typing import Optional
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
+import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
+from snowflake.snowpark import Column
+from snowflake.snowpark._internal.analyzer.unary_expression import Alias
 from snowflake.snowpark.types import DataType
 from snowflake.snowpark_connect.column_name_handler import (
     make_column_names_snowpark_compatible,
@@ -21,6 +25,7 @@ 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 (
+    get_is_evaluating_sql,
     set_current_grouping_columns,
     temporary_pivot_expression,
 )
@@ -131,19 +136,57 @@ def map_pivot_aggregate(
         get_literal_field_and_name(lit)[0] for lit in rel.aggregate.pivot.values
     ]
 
+    used_columns = {pivot_column[1].col._expression.name}
+    if get_is_evaluating_sql():
+        # When evaluating SQL spark doesn't trim columns from the result
+        used_columns = {"*"}
+    else:
+        for expression in rel.aggregate.aggregate_expressions:
+            matched_identifiers = re.findall(
+                r'unparsed_identifier: "(.*)"', expression.__str__()
+            )
+            for identifier in matched_identifiers:
+                mapped_col = input_container.column_map.spark_to_col.get(
+                    identifier, None
+                )
+                if mapped_col:
+                    used_columns.add(mapped_col[0].snowpark_name)
+
     if len(columns.grouping_expressions()) == 0:
-        result = input_df_actual.pivot(
-            pivot_column[1].col, pivot_values if pivot_values else None
-        ).agg(*columns.aggregation_expressions())
+        result = (
+            input_df_actual.select(*used_columns)
+            .pivot(pivot_column[1].col, pivot_values if pivot_values else None)
+            .agg(*columns.aggregation_expressions(unalias=True))
+        )
     else:
         result = (
             input_df_actual.group_by(*columns.grouping_expressions())
             .pivot(pivot_column[1].col, pivot_values if pivot_values else None)
-            .agg(*columns.aggregation_expressions())
+            .agg(*columns.aggregation_expressions(unalias=True))
         )
 
+    agg_name_list = [c.spark_name for c in columns.grouping_columns]
+
+    # Calculate number of pivot values for proper Spark-compatible indexing
+    total_pivot_columns = len(result.columns) - len(agg_name_list)
+    num_pivot_values = (
+        total_pivot_columns // len(columns.aggregation_columns)
+        if len(columns.aggregation_columns) > 0
+        else 1
+    )
+
+    def _get_agg_exp_alias_for_col(col_index: int) -> Optional[str]:
+        if col_index < len(agg_name_list) or len(columns.aggregation_columns) <= 1:
+            return None
+        else:
+            index = (col_index - len(agg_name_list)) // num_pivot_values
+            return columns.aggregation_columns[index].spark_name
+
     spark_columns = []
-    for col in [string_parser(s) for s in result.columns]:
+    for col in [
+        pivot_column_name(c, _get_agg_exp_alias_for_col(i))
+        for i, c in enumerate(result.columns)
+    ]:
         spark_col = (
             input_container.column_map.get_spark_column_name_from_snowpark_column_name(
                 col, allow_non_exists=True
@@ -153,22 +196,57 @@ def map_pivot_aggregate(
         if spark_col is not None:
             spark_columns.append(spark_col)
         else:
-            spark_columns.append(col)
+            # Handle NULL column names to match Spark behavior (lowercase 'null')
+            if col == "NULL":
+                spark_columns.append(col.lower())
+            else:
+                spark_columns.append(col)
+
+    grouping_cols_count = len(agg_name_list)
+    pivot_cols = result.columns[grouping_cols_count:]
+    spark_pivot_cols = spark_columns[grouping_cols_count:]
+
+    num_agg_functions = len(columns.aggregation_columns)
+    num_pivot_values = len(pivot_cols) // num_agg_functions
+
+    reordered_snowpark_cols = []
+    reordered_spark_cols = []
+    column_indices = []  # 1-based indexing
+
+    for i in range(grouping_cols_count):
+        reordered_snowpark_cols.append(result.columns[i])
+        reordered_spark_cols.append(spark_columns[i])
+        column_indices.append(i + 1)
+
+    for pivot_idx in range(num_pivot_values):
+        for agg_idx in range(num_agg_functions):
+            current_pos = agg_idx * num_pivot_values + pivot_idx
+            if current_pos < len(pivot_cols):
+                reordered_snowpark_cols.append(pivot_cols[current_pos])
+                reordered_spark_cols.append(spark_pivot_cols[current_pos])
+                original_index = grouping_cols_count + current_pos
+                column_indices.append(original_index + 1)
+
+    reordered_result = result.select(
+        *[snowpark_fn.col(f"${idx}") for idx in column_indices]
+    )
 
-    agg_name_list = [c.spark_name for c in columns.grouping_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,
+        dataframe=reordered_result,
+        spark_column_names=reordered_spark_cols,
+        snowpark_column_names=[f"${idx}" for idx in column_indices],
         column_qualifiers=(
             columns.get_qualifiers()[: len(agg_name_list)]
-            + [[]] * (len(spark_columns) - len(agg_name_list))
+            + [[]] * (len(reordered_spark_cols) - len(agg_name_list))
         ),
         parent_column_name_map=input_container.column_map,
+        snowpark_column_types=[
+            result.schema.fields[idx - 1].datatype for idx in column_indices
+        ],
     )
 
 
-def string_parser(s):
+def pivot_column_name(snowpark_cname, opt_alias: Optional[str] = None) -> Optional[str]:
     # For values that are used as pivoted columns, the input and output are in the following format (outermost double quotes are part of the input):
 
     # 1. "'Java'" -> Java
@@ -183,7 +261,7 @@ def string_parser(s):
 
     try:
         # handling values that are used as pivoted columns
-        match = re.match(r'^"\'(.*)\'"$', s)
+        match = re.match(r'^"\'(.*)\'"$', snowpark_cname)
         # extract the content between the outermost double quote followed by a single quote "'
         content = match.group(1)
         # convert the escaped double quote to the actual double quote
@@ -195,10 +273,10 @@ def string_parser(s):
         content = re.sub(r"'", "", content)
         # replace the placeholder with the single quote which we want to preserve
         result = content.replace(escape_single_quote_placeholder, "'")
-        return result
+        return f"{result}_{opt_alias}" if opt_alias else result
     except Exception:
         # fallback to the original logic, handling aliased column names
-        double_quote_list = re.findall(r'"(.*?)"', s)
+        double_quote_list = re.findall(r'"(.*?)"', snowpark_cname)
         spark_string = ""
         for entry in list(filter(None, double_quote_list)):
             if "'" in entry:
@@ -210,7 +288,7 @@ def string_parser(s):
                 spark_string += entry
             else:
                 spark_string += '"' + entry + '"'
-        return s if spark_string == "" else spark_string
+        return snowpark_cname if spark_string == "" else spark_string
 
 
 @dataclass(frozen=True)
@@ -231,8 +309,14 @@ class _Columns:
     def grouping_expressions(self) -> list[snowpark.Column]:
         return [col.expression for col in self.grouping_columns]
 
-    def aggregation_expressions(self) -> list[snowpark.Column]:
-        return [col.expression for col in self.aggregation_columns]
+    def aggregation_expressions(self, unalias: bool = False) -> list[snowpark.Column]:
+        def _unalias(col: snowpark.Column) -> snowpark.Column:
+            if unalias and hasattr(col, "_expr1") and isinstance(col._expr1, Alias):
+                return _unalias(Column(col._expr1.child))
+            else:
+                return col
+
+        return [_unalias(col.expression) for col in self.aggregation_columns]
 
     def expressions(self) -> list[snowpark.Column]:
         return self.grouping_expressions() + self.aggregation_expressions()

snowflake/snowpark_connect/relation/map_column_ops.py

@@ -666,10 +666,29 @@ def map_with_columns_renamed(
         )
 
     # Validate for naming conflicts
-    new_names_list = list(dict(rel.with_columns_renamed.rename_columns_map).values())
+    rename_map = dict(rel.with_columns_renamed.rename_columns_map)
+    new_names_list = list(rename_map.values())
     seen = set()
     for new_name in new_names_list:
-        if column_map.has_spark_column(new_name):
+        # Check if this new name conflicts with existing columns
+        # But allow renaming a column to a different case version of itself
+        is_case_insensitive_self_rename = False
+        if not global_config.spark_sql_caseSensitive:
+            # Find the source column(s) that map to this new name
+            source_columns = [
+                old_name
+                for old_name, new_name_candidate in rename_map.items()
+                if new_name_candidate == new_name
+            ]
+            # Check if any source column is the same as new name when case-insensitive
+            is_case_insensitive_self_rename = any(
+                source_col.lower() == new_name.lower() for source_col in source_columns
+            )
+
+        if (
+            column_map.has_spark_column(new_name)
+            and not is_case_insensitive_self_rename
+        ):
             # Spark doesn't allow reusing existing names, even if the result df will not contain duplicate columns
             raise _column_exists_error(new_name)
         if (global_config.spark_sql_caseSensitive and new_name in seen) or (

snowflake/snowpark_connect/relation/map_map_partitions.py

@@ -14,7 +14,6 @@ from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.type_mapping import proto_to_snowpark_type
 from snowflake.snowpark_connect.utils.pandas_udtf_utils import create_pandas_udtf
-from snowflake.snowpark_connect.utils.session import get_python_udxf_import_files
 from snowflake.snowpark_connect.utils.udf_helper import (
     SnowparkUDF,
     process_udf_in_sproc,
@@ -28,6 +27,9 @@ from snowflake.snowpark_connect.utils.udtf_helper import (
     create_pandas_udtf_in_sproc,
     require_creating_udtf_in_sproc,
 )
+from snowflake.snowpark_connect.utils.udxf_import_utils import (
+    get_python_udxf_import_files,
+)
 
 
 def map_map_partitions(

snowflake/snowpark_connect/relation/map_sql.py

@@ -14,7 +14,10 @@ import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 import sqlglot
 from google.protobuf.any_pb2 import Any
-from pyspark.errors.exceptions.base import AnalysisException
+from pyspark.errors.exceptions.base import (
+    AnalysisException,
+    UnsupportedOperationException,
+)
 from sqlglot.expressions import ColumnDef, DataType, FileFormatProperty, Identifier
 
 import snowflake.snowpark.functions as snowpark_fn
@@ -27,7 +30,6 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
 )
 from snowflake.snowpark._internal.type_utils import convert_sp_to_sf_type
 from snowflake.snowpark._internal.utils import is_sql_select_statement, quote_name
-from snowflake.snowpark.functions import when_matched, when_not_matched
 from snowflake.snowpark_connect.config import (
     auto_uppercase_non_column_identifiers,
     get_boolean_session_config_param,
@@ -56,16 +58,15 @@ from snowflake.snowpark_connect.utils.context import (
     get_session_id,
     push_evaluating_sql_scope,
     push_sql_scope,
-    set_plan_id_map,
     set_sql_args,
     set_sql_plan_name,
 )
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
+    telemetry,
 )
 
-from .. import column_name_handler
 from ..expression.map_sql_expression import (
     _window_specs,
     as_java_list,
@@ -163,6 +164,7 @@ def parse_pos_args(
 
 def execute_logical_plan(logical_plan) -> DataFrameContainer:
     proto = map_logical_plan_relation(logical_plan)
+    telemetry.report_parsed_sql_plan(proto)
     with push_evaluating_sql_scope():
         return map_relation(proto)
 
@@ -712,197 +714,22 @@ def map_sql_to_pandas_df(
                     f"INSERT {overwrite_str} INTO {name} {cols_str} {final_query}",
                 ).collect()
             case "MergeIntoTable":
-
-                def _get_condition_from_action(action, column_mapping, typer):
-                    condition = None
-                    if action.condition().isDefined():
-                        (_, condition_typed_col,) = map_single_column_expression(
-                            map_logical_plan_expression(
-                                matched_action.condition().get()
-                            ),
-                            column_mapping,
-                            typer,
-                        )
-                        condition = condition_typed_col.col
-                    return condition
-
-                def _get_assignments_from_action(
-                    action,
-                    column_mapping_source,
-                    column_mapping_target,
-                    typer_source,
-                    typer_target,
-                ):
-                    assignments = dict()
-                    if (
-                        action.getClass().getSimpleName() == "InsertAction"
-                        or action.getClass().getSimpleName() == "UpdateAction"
-                    ):
-                        incoming_assignments = as_java_list(action.assignments())
-                        for assignment in incoming_assignments:
-                            (key_name, _) = map_single_column_expression(
-                                map_logical_plan_expression(assignment.key()),
-                                column_mapping=column_mapping_target,
-                                typer=typer_source,
-                            )
-
-                            (_, val_typ_col) = map_single_column_expression(
-                                map_logical_plan_expression(assignment.value()),
-                                column_mapping=column_mapping_source,
-                                typer=typer_target,
-                            )
-
-                            assignments[key_name] = val_typ_col.col
-                    elif (
-                        action.getClass().getSimpleName() == "InsertStarAction"
-                        or action.getClass().getSimpleName() == "UpdateStarAction"
-                    ):
-                        if len(column_mapping_source.columns) != len(
-                            column_mapping_target.columns
-                        ):
-                            raise ValueError(
-                                "source and target must have the same number of columns for InsertStarAction or UpdateStarAction"
-                            )
-                        for i, col in enumerate(column_mapping_target.columns):
-                            if assignments.get(col.snowpark_name) is not None:
-                                raise SnowparkConnectNotImplementedError(
-                                    "UpdateStarAction or InsertStarAction is not supported with duplicate columns."
-                                )
-                            assignments[col.snowpark_name] = snowpark_fn.col(
-                                column_mapping_source.columns[i].snowpark_name
-                            )
-                    return assignments
-
-                source_df_container = map_relation(
-                    map_logical_plan_relation(logical_plan.sourceTable())
-                )
-                source_df = source_df_container.dataframe
-                plan_id = gen_sql_plan_id()
-                target_df_container = map_relation(
-                    map_logical_plan_relation(logical_plan.targetTable(), plan_id)
-                )
-                target_df = target_df_container.dataframe
-
-                for col in target_df_container.column_map.columns:
-                    target_df = target_df.with_column_renamed(
-                        col.snowpark_name,
-                        spark_to_sf_single_id(col.spark_name, is_column=True),
-                    )
-                target_df_container = DataFrameContainer.create_with_column_mapping(
-                    dataframe=target_df,
-                    spark_column_names=target_df.columns,
-                    snowpark_column_names=target_df.columns,
-                )
-
-                set_plan_id_map(plan_id, target_df_container)
-
-                joined_df_before_condition: snowpark.DataFrame = source_df.join(
-                    target_df
-                )
-
-                column_mapping_for_conditions = column_name_handler.JoinColumnNameMap(
-                    source_df_container.column_map,
-                    target_df_container.column_map,
-                )
-                typer_for_expressions = ExpressionTyper(joined_df_before_condition)
-
-                (_, merge_condition_typed_col,) = map_single_column_expression(
-                    map_logical_plan_expression(logical_plan.mergeCondition()),
-                    column_mapping=column_mapping_for_conditions,
-                    typer=typer_for_expressions,
-                )
-
-                clauses = []
-
-                for matched_action in as_java_list(logical_plan.matchedActions()):
-                    condition = _get_condition_from_action(
-                        matched_action,
-                        column_mapping_for_conditions,
-                        typer_for_expressions,
-                    )
-                    if matched_action.getClass().getSimpleName() == "DeleteAction":
-                        clauses.append(when_matched(condition).delete())
-                    elif (
-                        matched_action.getClass().getSimpleName() == "UpdateAction"
-                        or matched_action.getClass().getSimpleName()
-                        == "UpdateStarAction"
-                    ):
-                        assignments = _get_assignments_from_action(
-                            matched_action,
-                            source_df_container.column_map,
-                            target_df_container.column_map,
-                            ExpressionTyper(source_df),
-                            ExpressionTyper(target_df),
-                        )
-                        clauses.append(when_matched(condition).update(assignments))
-
-                for not_matched_action in as_java_list(
-                    logical_plan.notMatchedActions()
-                ):
-                    condition = _get_condition_from_action(
-                        not_matched_action,
-                        column_mapping_for_conditions,
-                        typer_for_expressions,
-                    )
-                    if (
-                        not_matched_action.getClass().getSimpleName() == "InsertAction"
-                        or not_matched_action.getClass().getSimpleName()
-                        == "InsertStarAction"
-                    ):
-                        assignments = _get_assignments_from_action(
-                            not_matched_action,
-                            source_df_container.column_map,
-                            target_df_container.column_map,
-                            ExpressionTyper(source_df),
-                            ExpressionTyper(target_df),
-                        )
-                        clauses.append(when_not_matched(condition).insert(assignments))
-
-                if not as_java_list(logical_plan.notMatchedBySourceActions()).isEmpty():
-                    raise SnowparkConnectNotImplementedError(
-                        "Snowflake does not support 'not matched by source' actions in MERGE statements."
-                    )
-
-                if (
-                    logical_plan.targetTable().getClass().getSimpleName()
-                    == "UnresolvedRelation"
-                ):
-                    target_table_name = _spark_to_snowflake(
-                        logical_plan.targetTable().multipartIdentifier()
-                    )
-                else:
-                    target_table_name = _spark_to_snowflake(
-                        logical_plan.targetTable().child().multipartIdentifier()
-                    )
-                session.table(target_table_name).merge(
-                    source_df, merge_condition_typed_col.col, clauses
+                raise UnsupportedOperationException(
+                    "[UNSUPPORTED_SQL_EXTENSION] The MERGE INTO command failed.\n"
+                    + "Reason: This command is a platform-specific SQL extension and is not part of the standard Apache Spark specification that this interface uses."
                 )
             case "DeleteFromTable":
-                df_container = map_relation(
-                    map_logical_plan_relation(logical_plan.table())
-                )
-                df = df_container.dataframe
-                for col in df_container.column_map.columns:
-                    df = df.with_column_renamed(
-                        col.snowpark_name,
-                        spark_to_sf_single_id(col.spark_name, is_column=True),
-                    )
-                df_container = column_name_handler.create_with_column_mapping(
-                    dataframe=df,
-                    spark_column_names=df.columns,
-                    snowpark_column_names=df.columns,
+                raise UnsupportedOperationException(
+                    "[UNSUPPORTED_SQL_EXTENSION] The DELETE FROM command failed.\n"
+                    + "Reason: This command is a platform-specific SQL extension and is not part of the standard Apache Spark specification that this interface uses."
                 )
-                df = df_container.dataframe
-                name = get_relation_identifier_name(logical_plan.table(), True)
-                (
-                    condition_column_name,
-                    condition_typed_col,
-                ) = map_single_column_expression(
-                    map_logical_plan_expression(logical_plan.condition()),
-                    df_container.column_map,
-                    ExpressionTyper(df),
+            case "UpdateTable":
+                # Databricks/Delta-specific extension not supported by SAS.
+                # Provide an actionable, clear error.
+                raise UnsupportedOperationException(
+                    "[UNSUPPORTED_SQL_EXTENSION] The UPDATE TABLE command failed.\n"
+                    + "Reason: This command is a platform-specific SQL extension and is not part of the standard Apache Spark specification that this interface uses."
                 )
-                session.table(name).delete(condition_typed_col.col)
             case "RenameColumn":
                 table_name = get_relation_identifier_name(logical_plan.table(), True)
                 column_obj = logical_plan.column()

snowflake/snowpark_connect/relation/map_udtf.py

@@ -31,10 +31,7 @@ from snowflake.snowpark_connect.type_mapping import (
     proto_to_snowpark_type,
 )
 from snowflake.snowpark_connect.utils.context import push_udtf_context
-from snowflake.snowpark_connect.utils.session import (
-    get_or_create_snowpark_session,
-    get_python_udxf_import_files,
-)
+from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.udtf_helper import (
     SnowparkUDTF,
     create_udtf_in_sproc,
@@ -42,6 +39,9 @@ from snowflake.snowpark_connect.utils.udtf_helper import (
     udtf_check,
 )
 from snowflake.snowpark_connect.utils.udtf_utils import create_udtf
+from snowflake.snowpark_connect.utils.udxf_import_utils import (
+    get_python_udxf_import_files,
+)
 
 
 def build_expected_types_from_parsed(

snowflake/snowpark_connect/relation/read/map_read_json.py

@@ -6,6 +6,7 @@ import copy
 import json
 import typing
 from contextlib import suppress
+from datetime import datetime
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
@@ -21,6 +22,7 @@ from snowflake.snowpark.types import (
     StringType,
     StructField,
     StructType,
+    TimestampType,
 )
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read.map_read import JsonReaderConfig
@@ -204,6 +206,8 @@ def merge_row_schema(
             next_level_content = row[col_name]
             if next_level_content is not None:
                 with suppress(json.JSONDecodeError):
+                    if isinstance(next_level_content, datetime):
+                        next_level_content = str(next_level_content)
                     next_level_content = json.loads(next_level_content)
                 if isinstance(next_level_content, dict):
                     sf.datatype = merge_json_schema(
@@ -235,6 +239,9 @@ def merge_row_schema(
                             sf.datatype.element_type,
                             dropFieldIfAllNull,
                         )
+        elif isinstance(sf.datatype, TimestampType):
+            sf.datatype = StringType()
+            columns_with_valid_contents.add(col_name)
         elif row[col_name] is not None:
             columns_with_valid_contents.add(col_name)
 
@@ -265,7 +272,7 @@ def construct_dataframe_by_schema(
     rows: typing.Iterator[Row],
     session: snowpark.Session,
     snowpark_options: dict,
-    batch_size: int = 100,
+    batch_size: int = 1000,
 ) -> snowpark.DataFrame:
     result = None
 
@@ -280,6 +287,8 @@ def construct_dataframe_by_schema(
                 session,
             )
 
+            current_data = []
+
     if len(current_data) > 0:
         result = union_data_into_df(
             result,
@@ -288,6 +297,8 @@ def construct_dataframe_by_schema(
             session,
         )
 
+        current_data = []
+
     if result is None:
         raise ValueError("Dataframe cannot be empty")
     return result

snowflake/snowpark_connect/relation/read/reader_config.py

@@ -346,6 +346,7 @@ class JsonReaderConfig(ReaderWriterConfig):
                     "compression",
                     # "ignoreNullFields",
                     "rowsToInferSchema",
+                    # "inferTimestamp",
                 },
                 boolean_config_list=[
                     "multiLine",