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

snowflake/snowpark_connect/relation/write/map_write.py

@@ -234,6 +234,10 @@ 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"
+                        )
                     _validate_schema_and_get_writer(
                         input_df, "append", snowpark_table_name
                     ).saveAsTable(
@@ -260,7 +264,10 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         )
                 case "overwrite":
                     if check_snowflake_table_existence(snowpark_table_name, session):
-                        session.sql(f"DELETE FROM {snowpark_table_name}").collect()
+                        if check_table_type(snowpark_table_name, session) != "ICEBERG":
+                            raise AnalysisException(
+                                f"Table {snowpark_table_name} is not an iceberg table"
+                            )
                     else:
                         create_iceberg_table(
                             snowpark_table_name=snowpark_table_name,
@@ -269,10 +276,10 @@ def map_write(request: proto_base.ExecutePlanRequest):
                             snowpark_session=session,
                         )
                     _validate_schema_and_get_writer(
-                        input_df, "append", snowpark_table_name
+                        input_df, "truncate", snowpark_table_name
                     ).saveAsTable(
                         table_name=snowpark_table_name,
-                        mode="append",
+                        mode="truncate",
                         column_order=_column_order_for_write,
                     )
                 case _:
@@ -286,13 +293,46 @@ def map_write(request: proto_base.ExecutePlanRequest):
                 write_op.table.save_method
                 == commands_proto.WriteOperation.SaveTable.TableSaveMethod.TABLE_SAVE_METHOD_SAVE_AS_TABLE
             ):
-                _validate_schema_and_get_writer(
-                    input_df, write_mode, snowpark_table_name
-                ).saveAsTable(
-                    table_name=snowpark_table_name,
-                    mode=write_mode,
-                    column_order=_column_order_for_write,
-                )
+                match write_mode:
+                    case "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"
+                                )
+                            write_mode = "truncate"
+                        _validate_schema_and_get_writer(
+                            input_df, write_mode, snowpark_table_name
+                        ).saveAsTable(
+                            table_name=snowpark_table_name,
+                            mode=write_mode,
+                            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"
+                            )
+                        _validate_schema_and_get_writer(
+                            input_df, write_mode, snowpark_table_name
+                        ).saveAsTable(
+                            table_name=snowpark_table_name,
+                            mode=write_mode,
+                            column_order=_column_order_for_write,
+                        )
+                    case _:
+                        _validate_schema_and_get_writer(
+                            input_df, write_mode, snowpark_table_name
+                        ).saveAsTable(
+                            table_name=snowpark_table_name,
+                            mode=write_mode,
+                            column_order=_column_order_for_write,
+                        )
             elif (
                 write_op.table.save_method
                 == commands_proto.WriteOperation.SaveTable.TableSaveMethod.TABLE_SAVE_METHOD_INSERT_INTO
@@ -312,21 +352,6 @@ def map_write(request: proto_base.ExecutePlanRequest):
 
 def map_write_v2(request: proto_base.ExecutePlanRequest):
     write_op = request.plan.command.write_operation_v2
-    match write_op.mode:
-        case commands_proto.WriteOperationV2.MODE_APPEND:
-            write_mode = "append"
-        case commands_proto.WriteOperationV2.MODE_CREATE:
-            write_mode = "errorifexists"
-        case commands_proto.WriteOperationV2.MODE_OVERWRITE:
-            write_mode = "overwrite"
-        case commands_proto.WriteOperationV2.MODE_REPLACE:
-            write_mode = "overwrite"
-        case commands_proto.WriteOperationV2.MODE_CREATE_OR_REPLACE:
-            write_mode = "overwrite"
-        case _:
-            raise SnowparkConnectNotImplementedError(
-                f"Write operation {write_op.mode} not implemented."
-            )
 
     snowpark_table_name = _spark_to_snowflake(write_op.table_name)
     result = map_relation(write_op.input)
@@ -338,55 +363,176 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
             "Write operation V2 only support table writing now"
         )
 
-    # For OVERWRITE and APPEND modes, check if table exists first - Spark requires table to exist for these operations
-    if write_op.mode in (
-        commands_proto.WriteOperationV2.MODE_OVERWRITE,
-        commands_proto.WriteOperationV2.MODE_APPEND,
-    ):
-        if not check_snowflake_table_existence(snowpark_table_name, session):
-            raise AnalysisException(
-                f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found. "
-                f"Verify the spelling and correctness of the schema and catalog.\n"
-            )
-
     if write_op.provider.lower() == "iceberg":
-        if write_mode == "overwrite" and check_snowflake_table_existence(
-            snowpark_table_name, session
-        ):
-            session.sql(f"DELETE FROM {snowpark_table_name}").collect()
-            write_mode = "append"
-
-        if write_mode in (
-            "errorifexists",
-            "overwrite",
-        ):
-            create_iceberg_table(
-                snowpark_table_name=snowpark_table_name,
-                location=write_op.table_properties.get("location"),
-                schema=input_df.schema,
-                snowpark_session=session,
-            )
-        _validate_schema_and_get_writer(
-            input_df, write_mode, snowpark_table_name
-        ).saveAsTable(
-            table_name=snowpark_table_name,
-            mode="append",
-            column_order=_column_order_for_write,
-        )
+        match write_op.mode:
+            case commands_proto.WriteOperationV2.MODE_CREATE:
+                if check_snowflake_table_existence(snowpark_table_name, session):
+                    raise AnalysisException(
+                        f"Table {snowpark_table_name} already exists"
+                    )
+                create_iceberg_table(
+                    snowpark_table_name=snowpark_table_name,
+                    location=write_op.table_properties.get("location"),
+                    schema=input_df.schema,
+                    snowpark_session=session,
+                )
+                _validate_schema_and_get_writer(
+                    input_df, "append", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="append",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_APPEND:
+                if not check_snowflake_table_existence(snowpark_table_name, session):
+                    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"
+                    )
+                _validate_schema_and_get_writer(
+                    input_df, "append", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="append",
+                    column_order=_column_order_for_write,
+                )
+            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"
+                        )
+                else:
+                    raise AnalysisException(
+                        f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
+                    )
+                _validate_schema_and_get_writer(
+                    input_df, "truncate", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="truncate",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_REPLACE:
+                if check_snowflake_table_existence(snowpark_table_name, session):
+                    create_iceberg_table(
+                        snowpark_table_name=snowpark_table_name,
+                        location=write_op.table_properties.get("location"),
+                        schema=input_df.schema,
+                        snowpark_session=session,
+                        mode="replace",
+                    )
+                else:
+                    raise AnalysisException(
+                        f"Table {snowpark_table_name} does not exist"
+                    )
+                _validate_schema_and_get_writer(
+                    input_df, "replace", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="append",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_CREATE_OR_REPLACE:
+                create_iceberg_table(
+                    snowpark_table_name=snowpark_table_name,
+                    location=write_op.table_properties.get("location"),
+                    schema=input_df.schema,
+                    snowpark_session=session,
+                    mode="create_or_replace",
+                )
+                _validate_schema_and_get_writer(
+                    input_df, "create_or_replace", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="append",
+                    column_order=_column_order_for_write,
+                )
+            case _:
+                raise SnowparkConnectNotImplementedError(
+                    f"Write mode {commands_proto.WriteOperationV2.Mode.Name(write_op.mode)} is not supported"
+                )
     else:
-        _validate_schema_and_get_writer(
-            input_df, write_mode, snowpark_table_name
-        ).saveAsTable(
-            table_name=snowpark_table_name,
-            mode=write_mode,
-            column_order=_column_order_for_write,
-        )
+        match write_op.mode:
+            case commands_proto.WriteOperationV2.MODE_CREATE:
+                _validate_schema_and_get_writer(
+                    input_df, "errorifexists", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="errorifexists",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_APPEND:
+                if not check_snowflake_table_existence(snowpark_table_name, session):
+                    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"
+                    )
+                _validate_schema_and_get_writer(
+                    input_df, "append", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="append",
+                    column_order=_column_order_for_write,
+                )
+            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"
+                        )
+                else:
+                    raise AnalysisException(
+                        f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
+                    )
+                _validate_schema_and_get_writer(
+                    input_df, "truncate", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="truncate",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_REPLACE:
+                if not check_snowflake_table_existence(snowpark_table_name, session):
+                    raise AnalysisException(
+                        f"Table {snowpark_table_name} does not exist"
+                    )
+                _validate_schema_and_get_writer(
+                    input_df, "replace", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="overwrite",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_CREATE_OR_REPLACE:
+                _validate_schema_and_get_writer(
+                    input_df, "create_or_replace", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="overwrite",
+                    column_order=_column_order_for_write,
+                )
+            case _:
+                raise SnowparkConnectNotImplementedError(
+                    f"Write mode {commands_proto.WriteOperationV2.Mode.Name(write_op.mode)} is not supported"
+                )
 
 
 def _validate_schema_and_get_writer(
     input_df: snowpark.DataFrame, write_mode: str, snowpark_table_name: str
 ) -> snowpark.DataFrameWriter:
-    if write_mode == "overwrite":
+    if write_mode is not None and write_mode.lower() in (
+        "replace",
+        "create_or_replace",
+    ):
         return input_df.write
 
     table_schema = None
@@ -528,6 +674,7 @@ def create_iceberg_table(
     location: str,
     schema: StructType,
     snowpark_session: snowpark.Session,
+    mode: str = "create",
 ):
     table_schema = [
         f"{spark_to_sf_single_id(unquote_if_quoted(field.name), is_column = True)} {snowpark_to_iceberg_type(field.datatype)}"
@@ -550,8 +697,20 @@ def create_iceberg_table(
         else f"EXTERNAL_VOLUME = '{config_external_volume}'"
     )
 
+    match mode:
+        case "create":
+            create_sql = "CREATE"
+        case "replace":
+            # There's no replace for iceberg table, so we use create or replace
+            create_sql = "CREATE OR REPLACE"
+        case "create_or_replace":
+            create_sql = "CREATE OR REPLACE"
+        case _:
+            raise SnowparkConnectNotImplementedError(
+                f"Write mode {mode} is not supported for iceberg table"
+            )
     sql = f"""
-        CREATE ICEBERG TABLE {snowpark_table_name} ({",".join(table_schema)})
+        {create_sql} ICEBERG TABLE {snowpark_table_name} ({",".join(table_schema)})
         CATALOG = 'SNOWFLAKE'
         {external_volume}
         {base_location};
@@ -637,3 +796,19 @@ def check_snowflake_table_existence(
         return True
     except Exception:
         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"

snowflake/snowpark_connect/server.py

@@ -83,6 +83,7 @@ from snowflake.snowpark_connect.utils.context import (
     set_session_id,
     set_spark_version,
 )
+from snowflake.snowpark_connect.utils.env_utils import get_int_from_env
 from snowflake.snowpark_connect.utils.interrupt import (
     interrupt_all_queries,
     interrupt_queries_with_tag,
@@ -700,11 +701,27 @@ def _serve(
             return
 
         server_options = [
-            ("grpc.max_receive_message_length", _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE),
-            ("grpc.max_metadata_size", _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE),
+            (
+                "grpc.max_receive_message_length",
+                get_int_from_env(
+                    "SNOWFLAKE_GRPC_MAX_MESSAGE_SIZE",
+                    _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE,
+                ),
+            ),
+            (
+                "grpc.max_metadata_size",
+                get_int_from_env(
+                    "SNOWFLAKE_GRPC_MAX_METADATA_SIZE",
+                    _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE,
+                ),
+            ),
             (
                 "grpc.absolute_max_metadata_size",
-                _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE * 2,
+                get_int_from_env(
+                    "SNOWFLAKE_GRPC_MAX_METADATA_SIZE",
+                    _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE,
+                )
+                * 2,
             ),
         ]
         server = grpc.server(
@@ -812,8 +829,11 @@ class UnixDomainSocketChannelBuilder(ChannelBuilder):
     Spark Connect gRPC channel builder for Unix domain sockets
     """
 
-    def __init__(self, channelOptions: Optional[List[Tuple[str, Any]]] = None) -> None:
-        url: str = get_client_url()
+    def __init__(
+        self, url: str = None, channelOptions: Optional[List[Tuple[str, Any]]] = None
+    ) -> None:
+        if url is None:
+            url = get_client_url()
         if url[:6] != "unix:/" or len(url) < 7:
             raise PySparkValueError(
                 error_class="INVALID_CONNECT_URL",

snowflake/snowpark_connect/type_mapping.py

@@ -59,7 +59,7 @@ def _get_struct_type_class():
 
 
 @cache
-def _get_python_sql_utils_class():
+def get_python_sql_utils_class():
     return jpype.JClass("org.apache.spark.sql.api.python.PythonSQLUtils")
 
 
@@ -70,7 +70,7 @@ def parse_ddl_with_spark_scala(ddl_string: str) -> pyspark.sql.types.DataType:
     This mimics pysparks.ddl parsing logic pyspark.sql.types._py_parse_datatype_string
     """
     struct_type_class = _get_struct_type_class()
-    python_sql_utils = _get_python_sql_utils_class()
+    python_sql_utils = get_python_sql_utils_class()
 
     try:
         # DDL format, "fieldname datatype, fieldname datatype".

snowflake/snowpark_connect/utils/env_utils.py

@@ -0,0 +1,55 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+"""
+Environment variable utilities for Snowpark Connect.
+"""
+
+import os
+
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
+
+
+def get_int_from_env(env_var: str, default: int) -> int:
+    """
+    Safely get integer value from environment variable with fallback to default.
+
+    Args:
+        env_var: Environment variable name
+        default: Default integer value if env var is not set or invalid
+
+    Returns:
+        Integer value from environment variable or default
+
+    Raises:
+        TypeError: If default is not an integer
+
+    Examples:
+        >>> get_int_from_env("MAX_WORKERS", 10)
+        10
+        >>> os.environ["MAX_WORKERS"] = "20"
+        >>> get_int_from_env("MAX_WORKERS", 10)
+        20
+        >>> os.environ["MAX_WORKERS"] = "invalid"
+        >>> get_int_from_env("MAX_WORKERS", 10)  # logs warning, returns 10
+        10
+    """
+    # Validate that default is actually an integer
+    if not isinstance(default, int):
+        raise TypeError(
+            f"Default value must be an integer, got {type(default).__name__}: {default}"
+        )
+
+    value = os.getenv(env_var)
+    if value is None:
+        return default
+
+    try:
+        return int(value)
+    except ValueError:
+        logger.warning(
+            f"Invalid integer value for environment variable {env_var}: '{value}', "
+            f"using default: {default}"
+        )
+        return default

snowflake/snowpark_connect/utils/session.py

@@ -73,6 +73,27 @@ def configure_snowpark_session(session: snowpark.Session):
     session.connection.arrow_number_to_decimal_setter = True
     session.custom_package_usage_config["enabled"] = True
 
+    default_fallback_timezone = "UTC"
+    if global_config.spark_sql_session_timeZone is None:
+        try:
+            result = session.sql("SHOW PARAMETERS LIKE 'TIMEZONE'").collect()
+            if result and len(result) > 0:
+                value = result[0]["value"]
+                logger.warning(
+                    f"Using Snowflake session timezone parameter as fallback: {value}"
+                )
+            else:
+                value = default_fallback_timezone
+                logger.warning(
+                    f"Could not determine timezone from parameters, defaulting to {default_fallback_timezone}"
+                )
+        except Exception as e:
+            value = default_fallback_timezone
+            logger.warning(
+                f"Could not query Snowflake timezone parameter ({e}), defaulting to {default_fallback_timezone}"
+            )
+        global_config.spark_sql_session_timeZone = value
+
     session_params = {
         "TIMESTAMP_TYPE_MAPPING": "TIMESTAMP_LTZ",
         "TIMEZONE": f"'{global_config.spark_sql_session_timeZone}'",

snowflake/snowpark_connect/version.py

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

snowflake/snowpark_decoder/spark_decoder.py

@@ -72,5 +72,5 @@ class Context:
     def __init__(self) -> None:
         pass
 
-    def abort_with_status(status):
+    def abort_with_status(self, status):
         raise NotImplementedError("abort_with_status is not implemented")

{snowpark_connect-0.22.1.dist-info → snowpark_connect-0.24.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-connect
-Version: 0.22.1
+Version: 0.24.0
 Summary: Snowpark Connect for Spark
 Author: Snowflake, Inc
 License: Apache License, Version 2.0
@@ -16,7 +16,7 @@ Requires-Dist: jpype1
 Requires-Dist: protobuf<5.0,>=4.25.3
 Requires-Dist: s3fs>=2025.3.0
 Requires-Dist: snowflake.core<2,>=1.0.5
-Requires-Dist: snowflake-snowpark-python[pandas]!=1.35.0,!=1.36.0,>=1.34.0
+Requires-Dist: snowflake-snowpark-python[pandas]<1.38.0,>=1.37.0
 Requires-Dist: sqlglot>=26.3.8
 Requires-Dist: jaydebeapi
 Requires-Dist: aiobotocore~=2.23.0