snowpark-connect 0.30.0__py3-none-any.whl → 0.31.0__py3-none-any.whl

snowflake/snowpark_connect/relation/write/map_write.py

@@ -28,16 +28,18 @@ from snowflake.snowpark.types import (
     _NumericType,
 )
 from snowflake.snowpark_connect.config import (
+    auto_uppercase_column_identifiers,
     global_config,
     sessions_config,
     str_to_bool,
 )
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 from snowflake.snowpark_connect.relation.io_utils import (
     convert_file_prefix_path,
+    get_compression_for_source_and_options,
     is_cloud_path,
-    is_supported_compression,
-    supported_compressions_for_format,
 )
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.relation.read.metadata_utils import (
@@ -147,6 +149,11 @@ def map_write(request: proto_base.ExecutePlanRequest):
 
     session: snowpark.Session = get_or_create_snowpark_session()
 
+    # Check for partition hint early to determine precedence over single option
+    partition_hint = (
+        result.partition_hint if hasattr(result, "partition_hint") else None
+    )
+
     # Snowflake saveAsTable doesn't support format
     if (
         write_op.HasField("table")
@@ -176,8 +183,11 @@ def map_write(request: proto_base.ExecutePlanRequest):
             # Generate Spark-compatible filename with proper extension
             extension = write_op.source if write_op.source != "text" else "txt"
 
-            # Get compression from options for proper filename generation
-            compression_option = write_op.options.get("compression", "none")
+            compression = get_compression_for_source_and_options(
+                write_op.source, write_op.options, from_read=False
+            )
+            if compression is not None:
+                write_op.options["compression"] = compression
 
             # Generate Spark-compatible filename or prefix
             # we need a random prefix to support "append" mode
@@ -203,12 +213,12 @@ def map_write(request: proto_base.ExecutePlanRequest):
                 except Exception as e:
                     logger.warning(f"Could not clear directory {write_path}: {e}")
 
-            if should_write_to_single_file:
+            if should_write_to_single_file and partition_hint is None:
                 # Single file: generate complete filename with extension
                 spark_filename = generate_spark_compatible_filename(
                     task_id=0,
                     attempt_number=0,
-                    compression=compression_option,
+                    compression=compression,
                     format_ext=extension,
                 )
                 temp_file_prefix_on_stage = f"{write_path}/{spark_filename}"
@@ -217,29 +227,11 @@ def map_write(request: proto_base.ExecutePlanRequest):
                 spark_filename_prefix = generate_spark_compatible_filename(
                     task_id=0,
                     attempt_number=0,
-                    compression=compression_option,
+                    compression=None,
                     format_ext="",  # No extension for prefix
                 )
                 temp_file_prefix_on_stage = f"{write_path}/{spark_filename_prefix}"
 
-            default_compression = "NONE" if write_op.source != "parquet" else "snappy"
-            compression = write_op.options.get(
-                "compression", default_compression
-            ).upper()
-
-            if not is_supported_compression(write_op.source, compression):
-                supported_compressions = supported_compressions_for_format(
-                    write_op.source
-                )
-                raise AnalysisException(
-                    f"Compression {compression} is not supported for {write_op.source} format. "
-                    + (
-                        f"Supported compressions: {sorted(supported_compressions)}"
-                        if supported_compressions
-                        else "No compression supported for this format."
-                    )
-                )
-
             parameters = {
                 "location": temp_file_prefix_on_stage,
                 "file_format_type": write_op.source
@@ -254,9 +246,6 @@ def map_write(request: proto_base.ExecutePlanRequest):
             # Using the base avoids coupling to exact filenames/prefixes.
             download_stage_path = write_path
 
-            # Check for partition hint early to determine precedence over single option
-            partition_hint = result.partition_hint
-
             # Apply max_file_size for both single and multi-file scenarios
             # This helps control when Snowflake splits files into multiple parts
             if max_file_size:
@@ -268,9 +257,13 @@ def map_write(request: proto_base.ExecutePlanRequest):
             get_param_from_options(parameters, write_op.options, write_op.source)
             if write_op.partitioning_columns:
                 if write_op.source != "parquet":
-                    raise SnowparkConnectNotImplementedError(
+                    exception = SnowparkConnectNotImplementedError(
                         "Partitioning is only supported for parquet format"
                     )
+                    attach_custom_error_code(
+                        exception, ErrorCodes.UNSUPPORTED_OPERATION
+                    )
+                    raise exception
                 # Build Spark-style directory structure: col1=value1/col2=value2/...
                 # Example produced expression (Snowflake SQL):
                 #   'department=' || TO_VARCHAR("department") || '/' || 'region=' || TO_VARCHAR("region")
@@ -314,7 +307,7 @@ def map_write(request: proto_base.ExecutePlanRequest):
                     per_part_prefix = generate_spark_compatible_filename(
                         task_id=part_idx,
                         attempt_number=0,
-                        compression=compression_option,
+                        compression=None,
                         format_ext="",  # prefix only; Snowflake appends extension/counters
                     )
                     part_params["location"] = f"{write_path}/{per_part_prefix}"
@@ -355,9 +348,13 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         snowpark_table_name, session
                     )
                     if isinstance(table_schema_or_error, DataType):  # Table exists
-                        raise AnalysisException(
+                        exception = AnalysisException(
                             f"Table {snowpark_table_name} already exists"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_OPERATION
+                        )
+                        raise exception
                     create_iceberg_table(
                         snowpark_table_name=snowpark_table_name,
                         location=write_op.options.get("location", None),
@@ -380,9 +377,13 @@ def map_write(request: proto_base.ExecutePlanRequest):
                             "ICEBERG",
                             "TABLE",
                         ):
-                            raise AnalysisException(
+                            exception = AnalysisException(
                                 f"Table {snowpark_table_name} is not an iceberg table"
                             )
+                            attach_custom_error_code(
+                                exception, ErrorCodes.INVALID_OPERATION
+                            )
+                            raise exception
                     else:
                         create_iceberg_table(
                             snowpark_table_name=snowpark_table_name,
@@ -426,9 +427,13 @@ def map_write(request: proto_base.ExecutePlanRequest):
                             "ICEBERG",
                             "TABLE",
                         ):
-                            raise AnalysisException(
+                            exception = AnalysisException(
                                 f"Table {snowpark_table_name} is not an iceberg table"
                             )
+                            attach_custom_error_code(
+                                exception, ErrorCodes.INVALID_OPERATION
+                            )
+                            raise exception
                     else:
                         create_iceberg_table(
                             snowpark_table_name=snowpark_table_name,
@@ -444,9 +449,13 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         column_order=_column_order_for_write,
                     )
                 case _:
-                    raise SnowparkConnectNotImplementedError(
+                    exception = SnowparkConnectNotImplementedError(
                         f"Write mode {write_mode} is not supported"
                     )
+                    attach_custom_error_code(
+                        exception, ErrorCodes.UNSUPPORTED_OPERATION
+                    )
+                    raise exception
         case _:
             snowpark_table_name = _spark_to_snowflake(write_op.table.table_name)
             save_method = write_op.table.save_method
@@ -462,9 +471,13 @@ def map_write(request: proto_base.ExecutePlanRequest):
                 if len(write_op.table.table_name) == 0:
                     dbtable_name = write_op.options.get("dbtable", "")
                     if len(dbtable_name) == 0:
-                        raise SnowparkConnectNotImplementedError(
+                        exception = SnowparkConnectNotImplementedError(
                             "Save command is not supported without a table name"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.UNSUPPORTED_OPERATION
+                        )
+                        raise exception
                     else:
                         snowpark_table_name = _spark_to_snowflake(dbtable_name)
 
@@ -482,9 +495,13 @@ def map_write(request: proto_base.ExecutePlanRequest):
                                 "NORMAL",
                                 "TABLE",
                             ):
-                                raise AnalysisException(
+                                exception = AnalysisException(
                                     f"Table {snowpark_table_name} is not a FDN table"
                                 )
+                                attach_custom_error_code(
+                                    exception, ErrorCodes.INVALID_OPERATION
+                                )
+                                raise exception
                             write_mode = "truncate"
                         _validate_schema_and_get_writer(
                             input_df,
@@ -508,9 +525,13 @@ def map_write(request: proto_base.ExecutePlanRequest):
                             "NORMAL",
                             "TABLE",
                         ):
-                            raise AnalysisException(
+                            exception = AnalysisException(
                                 f"Table {snowpark_table_name} is not a FDN table"
                             )
+                            attach_custom_error_code(
+                                exception, ErrorCodes.INVALID_OPERATION
+                            )
+                            raise exception
 
                         _validate_schema_and_get_writer(
                             input_df,
@@ -542,9 +563,11 @@ def map_write(request: proto_base.ExecutePlanRequest):
                     column_order=_column_order_for_write,
                 )
             else:
-                raise SnowparkConnectNotImplementedError(
+                exception = SnowparkConnectNotImplementedError(
                     f"Save command not supported: {save_method}"
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+                raise exception
 
 
 def map_write_v2(request: proto_base.ExecutePlanRequest):
@@ -569,9 +592,11 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
     session: snowpark.Session = get_or_create_snowpark_session()
 
     if write_op.table_name is None or write_op.table_name == "":
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "Write operation V2 only support table writing now"
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     if write_op.provider.lower() == "iceberg":
         match write_op.mode:
@@ -580,9 +605,11 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     snowpark_table_name, session
                 )
                 if isinstance(table_schema_or_error, DataType):  # Table exists
-                    raise AnalysisException(
+                    exception = AnalysisException(
                         f"Table {snowpark_table_name} already exists"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+                    raise exception
                 create_iceberg_table(
                     snowpark_table_name=snowpark_table_name,
                     location=write_op.table_properties.get("location"),
@@ -601,16 +628,20 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     snowpark_table_name, session
                 )
                 if not isinstance(table_schema_or_error, DataType):  # Table not exists
-                    raise AnalysisException(
+                    exception = AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found."
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+                    raise exception
                 if get_table_type(snowpark_table_name, session) not in (
                     "ICEBERG",
                     "TABLE",
                 ):
-                    raise AnalysisException(
+                    exception = AnalysisException(
                         f"Table {snowpark_table_name} is not an iceberg table"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+                    raise exception
                 _validate_schema_and_get_writer(
                     input_df, "append", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
@@ -628,13 +659,19 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                         "ICEBERG",
                         "TABLE",
                     ):
-                        raise AnalysisException(
+                        exception = AnalysisException(
                             f"Table {snowpark_table_name} is not an iceberg table"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_OPERATION
+                        )
+                        raise exception
                 else:
-                    raise AnalysisException(
+                    exception = AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+                    raise exception
                 _validate_schema_and_get_writer(
                     input_df, "truncate", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
@@ -655,9 +692,11 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                         mode="replace",
                     )
                 else:
-                    raise AnalysisException(
+                    exception = AnalysisException(
                         f"Table {snowpark_table_name} does not exist"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+                    raise exception
                 _validate_schema_and_get_writer(
                     input_df, "replace", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
@@ -681,9 +720,11 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     column_order=_column_order_for_write,
                 )
             case _:
-                raise SnowparkConnectNotImplementedError(
+                exception = SnowparkConnectNotImplementedError(
                     f"Write mode {commands_proto.WriteOperationV2.Mode.Name(write_op.mode)} is not supported"
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+                raise exception
     else:
         match write_op.mode:
             case commands_proto.WriteOperationV2.MODE_CREATE:
@@ -699,16 +740,20 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     snowpark_table_name, session
                 )
                 if not isinstance(table_schema_or_error, DataType):  # Table not exists
-                    raise AnalysisException(
+                    exception = AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found."
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+                    raise exception
                 if get_table_type(snowpark_table_name, session) not in (
                     "NORMAL",
                     "TABLE",
                 ):
-                    raise AnalysisException(
+                    exception = AnalysisException(
                         f"Table {snowpark_table_name} is not a FDN table"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+                    raise exception
                 _validate_schema_and_get_writer(
                     input_df, "append", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
@@ -726,13 +771,19 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                         "NORMAL",
                         "TABLE",
                     ):
-                        raise AnalysisException(
+                        exception = AnalysisException(
                             f"Table {snowpark_table_name} is not a FDN table"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_OPERATION
+                        )
+                        raise exception
                 else:
-                    raise AnalysisException(
+                    exception = AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.TABLE_NOT_FOUND)
+                    raise exception
                 _validate_schema_and_get_writer(
                     input_df, "truncate", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
@@ -745,9 +796,11 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     snowpark_table_name, session
                 )
                 if not isinstance(table_schema_or_error, DataType):  # Table not exists
-                    raise AnalysisException(
+                    exception = AnalysisException(
                         f"Table {snowpark_table_name} does not exist"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.TABLE_NOT_FOUND)
+                    raise exception
                 _validate_schema_and_get_writer(
                     input_df, "replace", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
@@ -764,9 +817,11 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     column_order=_column_order_for_write,
                 )
             case _:
-                raise SnowparkConnectNotImplementedError(
+                exception = SnowparkConnectNotImplementedError(
                     f"Write mode {commands_proto.WriteOperationV2.Mode.Name(write_op.mode)} is not supported"
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+                raise exception
 
 
 def _get_table_schema_or_error(
@@ -778,6 +833,20 @@ def _get_table_schema_or_error(
         return e
 
 
+def _get_writer_for_table_creation(df: snowpark.DataFrame) -> snowpark.DataFrameWriter:
+    # When creating a new table, if case sensitivity is not enabled, we need to rename the columns
+    # to upper case so they are case-insensitive in Snowflake.
+    if auto_uppercase_column_identifiers():
+        for field in df.schema.fields:
+            col_name = field.name
+            # Uppercasing is fine, regardless of whether the original name was quoted or not.
+            # In Snowflake these are equivalent "COL" == COL == col == coL
+            uppercased_name = col_name.upper()
+            if col_name != uppercased_name:
+                df = df.withColumnRenamed(col_name, uppercased_name)
+    return df.write
+
+
 def _validate_schema_and_get_writer(
     input_df: snowpark.DataFrame,
     write_mode: str,
@@ -788,7 +857,7 @@ def _validate_schema_and_get_writer(
         "replace",
         "create_or_replace",
     ):
-        return input_df.write
+        return _get_writer_for_table_creation(input_df)
 
     table_schema = None
     if table_schema_or_error is not None:
@@ -797,6 +866,9 @@ def _validate_schema_and_get_writer(
             if "SQL compilation error" in msg and "does not exist" in msg:
                 pass
             else:
+                attach_custom_error_code(
+                    table_schema_or_error, ErrorCodes.INTERNAL_ERROR
+                )
                 raise table_schema_or_error
         elif isinstance(table_schema_or_error, DataType):
             table_schema = table_schema_or_error
@@ -810,16 +882,17 @@ def _validate_schema_and_get_writer(
             if "SQL compilation error" in msg and "does not exist" in msg:
                 pass
             else:
+                attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
                 raise e
 
     if table_schema is None:
         # If table does not exist, we can skip the schema validation
-        return input_df.write
+        return _get_writer_for_table_creation(input_df)
 
     _validate_schema_for_append(table_schema, input_df.schema, snowpark_table_name)
 
     # if table exists and case sensitivity is not enabled, we need to rename the columns to match existing table schema
-    if not global_config.spark_sql_caseSensitive:
+    if auto_uppercase_column_identifiers():
 
         for field in input_df.schema.fields:
             # Find the matching field in the table schema (case-insensitive)
@@ -829,8 +902,8 @@ def _validate_schema_and_get_writer(
                 (
                     f
                     for f in table_schema.fields
-                    if unquote_if_quoted(f.name).lower()
-                    == unquote_if_quoted(col_name).lower()
+                    if unquote_if_quoted(f.name).upper()
+                    == unquote_if_quoted(col_name).upper()
                 ),
                 None,
             )
@@ -865,21 +938,25 @@ def _validate_schema_for_append(
         case (StructType() as table_struct, StructType() as data_struct):
 
             def _comparable_col_name(col: str) -> str:
-                name = col if global_config.spark_sql_caseSensitive else col.lower()
+                name = col.upper() if auto_uppercase_column_identifiers() else col
                 if compare_structs:
                     return name
                 else:
                     return unquote_if_quoted(name)
 
             def invalid_struct_schema():
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Cannot resolve columns for the existing table {snowpark_table_name} ({table_schema.simple_string()}) with the data schema ({data_schema.simple_string()})."
                 )
+                attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+                raise exception
 
             if len(table_struct.fields) != len(data_struct.fields):
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"The column number of the existing table {snowpark_table_name} ({table_schema.simple_string()}) doesn't match the data schema ({data_schema.simple_string()}).)"
                 )
+                attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+                raise exception
 
             table_field_names = {
                 _comparable_col_name(field.name) for field in table_struct.fields
@@ -942,9 +1019,11 @@ def _validate_schema_for_append(
         case (DateType(), _) if isinstance(data_schema, (DateType, TimestampType)):
             return
         case (_, _):
-            raise AnalysisException(
+            exception = AnalysisException(
                 f"[INCOMPATIBLE_DATA_FOR_TABLE.CANNOT_SAFELY_CAST] Cannot write incompatible data for the table {snowpark_table_name}: Cannot safely cast {data_schema.simple_string()} to {table_schema.simple_string()}"
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+            raise exception
 
 
 def create_iceberg_table(
@@ -984,9 +1063,11 @@ def create_iceberg_table(
         case "create_or_replace":
             create_sql = "CREATE OR REPLACE"
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Write mode {mode} is not supported for iceberg table"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
     sql = f"""
         {create_sql} ICEBERG TABLE {snowpark_table_name} ({",".join(table_schema)})
         CATALOG = 'SNOWFLAKE'
@@ -1093,9 +1174,11 @@ def store_files_locally(
 
 def _truncate_directory(directory_path: Path) -> None:
     if not directory_path.exists():
-        raise FileNotFoundError(
+        exception = FileNotFoundError(
             f"The specified directory {directory_path} does not exist."
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+        raise exception
     # Iterate over all the files and directories in the specified directory
     for file in directory_path.iterdir():
         # Check if it is a file or directory and remove it

snowflake/snowpark_connect/relation/write/map_write_jdbc.py

@@ -4,6 +4,8 @@
 
 from snowflake import snowpark
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 from snowflake.snowpark_connect.relation.read.map_read_jdbc import (
     close_connection,
     create_connection,
@@ -35,7 +37,9 @@ def map_write_jdbc(
         dbtable = None
 
     if dbtable is None:
-        raise ValueError("Include dbtable is required option")
+        exception = ValueError("Include dbtable is required option")
+        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+        raise exception
 
     try:
         JdbcDataFrameWriter(session, jdbc_options).jdbc_write_dbapi(
@@ -46,4 +50,6 @@ def map_write_jdbc(
             write_mode=write_mode,
         )
     except Exception as e:
-        raise Exception(f"Error accessing JDBC datasource for write: {e}")
+        exception = Exception(f"Error accessing JDBC datasource for write: {e}")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception

snowflake/snowpark_connect/resources_initializer.py

@@ -5,6 +5,8 @@ import pathlib
 import threading
 import time
 
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 
@@ -119,9 +121,11 @@ def wait_for_resource_initialization() -> None:
         logger.error(
             "Resource initialization failed - initializer thread has been running for over 300 seconds."
         )
-        raise RuntimeError(
+        exception = RuntimeError(
             "Resource initialization failed - initializer thread has been running for over 300 seconds."
         )
+        attach_custom_error_code(exception, ErrorCodes.RESOURCE_INITIALIZATION_FAILED)
+        raise exception
 
 
 def set_upload_jars(upload: bool) -> None: