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

snowflake/snowpark_connect/relation/write/map_write.py

@@ -50,7 +50,9 @@ from snowflake.snowpark_connect.utils.identifiers import (
     spark_to_sf_single_id,
     split_fully_qualified_spark_name,
 )
+from snowflake.snowpark_connect.utils.io_utils import get_table_type
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
     telemetry,
@@ -160,6 +162,29 @@ def map_write(request: proto_base.ExecutePlanRequest):
             compression_option = write_op.options.get("compression", "none")
 
             # Generate Spark-compatible filename or prefix
+            # we need a random prefix to support "append" mode
+            # otherwise copy into with overwrite=False will fail if the file already exists
+            overwrite = (
+                write_op.mode
+                == commands_proto.WriteOperation.SaveMode.SAVE_MODE_OVERWRITE
+            )
+
+            if overwrite:
+                try:
+                    path_after_stage = (
+                        write_path.split("/", 1)[1] if "/" in write_path else ""
+                    )
+                    if not path_after_stage or path_after_stage == "/":
+                        logger.warning(
+                            f"Skipping REMOVE for root path {write_path} - too broad scope"
+                        )
+                    else:
+                        remove_command = f"REMOVE {write_path}/"
+                        session.sql(remove_command).collect()
+                        logger.info(f"Successfully cleared directory: {write_path}")
+                except Exception as e:
+                    logger.warning(f"Could not clear directory {write_path}: {e}")
+
             if should_write_to_single_file:
                 # Single file: generate complete filename with extension
                 spark_filename = generate_spark_compatible_filename(
@@ -178,10 +203,6 @@ def map_write(request: proto_base.ExecutePlanRequest):
                     format_ext="",  # No extension for prefix
                 )
                 temp_file_prefix_on_stage = f"{write_path}/{spark_filename_prefix}"
-            overwrite = (
-                write_op.mode
-                == commands_proto.WriteOperation.SaveMode.SAVE_MODE_OVERWRITE
-            )
 
             default_compression = "NONE" if write_op.source != "parquet" else "snappy"
             compression = write_op.options.get(
@@ -291,7 +312,10 @@ def map_write(request: proto_base.ExecutePlanRequest):
 
             match write_mode:
                 case None | "error" | "errorifexists":
-                    if check_snowflake_table_existence(snowpark_table_name, session):
+                    table_schema_or_error = _get_table_schema_or_error(
+                        snowpark_table_name, session
+                    )
+                    if isinstance(table_schema_or_error, DataType):  # Table exists
                         raise AnalysisException(
                             f"Table {snowpark_table_name} already exists"
                         )
@@ -302,29 +326,45 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         snowpark_session=session,
                     )
                     _validate_schema_and_get_writer(
-                        input_df, "append", snowpark_table_name
+                        input_df, "append", snowpark_table_name, table_schema_or_error
                     ).saveAsTable(
                         table_name=snowpark_table_name,
                         mode="append",
                         column_order=_column_order_for_write,
                     )
                 case "append":
-                    # TODO: SNOW-2299414 Fix the implementation of table type check
-                    # if check_table_type(snowpark_table_name, session) != "ICEBERG":
-                    #     raise AnalysisException(
-                    #         f"Table {snowpark_table_name} is not an iceberg table"
-                    #     )
+                    table_schema_or_error = _get_table_schema_or_error(
+                        snowpark_table_name, session
+                    )
+                    if isinstance(table_schema_or_error, DataType):  # Table exists
+                        if get_table_type(snowpark_table_name, session) not in (
+                            "ICEBERG",
+                            "TABLE",
+                        ):
+                            raise AnalysisException(
+                                f"Table {snowpark_table_name} is not an iceberg table"
+                            )
+                    else:
+                        create_iceberg_table(
+                            snowpark_table_name=snowpark_table_name,
+                            location=write_op.options.get("location", None),
+                            schema=input_df.schema,
+                            snowpark_session=session,
+                        )
                     _validate_schema_and_get_writer(
-                        input_df, "append", snowpark_table_name
+                        input_df, "append", snowpark_table_name, table_schema_or_error
                     ).saveAsTable(
                         table_name=snowpark_table_name,
                         mode="append",
                         column_order=_column_order_for_write,
                     )
                 case "ignore":
-                    if not check_snowflake_table_existence(
+                    table_schema_or_error = _get_table_schema_or_error(
                         snowpark_table_name, session
-                    ):
+                    )
+                    if not isinstance(
+                        table_schema_or_error, DataType
+                    ):  # Table not exists
                         create_iceberg_table(
                             snowpark_table_name=snowpark_table_name,
                             location=write_op.options.get("location", None),
@@ -339,13 +379,17 @@ def map_write(request: proto_base.ExecutePlanRequest):
                             column_order=_column_order_for_write,
                         )
                 case "overwrite":
