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

snowflake/snowpark_connect/relation/write/map_write.py

@@ -50,6 +50,7 @@ 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 (
@@ -311,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"
                         )
@@ -322,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),
@@ -359,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,
@@ -374,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",
@@ -393,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,
@@ -466,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"
                     )
@@ -477,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",
@@ -502,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"),
@@ -534,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",
@@ -570,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",
@@ -588,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",
@@ -632,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",
@@ -642,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
@@ -904,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

@@ -725,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
         )
@@ -1053,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/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