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

snowflake/snowpark_connect/relation/map_extension.py

@@ -374,23 +374,31 @@ def map_aggregate(
     snowpark_columns: list[str] = []
     snowpark_column_types: list[snowpark_types.DataType] = []
 
-    def _add_column(spark_name: str, snowpark_column: TypedColumn) -> snowpark.Column:
-        alias = make_column_names_snowpark_compatible(
-            [spark_name], plan_id, len(spark_columns)
-        )[0]
+    # Use grouping columns directly without aliases
+    groupings = [col.col for _, col in raw_groupings]
+
+    # Create aliases only for aggregation columns
+    aggregations = []
+    for i, (spark_name, snowpark_column) in enumerate(raw_aggregations):
+        alias = make_column_names_snowpark_compatible([spark_name], plan_id, i)[0]
 
         spark_columns.append(spark_name)
         snowpark_columns.append(alias)
         snowpark_column_types.append(snowpark_column.typ)
 
-        return snowpark_column.col.alias(alias)
-
-    groupings = [_add_column(name, col) for name, col in raw_groupings]
-    aggregations = [_add_column(name, col) for name, col in raw_aggregations]
+        aggregations.append(snowpark_column.col.alias(alias))
 
     match aggregate.group_type:
         case snowflake_proto.Aggregate.GROUP_TYPE_GROUPBY:
-            result = input_df.group_by(groupings)
+            if groupings:
+                # Normal GROUP BY with explicit grouping columns
+                result = input_df.group_by(groupings)
+            else:
+                # No explicit GROUP BY - this is an aggregate over the entire table
+                # Use a dummy constant that will be excluded from the final result
+                result = input_df.with_column(
+                    "__dummy_group__", snowpark_fn.lit(1)
+                ).group_by("__dummy_group__")
         case snowflake_proto.Aggregate.GROUP_TYPE_ROLLUP:
             result = input_df.rollup(groupings)
         case snowflake_proto.Aggregate.GROUP_TYPE_CUBE:
@@ -410,28 +418,54 @@ def map_aggregate(
                 f"Unsupported GROUP BY type: {other}"
             )
 
-    result = result.agg(*aggregations)
+    result = result.agg(*aggregations, exclude_grouping_columns=True)
+
+    # If we added a dummy grouping column, make sure it's excluded
+    if not groupings and "__dummy_group__" in result.columns:
+        result = result.drop("__dummy_group__")
+
+    # Apply HAVING condition if present
+    if aggregate.HasField("having_condition"):
+        from snowflake.snowpark_connect.expression.hybrid_column_map import (
+            create_hybrid_column_map_for_having,
+        )
+
+        # Create aggregated DataFrame column map
+        aggregated_column_map = DataFrameContainer.create_with_column_mapping(
+            dataframe=result,
+            spark_column_names=spark_columns,
+            snowpark_column_names=snowpark_columns,
+            snowpark_column_types=snowpark_column_types,
+        ).column_map
+
+        # Create hybrid column map that can resolve both input and aggregate contexts
+        hybrid_map = create_hybrid_column_map_for_having(
+            input_df=input_df,
+            input_column_map=input_container.column_map,
+            aggregated_df=result,
+            aggregated_column_map=aggregated_column_map,
+            aggregate_expressions=list(aggregate.aggregate_expressions),
+            grouping_expressions=list(aggregate.grouping_expressions),
+            spark_columns=spark_columns,
+            raw_aggregations=raw_aggregations,
+        )
+
+        # Map the HAVING condition using hybrid resolution
+        _, having_column = hybrid_map.resolve_expression(aggregate.having_condition)
+
+        # Apply the HAVING filter
+        result = result.filter(having_column.col)
 
     if aggregate.group_type == snowflake_proto.Aggregate.GROUP_TYPE_GROUPING_SETS:
         # Immediately drop extra columns. Unlike other GROUP BY operations,
         # grouping sets don't allow ORDER BY with columns that aren't in the aggregate list.
-        result = result.select(result.columns[-len(spark_columns) :])
+        result = result.select(result.columns[-len(aggregations) :])
 
-    # Build a parent column map that includes groupings.
-    result_container = DataFrameContainer.create_with_column_mapping(
+    # Return only aggregation columns in the column map
+    return DataFrameContainer.create_with_column_mapping(
         dataframe=result,
         spark_column_names=spark_columns,
         snowpark_column_names=snowpark_columns,
         snowpark_column_types=snowpark_column_types,
-    )
-
-    # Drop the groupings.
-    grouping_count = len(groupings)
-
-    return DataFrameContainer.create_with_column_mapping(
-        result.drop(snowpark_columns[:grouping_count]),
-        spark_columns[grouping_count:],
-        snowpark_columns[grouping_count:],
-        snowpark_column_types[grouping_count:],
-        parent_column_name_map=result_container.column_map,
+        parent_column_name_map=input_df._column_map,
     )

snowflake/snowpark_connect/relation/map_local_relation.py

@@ -4,6 +4,7 @@
 
 import json
 import re
+from json import JSONDecodeError
 
 import numpy as np
 import pyarrow as pa
@@ -19,6 +20,7 @@ from snowflake.snowpark_connect.column_name_handler import (
 )
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.type_mapping import (
+    get_python_sql_utils_class,
     map_json_schema_to_snowpark,
     map_pyarrow_to_snowpark_types,
     map_simple_types,
@@ -34,7 +36,12 @@ def parse_local_relation_schema_string(rel: relation_proto.Relation):
     # schema_str can be a dict, or just a type string, e.g. INTEGER.
     schema_str = rel.local_relation.schema
     assert schema_str
-    schema_dict = json.loads(schema_str)
+    try:
+        schema_dict = json.loads(schema_str)
+    except JSONDecodeError:
+        # Legacy scala clients sends unparsed struct type strings like "struct<id:bigint,a:int,b:double>"
+        spark_datatype = get_python_sql_utils_class().parseDataType(schema_str)
+        schema_dict = json.loads(spark_datatype.json())
 
     column_metadata = {}
     if isinstance(schema_dict, dict):

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_row_ops.py

@@ -1,6 +1,7 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
+from copy import copy
 
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
@@ -8,6 +9,7 @@ from pyspark.errors.exceptions.base import AnalysisException, IllegalArgumentExc
 
 import snowflake.snowpark_connect.relation.utils as utils
 from snowflake import snowpark
+from snowflake.snowpark._internal.analyzer.binary_expression import And
 from snowflake.snowpark.functions import col, expr as snowpark_expr
 from snowflake.snowpark.types import (
     BooleanType,
@@ -29,6 +31,7 @@ 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.relation.utils import can_filter_be_flattened
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -551,7 +554,33 @@ def map_filter(
     _, condition = map_single_column_expression(
         rel.filter.condition, input_container.column_map, typer
     )
-    result = input_df.filter(condition.col)
+
+    select_statement = getattr(input_df, "_select_statement", None)
+    condition_exp = condition.col._expression
+    if (
+        can_filter_be_flattened(select_statement, condition_exp)
+        and input_df._ops_after_agg is None
+    ):
+        new = copy(select_statement)
+        new.from_ = select_statement.from_.to_subqueryable()
+        new.pre_actions = new.from_.pre_actions
+        new.post_actions = new.from_.post_actions
+        new.column_states = select_statement.column_states
+        new.where = (
+            And(select_statement.where, condition_exp)
+            if select_statement.where is not None
+            else condition_exp
+        )
+        new._merge_projection_complexity_with_subquery = False
+        new.df_ast_ids = (
+            select_statement.df_ast_ids.copy()
+            if select_statement.df_ast_ids is not None
+            else None
+        )
+        new.attributes = select_statement.attributes
+        result = input_df._with_plan(new)
+    else:
+        result = input_df.filter(condition.col)
 
     return DataFrameContainer(
         result,

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,
@@ -76,6 +77,9 @@ from ..expression.map_sql_expression import (
 from ..utils.identifiers import spark_to_sf_single_id
 
 _ctes = ContextVar[dict[str, relation_proto.Relation]]("_ctes", default={})
+_having_condition = ContextVar[expressions_proto.Expression | None](
+    "_having_condition", default=None
+)
 
 
 def _is_sql_select_statement_helper(sql_string: str) -> bool:
@@ -163,6 +167,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 +717,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()
@@ -1319,6 +1149,7 @@ def map_logical_plan_relation(
                             grouping_expressions=grouping_expressions,
                             aggregate_expressions=aggregate_expressions,
                             grouping_sets=grouping_sets,
+                            having_condition=_having_condition.get(),
                         )
                     )
                 )
@@ -1562,12 +1393,25 @@ def map_logical_plan_relation(
                 )
             )
         case "UnresolvedHaving":
-            proto = relation_proto.Relation(
-                filter=relation_proto.Filter(
-                    input=map_logical_plan_relation(rel.child()),
-                    condition=map_logical_plan_expression(rel.havingCondition()),
+            # Store the having condition in context and process the child aggregate
+            child_relation = rel.child()
+            if str(child_relation.getClass().getSimpleName()) != "Aggregate":
+                raise SnowparkConnectNotImplementedError(
+                    "UnresolvedHaving can only be applied to Aggregate relations"
                 )
-            )
+
+            # Store having condition in a context variable for the Aggregate case to pick up
+            having_condition = map_logical_plan_expression(rel.havingCondition())
+
+            # Store in thread-local context (similar to how _ctes works)
+            token = _having_condition.set(having_condition)
+
+            try:
+                # Recursively call map_logical_plan_relation on the child Aggregate
+                # The Aggregate case will pick up the having condition from context
+                proto = map_logical_plan_relation(child_relation, plan_id)
+            finally:
+                _having_condition.reset(token)
         case "UnresolvedHint":
             proto = relation_proto.Relation(
                 hint=relation_proto.Hint(

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.py

@@ -95,7 +95,8 @@ def map_read(
             if len(rel.read.data_source.paths) > 0:
                 # Normalize paths to ensure consistent behavior
                 clean_source_paths = [
-                    str(Path(path)) for path in rel.read.data_source.paths
+                    path.rstrip("/") if is_cloud_path(path) else str(Path(path))
+                    for path in rel.read.data_source.paths
                 ]
 
                 result = _read_file(

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/map_read_parquet.py

@@ -54,10 +54,17 @@ def map_read_parquet(
     if len(paths) == 1:
         df = _read_parquet_with_partitions(session, reader, paths[0])
     else:
+        is_merge_schema = options.config.get("mergeschema")
         df = _read_parquet_with_partitions(session, reader, paths[0])
+        schema_cols = df.columns
         for p in paths[1:]:
             reader._user_schema = None
-            df = df.union_all(_read_parquet_with_partitions(session, reader, p))
+            df = df.union_all_by_name(
+                _read_parquet_with_partitions(session, reader, p),
+                allow_missing_columns=True,
+            )
+        if not is_merge_schema:
+            df = df.select(*schema_cols)
 
     renamed_df, snowpark_column_names = rename_columns_as_snowflake_standard(
         df, rel.common.plan_id

snowflake/snowpark_connect/relation/read/reader_config.py

@@ -346,6 +346,7 @@ class JsonReaderConfig(ReaderWriterConfig):
                     "compression",
                     # "ignoreNullFields",
                     "rowsToInferSchema",
+                    # "inferTimestamp",
                 },
                 boolean_config_list=[
                     "multiLine",
@@ -397,3 +398,12 @@ class ParquetReaderConfig(ReaderWriterConfig):
             ),
             options,
         )
+
+    def convert_to_snowpark_args(self) -> dict[str, Any]:
+        snowpark_args = super().convert_to_snowpark_args()
+
+        # Should be determined by spark.sql.parquet.binaryAsString, but currently Snowpark Connect only supports
+        # the default value (false). TODO: Add support for spark.sql.parquet.binaryAsString equal to "true".
+        snowpark_args["BINARY_AS_TEXT"] = False
+
+        return snowpark_args