-                    if check_snowflake_table_existence(snowpark_table_name, session):
-                        # TODO: SNOW-2299414 Fix the implementation of table type check
-                        # if check_table_type(snowpark_table_name, session) != "ICEBERG":
-                        #     raise AnalysisException(
-                        #         f"Table {snowpark_table_name} is not an iceberg table"
-                        #     )
-                        pass
+                    table_schema_or_error = _get_table_schema_or_error(
+                        snowpark_table_name, session
+                    )
+                    if isinstance(table_schema_or_error, DataType):  # Table exists
+                        if get_table_type(snowpark_table_name, session) not in (
+                            "ICEBERG",
+                            "TABLE",
+                        ):
+                            raise AnalysisException(
+                                f"Table {snowpark_table_name} is not an iceberg table"
+                            )
                     else:
                         create_iceberg_table(
                             snowpark_table_name=snowpark_table_name,
@@ -354,7 +398,7 @@ def map_write(request: proto_base.ExecutePlanRequest):
                             snowpark_session=session,
                         )
                     _validate_schema_and_get_writer(
-                        input_df, "truncate", snowpark_table_name
+                        input_df, "truncate", snowpark_table_name, table_schema_or_error
                     ).saveAsTable(
                         table_name=snowpark_table_name,
                         mode="truncate",
@@ -373,33 +417,49 @@ def map_write(request: proto_base.ExecutePlanRequest):
             ):
                 match write_mode:
                     case "overwrite":
-                        if check_snowflake_table_existence(
+                        table_schema_or_error = _get_table_schema_or_error(
                             snowpark_table_name, session
-                        ):
-                            # TODO: SNOW-2299414 Fix the implementation of table type check
-                            # if (
-                            #     check_table_type(snowpark_table_name, session)
-                            #     != "TABLE"
-                            # ):
-                            #     raise AnalysisException(
-                            #         f"Table {snowpark_table_name} is not a FDN table"
-                            #     )
+                        )
+                        if isinstance(table_schema_or_error, DataType):  # Table exists
+                            if get_table_type(snowpark_table_name, session) not in (
+                                "NORMAL",
+                                "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
+                            input_df,
+                            write_mode,
+                            snowpark_table_name,
+                            table_schema_or_error,
                         ).saveAsTable(
                             table_name=snowpark_table_name,
                             mode=write_mode,
                             column_order=_column_order_for_write,
                         )
                     case "append":
-                        # TODO: SNOW-2299414 Fix the implementation of table type check
-                        # if check_table_type(snowpark_table_name, session) != "TABLE":
-                        #     raise AnalysisException(
-                        #         f"Table {snowpark_table_name} is not a FDN table"
-                        #     )
+                        table_schema_or_error = _get_table_schema_or_error(
+                            snowpark_table_name, session
+                        )
+                        if isinstance(
+                            table_schema_or_error, DataType
+                        ) and get_table_type(  # Table exists
+                            snowpark_table_name, session
+                        ) not in (
+                            "NORMAL",
+                            "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
+                            input_df,
+                            write_mode,
+                            snowpark_table_name,
+                            table_schema_or_error,
                         ).saveAsTable(
                             table_name=snowpark_table_name,
                             mode=write_mode,
@@ -446,7 +506,10 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
     if write_op.provider.lower() == "iceberg":
         match write_op.mode:
             case commands_proto.WriteOperationV2.MODE_CREATE:
-                if check_snowflake_table_existence(snowpark_table_name, session):
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if isinstance(table_schema_or_error, DataType):  # Table exists
                     raise AnalysisException(
                         f"Table {snowpark_table_name} already exists"
                     )
@@ -457,24 +520,29 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     snowpark_session=session,
                 )
                 _validate_schema_and_get_writer(
-                    input_df, "append", snowpark_table_name
+                    input_df, "append", snowpark_table_name, table_schema_or_error
                 ).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):
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if not isinstance(table_schema_or_error, DataType):  # Table not exists
                     raise AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found."
                     )
