snowpark-connect 0.24.0__py3-none-any.whl → 0.26.0__py3-none-any.whl

snowflake/snowpark_connect/relation/write/map_write.py

@@ -16,7 +16,7 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     unquote_if_quoted,
 )
 from snowflake.snowpark.exceptions import SnowparkSQLException
-from snowflake.snowpark.functions import col, lit, object_construct
+from snowflake.snowpark.functions import col, lit, object_construct, sql_expr
 from snowflake.snowpark.types import (
     ArrayType,
     DataType,
@@ -40,7 +40,10 @@ from snowflake.snowpark_connect.relation.io_utils import (
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.relation.read.reader_config import CsvWriterConfig
 from snowflake.snowpark_connect.relation.stage_locator import get_paths_from_stage
-from snowflake.snowpark_connect.relation.utils import random_string
+from snowflake.snowpark_connect.relation.utils import (
+    generate_spark_compatible_filename,
+    random_string,
+)
 from snowflake.snowpark_connect.type_mapping import snowpark_to_iceberg_type
 from snowflake.snowpark_connect.utils.context import get_session_id
 from snowflake.snowpark_connect.utils.identifiers import (
@@ -107,10 +110,7 @@ def _spark_to_snowflake(multipart_id: str) -> str:
 
 def map_write(request: proto_base.ExecutePlanRequest):
     write_op = request.plan.command.write_operation
-    if write_op.options is not None:
-        telemetry.report_io_write(write_op.source, dict(write_op.options))
-    else:
-        telemetry.report_io_write(write_op.source)
+    telemetry.report_io_write(write_op.source)
 
     write_mode = None
     match write_op.mode:
@@ -136,45 +136,80 @@ def map_write(request: proto_base.ExecutePlanRequest):
         write_op.source = ""
 
     should_write_to_single_file = str_to_bool(write_op.options.get("single", "false"))
-    if should_write_to_single_file:
-        # providing default size as 1GB
-        max_file_size = int(
-            write_op.options.get("snowflake_max_file_size", "1073741824")
-        )
+
+    # Support Snowflake-specific snowflake_max_file_size option. This is NOT a spark option.
+    max_file_size = None
+    if (
+        "snowflake_max_file_size" in write_op.options
+        and int(write_op.options["snowflake_max_file_size"]) > 0
+    ):
+        max_file_size = int(write_op.options["snowflake_max_file_size"])
+    elif should_write_to_single_file:
+        # providing default size as 1GB for single file write
+        max_file_size = 1073741824
     match write_op.source:
         case "csv" | "parquet" | "json" | "text":
             write_path = get_paths_from_stage(
                 [write_op.path],
                 session=session,
             )[0]
-            # we need a random prefix to support "append" mode
-            # otherwise copy into with overwrite=False will fail if the file already exists
+            # Generate Spark-compatible filename with proper extension
+            extension = write_op.source if write_op.source != "text" else "txt"
+
+            # Get compression from options for proper filename generation
+            compression_option = write_op.options.get("compression", "none")
+
+            # Generate Spark-compatible filename or prefix
             if should_write_to_single_file:
-                extention = write_op.source if write_op.source != "text" else "txt"
-                temp_file_prefix_on_stage = (
-                    f"{write_path}/{random_string(10, 'sas_file_')}.{extention}"
+                # Single file: generate complete filename with extension
+                spark_filename = generate_spark_compatible_filename(
+                    task_id=0,
+                    attempt_number=0,
+                    compression=compression_option,
+                    format_ext=extension,
                 )
+                temp_file_prefix_on_stage = f"{write_path}/{spark_filename}"
             else:
-                temp_file_prefix_on_stage = (
-                    f"{write_path}/{random_string(10, 'sas_file_')}"
+                # Multiple files: generate prefix without extension (Snowflake will add extensions)
+                spark_filename_prefix = generate_spark_compatible_filename(
+                    task_id=0,
+                    attempt_number=0,
+                    compression=compression_option,
+                    format_ext="",  # No extension for prefix
                 )
+                temp_file_prefix_on_stage = f"{write_path}/{spark_filename_prefix}"
             overwrite = (
                 write_op.mode
                 == commands_proto.WriteOperation.SaveMode.SAVE_MODE_OVERWRITE
             )
+
+            default_compression = "NONE" if write_op.source != "parquet" else "snappy"
+            compression = write_op.options.get(
+                "compression", default_compression
+            ).upper()
             parameters = {
                 "location": temp_file_prefix_on_stage,
                 "file_format_type": write_op.source
                 if write_op.source != "text"
                 else "csv",
                 "format_type_options": {
-                    "COMPRESSION": "NONE",
+                    "COMPRESSION": compression,
                 },
                 "overwrite": overwrite,
             }
-            if should_write_to_single_file:
-                parameters["single"] = True
+            # By default, download from the same prefix we wrote to.
+            download_stage_path = temp_file_prefix_on_stage
+
+            # Check for partition hint early to determine precedence over single option
+            partition_hint = result.partition_hint
+
+            # Apply max_file_size for both single and multi-file scenarios
+            # This helps control when Snowflake splits files into multiple parts
+            if max_file_size:
                 parameters["max_file_size"] = max_file_size
+            # Only apply single option if no partition hint is present (partition hint takes precedence)
+            if should_write_to_single_file and partition_hint is None:
+                parameters["single"] = True
             rewritten_df: snowpark.DataFrame = rewrite_df(input_df, write_op.source)
             get_param_from_options(parameters, write_op.options, write_op.source)
             if write_op.partitioning_columns:
@@ -189,10 +224,50 @@ def map_write(request: proto_base.ExecutePlanRequest):
                     )
                 else:
                     parameters["partition_by"] = partitioning_columns[0]
-            rewritten_df.write.copy_into_location(**parameters)
+
+            # If a partition hint is present (from DataFrame.repartition(n)), optionally split the
+            # write into n COPY INTO calls by assigning a synthetic partition id. Controlled by config.
+            # Note: This affects only the number of output files, not computation semantics.
+            # Partition hints take precedence over single option (matches Spark behavior) when enabled.
+            repartition_for_writes_enabled = (
+                global_config.snowflake_repartition_for_writes
+            )
+            if repartition_for_writes_enabled and partition_hint and partition_hint > 0:
+                # Create a stable synthetic file number per row using ROW_NUMBER() over a
+                # randomized order, then modulo partition_hint. We rely on sql_expr to avoid
+                # adding new helpers.
+                file_num_col = "_sas_file_num"
+                partitioned_df = rewritten_df.withColumn(
+                    file_num_col,
+                    sql_expr(
+                        f"(ROW_NUMBER() OVER (ORDER BY RANDOM())) % {partition_hint}"
+                    ),
+                )
+
+                # Execute multiple COPY INTO operations, one per target file.
+                # Since we write per-partition with distinct prefixes, download from the base write path.
+                download_stage_path = write_path
+                for part_idx in range(partition_hint):
+                    part_params = dict(parameters)
+                    # Preserve Spark-like filename prefix per partition so downloaded basenames
+                    # match the expected Spark pattern (with possible Snowflake counters appended).
+                    per_part_prefix = generate_spark_compatible_filename(
+                        task_id=part_idx,
+                        attempt_number=0,
+                        compression=compression_option,
+                        format_ext="",  # prefix only; Snowflake appends extension/counters
+                    )
+                    part_params["location"] = f"{write_path}/{per_part_prefix}"
+                    (
+                        partitioned_df.filter(col(file_num_col) == lit(part_idx))
+                        .drop(file_num_col)
+                        .write.copy_into_location(**part_params)
+                    )
+            else:
+                rewritten_df.write.copy_into_location(**parameters)
             if not is_cloud_path(write_op.path):
                 store_files_locally(
-                    temp_file_prefix_on_stage,
+                    download_stage_path,
                     write_op.path,
                     overwrite,
                     session,
@@ -234,10 +309,11 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         column_order=_column_order_for_write,
                     )
                 case "append":
-                    if check_table_type(snowpark_table_name, session) != "ICEBERG":
-                        raise AnalysisException(
-                            f"Table {snowpark_table_name} is not an iceberg table"
-                        )
+                    # TODO: SNOW-2299414 Fix the implementation of table type check
+                    # if check_table_type(snowpark_table_name, session) != "ICEBERG":
+                    #     raise AnalysisException(
+                    #         f"Table {snowpark_table_name} is not an iceberg table"
+                    #     )
                     _validate_schema_and_get_writer(
                         input_df, "append", snowpark_table_name
                     ).saveAsTable(
@@ -264,10 +340,12 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         )
                 case "overwrite":
                     if check_snowflake_table_existence(snowpark_table_name, session):
-                        if check_table_type(snowpark_table_name, session) != "ICEBERG":
-                            raise AnalysisException(
-                                f"Table {snowpark_table_name} is not an iceberg table"
-                            )
+                        # TODO: SNOW-2299414 Fix the implementation of table type check
+                        # if check_table_type(snowpark_table_name, session) != "ICEBERG":
+                        #     raise AnalysisException(
+                        #         f"Table {snowpark_table_name} is not an iceberg table"
+                        #     )
+                        pass
                     else:
                         create_iceberg_table(
                             snowpark_table_name=snowpark_table_name,
@@ -298,13 +376,14 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         if check_snowflake_table_existence(
                             snowpark_table_name, session
                         ):
-                            if (
-                                check_table_type(snowpark_table_name, session)
-                                != "TABLE"
-                            ):
-                                raise AnalysisException(
-                                    f"Table {snowpark_table_name} is not a FDN table"
-                                )
+                            # TODO: SNOW-2299414 Fix the implementation of table type check
+                            # if (
+                            #     check_table_type(snowpark_table_name, session)
+                            #     != "TABLE"
+                            # ):
+                            #     raise AnalysisException(
+                            #         f"Table {snowpark_table_name} is not a FDN table"
+                            #     )
                             write_mode = "truncate"
                         _validate_schema_and_get_writer(
                             input_df, write_mode, snowpark_table_name
@@ -314,10 +393,11 @@ def map_write(request: proto_base.ExecutePlanRequest):
                             column_order=_column_order_for_write,
                         )
                     case "append":
-                        if check_table_type(snowpark_table_name, session) != "TABLE":
-                            raise AnalysisException(
-                                f"Table {snowpark_table_name} is not a FDN table"
-                            )
+                        # TODO: SNOW-2299414 Fix the implementation of table type check
+                        # if check_table_type(snowpark_table_name, session) != "TABLE":
+                        #     raise AnalysisException(
+                        #         f"Table {snowpark_table_name} is not a FDN table"
+                        #     )
                         _validate_schema_and_get_writer(
                             input_df, write_mode, snowpark_table_name
                         ).saveAsTable(
@@ -388,10 +468,11 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     raise AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found."
                     )
-                if check_table_type(snowpark_table_name, session) != "ICEBERG":
-                    raise AnalysisException(
-                        f"Table {snowpark_table_name} is not an iceberg table"
-                    )
+                # TODO: SNOW-2299414 Fix the implementation of table type check
+                # if check_table_type(snowpark_table_name, session) != "ICEBERG":
+                #     raise AnalysisException(
+                #         f"Table {snowpark_table_name} is not an iceberg table"
+                #     )
                 _validate_schema_and_get_writer(
                     input_df, "append", snowpark_table_name
                 ).saveAsTable(
@@ -402,10 +483,12 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
             case commands_proto.WriteOperationV2.MODE_OVERWRITE | commands_proto.WriteOperationV2.MODE_OVERWRITE_PARTITIONS:
                 # TODO: handle the filter condition for MODE_OVERWRITE
                 if check_snowflake_table_existence(snowpark_table_name, session):
-                    if check_table_type(snowpark_table_name, session) != "ICEBERG":
-                        raise AnalysisException(
-                            f"Table {snowpark_table_name} is not an iceberg table"
-                        )
+                    # TODO: SNOW-2299414 Fix the implementation of table type check
+                    # if check_table_type(snowpark_table_name, session) != "ICEBERG":
+                    #     raise AnalysisException(
+                    #         f"Table {snowpark_table_name} is not an iceberg table"
+                    #     )
+                    pass
                 else:
                     raise AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
@@ -471,10 +554,11 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     raise AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found."
                     )
-                if check_table_type(snowpark_table_name, session) != "TABLE":
-                    raise AnalysisException(
-                        f"Table {snowpark_table_name} is not a FDN table"
-                    )
+                # TODO: SNOW-2299414 Fix the implementation of table type check
+                # if check_table_type(snowpark_table_name, session) != "TABLE":
+                #     raise AnalysisException(
+                #         f"Table {snowpark_table_name} is not a FDN table"
+                #     )
                 _validate_schema_and_get_writer(
                     input_df, "append", snowpark_table_name
                 ).saveAsTable(
@@ -485,10 +569,12 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
             case commands_proto.WriteOperationV2.MODE_OVERWRITE | commands_proto.WriteOperationV2.MODE_OVERWRITE_PARTITIONS:
                 # TODO: handle the filter condition for MODE_OVERWRITE
                 if check_snowflake_table_existence(snowpark_table_name, session):
-                    if check_table_type(snowpark_table_name, session) != "TABLE":
-                        raise AnalysisException(
-                            f"Table {snowpark_table_name} is not a FDN table"
-                        )
+                    # TODO: SNOW-2299414 Fix the implementation of table type check
+                    # if check_table_type(snowpark_table_name, session) != "TABLE":
+                    #     raise AnalysisException(
+                    #         f"Table {snowpark_table_name} is not a FDN table"
+                    #     )
+                    pass
                 else:
                     raise AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
@@ -561,7 +647,12 @@ def _validate_schema_and_get_writer(
             col_name = field.name
             renamed = col_name
             matching_field = next(
-                (f for f in table_schema.fields if f.name.lower() == col_name.lower()),
+                (
+                    f
+                    for f in table_schema.fields
+                    if unquote_if_quoted(f.name).lower()
+                    == unquote_if_quoted(col_name).lower()
+                ),
                 None,
             )
             if matching_field is not None and matching_field != col_name:
@@ -583,7 +674,10 @@ def _validate_schema_and_get_writer(
 
 
 def _validate_schema_for_append(
-    table_schema: DataType, data_schema: DataType, snowpark_table_name: str
+    table_schema: DataType,
+    data_schema: DataType,
+    snowpark_table_name: str,
+    compare_structs: bool = False,
 ):
     match (table_schema, data_schema):
         case (_, _) if table_schema == data_schema:
@@ -592,7 +686,11 @@ def _validate_schema_for_append(
         case (StructType() as table_struct, StructType() as data_struct):
 
             def _comparable_col_name(col: str) -> str:
-                return col if global_config.spark_sql_caseSensitive else col.lower()
+                name = col if global_config.spark_sql_caseSensitive else col.lower()
+                if compare_structs:
+                    return name
+                else:
+                    return unquote_if_quoted(name)
 
             def invalid_struct_schema():
                 raise AnalysisException(
@@ -632,6 +730,7 @@ def _validate_schema_for_append(
                         matching_table_field.datatype,
                         data_field.datatype,
                         snowpark_table_name,
+                        compare_structs=True,
                     )
 
             return
@@ -798,17 +897,18 @@ def check_snowflake_table_existence(
         return False
 
 
-def check_table_type(
-    snowpark_table_name: str,
-    snowpark_session: snowpark.Session,
-) -> str:
-    # currently we only support iceberg table and FDN table
-    metadata = snowpark_session.sql(
-        f"SHOW TABLES LIKE '{unquote_if_quoted(snowpark_table_name)}';"
-    ).collect()
-    if metadata is None or len(metadata) == 0:
-        raise AnalysisException(f"Table {snowpark_table_name} does not exist")
-    metadata = metadata[0]
-    if metadata.as_dict().get("is_iceberg") == "Y":
-        return "ICEBERG"
-    return "TABLE"
+# TODO: SNOW-2299414 Fix the implementation of table type check
+# def check_table_type(
+#     snowpark_table_name: str,
+#     snowpark_session: snowpark.Session,
+# ) -> str:
+#     # currently we only support iceberg table and FDN table
+#     metadata = snowpark_session.sql(
+#         f"SHOW TABLES LIKE '{unquote_if_quoted(snowpark_table_name)}';"
+#     ).collect()
+#     if metadata is None or len(metadata) == 0:
+#         raise AnalysisException(f"Table {snowpark_table_name} does not exist")
+#     metadata = metadata[0]
+#     if metadata.as_dict().get("is_iceberg") == "Y":
+#         return "ICEBERG"
+#     return "TABLE"

snowflake/snowpark_connect/resources_initializer.py

@@ -1,7 +1,7 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-
+import pathlib
 import threading
 import time
 
@@ -9,6 +9,9 @@ from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_sess
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 
 _resources_initialized = threading.Event()
+_initializer_lock = threading.Lock()
+SPARK_VERSION = "3.5.6"
+RESOURCE_PATH = "/snowflake/snowpark_connect/resources"
 
 
 def initialize_resources() -> None:
@@ -41,6 +44,34 @@ def initialize_resources() -> None:
 
         session.sql("select 1 as sf_connection_warm_up").collect()
 
+    def upload_scala_udf_jars() -> None:
+        """Upload Spark jar files required for creating Scala UDFs."""
+        stage = session.get_session_stage()
+        resource_path = stage + RESOURCE_PATH
+        import snowflake
+
+        pyspark_jars = (
+            pathlib.Path(snowflake.snowpark_connect.__file__).parent / "includes/jars"
+        )
+        jar_files = [
+            f"spark-sql_2.12-{SPARK_VERSION}.jar",
+            f"spark-connect-client-jvm_2.12-{SPARK_VERSION}.jar",
+            f"spark-common-utils_2.12-{SPARK_VERSION}.jar",
+            "json4s-ast_2.12-3.7.0-M11.jar",
+            "json4s-native_2.12-3.7.0-M11.jar",
+            "json4s-core_2.12-3.7.0-M11.jar",
+            "paranamer-2.8.3.jar",
+        ]
+
+        for jar in jar_files:
+            session.file.put(
+                str(pyspark_jars) + "/" + jar,
+                resource_path,
+                auto_compress=False,
+                overwrite=False,
+                source_compression="NONE",
+            )
+
     start_time = time.time()
 
     resources = [
@@ -49,6 +80,7 @@ def initialize_resources() -> None:
         ("Initialize Session Stage", initialize_session_stage),  # Takes about 0.3s
         ("Initialize Session Catalog", initialize_catalog),  # Takes about 1.2s
         ("Snowflake Connection Warm Up", warm_up_sf_connection),  # Takes about 1s
+        ("Upload Scala UDF Jars", upload_scala_udf_jars),
     ]
 
     for name, resource_func in resources:
@@ -66,10 +98,19 @@ def initialize_resources() -> None:
     logger.info(f"All resources initialized in {time.time() - start_time:.2f}s")
 
 
+_resource_initializer = threading.Thread(
+    target=initialize_resources, name="ResourceInitializer"
+)
+
+
 def initialize_resources_async() -> threading.Thread:
     """Start resource initialization in background."""
-    thread = threading.Thread(
-        target=initialize_resources, name="ResourceInitializer", daemon=True
-    )
-    thread.start()
-    return thread
+    with _initializer_lock:
+        if not _resource_initializer.is_alive() and _resource_initializer.ident is None:
+            _resource_initializer.start()
+        return _resource_initializer
+
+
+def wait_for_resource_initialization() -> None:
+    with _initializer_lock:
+        _resource_initializer.join()

snowflake/snowpark_connect/server.py

@@ -68,10 +68,14 @@ from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.relation.utils import get_semantic_string
 from snowflake.snowpark_connect.resources_initializer import initialize_resources_async
 from snowflake.snowpark_connect.type_mapping import (
-    parse_ddl_string,
+    map_type_string_to_proto,
     snowpark_to_proto_type,
 )
-from snowflake.snowpark_connect.utils.artifacts import check_checksum, write_artifact
+from snowflake.snowpark_connect.utils.artifacts import (
+    check_checksum,
+    write_artifact,
+    write_class_files_to_stage,
+)
 from snowflake.snowpark_connect.utils.cache import (
     df_cache_map_get,
     df_cache_map_pop,
@@ -84,6 +88,9 @@ from snowflake.snowpark_connect.utils.context import (
     set_spark_version,
 )
 from snowflake.snowpark_connect.utils.env_utils import get_int_from_env
+from snowflake.snowpark_connect.utils.external_udxf_cache import (
+    clear_external_udxf_cache,
+)
 from snowflake.snowpark_connect.utils.interrupt import (
     interrupt_all_queries,
     interrupt_queries_with_tag,
@@ -249,7 +256,9 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
                     return proto_base.AnalyzePlanResponse(
                         session_id=request.session_id,
                         ddl_parse=proto_base.AnalyzePlanResponse.DDLParse(
-                            parsed=parse_ddl_string(request.ddl_parse.ddl_string)
+                            parsed=map_type_string_to_proto(
+                                request.ddl_parse.ddl_string
+                            )
                         ),
                     )
                 case "get_storage_level":
@@ -430,7 +439,8 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
                         lambda: map_local_relation(relation),  # noqa: B023
                         materialize=True,
                     )
-                except Exception:
+                except Exception as e:
+                    logger.warning("Failed to put df into cache: %s", str(e))
                     # fallback - treat as regular artifact
                     _handle_regular_artifact()
             else:
@@ -516,7 +526,13 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
         for name, data in cache_data.items():
             _try_handle_local_relation(name, bytes(data))
 
+        class_files: dict[str, str] = {}
         for (name, filepath) in filenames.items():
+            if name.endswith(".class"):
+                # name is <dir>/<package>/<class_name>
+                # we don't need the dir name, but require the package, so only remove dir
+                class_files[name.split("/", 1)[-1]] = filepath
+                continue
             session.file.put(
                 filepath,
                 session.get_session_stage(),
@@ -541,6 +557,12 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
             if not name.startswith("pyfiles"):
                 session._import_files.add(cached_name)
 
+        if class_files:
+            write_class_files_to_stage(session, class_files)
+
+        if any(not name.startswith("cache") for name in filenames.keys()):
+            clear_external_udxf_cache(session)
+
         return proto_base.AddArtifactsResponse(artifacts=list(response.values()))
 
     def ArtifactStatus(self, request, context):

snowflake/snowpark_connect/type_mapping.py

@@ -19,8 +19,9 @@ from pyspark.sql.connect.proto import expressions_pb2
 
 from snowflake import snowpark
 from snowflake.snowpark._internal.utils import quote_name
-from snowflake.snowpark.types import TimestampTimeZone
+from snowflake.snowpark.types import TimestampTimeZone, TimestampType
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+from snowflake.snowpark_connect.config import get_timestamp_type
 from snowflake.snowpark_connect.constants import (
     COLUMN_METADATA_COLLISION_KEY,
     STRUCTURED_TYPES_ENABLED,
@@ -29,6 +30,7 @@ from snowflake.snowpark_connect.date_time_format_mapping import (
     convert_spark_format_to_snowflake,
 )
 from snowflake.snowpark_connect.expression.literal import get_literal_field_and_name
+from snowflake.snowpark_connect.utils.context import get_is_evaluating_sql
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
@@ -63,7 +65,7 @@ def get_python_sql_utils_class():
     return jpype.JClass("org.apache.spark.sql.api.python.PythonSQLUtils")
 
 
-def parse_ddl_with_spark_scala(ddl_string: str) -> pyspark.sql.types.DataType:
+def _parse_ddl_with_spark_scala(ddl_string: str) -> pyspark.sql.types.DataType:
     """
     Parse DDL string using PySpark's Scala StructType.fromDDL() method.
 
@@ -869,51 +871,53 @@ def map_json_schema_to_snowpark(
             return map_simple_types(schema["type"])
 
 
-def map_type_string_to_snowpark_type(type_string: str) -> snowpark.types.DataType:
+def _map_type_string(
+    type_string: str,
+) -> tuple[pyspark.sql.types.DataType, snowpark.types.DataType]:
     """
-    Converts a pyspark type string like x: int or struct<x: int, y: string> etc. to a snowpark type.
+    Map a ddl-string type to a tuple of PySpark data type and a Snowpark data type.
     """
-    pyspark_type = parse_ddl_with_spark_scala(type_string)
-    return map_pyspark_types_to_snowpark_types(pyspark_type)
+
+    pyspark_type = _parse_ddl_with_spark_scala(type_string)
+    match type_string:
+        case "timestamp" if not get_is_evaluating_sql():
+            return pyspark_type, get_timestamp_type()
+        case "timestamp" | "timestamp_ltz":
+            return pyspark_type, TimestampType(TimestampTimeZone.LTZ)
+        case _:
+            return pyspark_type, map_pyspark_types_to_snowpark_types(pyspark_type)
 
 
-def preprocess_tinyint_to_byte(ddl_string):
-    if ddl_string.strip().lower() == "tinyint":
-        # If the ddl_string is just "tinyint", we can directly return "byte"
-        return "byte"
-    # Replace any occurrence of ": tinyint" (with optional spaces) with ": byte"
-    return re.sub(r"(:\s*)tinyint\b", r"\1byte", ddl_string, flags=re.IGNORECASE)
+def map_type_string_to_snowpark_type(type_string: str) -> snowpark.types.DataType:
+    _, snowpark_type = _map_type_string(type_string)
+    return snowpark_type
 
 
-def parse_ddl_string(
-    ddl_string: str, return_raw_type: bool = False
+def map_type_string_to_proto(
+    ddl_string: str,
 ) -> types_proto.DataType | snowpark.types.DataType:
     """
     Parse a DDL string and return the Proto data type.
-
-    This function is used to parse a DDL string and return the Proto data type.
     """
-    pyspark_type = parse_ddl_with_spark_scala(ddl_string)
-    schema = map_pyspark_types_to_snowpark_types(pyspark_type)
-
-    if return_raw_type:
-        return schema
+    pyspark_type, snowpark_type = _map_type_string(ddl_string)
 
     column_name_map = None
     if isinstance(pyspark_type, pyspark.sql.types.StructType):
         spark_column_names = [field.name for field in pyspark_type.fields]
-        snowpark_column_names = [field.name for field in schema.fields]
+        snowpark_column_names = [field.name for field in snowpark_type.fields]
 
         if len(spark_column_names) == len(snowpark_column_names):
             column_name_map = ColumnNameMap(spark_column_names, snowpark_column_names)
 
-    struct_proto = list(snowpark_to_proto_type(schema, column_name_map).values())[0]
+    struct_proto = list(
+        snowpark_to_proto_type(snowpark_type, column_name_map).values()
+    )[0]
 
-    if isinstance(schema, snowpark.types.StructType):
+    if isinstance(snowpark_type, snowpark.types.StructType):
         return types_proto.DataType(struct=struct_proto)
     else:
         # Handle non-struct types (arrays, maps, primitives, etc.)
-        proto_type_mapping = snowpark_to_proto_type(schema, column_name_map)
+        proto_type_mapping = snowpark_to_proto_type(snowpark_type, column_name_map)
         proto_field_name = list(proto_type_mapping.keys())[0]
         return types_proto.DataType(**{proto_field_name: struct_proto})