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

snowflake/snowpark_connect/relation/map_sql.py

@@ -77,6 +77,7 @@ from ..expression.map_sql_expression import (
 from ..utils.identifiers import spark_to_sf_single_id
 
 _ctes = ContextVar[dict[str, relation_proto.Relation]]("_ctes", default={})
+_cte_definitions = ContextVar[dict[str, any]]("_cte_definitions", default={})
 _having_condition = ContextVar[expressions_proto.Expression | None](
     "_having_condition", default=None
 )
@@ -117,12 +118,15 @@ def _push_cte_scope():
     """
     Creates a new CTE scope when evaluating nested WITH clauses.
     """
-    cur = _ctes.get()
-    token = _ctes.set(cur.copy())
+    cur_ctes = _ctes.get()
+    cur_definitions = _cte_definitions.get()
+    cte_token = _ctes.set(cur_ctes.copy())
+    def_token = _cte_definitions.set(cur_definitions.copy())
     try:
         yield
     finally:
-        _ctes.reset(token)
+        _ctes.reset(cte_token)
+        _cte_definitions.reset(def_token)
 
 
 @contextmanager
@@ -405,15 +409,21 @@ def map_sql_to_pandas_df(
                 else:
                     # TODO: Unset the schema
                     pass
-            case "CreateTable":
+            case "CreateTable" | "ReplaceTable":
+                if class_name == "ReplaceTable":
+                    replace_table = " OR REPLACE "
+                    if_not_exists = ""
+                else:
+                    replace_table = ""
+                    if_not_exists = (
+                        "IF NOT EXISTS " if logical_plan.ignoreIfExists() else ""
+                    )
+
                 name = get_relation_identifier_name(logical_plan.name())
                 columns = ", ".join(
                     _spark_field_to_sql(f, True)
                     for f in logical_plan.tableSchema().fields()
                 )
-                if_not_exists = (
-                    "IF NOT EXISTS " if logical_plan.ignoreIfExists() else ""
-                )
                 comment_opt = logical_plan.tableSpec().comment()
                 comment = (
                     f"COMMENT = '{_escape_sql_comment(str(comment_opt.get()))}'"
@@ -422,7 +432,7 @@ def map_sql_to_pandas_df(
                 )
                 # NOTE: We are intentionally ignoring any FORMAT=... parameters here.
                 session.sql(
-                    f"CREATE TABLE {if_not_exists}{name} ({columns}) {comment}"
+                    f"CREATE {replace_table} TABLE {if_not_exists}{name} ({columns}) {comment}"
                 ).collect()
             case "CreateTableAsSelect":
                 mode = "ignore" if logical_plan.ignoreIfExists() else "errorifexists"
@@ -1351,8 +1361,28 @@ def map_logical_plan_relation(
             value_column_names = [e for e in as_java_list(rel.valueColumnNames())]
             variable_column_name = rel.variableColumnName()
 
+            # Check for multi-column UNPIVOT which Snowflake doesn't support
+            if len(value_column_names) > 1:
+                raise UnsupportedOperationException(
+                    f"Multi-column UNPIVOT is not supported. Snowflake SQL does not support unpivoting "
+                    f"multiple value columns ({', '.join(value_column_names)}) in a single operation. "
+                    f"Workaround: Use separate UNPIVOT operations for each value column and join the results, "
+                    f"or restructure your query to unpivot columns individually."
+                )
+
             values = []
-            for e1 in as_java_list(rel.values().get()):
+            values_groups = as_java_list(rel.values().get())
+
+            # Check if we have multi-column groups in the IN clause
+            if values_groups and len(as_java_list(values_groups[0])) > 1:
+                group_sizes = [len(as_java_list(group)) for group in values_groups]
+                raise UnsupportedOperationException(
+                    f"Multi-column UNPIVOT is not supported. Snowflake SQL does not support unpivoting "
+                    f"multiple columns together in groups. Found groups with {max(group_sizes)} columns. "
+                    f"Workaround: Unpivot each column separately and then join/union the results as needed."
+                )
+
+            for e1 in values_groups:
                 for e in as_java_list(e1):
                     values.append(map_logical_plan_expression(e))
 
@@ -1468,7 +1498,50 @@ def map_logical_plan_relation(
                 # The name corresponds to a `WITH` alias rather than a table.
                 # TODO: We currently evaluate the query each time its alias is used;
                 # we should eventually start using `WITH` in Snowflake SQL.
-                proto = cte_proto
+                # Each CTE reference should get completely fresh evaluation to prevent ambiguity
+                # when the same CTE is joined multiple times. Instead of reusing the same cte_proto,
+                # re-evaluate the CTE definition to get fresh column identifiers.
+
+                # Re-evaluate the CTE definition to get fresh column identifiers
+                cte_definition = _cte_definitions.get().get(name)
+                if cte_definition is not None:
+                    # Get the original column names for consistency across CTE references
+                    original_container = map_relation(cte_proto)
+                    original_spark_columns = (
+                        original_container.column_map.get_spark_columns()
+                    )
+
+                    # Re-evaluate the CTE definition with a fresh plan_id
+                    fresh_plan_id = gen_sql_plan_id()
+                    fresh_cte_proto = map_logical_plan_relation(
+                        cte_definition, fresh_plan_id
+                    )
+
+                    # Use SubqueryColumnAliases to ensure consistent column names across CTE references
+                    # This is crucial for CTEs that reference other CTEs
+                    any_proto = Any()
+                    any_proto.Pack(
+                        snowflake_proto.Extension(
+                            subquery_column_aliases=snowflake_proto.SubqueryColumnAliases(
+                                input=fresh_cte_proto,
+                                aliases=original_spark_columns,
+                            )
+                        )
+                    )
+                    column_aliased_proto = relation_proto.Relation(extension=any_proto)
+                    column_aliased_proto.common.plan_id = gen_sql_plan_id()
+
+                    # Wrap in SubqueryAlias with the CTE name
+                    proto = relation_proto.Relation(
+                        subquery_alias=relation_proto.SubqueryAlias(
+                            input=column_aliased_proto,
+                            alias=name,
+                        )
+                    )
+                    proto.common.plan_id = gen_sql_plan_id()
+                else:
+                    # Fallback to stored CTE if definition not found
+                    proto = cte_proto
             else:
                 tmp_views = _get_current_temp_objects()
                 current_schema = session.connection.schema
@@ -1600,7 +1673,11 @@ def map_logical_plan_relation(
             with _push_cte_scope():
                 for cte in as_java_list(rel.cteRelations()):
                     name = str(cte._1())
-                    cte_proto = map_logical_plan_relation(cte._2())
+                    # Store the original CTE definition for re-evaluation
+                    _cte_definitions.get()[name] = cte._2()
+                    # Process CTE definition with a unique plan_id to ensure proper column naming
+                    cte_plan_id = gen_sql_plan_id()
+                    cte_proto = map_logical_plan_relation(cte._2(), cte_plan_id)
                     _ctes.get()[name] = cte_proto
 
                 proto = map_logical_plan_relation(rel.child())

snowflake/snowpark_connect/relation/map_udtf.py

@@ -27,7 +27,7 @@ from snowflake.snowpark_connect.expression.map_expression import (
 )
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.type_mapping import (
-    parse_ddl_string,
+    map_type_string_to_proto,
     proto_to_snowpark_type,
 )
 from snowflake.snowpark_connect.utils.context import push_udtf_context
@@ -125,7 +125,9 @@ def process_return_type(
 ) -> tuple[list[tuple[str, Any]], DataType, StructType, list[str]]:
     try:
         if return_type.HasField("unparsed"):
-            parsed_return = parse_ddl_string(return_type.unparsed.data_type_string)
+            parsed_return = map_type_string_to_proto(
+                return_type.unparsed.data_type_string
+            )
         else:
             parsed_return = return_type
     except ValueError as e:

snowflake/snowpark_connect/relation/read/map_read.py

@@ -65,7 +65,7 @@ def map_read(
                 read_format = "parquet"
 
             if read_format.lower() == "iceberg":
-                telemetry.report_io_read("iceberg", None)
+                telemetry.report_io_read("iceberg")
                 return map_read_table(rel)
 
             if rel.read.data_source.schema == "":
@@ -90,7 +90,7 @@ def map_read(
                         )
                 schema = map_json_schema_to_snowpark(parsed_schema)
             options = dict(rel.read.data_source.options)
-            telemetry.report_io_read(read_format, options)
+            telemetry.report_io_read(read_format)
             session: snowpark.Session = get_or_create_snowpark_session()
             if len(rel.read.data_source.paths) > 0:
                 # Normalize paths to ensure consistent behavior
@@ -175,7 +175,7 @@ def map_read_table_or_file(rel):
             rel.read.named_table.unparsed_identifier
         )
         options = {}
-        telemetry.report_io_read(read_format, options)
+        telemetry.report_io_read(read_format)
         session: snowpark.Session = get_or_create_snowpark_session()
 
         clean_source_paths = [

snowflake/snowpark_connect/relation/read/map_read_jdbc.py

@@ -94,7 +94,7 @@ def map_read_jdbc(
             num_partitions=num_partitions,
             predicates=predicates,
         )
-        true_names = list(map(lambda x: unquote_if_quoted(x).lower(), df.columns))
+        true_names = list(map(lambda x: unquote_if_quoted(x), df.columns))
         renamed_df, snowpark_cols = rename_columns_as_snowflake_standard(
             df, rel.common.plan_id
         )

snowflake/snowpark_connect/relation/read/map_read_json.py

@@ -34,6 +34,7 @@ from snowflake.snowpark_connect.type_mapping import (
     cast_to_match_snowpark_type,
     map_simple_types,
 )
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -66,6 +67,7 @@ def map_read_json(
 
         rows_to_infer_schema = snowpark_options.pop("rowstoinferschema", 1000)
         dropFieldIfAllNull = snowpark_options.pop("dropfieldifallnull", False)
+        batch_size = snowpark_options.pop("batchsize", 1000)
 
         reader = session.read.options(snowpark_options)
 
@@ -99,7 +101,7 @@ def map_read_json(
                 ]
 
         df = construct_dataframe_by_schema(
-            schema, df.to_local_iterator(), session, snowpark_options
+            schema, df.to_local_iterator(), session, snowpark_options, batch_size
         )
 
         spark_column_names = get_spark_column_names_from_snowpark_columns(df.columns)
@@ -277,9 +279,11 @@ def construct_dataframe_by_schema(
     result = None
 
     current_data = []
+    progress = 0
     for row in rows:
         current_data.append(construct_row_by_schema(row, schema, snowpark_options))
         if len(current_data) >= batch_size:
+            progress += len(current_data)
             result = union_data_into_df(
                 result,
                 current_data,
@@ -287,9 +291,11 @@ def construct_dataframe_by_schema(
                 session,
             )
 
+            logger.info(f"JSON reader: finished processing {progress} rows")
             current_data = []
 
     if len(current_data) > 0:
+        progress += len(current_data)
         result = union_data_into_df(
             result,
             current_data,
@@ -297,6 +303,7 @@ def construct_dataframe_by_schema(
             session,
         )
 
+        logger.info(f"JSON reader: finished processing {progress} rows")
         current_data = []
 
     if result is None:

snowflake/snowpark_connect/relation/read/map_read_table.py

@@ -11,7 +11,6 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     unquote_if_quoted,
 )
 from snowflake.snowpark.exceptions import SnowparkSQLException
-from snowflake.snowpark_connect.column_name_handler import ALREADY_QUOTED
 from snowflake.snowpark_connect.config import auto_uppercase_non_column_identifiers
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read.utils import (
@@ -29,15 +28,8 @@ from snowflake.snowpark_connect.utils.telemetry import (
 def post_process_df(
     df: snowpark.DataFrame, plan_id: int, source_table_name: str = None
 ) -> DataFrameContainer:
-    def _lower_or_unquote(string):
-        return (
-            string[1:-1].replace('""', '"')
-            if ALREADY_QUOTED.match(string)
-            else string.lower()
-        )
-
     try:
-        true_names = list(map(lambda x: _lower_or_unquote(x), df.columns))
+        true_names = list(map(lambda x: unquote_if_quoted(x), df.columns))
         renamed_df, snowpark_column_names = rename_columns_as_snowflake_standard(
             df, plan_id
         )

snowflake/snowpark_connect/relation/read/reader_config.py

@@ -317,6 +317,7 @@ class JsonReaderConfig(ReaderWriterConfig):
                     # TODO: modifiedAfter: Union[bool, str, None] = None,
                     # TODO: allowNonNumericNumbers: Union[bool, str, None] = None,
                     "rowsToInferSchema": 1000,
+                    "batchSize": 1000,
                 },
                 supported_options={
                     "schema",
@@ -347,12 +348,13 @@ class JsonReaderConfig(ReaderWriterConfig):
                     # "ignoreNullFields",
                     "rowsToInferSchema",
                     # "inferTimestamp",
+                    "batchSize",
                 },
                 boolean_config_list=[
                     "multiLine",
                     "dropFieldIfAllNull",
                 ],
-                int_config_list=["rowsToInferSchema"],
+                int_config_list=["rowsToInferSchema", "batchSize"],
                 float_config_list=["samplingRatio"],
             ),
             options,

snowflake/snowpark_connect/relation/write/map_write.py

@@ -107,10 +107,7 @@ def _spark_to_snowflake(multipart_id: str) -> str:
 
 def map_write(request: proto_base.ExecutePlanRequest):
     write_op = request.plan.command.write_operation
-    if write_op.options is not None:
-        telemetry.report_io_write(write_op.source, dict(write_op.options))
-    else:
-        telemetry.report_io_write(write_op.source)
+    telemetry.report_io_write(write_op.source)
 
     write_mode = None
     match write_op.mode:
@@ -234,10 +231,11 @@ 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"
-                        )
+                    # 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"
+                    #     )
                     _validate_schema_and_get_writer(
                         input_df, "append", snowpark_table_name
                     ).saveAsTable(
@@ -264,10 +262,12 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         )
                 case "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"
-                            )
+                        # 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
                     else:
                         create_iceberg_table(
                             snowpark_table_name=snowpark_table_name,
@@ -298,13 +298,14 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         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"
-                                )
+                            # 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"
+                            #     )
                             write_mode = "truncate"
                         _validate_schema_and_get_writer(
                             input_df, write_mode, snowpark_table_name
@@ -314,10 +315,11 @@ def map_write(request: proto_base.ExecutePlanRequest):
                             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"
-                            )
+                        # 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"
+                        #     )
                         _validate_schema_and_get_writer(
                             input_df, write_mode, snowpark_table_name
                         ).saveAsTable(
@@ -388,10 +390,11 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     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"
-                    )
+                # 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"
+                #     )
                 _validate_schema_and_get_writer(
                     input_df, "append", snowpark_table_name
                 ).saveAsTable(
@@ -402,10 +405,12 @@ 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):
-                    if check_table_type(snowpark_table_name, session) != "ICEBERG":
-                        raise AnalysisException(
-                            f"Table {snowpark_table_name} is not an iceberg table"
-                        )
+                    # 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
                 else:
                     raise AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
@@ -471,10 +476,11 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     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"
-                    )
+                # 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"
+                #     )
                 _validate_schema_and_get_writer(
                     input_df, "append", snowpark_table_name
                 ).saveAsTable(
@@ -485,10 +491,12 @@ 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):
-                    if check_table_type(snowpark_table_name, session) != "TABLE":
-                        raise AnalysisException(
-                            f"Table {snowpark_table_name} is not a FDN table"
-                        )
+                    # 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
                 else:
                     raise AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
@@ -798,17 +806,18 @@ def check_snowflake_table_existence(
         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"
+# 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

@@ -1,7 +1,7 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-
+import pathlib
 import threading
 import time
 
@@ -9,6 +9,8 @@ from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_sess
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 
 _resources_initialized = threading.Event()
+SPARK_VERSION = "3.5.6"
+RESOURCE_PATH = "/snowflake/snowpark_connect/resources"
 
 
 def initialize_resources() -> None:
@@ -41,6 +43,31 @@ def initialize_resources() -> None:
 
         session.sql("select 1 as sf_connection_warm_up").collect()
 
+    def upload_scala_udf_jars() -> None:
+        """Upload Spark jar files required for creating Scala UDFs."""
+        stage = session.get_session_stage()
+        resource_path = stage + RESOURCE_PATH
+        import snowflake
+
+        pyspark_jars = (
+            pathlib.Path(snowflake.snowpark_connect.__file__).parent / "includes/jars"
+        )
+        jar_files = [
+            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",
+            "json4s-ast_2.12-3.7.0-M11.jar",
+        ]
+
+        for jar in jar_files:
+            session.file.put(
+                str(pyspark_jars) + "/" + jar,
+                resource_path,
+                auto_compress=False,
+                overwrite=False,
+                source_compression="NONE",
+            )
+
     start_time = time.time()
 
     resources = [
@@ -49,6 +76,7 @@ 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),
     ]
 
     for name, resource_func in resources:

snowflake/snowpark_connect/server.py

@@ -68,10 +68,14 @@ from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.relation.utils import get_semantic_string
 from snowflake.snowpark_connect.resources_initializer import initialize_resources_async
 from snowflake.snowpark_connect.type_mapping import (
-    parse_ddl_string,
+    map_type_string_to_proto,
     snowpark_to_proto_type,
 )
-from snowflake.snowpark_connect.utils.artifacts import check_checksum, write_artifact
+from snowflake.snowpark_connect.utils.artifacts import (
+    check_checksum,
+    write_artifact,
+    write_class_files_to_stage,
+)
 from snowflake.snowpark_connect.utils.cache import (
     df_cache_map_get,
     df_cache_map_pop,
@@ -249,7 +253,9 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
                     return proto_base.AnalyzePlanResponse(
                         session_id=request.session_id,
                         ddl_parse=proto_base.AnalyzePlanResponse.DDLParse(
-                            parsed=parse_ddl_string(request.ddl_parse.ddl_string)
+                            parsed=map_type_string_to_proto(
+                                request.ddl_parse.ddl_string
+                            )
                         ),
                     )
                 case "get_storage_level":
@@ -516,7 +522,13 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
         for name, data in cache_data.items():
             _try_handle_local_relation(name, bytes(data))
 
+        class_files: dict[str, str] = {}
         for (name, filepath) in filenames.items():
+            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
+                continue
             session.file.put(
                 filepath,
                 session.get_session_stage(),
@@ -541,6 +553,9 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
             if not name.startswith("pyfiles"):
                 session._import_files.add(cached_name)
 
+        if class_files:
+            write_class_files_to_stage(session, class_files)
+
         return proto_base.AddArtifactsResponse(artifacts=list(response.values()))
 
     def ArtifactStatus(self, request, context):