-                # TODO: SNOW-2299414 Fix the implementation of table type check
-                # if check_table_type(snowpark_table_name, session) != "ICEBERG":
-                #     raise AnalysisException(
-                #         f"Table {snowpark_table_name} is not an iceberg table"
-                #     )
+                if get_table_type(snowpark_table_name, session) not in (
+                    "ICEBERG",
+                    "TABLE",
+                ):
+                    raise AnalysisException(
+                        f"Table {snowpark_table_name} is not an iceberg table"
+                    )
                 _validate_schema_and_get_writer(
-                    input_df, "append", snowpark_table_name
+                    input_df, "append", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
                     table_name=snowpark_table_name,
                     mode="append",
@@ -482,26 +550,33 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                 )
             case commands_proto.WriteOperationV2.MODE_OVERWRITE | commands_proto.WriteOperationV2.MODE_OVERWRITE_PARTITIONS:
                 # TODO: handle the filter condition for MODE_OVERWRITE
-                if check_snowflake_table_existence(snowpark_table_name, session):
-                    # TODO: SNOW-2299414 Fix the implementation of table type check
-                    # if check_table_type(snowpark_table_name, session) != "ICEBERG":
-                    #     raise AnalysisException(
-                    #         f"Table {snowpark_table_name} is not an iceberg table"
-                    #     )
-                    pass
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if isinstance(table_schema_or_error, DataType):  # Table exists
+                    if get_table_type(snowpark_table_name, session) not in (
+                        "ICEBERG",
+                        "TABLE",
+                    ):
+                        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
+                    input_df, "truncate", snowpark_table_name, table_schema_or_error
                 ).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):
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if isinstance(table_schema_or_error, DataType):  # Table exists
                     create_iceberg_table(
                         snowpark_table_name=snowpark_table_name,
                         location=write_op.table_properties.get("location"),
@@ -514,7 +589,7 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                         f"Table {snowpark_table_name} does not exist"
                     )
                 _validate_schema_and_get_writer(
-                    input_df, "replace", snowpark_table_name
+                    input_df, "replace", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
                     table_name=snowpark_table_name,
                     mode="append",
@@ -550,17 +625,22 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     column_order=_column_order_for_write,
                 )
             case commands_proto.WriteOperationV2.MODE_APPEND:
-                if not check_snowflake_table_existence(snowpark_table_name, session):
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if not isinstance(table_schema_or_error, DataType):  # Table not exists
                     raise AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found."
                     )
-                # TODO: SNOW-2299414 Fix the implementation of table type check
-                # if check_table_type(snowpark_table_name, session) != "TABLE":
-                #     raise AnalysisException(
-                #         f"Table {snowpark_table_name} is not a FDN table"
-                #     )
+                if get_table_type(snowpark_table_name, session) not in (
+                    "NORMAL",
+                    "TABLE",
+                ):
+                    raise AnalysisException(
+                        f"Table {snowpark_table_name} is not a FDN table"
+                    )
                 _validate_schema_and_get_writer(
-                    input_df, "append", snowpark_table_name
+                    input_df, "append", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
                     table_name=snowpark_table_name,
                     mode="append",
@@ -568,31 +648,38 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                 )
             case commands_proto.WriteOperationV2.MODE_OVERWRITE | commands_proto.WriteOperationV2.MODE_OVERWRITE_PARTITIONS:
                 # TODO: handle the filter condition for MODE_OVERWRITE
