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

snowflake/snowpark_connect/relation/map_aggregate.py

@@ -11,9 +11,10 @@ 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.types import DataType, StructType
 from snowflake.snowpark_connect.column_name_handler import (
     make_column_names_snowpark_compatible,
+    make_unique_snowpark_name,
 )
 from snowflake.snowpark_connect.column_qualifier import ColumnQualifier
 from snowflake.snowpark_connect.dataframe_container import (
@@ -57,6 +58,47 @@ def map_group_by_aggregate(
             *columns.aggregation_expressions()
         )
 
+    for rel_aggregate_expression, aggregate_original_column in zip(
+        rel.aggregate.aggregate_expressions, columns.aggregation_columns
+    ):
+        aggregate_original_data_type = aggregate_original_column.data_type
+
+        if not (
+            rel_aggregate_expression.HasField("unresolved_function")
+            and rel_aggregate_expression.unresolved_function.function_name == "reduce"
+        ) or not isinstance(aggregate_original_data_type, StructType):
+            continue
+
+        cols = []
+        new_snowpark_column_names = []
+        new_snowpark_column_types = [
+            field.datatype for field in aggregate_original_data_type.fields
+        ]
+
+        if not result.columns or len(result.columns) != 1:
+            raise ValueError(
+                "Expected result DataFrame to have exactly one column for reduce(StructType)"
+            )
+        aggregate_col = snowpark_fn.col(result.columns[0])
+
+        # Extract each field from the StructType result after aggregation to create separate columns
+        for spark_col_name in input_df_container.column_map.get_spark_columns():
+            unique_snowpark_name = make_unique_snowpark_name(spark_col_name)
+            cols.append(
+                snowpark_fn.get(aggregate_col, snowpark_fn.lit(spark_col_name)).alias(
+                    unique_snowpark_name
+                )
+            )
+            new_snowpark_column_names.append(unique_snowpark_name)
+
+        result = result.select(*cols)
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=result,
+            spark_column_names=input_df_container.column_map.get_spark_columns(),
+            snowpark_column_names=new_snowpark_column_names,
+            snowpark_column_types=new_snowpark_column_types,
+        )
+
     # Store aggregate metadata for ORDER BY resolution
     aggregate_metadata = AggregateMetadata(
         input_column_map=input_df_container.column_map,

snowflake/snowpark_connect/relation/read/map_read_csv.py

@@ -11,7 +11,18 @@ from pyspark.errors.exceptions.base import AnalysisException
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
 from snowflake.snowpark.dataframe_reader import DataFrameReader
-from snowflake.snowpark.types import StringType, StructField, StructType
+from snowflake.snowpark.types import (
+    DataType,
+    DecimalType,
+    DoubleType,
+    IntegerType,
+    LongType,
+    StringType,
+    StructField,
+    StructType,
+    _FractionalType,
+    _IntegralType,
+)
 from snowflake.snowpark_connect.config import global_config, str_to_bool
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.error.error_codes import ErrorCodes
@@ -26,6 +37,10 @@ from snowflake.snowpark_connect.relation.read.utils import (
     get_spark_column_names_from_snowpark_columns,
     rename_columns_as_snowflake_standard,
 )
+from snowflake.snowpark_connect.type_support import (
+    _integral_types_conversion_enabled,
+    emulate_integral_types,
+)
 from snowflake.snowpark_connect.utils.io_utils import cached_file_format
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
@@ -108,7 +123,7 @@ def map_read_csv(
                 df = df.union_all(reader.csv(p))
 
         if schema is None and not str_to_bool(
-            str(raw_options.get("inferSchema", "false"))
+            str(raw_options.get("inferSchema", raw_options.get("inferschema", "false")))
         ):
             df = df.select(
                 [snowpark_fn.col(c).cast("STRING").alias(c) for c in df.schema.names]
@@ -123,7 +138,9 @@ def map_read_csv(
             dataframe=renamed_df,
             spark_column_names=spark_column_names,
             snowpark_column_names=snowpark_column_names,
-            snowpark_column_types=[f.datatype for f in df.schema.fields],
+            snowpark_column_types=[
+                _emulate_integral_types_for_csv(f.datatype) for f in df.schema.fields
+            ],
         )
 
 
@@ -320,7 +337,13 @@ def read_data(
     # Create schema with the column names and read CSV
     if len(headers) > 0:
         if (
-            not str_to_bool(str(raw_options.get("inferSchema", "false")))
+            not str_to_bool(
+                str(
+                    raw_options.get(
+                        "inferSchema", raw_options.get("inferschema", "false")
+                    )
+                )
+            )
             and schema is None
         ):
             inferred_schema = StructType(
@@ -350,3 +373,49 @@ def read_data(
 
     # Fallback: no headers, shouldn't reach here
     return reader.csv(path)
+
+
+def _emulate_integral_types_for_csv(t: DataType) -> DataType:
+    """
+    CSV requires different type handling to match OSS Spark CSV schema inference.
+
+    After applying emulate_integral_types, converts to Spark CSV types:
+    - IntegerType, ShortType, ByteType -> IntegerType
+    - LongType -> LongType
+    - DecimalType with scale > 0 -> DoubleType
+    - DecimalType with precision > 18 -> DecimalType (too big for long)
+    - DecimalType with precision > 9 -> LongType
+    - DecimalType with precision <= 9 -> IntegerType
+    - FloatType, DoubleType -> DoubleType
+    """
+    if not _integral_types_conversion_enabled:
+        return t
+
+    # First apply standard integral type conversion
+    t = emulate_integral_types(t)
+
+    if isinstance(t, LongType):
+        return LongType()
+
+    elif isinstance(t, _IntegralType):
+        # ByteType, ShortType, IntegerType -> IntegerType
+        return IntegerType()
+
+    elif isinstance(t, DecimalType):
+        # DecimalType with scale > 0 means it has decimal places -> DoubleType
+        if t.scale > 0:
+            return DoubleType()
+        # DecimalType with scale = 0 is integral
+        if t.precision > 18:
+            # Too big for long, keep as DecimalType
+            return DecimalType(t.precision, 0)
+        elif t.precision > 9:
+            return LongType()
+        else:
+            return IntegerType()
+
+    elif isinstance(t, _FractionalType):
+        # FloatType, DoubleType -> DoubleType
+        return DoubleType()
+
+    return t

snowflake/snowpark_connect/relation/read/map_read_jdbc.py

@@ -16,6 +16,7 @@ from snowflake.snowpark_connect.relation.read.utils import (
     Connection,
     rename_columns_as_snowflake_standard,
 )
+from snowflake.snowpark_connect.type_support import emulate_integral_types
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 
 
@@ -112,7 +113,9 @@ def map_read_jdbc(
             dataframe=renamed_df,
             spark_column_names=true_names,
             snowpark_column_names=snowpark_cols,
-            snowpark_column_types=[f.datatype for f in df.schema.fields],
+            snowpark_column_types=[
+                emulate_integral_types(f.datatype) for f in df.schema.fields
+            ],
         )
     except Exception as e:
         exception = Exception(f"Error accessing JDBC datasource for read: {e}")

snowflake/snowpark_connect/relation/read/map_read_json.py

@@ -45,6 +45,7 @@ from snowflake.snowpark_connect.type_mapping import (
     map_simple_types,
     merge_different_types,
 )
+from snowflake.snowpark_connect.type_support import emulate_integral_types
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
@@ -148,7 +149,9 @@ def map_read_json(
             dataframe=renamed_df,
             spark_column_names=spark_column_names,
             snowpark_column_names=snowpark_column_names,
-            snowpark_column_types=[f.datatype for f in df.schema.fields],
+            snowpark_column_types=[
+                emulate_integral_types(f.datatype) for f in df.schema.fields
+            ],
         )
 
 

snowflake/snowpark_connect/relation/read/map_read_parquet.py

@@ -44,6 +44,7 @@ from snowflake.snowpark_connect.relation.read.utils import (
     apply_metadata_exclusion_pattern,
     rename_columns_as_snowflake_standard,
 )
+from snowflake.snowpark_connect.type_support import emulate_integral_types
 from snowflake.snowpark_connect.utils.io_utils import cached_file_format
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
@@ -126,7 +127,9 @@ def map_read_parquet(
         dataframe=renamed_df,
         spark_column_names=[analyzer_utils.unquote_if_quoted(c) for c in df.columns],
         snowpark_column_names=snowpark_column_names,
-        snowpark_column_types=[f.datatype for f in df.schema.fields],
+        snowpark_column_types=[
+            emulate_integral_types(f.datatype) for f in df.schema.fields
+        ],
         can_be_cached=can_be_cached,
     )
 

snowflake/snowpark_connect/relation/read/map_read_socket.py

@@ -11,6 +11,7 @@ from snowflake import snowpark
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 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.type_support import emulate_integral_types
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -58,6 +59,9 @@ def map_read_socket(
                     dataframe=df,
                     spark_column_names=[spark_cname],
                     snowpark_column_names=[snowpark_cname],
+                    snowpark_column_types=[
+                        emulate_integral_types(f.datatype) for f in df.schema.fields
+                    ],
                 )
             except OSError as e:
                 exception = Exception(f"Error connecting to {host}:{port} - {e}")

snowflake/snowpark_connect/relation/read/map_read_table.py

@@ -24,6 +24,7 @@ from snowflake.snowpark_connect.error.error_utils import attach_custom_error_cod
 from snowflake.snowpark_connect.relation.read.utils import (
     rename_columns_as_snowflake_standard,
 )
+from snowflake.snowpark_connect.type_support import emulate_integral_types
 from snowflake.snowpark_connect.utils.context import get_processed_views
 from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
@@ -58,7 +59,9 @@ def post_process_df(
             dataframe=renamed_df,
             spark_column_names=true_names,
             snowpark_column_names=snowpark_column_names,
-            snowpark_column_types=[f.datatype for f in df.schema.fields],
+            snowpark_column_types=[
+                emulate_integral_types(f.datatype) for f in df.schema.fields
+            ],
             column_qualifiers=[{ColumnQualifier(tuple(name_parts))} for _ in true_names]
             if source_table_name
             else None,

snowflake/snowpark_connect/relation/read/map_read_text.py

@@ -14,6 +14,7 @@ from snowflake.snowpark_connect.relation.read.utils import (
     get_spark_column_names_from_snowpark_columns,
     rename_columns_as_snowflake_standard,
 )
+from snowflake.snowpark_connect.type_support import emulate_integral_types
 from snowflake.snowpark_connect.utils.io_utils import file_format
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
@@ -117,5 +118,7 @@ def map_read_text(
         dataframe=renamed_df,
         spark_column_names=spark_column_names,
         snowpark_column_names=snowpark_column_names,
-        snowpark_column_types=[f.datatype for f in df.schema.fields],
+        snowpark_column_types=[
+            emulate_integral_types(f.datatype) for f in df.schema.fields
+        ],
     )

snowflake/snowpark_connect/relation/read/reader_config.py

@@ -441,4 +441,10 @@ class ParquetReaderConfig(ReaderWriterConfig):
             "snowpark.connect.parquet.useVectorizedScanner"
         )
 
+        # Set USE_LOGICAL_TYPE from global config to properly handle Parquet logical types like TIMESTAMP.
+        # Without this, Parquet TIMESTAMP (INT64 physical) is incorrectly read as NUMBER(38,0).
+        snowpark_args["USE_LOGICAL_TYPE"] = global_config._get_config_setting(
+            "snowpark.connect.parquet.useLogicalType"
+        )
+
         return snowpark_args

snowflake/snowpark_connect/resources_initializer.py

@@ -3,7 +3,12 @@
 #
 import threading
 import time
+from collections.abc import Callable
+from pathlib import Path
 
+from snowflake.snowpark_connect.client.error_utils import attach_custom_error_code
+from snowflake.snowpark_connect.config import get_scala_version
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 
@@ -11,22 +16,31 @@ SPARK_VERSION = "3.5.6"
 RESOURCE_PATH = "/snowflake/snowpark_connect/resources"
 
 # On demand Scala UDF jar upload state - separate from general resource initialization
-_scala_jars_uploaded = threading.Event()
-_scala_jars_lock = threading.Lock()
+_scala_2_12_jars_uploaded = threading.Event()
+_scala_2_12_jars_lock = threading.Lock()
+_scala_2_13_jars_uploaded = threading.Event()
+_scala_2_13_jars_lock = threading.Lock()
 
 # Define Scala resource names
-SPARK_SQL_JAR = f"spark-sql_2.12-{SPARK_VERSION}.jar"
-SPARK_CONNECT_CLIENT_JAR = f"spark-connect-client-jvm_2.12-{SPARK_VERSION}.jar"
-SPARK_COMMON_UTILS_JAR = f"spark-common-utils_2.12-{SPARK_VERSION}.jar"
-SAS_SCALA_UDF_JAR = "sas-scala-udf_2.12-0.2.0.jar"
-JSON_4S_JAR = "json4s-ast_2.12-3.7.0-M11.jar"
-SCALA_REFLECT_JAR = "scala-reflect-2.12.18.jar"
-
-
-def _upload_scala_udf_jars_impl() -> None:
+SPARK_SQL_JAR_212 = f"spark-sql_2.12-{SPARK_VERSION}.jar"
+SPARK_CONNECT_CLIENT_JAR_212 = f"spark-connect-client-jvm_2.12-{SPARK_VERSION}.jar"
+SPARK_COMMON_UTILS_JAR_212 = f"spark-common-utils_2.12-{SPARK_VERSION}.jar"
+SAS_SCALA_UDF_JAR_212 = "sas-scala-udf_2.12-0.2.0.jar"
+JSON_4S_JAR_212 = "json4s-ast_2.12-3.7.0-M11.jar"
+SCALA_REFLECT_JAR_212 = "scala-reflect-2.12.18.jar"
+
+# Static dependencies for Scala 2.13
+SPARK_SQL_JAR_213 = f"spark-sql_2.13-{SPARK_VERSION}.jar"
+SPARK_CONNECT_CLIENT_JAR_213 = f"spark-connect-client-jvm_2.13-{SPARK_VERSION}.jar"
+SPARK_COMMON_UTILS_JAR_213 = f"spark-common-utils_2.13-{SPARK_VERSION}.jar"
+SAS_SCALA_UDF_JAR_213 = "sas-scala-udf_2.13-0.2.0.jar"
+JSON_4S_JAR_213 = "json4s-ast_2.13-3.7.0-M11.jar"
+SCALA_REFLECT_JAR_213 = "scala-reflect-2.13.16.jar"
+
+
+def _upload_scala_udf_jars(jar_files: list[str]) -> None:
     """Upload Spark jar files required for creating Scala UDFs.
     This is the internal implementation - use ensure_scala_udf_jars_uploaded() for thread-safe lazy loading."""
-    from pathlib import Path
 
     session = get_or_create_snowpark_session()
     stage = session.get_session_stage()
@@ -34,15 +48,6 @@ def _upload_scala_udf_jars_impl() -> None:
     import snowpark_connect_deps_1
     import snowpark_connect_deps_2
 
-    jar_files = [
-        SPARK_SQL_JAR,
-        SPARK_CONNECT_CLIENT_JAR,
-        SPARK_COMMON_UTILS_JAR,
-        SAS_SCALA_UDF_JAR,
-        JSON_4S_JAR,
-        SCALA_REFLECT_JAR,  # Required for deserializing Scala lambdas
-    ]
-
     # Path to includes/jars directory
     includes_jars_dir = Path(__file__).parent / "includes" / "jars"
 
@@ -78,31 +83,87 @@ def _upload_scala_udf_jars_impl() -> None:
             raise RuntimeError(f"Failed to upload JAR {jar_name}: {e}")
 
 
-def ensure_scala_udf_jars_uploaded() -> None:
-    """Ensure Scala UDF jars are uploaded to Snowflake, uploading them lazily if not already done.
-    This function is thread-safe and will only upload once even if called from multiple threads."""
+def _upload_scala_2_12_jars() -> None:
+    scala_2_12_jars = [
+        SPARK_SQL_JAR_212,
+        SPARK_CONNECT_CLIENT_JAR_212,
+        SPARK_COMMON_UTILS_JAR_212,
+        SAS_SCALA_UDF_JAR_212,
+        JSON_4S_JAR_212,
+        SCALA_REFLECT_JAR_212,  # Required for deserializing Scala lambdas
+    ]
+    _upload_scala_udf_jars(scala_2_12_jars)
 
+
+def _upload_scala_2_13_jars() -> None:
+    scala_2_13_jars = [
+        SPARK_SQL_JAR_213,
+        SPARK_CONNECT_CLIENT_JAR_213,
+        SPARK_COMMON_UTILS_JAR_213,
+        SAS_SCALA_UDF_JAR_213,
+        JSON_4S_JAR_213,
+        SCALA_REFLECT_JAR_213,
+    ]
+    _upload_scala_udf_jars(scala_2_13_jars)
+
+
+def _ensure_configured_scala_jars_uploaded(
+    jars_uploaded: threading.Event, lock: threading.Lock, upload_fn: Callable[[], None]
+) -> None:
+    """
+    Ensure Scala UDF jars are uploaded to Snowflake, uploading them lazily if not already done.
+    This function is thread-safe and will only upload once even if called from multiple threads.
+
+    Uses the given upload_fn to upload Scala jars if the jars_uploaded event is not set yet.
+    """
     # Fast path: if already uploaded, return immediately without acquiring lock
-    if _scala_jars_uploaded.is_set():
+    if jars_uploaded.is_set():
         return
 
     # Slow path: need to upload, acquire lock to ensure only one thread does it
-    with _scala_jars_lock:
+    with lock:
         # Double-check pattern: another thread might have uploaded while we waited for the lock
-        if _scala_jars_uploaded.is_set():
+        if jars_uploaded.is_set():
             return
 
         try:
             start_time = time.time()
             logger.info("Uploading Scala UDF jars on-demand...")
-            _upload_scala_udf_jars_impl()
-            _scala_jars_uploaded.set()
+            upload_fn()
+            jars_uploaded.set()
             logger.info(f"Scala UDF jars uploaded in {time.time() - start_time:.2f}s")
         except Exception as e:
             logger.error(f"Failed to upload Scala UDF jars: {e}")
             raise
 
 
+def ensure_scala_udf_jars_uploaded() -> None:
+    """
+    Public function to make sure Scala jars are uploaded and available for imports.
+    """
+    scala_version = get_scala_version()
+
+    match scala_version:
+        case "2.12":
+            _ensure_configured_scala_jars_uploaded(
+                _scala_2_12_jars_uploaded,
+                _scala_2_12_jars_lock,
+                _upload_scala_2_12_jars,
+            )
+        case "2.13":
+            _ensure_configured_scala_jars_uploaded(
+                _scala_2_13_jars_uploaded,
+                _scala_2_13_jars_lock,
+                _upload_scala_2_13_jars,
+            )
+        case _:
+            exception = ValueError(
+                f"Unsupported Scala version: {scala_version}. Snowpark Connect supports Scala 2.12 and 2.13"
+            )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_CONFIG_VALUE)
+            raise exception
+
+
 def initialize_resources() -> None:
     """Initialize all expensive resources. We should initialize what we can here, so that actual rpc calls like
     ExecutePlan are as fast as possible."""

snowflake/snowpark_connect/server.py

@@ -45,7 +45,10 @@ import snowflake.snowpark_connect.proto.control_pb2_grpc as control_grpc
 import snowflake.snowpark_connect.tcm as tcm
 from snowflake import snowpark
 from snowflake.snowpark_connect.analyze_plan.map_tree_string import map_tree_string
-from snowflake.snowpark_connect.config import route_config_proto
+from snowflake.snowpark_connect.config import (
+    route_config_proto,
+    set_java_udf_creator_initialized_state,
+)
 from snowflake.snowpark_connect.constants import SERVER_SIDE_SESSION_ID
 from snowflake.snowpark_connect.control_server import ControlServicer
 from snowflake.snowpark_connect.error.error_codes import ErrorCodes
@@ -112,9 +115,6 @@ from snowflake.snowpark_connect.utils.interrupt import (
     interrupt_queries_with_tag,
     interrupt_query,
 )
-from snowflake.snowpark_connect.utils.java_stored_procedure import (
-    set_java_udf_creator_initialized_state,
-)
 from snowflake.snowpark_connect.utils.open_telemetry import (
     is_telemetry_enabled,
     otel_attach_context,
@@ -1126,51 +1126,16 @@ def start_jvm():
         attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
         raise exception
 
-    import pathlib
-    import zipfile
-
-    import snowflake.snowpark_connect
-
     # Import both JAR dependency packages
     import snowpark_connect_deps_1
     import snowpark_connect_deps_2
 
-    # First, add JARs from includes/jars directory
-    pyspark_jars = (
-        pathlib.Path(snowflake.snowpark_connect.__file__).parent / "includes" / "jars"
-    )
-
-    if "dataframe_processor.zip" in str(pyspark_jars):
-        # importlib.resource doesn't work when local stage package is used in TCM
-        zip_path = pathlib.Path(
-            snowflake.snowpark_connect.__file__
-        ).parent.parent.parent
-        temp_dir = tempfile.gettempdir()
-        extract_folder = "snowflake/snowpark_connect/includes/jars/"  # Folder to extract (must end with '/')
-
-        with zipfile.ZipFile(zip_path, "r") as zip_ref:
-            for member in zip_ref.namelist():
-                if member.startswith(extract_folder):
-                    zip_ref.extract(member, path=temp_dir)
-        pyspark_jars = pathlib.Path(temp_dir) / extract_folder
-
-    included_jar_names = set()
-
-    if pyspark_jars.exists():
-        for jar_path in pyspark_jars.glob(
-            "**/*.jar"
-        ):  # Use **/*.jar to handle nested paths in TCM
-            jpype.addClassPath(str(jar_path))
-            included_jar_names.add(jar_path.name)
-
-    # Load jar files from both packages, skipping those already loaded from includes/jars
+    # Load all the jar files from both packages
     jar_path_list = (
         snowpark_connect_deps_1.list_jars() + snowpark_connect_deps_2.list_jars()
     )
     for jar_path in jar_path_list:
-        # Skip if this JAR was already loaded from includes/jars
-        if jar_path.name not in included_jar_names:
-            jpype.addClassPath(jar_path)
+        jpype.addClassPath(jar_path)
 
     # TODO: Should remove convertStrings, but it breaks the JDBC code.
     jvm_settings: list[str] = list(

snowflake/snowpark_connect/server_common/__init__.py

@@ -369,7 +369,10 @@ def _setup_spark_environment(setup_java_home: bool = True) -> None:
                         lightweight client servers that don't need JVM.
     """
     if setup_java_home:
-        if os.environ.get("JAVA_HOME") is None:
+        if (
+            os.environ.get("JAVA_HOME") is None
+            or str(os.environ.get("JAVA_HOME")).strip() == ""
+        ):
             try:
                 # For Notebooks on SPCS
                 from jdk4py import JAVA_HOME