snowpark-connect 0.27.0__py3-none-any.whl → 1.7.0__py3-none-any.whl

snowflake/snowpark_connect/type_mapping.py

@@ -8,6 +8,7 @@ import typing
 from contextlib import suppress
 from datetime import datetime
 from functools import cache
+from typing import Union
 
 import jpype
 import pyarrow as pa
@@ -18,6 +19,7 @@ from pyspark.errors.exceptions.base import AnalysisException
 from pyspark.sql.connect.proto import expressions_pb2
 
 from snowflake import snowpark
+from snowflake.snowpark import types as snowpark_type
 from snowflake.snowpark._internal.utils import quote_name
 from snowflake.snowpark.types import TimestampTimeZone, TimestampType
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
@@ -29,8 +31,17 @@ from snowflake.snowpark_connect.constants import (
 from snowflake.snowpark_connect.date_time_format_mapping import (
     convert_spark_format_to_snowflake,
 )
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 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.expression.map_sql_expression import (
+    _INTERVAL_DAYTIME_PATTERN_RE,
+    _INTERVAL_YEARMONTH_PATTERN_RE,
+)
+from snowflake.snowpark_connect.utils.context import (
+    get_is_evaluating_sql,
+    get_jpype_jclass_lock,
+)
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
@@ -57,12 +68,14 @@ SNOWPARK_TYPE_NAME_TO_PYSPARK_TYPE_NAME = {
 
 @cache
 def _get_struct_type_class():
-    return jpype.JClass("org.apache.spark.sql.types.StructType")
+    with get_jpype_jclass_lock():
+        return jpype.JClass("org.apache.spark.sql.types.StructType")
 
 
 @cache
 def get_python_sql_utils_class():
-    return jpype.JClass("org.apache.spark.sql.api.python.PythonSQLUtils")
+    with get_jpype_jclass_lock():
+        return jpype.JClass("org.apache.spark.sql.api.python.PythonSQLUtils")
 
 
 def _parse_ddl_with_spark_scala(ddl_string: str) -> pyspark.sql.types.DataType:
@@ -93,7 +106,7 @@ def _parse_ddl_with_spark_scala(ddl_string: str) -> pyspark.sql.types.DataType:
 
 def snowpark_to_proto_type(
     data_type: snowpark.types.DataType,
-    column_name_map: ColumnNameMap,
+    column_name_map: ColumnNameMap | None = None,
     df: snowpark.DataFrame = None,  # remove this param after SNOW-1857090
     depth: int = 0,
 ) -> dict[str, types_proto.DataType]:
@@ -193,7 +206,7 @@ def snowpark_to_proto_type(
                 # For attributes inside struct type (depth > 0), they don't get renamed as normal dataframe column names. Thus no need to do the conversion from snowpark column name to spark column name.
                 spark_name = (
                     column_name_map.get_spark_column_name(index)
-                    if depth == 0
+                    if depth == 0 and column_name_map
                     else field.name
                 )
 
@@ -274,10 +287,24 @@ def snowpark_to_proto_type(
         case snowpark.types.VariantType:
             # For now we are returning a string type for variant types.
             return {"string": types_proto.DataType.String()}
+        case snowpark.types.YearMonthIntervalType:
+            return {
+                "year_month_interval": types_proto.DataType.YearMonthInterval(
+                    start_field=data_type.start_field, end_field=data_type.end_field
+                )
+            }
+        case snowpark.types.DayTimeIntervalType:
+            return {
+                "day_time_interval": types_proto.DataType.DayTimeInterval(
+                    start_field=data_type.start_field, end_field=data_type.end_field
+                )
+            }
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type: {data_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def cast_to_match_snowpark_type(
@@ -317,7 +344,9 @@ def cast_to_match_snowpark_type(
                 with suppress(TypeError):
                     date = datetime.strptime(content, format)
                     return date
-            raise ValueError(f"Date casting error for {str(content)}")
+            exception = ValueError(f"Date casting error for {str(content)}")
+            attach_custom_error_code(exception, ErrorCodes.INVALID_CAST)
+            raise exception
         case snowpark.types.ShortType:
             return int(content)
         case snowpark.types.StringType:
@@ -328,10 +357,32 @@ def cast_to_match_snowpark_type(
             return str(content)
         case snowpark.types.TimestampType:
             return str(content)
+        case snowpark.types.YearMonthIntervalType:
+            if isinstance(content, (int, float)):
+                total_months = int(content)
+                years = total_months // 12
+                months = total_months % 12
+                return f"INTERVAL '{years}-{months}' YEAR TO MONTH"
+            elif isinstance(content, str) and content.startswith(("+", "-")):
+                # Handle Snowflake's native interval format (e.g., "+11-08" or "-2-3")
+                # Convert to Spark's format: "INTERVAL 'Y-M' YEAR TO MONTH"
+                sign = content[0]
+                interval_part = content[1:]  # Remove sign
+                if sign == "-":
+                    return f"INTERVAL '-{interval_part}' YEAR TO MONTH"
+                else:
+                    return f"INTERVAL '{interval_part}' YEAR TO MONTH"
+            return str(content)
+        case snowpark.types.DayTimeIntervalType:
+            return str(content)
+        case snowpark.types.MapType:
+            return content
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type in casting: {data_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def snowpark_to_iceberg_type(data_type: snowpark.types.DataType) -> str:
@@ -364,9 +415,11 @@ def snowpark_to_iceberg_type(data_type: snowpark.types.DataType) -> str:
         case snowpark.types.TimestampType:
             return "timestamp"
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type for iceber: {data_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def proto_to_snowpark_type(
@@ -411,6 +464,18 @@ def proto_to_snowpark_type(
             # For UDT types, return the underlying SQL type
             logger.debug("Returning underlying sql type for udt")
             return proto_to_snowpark_type(data_type.udt.sql_type)
+        case "year_month_interval":
+            # Preserve start_field and end_field from protobuf
+            return snowpark.types.YearMonthIntervalType(
+                start_field=data_type.year_month_interval.start_field,
+                end_field=data_type.year_month_interval.end_field,
+            )
+        case "day_time_interval":
+            # Preserve start_field and end_field from protobuf
+            return snowpark.types.DayTimeIntervalType(
+                start_field=data_type.day_time_interval.start_field,
+                end_field=data_type.day_time_interval.end_field,
+            )
         case _:
             return map_simple_types(data_type.WhichOneof("kind"))
 
@@ -441,9 +506,11 @@ def map_snowpark_types_to_pyarrow_types(
                     )
                 )
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Unsupported arrow type {pa_type} for snowpark ArrayType."
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+                raise exception
         case snowpark.types.BinaryType:
             return pa.binary()
         case snowpark.types.BooleanType:
@@ -484,9 +551,11 @@ def map_snowpark_types_to_pyarrow_types(
                     ),
                 )
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Unsupported arrow type {pa_type} for snowpark MapType."
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+                raise exception
         case snowpark.types.NullType:
             return pa.string()
         case snowpark.types.ShortType:
@@ -503,30 +572,44 @@ def map_snowpark_types_to_pyarrow_types(
                         pa.field(
                             field.name if not rename_struct_columns else str(i),
                             map_snowpark_types_to_pyarrow_types(
-                                field.datatype, pa_type[i].type
+                                field.datatype,
+                                pa_type[i].type,
                             ),
-                            nullable=field.nullable,
+                            nullable=True,
                         )
                         for i, field in enumerate(snowpark_type.fields)
                     ]
                 )
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Unsupported arrow type {pa_type} for snowpark StructType."
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+                raise exception
         case snowpark.types.TimestampType:
-            unit = pa_type.unit
-            tz = pa_type.tz
+            # Check if pa_type has unit attribute (it should be a timestamp type)
+            unit = pa_type.unit if hasattr(pa_type, "unit") else "us"
+            tz = pa_type.tz if hasattr(pa_type, "tz") else None
+
+            # Spark truncates nanosecond precision to microseconds
             if unit == "ns":
-                # Spark truncates nanosecond precision to microseconds
                 unit = "us"
+
             return pa.timestamp(unit, tz=tz)
         case snowpark.types.VariantType:
             return pa.string()
+        case snowpark.types.YearMonthIntervalType:
+            # Return string type so formatted intervals are preserved in display
+            return pa.string()
+        case snowpark.types.DayTimeIntervalType:
+            # Return string type so formatted intervals are preserved in display
+            return pa.string()
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type: {snowpark_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def map_pyarrow_to_snowpark_types(pa_type: pa.DataType) -> snowpark.types.DataType:
@@ -595,10 +678,15 @@ def map_pyarrow_to_snowpark_types(pa_type: pa.DataType) -> snowpark.types.DataTy
         return snowpark.types.TimestampType()
     elif pa.types.is_null(pa_type):
         return snowpark.types.NullType()
+    elif pa.types.is_duration(pa_type):
+        # Map PyArrow duration[us] to DayTimeIntervalType
+        return snowpark.types.DayTimeIntervalType()
     else:
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             f"Unsupported PyArrow data type: {pa_type}"
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
 
 def map_pyspark_types_to_snowpark_types(
@@ -676,9 +764,19 @@ def map_pyspark_types_to_snowpark_types(
         return snowpark.types.TimestampType()
     if isinstance(type_to_map, pyspark.sql.types.TimestampNTZType):
         return snowpark.types.TimestampType(timezone=TimestampTimeZone.NTZ)
-    raise SnowparkConnectNotImplementedError(
+    if isinstance(type_to_map, pyspark.sql.types.YearMonthIntervalType):
+        return snowpark.types.YearMonthIntervalType(
+            type_to_map.startField, type_to_map.endField
+        )
+    if isinstance(type_to_map, pyspark.sql.types.DayTimeIntervalType):
+        return snowpark.types.DayTimeIntervalType(
+            type_to_map.startField, type_to_map.endField
+        )
+    exception = SnowparkConnectNotImplementedError(
         f"Unsupported spark data type: {type_to_map}"
     )
+    attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+    raise exception
 
 
 def map_snowpark_to_pyspark_types(
@@ -743,7 +841,81 @@ def map_snowpark_to_pyspark_types(
         if type_to_map.tz == snowpark.types.TimestampTimeZone.NTZ:
             return pyspark.sql.types.TimestampNTZType()
         return pyspark.sql.types.TimestampType()
-    raise SnowparkConnectNotImplementedError(f"Unsupported data type: {type_to_map}")
+    if isinstance(type_to_map, snowpark.types.YearMonthIntervalType):
+        return pyspark.sql.types.YearMonthIntervalType(
+            type_to_map.start_field, type_to_map.end_field
+        )
+    if isinstance(type_to_map, snowpark.types.DayTimeIntervalType):
+        return pyspark.sql.types.DayTimeIntervalType(
+            type_to_map.start_field, type_to_map.end_field
+        )
+    exception = SnowparkConnectNotImplementedError(
+        f"Unsupported data type: {type_to_map}"
+    )
+    attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+    raise exception
+
+
+def map_pyspark_types_to_pyarrow_types(
+    pyspark_type: pyspark.sql.types.DataType,
+) -> pa.DataType:
+    """
+    Map a PySpark data type to a PyArrow data type.
+
+    This function converts PySpark types to PyArrow types for generating
+    Parquet metadata files with correct schema structure.
+
+    Args:
+        pyspark_type: PySpark data type to convert
+
+    Returns:
+        Corresponding PyArrow data type
+    """
+    if isinstance(pyspark_type, pyspark.sql.types.StringType):
+        return pa.string()
+    elif isinstance(pyspark_type, pyspark.sql.types.LongType):
+        return pa.int64()
+    elif isinstance(pyspark_type, pyspark.sql.types.IntegerType):
+        return pa.int32()
+    elif isinstance(pyspark_type, pyspark.sql.types.ShortType):
+        return pa.int16()
+    elif isinstance(pyspark_type, pyspark.sql.types.ByteType):
+        return pa.int8()
+    elif isinstance(pyspark_type, pyspark.sql.types.DoubleType):
+        return pa.float64()
+    elif isinstance(pyspark_type, pyspark.sql.types.FloatType):
+        return pa.float32()
+    elif isinstance(pyspark_type, pyspark.sql.types.BooleanType):
+        return pa.bool_()
+    elif isinstance(pyspark_type, pyspark.sql.types.DateType):
+        return pa.date32()
+    elif isinstance(pyspark_type, pyspark.sql.types.TimestampType):
+        return pa.timestamp("us")
+    elif isinstance(pyspark_type, pyspark.sql.types.TimestampNTZType):
+        return pa.timestamp("us")
+    elif isinstance(pyspark_type, pyspark.sql.types.BinaryType):
+        return pa.binary()
+    elif isinstance(pyspark_type, pyspark.sql.types.DecimalType):
+        return pa.decimal128(pyspark_type.precision, pyspark_type.scale)
+    elif isinstance(pyspark_type, pyspark.sql.types.ArrayType):
+        element_type = map_pyspark_types_to_pyarrow_types(pyspark_type.elementType)
+        return pa.list_(element_type)
+    elif isinstance(pyspark_type, pyspark.sql.types.MapType):
+        key_type = map_pyspark_types_to_pyarrow_types(pyspark_type.keyType)
+        value_type = map_pyspark_types_to_pyarrow_types(pyspark_type.valueType)
+        return pa.map_(key_type, value_type)
+    elif isinstance(pyspark_type, pyspark.sql.types.StructType):
+        fields = [
+            pa.field(
+                f.name,
+                map_pyspark_types_to_pyarrow_types(f.dataType),
+                nullable=f.nullable,
+            )
+            for f in pyspark_type.fields
+        ]
+        return pa.struct(fields)
+    else:
+        return pa.string()  # Default fallback
 
 
 def map_simple_types(simple_type: str) -> snowpark.types.DataType:
@@ -785,18 +957,51 @@ def map_simple_types(simple_type: str) -> snowpark.types.DataType:
             return snowpark.types.TimestampType(snowpark.types.TimestampTimeZone.NTZ)
         case "timestamp_ltz":
             return snowpark.types.TimestampType(snowpark.types.TimestampTimeZone.LTZ)
+        case "year_month_interval":
+            return snowpark.types.YearMonthIntervalType()
         case "day_time_interval":
-            # this is not a column type in snowflake so there won't be a dataframe column
-            # with this, for now this type won't make any sense
-            return snowpark.types.StringType()
+            return snowpark.types.DayTimeIntervalType()
+        case type_name if _INTERVAL_YEARMONTH_PATTERN_RE.match(type_name):
+            return snowpark.types.YearMonthIntervalType()
+        case type_name if _INTERVAL_DAYTIME_PATTERN_RE.match(type_name):
+            return snowpark.types.DayTimeIntervalType()
+        # Year-Month interval cases
+        case "interval year":
+            return snowpark.types.YearMonthIntervalType(0)
+        case "interval month":
+            return snowpark.types.YearMonthIntervalType(1)
+        case "interval year to month":
+            return snowpark.types.YearMonthIntervalType(0, 1)
+        case "interval day":
+            return snowpark.types.DayTimeIntervalType(0)
+        case "interval hour":
+            return snowpark.types.DayTimeIntervalType(1)
+        case "interval minute":
+            return snowpark.types.DayTimeIntervalType(2)
+        case "interval second":
+            return snowpark.types.DayTimeIntervalType(3)
+        case "interval day to hour":
+            return snowpark.types.DayTimeIntervalType(0, 1)
+        case "interval day to minute":
+            return snowpark.types.DayTimeIntervalType(0, 2)
+        case "interval day to second":
+            return snowpark.types.DayTimeIntervalType(0, 3)
+        case "interval hour to minute":
+            return snowpark.types.DayTimeIntervalType(1, 2)
+        case "interval hour to second":
+            return snowpark.types.DayTimeIntervalType(1, 3)
+        case "interval minute to second":
+            return snowpark.types.DayTimeIntervalType(2, 3)
         case _:
             if simple_type.startswith("decimal"):
                 precision = int(simple_type.split("(")[1].split(",")[0])
                 scale = int(simple_type.split(",")[1].split(")")[0])
                 return snowpark.types.DecimalType(precision, scale)
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported simple type: {simple_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def map_json_schema_to_snowpark(
@@ -937,9 +1142,11 @@ def map_spark_timestamp_format_expression(
             lit_value, _ = get_literal_field_and_name(arguments.literal)
             return convert_spark_format_to_snowflake(lit_value, timestamp_input_type)
         case other:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported expression type {other} in timestamp format argument"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def map_spark_number_format_expression(
@@ -958,9 +1165,11 @@ def map_spark_number_format_expression(
             case "literal":
                 lit_value, _ = get_literal_field_and_name(arguments.literal)
             case other:
-                raise SnowparkConnectNotImplementedError(
+                exception = SnowparkConnectNotImplementedError(
                     f"Unsupported expression type {other} in number format argument"
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+                raise exception
 
     return _map_spark_to_snowflake_number_format(lit_value)
 
@@ -988,3 +1197,116 @@ def _map_spark_to_snowflake_number_format(spark_format: str) -> str:
         return ret
 
     return NUMBER_FORMAT_RE.sub(_replace, spark_format)
+
+
+def map_type_to_snowflake_type(
+    t: Union[snowpark_type.DataType, types_proto.DataType]
+) -> str:
+    """Maps a Snowpark or Spark protobuf type to a Snowflake type string."""
+    if not t:
+        return "VARCHAR"
+    is_snowpark_type = isinstance(t, snowpark_type.DataType)
+    condition = type(t) if is_snowpark_type else t.WhichOneof("kind")
+    match condition:
+        case snowpark_type.ArrayType | "array":
+            return (
+                f"ARRAY({map_type_to_snowflake_type(t.element_type)})"
+                if is_snowpark_type
+                else f"ARRAY({map_type_to_snowflake_type(t.array.element_type)})"
+            )
+        case snowpark_type.BinaryType | "binary":
+            return "BINARY"
+        case snowpark_type.BooleanType | "boolean":
+            return "BOOLEAN"
+        case snowpark_type.ByteType | "byte":
+            return "TINYINT"
+        case snowpark_type.DateType | "date":
+            return "DATE"
+        case snowpark_type.DecimalType | "decimal":
+            return "NUMBER"
+        case snowpark_type.DoubleType | "double":
+            return "DOUBLE"
+        case snowpark_type.FloatType | "float":
+            return "FLOAT"
+        case snowpark_type.GeographyType:
+            return "GEOGRAPHY"
+        case snowpark_type.IntegerType | "integer":
+            return "INT"
+        case snowpark_type.LongType | "long":
+            return "BIGINT"
+        case snowpark_type.MapType | "map":
+            # Maps to OBJECT in Snowflake if key and value types are not specified.
+            key_type = (
+                map_type_to_snowflake_type(t.key_type)
+                if is_snowpark_type
+                else map_type_to_snowflake_type(t.map.key_type)
+            )
+            value_type = (
+                map_type_to_snowflake_type(t.value_type)
+                if is_snowpark_type
+                else map_type_to_snowflake_type(t.map.value_type)
+            )
+            return (
+                f"MAP({key_type}, {value_type})"
+                if key_type and value_type
+                else "OBJECT"
+            )
+        case snowpark_type.NullType | "null":
+            return "VARCHAR"
+        case snowpark_type.ShortType | "short":
+            return "SMALLINT"
+        case snowpark_type.StringType | "string" | "char" | "varchar":
+            return "VARCHAR"
+        case snowpark_type.TimestampType | "timestamp" | "timestamp_ntz":
+            return "TIMESTAMP"
+        case snowpark_type.StructType | "struct":
+            return "VARIANT"
+        case snowpark_type.VariantType:
+            return "VARIANT"
+        case _:
+            exception = ValueError(f"Unsupported Snowpark type: {t}")
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+            raise exception
+
+
+def merge_different_types(
+    type1: snowpark_type.DataType,
+    type2: snowpark_type.DataType,
+) -> snowpark_type.DataType:
+    """
+    Merge two different Snowpark data types.
+    """
+    # If one type is NullType, return the other
+    if isinstance(type1, snowpark_type.NullType) or type1 is None:
+        return type2
+    if isinstance(type2, snowpark_type.NullType) or type2 is None:
+        return type1
+
+    if type1 == type2:
+        return type2
+    # Define type hierarchy - from narrowest to widest scope
+    # Each set contains types that can be merged to a common type
+    numeric_type_hierarchy = [
+        # Numeric hierarchy: byte -> short -> int -> long -> decimal -> float -> double
+        snowpark_type.ByteType,
+        snowpark_type.ShortType,
+        snowpark_type.IntegerType,
+        snowpark_type.LongType,
+        snowpark_type.DecimalType,
+        snowpark_type.FloatType,
+        snowpark_type.DoubleType,
+    ]
+
+    type1_index = next(
+        (i for i, t in enumerate(numeric_type_hierarchy) if isinstance(type1, t)), -1
+    )
+    type2_index = next(
+        (i for i, t in enumerate(numeric_type_hierarchy) if isinstance(type2, t)), -1
+    )
+
+    if type1_index >= 0 and type2_index >= 0:
+        broader_index = max(type1_index, type2_index)
+        return numeric_type_hierarchy[broader_index]()
+
+    # No common type found, default to StringType
+    return snowpark_type.StringType()

snowflake/snowpark_connect/type_support.py

@@ -0,0 +1,130 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import threading
+
+from snowflake import snowpark
+from snowflake.snowpark.types import (
+    ArrayType,
+    ByteType,
+    DataType,
+    DecimalType,
+    IntegerType,
+    LongType,
+    MapType,
+    ShortType,
+    StructField,
+    StructType,
+    _IntegralType,
+)
+
+_integral_types_conversion_enabled: bool = False
+_client_mode_lock = threading.Lock()
+
+
+def set_integral_types_conversion(enabled: bool) -> None:
+    global _integral_types_conversion_enabled
+
+    with _client_mode_lock:
+        if _integral_types_conversion_enabled == enabled:
+            return
+
+        _integral_types_conversion_enabled = enabled
+
+        if enabled:
+            snowpark.context._integral_type_default_precision = {
+                LongType: 19,
+                IntegerType: 10,
+                ShortType: 5,
+                ByteType: 3,
+            }
+        else:
+            snowpark.context._integral_type_default_precision = {}
+
+
+def set_integral_types_for_client_default(is_python_client: bool) -> None:
+    """
+    Set integral types based on client type when config is 'client_default'.
+    """
+    from snowflake.snowpark_connect.config import global_config
+
+    config_key = "snowpark.connect.integralTypesEmulation"
+    if global_config.get(config_key) != "client_default":
+        return
+
+    # if client mode matches, no action needed (no lock overhead)
+    if _integral_types_conversion_enabled == (not is_python_client):
+        return
+
+    set_integral_types_conversion(not is_python_client)
+
+
+def emulate_integral_types(t: DataType) -> DataType:
+    """
+    Map LongType based on precision attribute to appropriate integral types.
+
+    Mappings:
+    - _IntegralType with precision=19 -> LongType
+    - _IntegralType with precision=10 -> IntegerType
+    - _IntegralType with precision=5 -> ShortType
+    - _IntegralType with precision=3 -> ByteType
+    - _IntegralType with other precision -> DecimalType(precision, 0)
+
+    This conversion is controlled by the 'snowpark.connect.integralTypesEmulation' config.
+    When disabled, the function returns the input type unchanged.
+
+    Args:
+        t: The DataType to transform
+
+    Returns:
+        The transformed DataType with integral type conversions applied based on precision.
+    """
+    global _integral_types_conversion_enabled
+
+    with _client_mode_lock:
+        enabled = _integral_types_conversion_enabled
+    if not enabled:
+        return t
+    if isinstance(t, _IntegralType):
+        precision = getattr(t, "_precision", None)
+
+        if precision is None:
+            return t
+        elif precision == 19:
+            return LongType()
+        elif precision == 10:
+            return IntegerType()
+        elif precision == 5:
+            return ShortType()
+        elif precision == 3:
+            return ByteType()
+        else:
+            return DecimalType(precision, 0)
+
+    elif isinstance(t, StructType):
+        new_fields = [
+            StructField(
+                field.name,
+                emulate_integral_types(field.datatype),
+                field.nullable,
+                _is_column=field._is_column,
+            )
+            for field in t.fields
+        ]
+        return StructType(new_fields)
+
+    elif isinstance(t, ArrayType):
+        return ArrayType(
+            emulate_integral_types(t.element_type),
+            t.contains_null,
+        )
+
+    elif isinstance(t, MapType):
+        return MapType(
+            emulate_integral_types(t.key_type),
+            emulate_integral_types(t.value_type),
+            t.value_contains_null,
+        )
+
+    return t