-                if check_snowflake_table_existence(snowpark_table_name, session):
-                    # TODO: SNOW-2299414 Fix the implementation of table type check
-                    # if check_table_type(snowpark_table_name, session) != "TABLE":
-                    #     raise AnalysisException(
-                    #         f"Table {snowpark_table_name} is not a FDN table"
-                    #     )
-                    pass
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if isinstance(table_schema_or_error, DataType):  # Table exists
+                    if get_table_type(snowpark_table_name, session) not in (
+                        "NORMAL",
+                        "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
+                    input_df, "truncate", snowpark_table_name, table_schema_or_error
                 ).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):
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if not isinstance(table_schema_or_error, DataType):  # Table not exists
                     raise AnalysisException(
                         f"Table {snowpark_table_name} does not exist"
                     )
                 _validate_schema_and_get_writer(
-                    input_df, "replace", snowpark_table_name
+                    input_df, "replace", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
                     table_name=snowpark_table_name,
                     mode="overwrite",
@@ -612,8 +699,20 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                 )
 
 
+def _get_table_schema_or_error(
+    snowpark_table_name: str, snowpark_session: snowpark.Session
+) -> DataType | SnowparkSQLException:
+    try:
+        return snowpark_session.table(snowpark_table_name).schema
+    except SnowparkSQLException as e:
+        return e
+
+
 def _validate_schema_and_get_writer(
-    input_df: snowpark.DataFrame, write_mode: str, snowpark_table_name: str
+    input_df: snowpark.DataFrame,
+    write_mode: str,
+    snowpark_table_name: str,
+    table_schema_or_error: DataType | SnowparkSQLException | None = None,
 ) -> snowpark.DataFrameWriter:
     if write_mode is not None and write_mode.lower() in (
         "replace",
@@ -622,16 +721,26 @@ def _validate_schema_and_get_writer(
         return input_df.write
 
     table_schema = None
-    try:
-        table_schema = (
-            get_or_create_snowpark_session().table(snowpark_table_name).schema
-        )
-    except SnowparkSQLException as e:
-        msg = e.message
-        if "SQL compilation error" in msg and "does not exist" in msg:
-            pass
-        else:
-            raise e
+    if table_schema_or_error is not None:
+        if isinstance(table_schema_or_error, SnowparkSQLException):
+            msg = table_schema_or_error.message
+            if "SQL compilation error" in msg and "does not exist" in msg:
+                pass
+            else:
+                raise table_schema_or_error
+        elif isinstance(table_schema_or_error, DataType):
+            table_schema = table_schema_or_error
+    else:
+        try:
+            table_schema = (
+                get_or_create_snowpark_session().table(snowpark_table_name).schema
+            )
+        except SnowparkSQLException as e:
+            msg = e.message
+            if "SQL compilation error" in msg and "does not exist" in msg:
+                pass
+            else:
+                raise e
 
     if table_schema is None:
         # If table does not exist, we can skip the schema validation
@@ -884,31 +993,3 @@ def _truncate_directory(directory_path: Path) -> None:
             file.unlink()
         elif file.is_dir():
             shutil.rmtree(file)
-
-
-def check_snowflake_table_existence(
-    snowpark_table_name: str,
-    snowpark_session: snowpark.Session,
-):
-    try:
-        snowpark_session.sql(f"SELECT 1 FROM {snowpark_table_name} LIMIT 1").collect()
-        return True
-    except Exception:
-        return False
-
-
-# TODO: SNOW-2299414 Fix the implementation of table type check
-# def check_table_type(
-#     snowpark_table_name: str,
-#     snowpark_session: snowpark.Session,
-# ) -> str:
-#     # currently we only support iceberg table and FDN table
-#     metadata = snowpark_session.sql(
-#         f"SHOW TABLES LIKE '{unquote_if_quoted(snowpark_table_name)}';"
-#     ).collect()
-#     if metadata is None or len(metadata) == 0:
-#         raise AnalysisException(f"Table {snowpark_table_name} does not exist")
-#     metadata = metadata[0]
-#     if metadata.as_dict().get("is_iceberg") == "Y":
-#         return "ICEBERG"
-#     return "TABLE"

snowflake/snowpark_connect/resources_initializer.py

@@ -12,6 +12,7 @@ _resources_initialized = threading.Event()
 _initializer_lock = threading.Lock()
 SPARK_VERSION = "3.5.6"
 RESOURCE_PATH = "/snowflake/snowpark_connect/resources"
+_upload_jars = True  # Flag to control whether to upload jars. Required for Scala UDFs.
 
 
 def initialize_resources() -> None:
@@ -57,10 +58,8 @@ def initialize_resources() -> None:
             f"spark-sql_2.12-{SPARK_VERSION}.jar",
             f"spark-connect-client-jvm_2.12-{SPARK_VERSION}.jar",
             f"spark-common-utils_2.12-{SPARK_VERSION}.jar",
+            "sas-scala-udf_2.12-0.1.0.jar",
             "json4s-ast_2.12-3.7.0-M11.jar",
-            "json4s-native_2.12-3.7.0-M11.jar",
-            "json4s-core_2.12-3.7.0-M11.jar",
-            "paranamer-2.8.3.jar",
         ]
 
         for jar in jar_files:
@@ -80,9 +79,11 @@ def initialize_resources() -> None:
         ("Initialize Session Stage", initialize_session_stage),  # Takes about 0.3s
         ("Initialize Session Catalog", initialize_catalog),  # Takes about 1.2s
         ("Snowflake Connection Warm Up", warm_up_sf_connection),  # Takes about 1s
-        ("Upload Scala UDF Jars", upload_scala_udf_jars),
     ]
 
+    if _upload_jars:
+        resources.append(("Upload Scala UDF Jars", upload_scala_udf_jars))
+
     for name, resource_func in resources:
         resource_start = time.time()
         try:
@@ -113,4 +114,18 @@ def initialize_resources_async() -> threading.Thread:
 
 def wait_for_resource_initialization() -> None:
     with _initializer_lock:
