snowpark-connect 0.30.0__tar.gz → 0.30.1__tar.gz

{snowpark_connect-0.30.0/src/snowpark_connect.egg-info → snowpark_connect-0.30.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-connect
-Version: 0.30.0
+Version: 0.30.1
 Summary: Snowpark Connect for Spark
 Author: Snowflake, Inc
 License: Apache License, Version 2.0

{snowpark_connect-0.30.0 → snowpark_connect-0.30.1}/src/snowflake/snowpark_connect/config.py

@@ -143,6 +143,9 @@ class GlobalConfig:
         # Control whether repartition(n) on a DataFrame forces splitting into n files during writes
         # This matches spark behavior more closely, but introduces overhead.
         "snowflake.repartition.for.writes": "false",
+        "snowpark.connect.structured_types.fix": str(
+            snowpark.context._enable_fix_2360274
+        ).lower(),
     }
 
     boolean_config_list = [

snowpark_connect-0.30.1/src/snowflake/snowpark_connect/proto/snowflake_expression_ext_pb2_grpc.py

@@ -0,0 +1,4 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+

snowpark_connect-0.30.1/src/snowflake/snowpark_connect/proto/snowflake_relation_ext_pb2_grpc.py

@@ -0,0 +1,4 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+

{snowpark_connect-0.30.0 → snowpark_connect-0.30.1}/src/snowflake/snowpark_connect/relation/io_utils.py

@@ -4,6 +4,8 @@
 
 from urllib.parse import urlparse
 
+from pyspark.errors.exceptions.base import AnalysisException
+
 CLOUD_PREFIX_TO_CLOUD = {
     "abfss": "azure",
     "wasbs": "azure",
@@ -12,10 +14,28 @@ CLOUD_PREFIX_TO_CLOUD = {
 }
 
 SUPPORTED_COMPRESSION_PER_FORMAT = {
-    "csv": {"AUTO", "GZIP", "BZ2", "BROTLI", "ZSTD", "DEFLATE", "RAW_DEFLATE", "NONE"},
-    "json": {"AUTO", "GZIP", "BZ2", "BROTLI", "ZSTD", "DEFLATE", "RAW_DEFLATE", "NONE"},
-    "parquet": {"AUTO", "LZO", "SNAPPY", "NONE"},
-    "text": {"NONE"},
+    "csv": {
+        "GZIP",
+        "BZ2",
+        "BROTLI",
+        "ZSTD",
+        "DEFLATE",
+        "RAW_DEFLATE",
+        "NONE",
+        "UNCOMPRESSED",
+    },
+    "json": {
+        "GZIP",
+        "BZ2",
+        "BROTLI",
+        "ZSTD",
+        "DEFLATE",
+        "RAW_DEFLATE",
+        "NONE",
+        "UNCOMPRESSED",
+    },
+    "parquet": {"LZO", "SNAPPY", "NONE", "UNCOMPRESSED"},
+    "text": {"NONE", "UNCOMPRESSED"},
 }
 
 
@@ -29,6 +49,43 @@ def is_supported_compression(format: str, compression: str | None) -> bool:
     return compression in supported_compressions_for_format(format)
 
 
+def get_compression_for_source_and_options(
+    source: str, options: dict[str, str], from_read: bool = False
+) -> str | None:
+    """
+    Determines the compression type to use for a given data source and options.
+    Args:
+        source (str): The data source format (e.g., "csv", "json", "parquet", "text").
+        options (dict[str, str]): A dictionary of options that may include a "compression" key.
+    Returns:
+        str: The compression type to use (e.g., "GZIP", "SNAPPY", "NONE").
+    Raises:
+        AnalysisException: If the specified compression is not supported for the given source format.
+    """
+    # From read, we don't have a default compression
+    if from_read and "compression" not in options:
+        return None
+
+    # Get compression from options for proper filename generation
+    default_compression = "NONE" if source != "parquet" else "snappy"
+    compression = options.get("compression", default_compression).upper()
+    if compression == "UNCOMPRESSED":
+        compression = "NONE"
+
+    if not is_supported_compression(source, compression):
+        supported_compressions = supported_compressions_for_format(source)
+        raise AnalysisException(
+            f"Compression {compression} is not supported for {source} format. "
+            + (
+                f"Supported compressions: {sorted(supported_compressions)}"
+                if supported_compressions
+                else "None compression supported for this format."
+            )
+        )
+
+    return compression
+
+
 def get_cloud_from_url(
     url: str,
 ):

{snowpark_connect-0.30.0 → snowpark_connect-0.30.1}/src/snowflake/snowpark_connect/relation/map_row_ops.py

@@ -35,6 +35,9 @@ 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.utils.identifiers import (
+    split_fully_qualified_spark_name,
+)
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -131,11 +134,19 @@ def map_fillna(
     input_df = input_container.dataframe
 
     if len(rel.fill_na.cols) > 0:
+        if rel.fill_na.cols == ["*"]:
+            # Expand "*" to all columns
+            spark_col_names = input_container.column_map.get_spark_columns()
+        else:
+            spark_col_names = list(rel.fill_na.cols)
+
+        # We don't validate the fully qualified spark name here as fillNa is no-op for structured type colums.
+        # It only works for scalar type columns like float, int, string or bool.
         columns: list[str] = [
             input_container.column_map.get_snowpark_column_name_from_spark_column_name(
-                c
+                split_fully_qualified_spark_name(c)[0]
             )
-            for c in rel.fill_na.cols
+            for c in spark_col_names
         ]
         values = [get_literal_field_and_name(v)[0] for v in rel.fill_na.values]
         if len(values) == 1:

{snowpark_connect-0.30.0 → snowpark_connect-0.30.1}/src/snowflake/snowpark_connect/relation/map_stats.py

@@ -309,9 +309,28 @@ def map_freq_items(rel: relation_proto.Relation) -> DataFrameContainer:
     cols = input_container.column_map.get_snowpark_column_names_from_spark_column_names(
         list(rel.freq_items.cols)
     )
+
+    # handle empty DataFrame case
+    row_count = input_df.count()
+
+    for sp_col_name in cols:
+        spark_col_names.append(
+            f"{input_container.column_map.get_spark_column_name_from_snowpark_column_name(sp_col_name)}_freqItems"
+        )
+
+    if row_count == 0:
+        # If DataFrame is empty, return empty arrays for each column
+        empty_values = [[] for _ in cols]
+        approx_top_k_df = session.createDataFrame([empty_values], spark_col_names)
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=approx_top_k_df,
+            spark_column_names=spark_col_names,
+            snowpark_column_names=spark_col_names,
+        )
+
     approx_top_k_df = input_df.select(
         *[
-            fn.function("approx_top_k")(fn.col(col), round(input_df.count() / support))
+            fn.function("approx_top_k")(fn.col(col), round(row_count / support))
             for col in cols
         ]
     )
@@ -330,10 +349,6 @@ def map_freq_items(rel: relation_proto.Relation) -> DataFrameContainer:
         for value in approx_top_k_values
     ]
 
-    for sp_col_name in cols:
-        spark_col_names.append(
-            f"{input_container.column_map.get_spark_column_name_from_snowpark_column_name(sp_col_name)}_freqItems"
-        )
     approx_top_k_df = session.createDataFrame([filtered_values], spark_col_names)
 
     return DataFrameContainer.create_with_column_mapping(

{snowpark_connect-0.30.0 → snowpark_connect-0.30.1}/src/snowflake/snowpark_connect/relation/read/map_read.py

@@ -17,6 +17,7 @@ from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.io_utils import (
     convert_file_prefix_path,
+    get_compression_for_source_and_options,
     is_cloud_path,
 )
 from snowflake.snowpark_connect.relation.read.map_read_table import map_read_table
@@ -237,6 +238,14 @@ def _read_file(
     )
     upload_files_if_needed(paths, clean_source_paths, session, read_format)
     paths = [_quote_stage_path(path) for path in paths]
+
+    if read_format in ("csv", "text", "json", "parquet"):
+        compression = get_compression_for_source_and_options(
+            read_format, options, from_read=True
+        )
+        if compression is not None:
+            options["compression"] = compression
+
     match read_format:
         case "csv":
             from snowflake.snowpark_connect.relation.read.map_read_csv import (

{snowpark_connect-0.30.0 → snowpark_connect-0.30.1}/src/snowflake/snowpark_connect/relation/utils.py

@@ -92,6 +92,21 @@ TYPE_MAP_FOR_TO_SCHEMA = {
 }
 
 
+# This mapping is used to map the compression type to the extension of the file.
+FILE_COMPRESSION_TO_EXTENSION = {
+    "GZIP": "gz",
+    "BZ2": "bz2",
+    "BROTLI": "br",
+    "ZSTD": "zst",
+    "DEFLATE": "deflate",
+    "RAW_DEFLATE": "raw_deflate",
+    "SNAPPY": "snappy",
+    "LZO": "lzo",
+    "LZ4": "lz4",
+    "BZIP2": "bz2",
+}
+
+
 def get_df_with_partition_row_number(
     container: DataFrameContainer,
     plan_id: int | None,
@@ -186,13 +201,15 @@ def generate_spark_compatible_filename(
 
     # Add compression if specified and not 'none'
     if compression and compression.lower() not in ("none", "uncompressed"):
-        compression_part = f".{compression.lower()}"
+        compression_part = f".{FILE_COMPRESSION_TO_EXTENSION.get(compression.upper(), compression.lower())}"
     else:
         compression_part = ""
 
     # Add format extension if specified
-    if format_ext:
+    if format_ext == "parquet":
         return f"{base_name}{compression_part}.{format_ext}"
+    elif format_ext is not None and format_ext != "":
+        return f"{base_name}.{format_ext}{compression_part}"
     else:
         return f"{base_name}{compression_part}"
 

{snowpark_connect-0.30.0 → snowpark_connect-0.30.1}/src/snowflake/snowpark_connect/relation/write/map_write.py

@@ -35,9 +35,8 @@ from snowflake.snowpark_connect.config import (
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.io_utils import (
     convert_file_prefix_path,
+    get_compression_for_source_and_options,
     is_cloud_path,
-    is_supported_compression,
-    supported_compressions_for_format,
 )
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.relation.read.metadata_utils import (
@@ -147,6 +146,11 @@ def map_write(request: proto_base.ExecutePlanRequest):
 
     session: snowpark.Session = get_or_create_snowpark_session()
 
+    # Check for partition hint early to determine precedence over single option
+    partition_hint = (
+        result.partition_hint if hasattr(result, "partition_hint") else None
+    )
+
     # Snowflake saveAsTable doesn't support format
     if (
         write_op.HasField("table")
@@ -176,8 +180,11 @@ def map_write(request: proto_base.ExecutePlanRequest):
             # 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")
+            compression = get_compression_for_source_and_options(
+                write_op.source, write_op.options, from_read=False
+            )
+            if compression is not None:
+                write_op.options["compression"] = compression
 
             # Generate Spark-compatible filename or prefix
             # we need a random prefix to support "append" mode
@@ -203,12 +210,12 @@ def map_write(request: proto_base.ExecutePlanRequest):
                 except Exception as e:
                     logger.warning(f"Could not clear directory {write_path}: {e}")
 
-            if should_write_to_single_file:
+            if should_write_to_single_file and partition_hint is None:
                 # Single file: generate complete filename with extension
                 spark_filename = generate_spark_compatible_filename(
                     task_id=0,
                     attempt_number=0,
-                    compression=compression_option,
+                    compression=compression,
                     format_ext=extension,
                 )
                 temp_file_prefix_on_stage = f"{write_path}/{spark_filename}"
@@ -217,29 +224,11 @@ def map_write(request: proto_base.ExecutePlanRequest):
                 spark_filename_prefix = generate_spark_compatible_filename(
                     task_id=0,
                     attempt_number=0,
-                    compression=compression_option,
+                    compression=None,
                     format_ext="",  # No extension for prefix
                 )
                 temp_file_prefix_on_stage = f"{write_path}/{spark_filename_prefix}"
 
-            default_compression = "NONE" if write_op.source != "parquet" else "snappy"
-            compression = write_op.options.get(
-                "compression", default_compression
-            ).upper()
-
-            if not is_supported_compression(write_op.source, compression):
-                supported_compressions = supported_compressions_for_format(
-                    write_op.source
-                )
-                raise AnalysisException(
-                    f"Compression {compression} is not supported for {write_op.source} format. "
-                    + (
-                        f"Supported compressions: {sorted(supported_compressions)}"
-                        if supported_compressions
-                        else "No compression supported for this format."
-                    )
-                )
-
             parameters = {
                 "location": temp_file_prefix_on_stage,
                 "file_format_type": write_op.source
@@ -254,9 +243,6 @@ def map_write(request: proto_base.ExecutePlanRequest):
             # Using the base avoids coupling to exact filenames/prefixes.
             download_stage_path = write_path
 
-            # 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:
@@ -314,7 +300,7 @@ def map_write(request: proto_base.ExecutePlanRequest):
                     per_part_prefix = generate_spark_compatible_filename(
                         task_id=part_idx,
                         attempt_number=0,
-                        compression=compression_option,
+                        compression=None,
                         format_ext="",  # prefix only; Snowflake appends extension/counters
                     )
                     part_params["location"] = f"{write_path}/{per_part_prefix}"

{snowpark_connect-0.30.0 → snowpark_connect-0.30.1}/src/snowflake/snowpark_connect/utils/session.py

@@ -120,7 +120,6 @@ def configure_snowpark_session(session: snowpark.Session):
         "PYTHON_SNOWPARK_USE_SCOPED_TEMP_OBJECTS": "false",  # this is required for creating udfs from sproc
         "ENABLE_STRUCTURED_TYPES_IN_SNOWPARK_CONNECT_RESPONSE": "true",
         "QUERY_TAG": f"'{query_tag}'",
-        "FEATURE_INTERVAL_TYPES": "enabled",
     }
 
     session.sql(

{snowpark_connect-0.30.0 → snowpark_connect-0.30.1}/src/snowflake/snowpark_connect/utils/telemetry.py

@@ -11,8 +11,8 @@ from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import Iterable
 from contextvars import ContextVar
+from dataclasses import dataclass
 from enum import Enum, unique
-from typing import Dict
 
 import google.protobuf.message
 import pyspark.sql.connect.proto.base_pb2 as proto_base
@@ -57,6 +57,7 @@ class TelemetryType(Enum):
 
 class EventType(Enum):
     SERVER_STARTED = "scos_server_started"
+    WARNING = "scos_warning"
 
 
 # global labels
@@ -106,7 +107,16 @@ REDACTED_PLAN_SUFFIXES = [
 ]
 
 
-def _basic_telemetry_data() -> Dict:
+@dataclass
+class TelemetryMessage:
+    """Container for telemetry messages in the processing queue."""
+
+    message: dict
+    timestamp: int
+    is_warning: bool
+
+
+def _basic_telemetry_data() -> dict:
     return {
         **STATIC_TELEMETRY_DATA,
         TelemetryField.KEY_EVENT_ID.value: str(uuid.uuid4()),
@@ -123,9 +133,11 @@ def safe(func):
     def wrap(*args, **kwargs):
         try:
             func(*args, **kwargs)
-        except Exception:
-            # We don't really care if telemetry fails, just want to be safe for the user
-            logger.warning(f"Telemetry operation failed: {func}", exc_info=True)
+        except Exception as e:
+            # report failed operation to telemetry
+            telemetry.send_warning_msg(
+                f"Telemetry operation {func} failed due to exception", e
+            )
 
     return wrap
 
@@ -528,8 +540,8 @@ class Telemetry:
     @safe
     def send_request_summary_telemetry(self):
         if self._not_in_request():
-            logger.warning(
-                "Truing to send request summary telemetry without initializing it"
+            self.send_warning_msg(
+                "Trying to send request summary telemetry without initializing it"
             )
             return
 
@@ -541,14 +553,56 @@ class Telemetry:
         }
         self._send(message)
 
-    def _send(self, msg: Dict) -> None:
+    def send_warning_msg(self, msg: str, e: Exception = None) -> None:
+        # using this within @safe decorator may result in recursive loop
+        try:
+            message = self._build_warning_message(msg, e)
+            if not message:
+                return
+
+            self._send(message, is_warning=True)
+        except Exception:
+            # if there's an exception here, there's nothing we can really do about it
+            pass
+
+    def _build_warning_message(self, warning_msg: str, e: Exception = None) -> dict:
+        try:
+            data = {"warning_message": warning_msg}
+            if e is not None:
+                data["exception"] = repr(e)
+
+            # add session and operation id if available
+            spark_session_id = self._request_summary.get().get("spark_session_id", None)
+            if spark_session_id is not None:
+                data["spark_session_id"] = spark_session_id
+
+            spark_operation_id = self._request_summary.get().get(
+                "spark_operation_id", None
+            )
+            if spark_operation_id is not None:
+                data["spark_operation_id"] = spark_operation_id
+
+            message = {
+                **_basic_telemetry_data(),
+                TelemetryField.KEY_TYPE.value: TelemetryType.TYPE_EVENT.value,
+                TelemetryType.EVENT_TYPE.value: EventType.WARNING.value,
+                TelemetryField.KEY_DATA.value: data,
+            }
+            return message
+        except Exception:
+            return {}
+
+    def _send(self, msg: dict, is_warning: bool = False) -> None:
         """Queue a telemetry message for asynchronous processing."""
         if not self._is_enabled:
             return
 
         timestamp = get_time_millis()
         try:
-            self._message_queue.put_nowait((msg, timestamp))
+            telemetry_msg = TelemetryMessage(
+                message=msg, timestamp=timestamp, is_warning=is_warning
+            )
+            self._message_queue.put_nowait(telemetry_msg)
         except queue.Full:
             # If queue is full, drop the message to avoid blocking
             logger.warning("Telemetry queue is full, dropping message")
@@ -566,13 +620,16 @@ class Telemetry:
         while True:
             try:
                 # block to allow the GIL to switch threads
-                message, timestamp = self._message_queue.get()
-                if timestamp is None and message is None:
-                    # shutdown
+                telemetry_msg = self._message_queue.get()
+                if telemetry_msg is None:
+                    # shutdown signal
                     break
-                self._sink.add_telemetry_data(message, timestamp)
-            except Exception:
-                logger.warning("Failed to add telemetry message to sink", exc_info=True)
+                self._sink.add_telemetry_data(
+                    telemetry_msg.message, telemetry_msg.timestamp
+                )
+            except Exception as e:
+                if not telemetry_msg.is_warning:
+                    self.send_warning_msg("Failed to add telemetry message to sink", e)
             finally:
                 self._message_queue.task_done()
 
@@ -585,7 +642,7 @@ class Telemetry:
             return
 
         try:
-            self._message_queue.put_nowait((None, None))
+            self._message_queue.put_nowait(None)
             # Wait for worker thread to finish
             self._worker_thread.join(timeout=3.0)
         except Exception:
@@ -594,7 +651,7 @@ class Telemetry:
             )
 
 
-def _error_location(e: Exception) -> Dict | None:
+def _error_location(e: Exception) -> dict | None:
     """
     Inspect the exception traceback and extract the file name, line number, and function name
     from the last frame (the one that raised the exception).
@@ -675,7 +732,7 @@ def _protobuf_to_json_with_redaction(
         """Recursively convert protobuf message to dict"""
 
         if not isinstance(msg, google.protobuf.message.Message):
-            logger.warning("Expected a protobuf message, got: %s", type(msg))
+            telemetry.send_warning_msg(f"Expected a protobuf message, got: {type(msg)}")
             return {}
 
         result = {}

{snowpark_connect-0.30.0 → snowpark_connect-0.30.1}/src/snowflake/snowpark_connect/version.py

@@ -2,4 +2,4 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-VERSION = (0,30,0)
+VERSION = (0,30,1)

{snowpark_connect-0.30.0 → snowpark_connect-0.30.1/src/snowpark_connect.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-connect
-Version: 0.30.0
+Version: 0.30.1
 Summary: Snowpark Connect for Spark
 Author: Snowflake, Inc
 License: Apache License, Version 2.0

{snowpark_connect-0.30.0 → snowpark_connect-0.30.1}/src/snowpark_connect.egg-info/SOURCES.txt

@@ -397,9 +397,11 @@ src/snowflake/snowpark_connect/proto/control_pb2.pyi
 src/snowflake/snowpark_connect/proto/control_pb2_grpc.py
 src/snowflake/snowpark_connect/proto/snowflake_expression_ext_pb2.py
 src/snowflake/snowpark_connect/proto/snowflake_expression_ext_pb2.pyi
+src/snowflake/snowpark_connect/proto/snowflake_expression_ext_pb2_grpc.py
 src/snowflake/snowpark_connect/proto/snowflake_rdd_pb2.pyi
 src/snowflake/snowpark_connect/proto/snowflake_relation_ext_pb2.py
 src/snowflake/snowpark_connect/proto/snowflake_relation_ext_pb2.pyi
+src/snowflake/snowpark_connect/proto/snowflake_relation_ext_pb2_grpc.py
 src/snowflake/snowpark_connect/relation/__init__.py
 src/snowflake/snowpark_connect/relation/io_utils.py
 src/snowflake/snowpark_connect/relation/map_aggregate.py