-        _resource_initializer.join()
+        _resource_initializer.join(timeout=300)  # wait at most 300 seconds
+    if _resource_initializer.is_alive():
+        logger.error(
+            "Resource initialization failed - initializer thread has been running for over 300 seconds."
+        )
+        raise RuntimeError(
+            "Resource initialization failed - initializer thread has been running for over 300 seconds."
+        )
+
+
+def set_upload_jars(upload: bool) -> None:
+    """Set whether to upload jars required for Scala UDFs. This should be set to False if Scala UDFs
+    are not used, to avoid the overhead of uploading jars."""
+    global _upload_jars
+    _upload_jars = upload

snowflake/snowpark_connect/server.py

@@ -531,7 +531,10 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
             if name.endswith(".class"):
                 # name is <dir>/<package>/<class_name>
                 # we don't need the dir name, but require the package, so only remove dir
-                class_files[name.split("/", 1)[-1]] = filepath
+                if os.name != "nt":
+                    class_files[name.split("/", 1)[-1]] = filepath
+                else:
+                    class_files[name.split("\\", 1)[-1]] = filepath
                 continue
             session.file.put(
                 filepath,
@@ -722,30 +725,33 @@ def _serve(
             # No need to start grpc server in TCM
             return
 
+        grpc_max_msg_size = 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,
+        )
         server_options = [
             (
                 "grpc.max_receive_message_length",
-                get_int_from_env(
-                    "SNOWFLAKE_GRPC_MAX_MESSAGE_SIZE",
-                    _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE,
-                ),
+                grpc_max_msg_size,
             ),
             (
                 "grpc.max_metadata_size",
-                get_int_from_env(
-                    "SNOWFLAKE_GRPC_MAX_METADATA_SIZE",
-                    _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE,
-                ),
+                grpc_max_metadata_size,
             ),
             (
                 "grpc.absolute_max_metadata_size",
-                get_int_from_env(
-                    "SNOWFLAKE_GRPC_MAX_METADATA_SIZE",
-                    _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE,
-                )
-                * 2,
+                grpc_max_metadata_size * 2,
             ),
         ]
+
+        from pyspark.sql.connect.client import ChannelBuilder
+
+        ChannelBuilder.MAX_MESSAGE_LENGTH = grpc_max_msg_size
+
         server = grpc.server(
             futures.ThreadPoolExecutor(max_workers=10), options=server_options
         )
@@ -1050,10 +1056,6 @@ def start_session(
         global _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE
         _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE = max_grpc_message_size
 
-        from pyspark.sql.connect.client import ChannelBuilder
-
-        ChannelBuilder.MAX_MESSAGE_LENGTH = max_grpc_message_size
-
         if os.environ.get("SPARK_ENV_LOADED"):
             raise RuntimeError(
                 "Snowpark Connect cannot be run inside of a Spark environment"

snowflake/snowpark_connect/utils/artifacts.py

@@ -39,7 +39,7 @@ def write_temporary_artifact(
     if os.name != "nt":
         filepath = f"/tmp/sas-{session.session_id}/{name}"
     else:
-        filepath = f"{tempfile.gettempdir()}/sas-{session.session_id}/{name}"
+        filepath = f"{tempfile.gettempdir()}\\sas-{session.session_id}\\{name}"
     # The name comes to us as a path (e.g. cache/<name>), so we need to create
     # the parent directory if it doesn't exist to avoid errors during writing.
     pathlib.Path(filepath).parent.mkdir(parents=True, exist_ok=True)
@@ -55,11 +55,10 @@ def write_class_files_to_stage(
 ) -> None:
     if os.name != "nt":
         filepath = f"/tmp/sas-{session.session_id}"
+        jar_name = f'{filepath}/{hashlib.sha256(str(files).encode("utf-8")).hexdigest()[:10]}.jar'
     else:
-        filepath = f"{tempfile.gettempdir()}/sas-{session.session_id}"
-    jar_name = (
-        f'{filepath}/{hashlib.sha256(str(files).encode("utf-8")).hexdigest()[:10]}.jar'
-    )
+        filepath = f"{tempfile.gettempdir()}\\sas-{session.session_id}"
+        jar_name = f'{filepath}\\{hashlib.sha256(str(files).encode("utf-8")).hexdigest()[:10]}.jar'
     with zipfile.ZipFile(jar_name, "w", zipfile.ZIP_DEFLATED) as jar:
         for name, path in files.items():
             jar.write(path, name)

snowflake/snowpark_connect/utils/concurrent.py

@@ -52,6 +52,10 @@ class SynchronizedDict(Mapping[K, V]):
         with self._lock.writer():
             self._dict[key] = value
 
+    def __delitem__(self, key: K) -> None:
+        with self._lock.writer():
+            del self._dict[key]
+
     def __contains__(self, key: K) -> bool:
         with self._lock.reader():
             return key in self._dict