snowpark-connect 0.28.1__py3-none-any.whl → 0.29.0__py3-none-any.whl

snowflake/snowpark_connect/config.py

@@ -264,16 +264,22 @@ SESSION_CONFIG_KEY_WHITELIST = {
     "snowpark.connect.udtf.compatibility_mode",
     "snowpark.connect.views.duplicate_column_names_handling_mode",
     "enable_snowflake_extension_behavior",
+    "spark.hadoop.fs.s3a.server-side-encryption.key",
+    "spark.hadoop.fs.s3a.assumed.role.arn",
 }
-AZURE_SAS_KEY = re.compile(
+AZURE_ACCOUNT_KEY = re.compile(
     r"^fs\.azure\.sas\.[^\.]+\.[^\.]+\.blob\.core\.windows\.net$"
 )
+AZURE_SAS_KEY = re.compile(
+    r"^fs\.azure\.sas\.fixed\.token\.[^\.]+\.dfs\.core\.windows\.net$"
+)
 
 
 def valid_session_config_key(key: str):
     return (
         key in SESSION_CONFIG_KEY_WHITELIST  # AWS session keys
         or AZURE_SAS_KEY.match(key)  # Azure session keys
+        or AZURE_ACCOUNT_KEY.match(key)  # Azure account keys
     )
 
 
@@ -578,7 +584,10 @@ def set_snowflake_parameters(
             cte_enabled = str_to_bool(value)
             snowpark_session.cte_optimization_enabled = cte_enabled
             logger.info(f"Updated snowpark session CTE optimization: {cte_enabled}")
-
+        case "snowpark.connect.structured_types.fix":
+            # TODO: SNOW-2367714 Remove this once the fix is automatically enabled in Snowpark
+            snowpark.context._enable_fix_2360274 = str_to_bool(value)
+            logger.info(f"Updated snowpark session structured types fix: {value}")
         case _:
             pass
 

snowflake/snowpark_connect/expression/map_unresolved_function.py

@@ -15,6 +15,7 @@ import tempfile
 import time
 import uuid
 from collections import defaultdict
+from collections.abc import Callable
 from contextlib import suppress
 from decimal import ROUND_HALF_EVEN, ROUND_HALF_UP, Context, Decimal
 from functools import partial, reduce
@@ -199,7 +200,7 @@ def _validate_numeric_args(
             case StringType():
                 # Cast strings to doubles following Spark
                 # https://github.com/apache/spark/blob/master/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala#L204
-                modified_args[i] = _try_cast_helper(snowpark_args[i], DoubleType())
+                modified_args[i] = snowpark_fn.try_cast(snowpark_args[i], DoubleType())
             case _:
                 raise TypeError(
                     f"Data type mismatch: {function_name} requires numeric types, but got {typed_args[0].typ} and {typed_args[1].typ}."
@@ -519,7 +520,7 @@ def map_unresolved_function(
                     DecimalType() as t,
                 ):
                     p1, s1 = _get_type_precision(t)
-                    result_type = _get_decimal_multiplication_result_type(
+                    result_type, _ = _get_decimal_multiplication_result_type(
                         p1, s1, p1, s1
                     )
                     result_exp = snowpark_fn.lit(None)
@@ -528,11 +529,17 @@ def map_unresolved_function(
                 ):
                     p1, s1 = _get_type_precision(snowpark_typed_args[0].typ)
                     p2, s2 = _get_type_precision(snowpark_typed_args[1].typ)
-                    result_type = _get_decimal_multiplication_result_type(
-                        p1, s1, p2, s2
-                    )
-                    result_exp = _get_decimal_multiplication_result_exp(
-                        result_type, t, snowpark_args
+                    (
+                        result_type,
+                        overflow_possible,
+                    ) = _get_decimal_multiplication_result_type(p1, s1, p2, s2)
+                    result_exp = _arithmetic_operation(
+                        snowpark_typed_args[0],
+                        snowpark_typed_args[1],
+                        lambda x, y: x * y,
+                        overflow_possible,
+                        global_config.spark_sql_ansi_enabled,
+                        result_type,
                     )
                 case (NullType(), NullType()):
                     result_type = DoubleType()
@@ -617,7 +624,7 @@ def map_unresolved_function(
             )
             match (snowpark_typed_args[0].typ, snowpark_typed_args[1].typ):
                 case (NullType(), _) | (_, NullType()):
-                    result_type = _get_add_sub_result_type(
+                    result_type, _ = _get_add_sub_result_type(
                         snowpark_typed_args[0].typ,
                         snowpark_typed_args[1].typ,
                         spark_function_name,
@@ -693,14 +700,21 @@ def map_unresolved_function(
                         f'[DATATYPE_MISMATCH.BINARY_OP_DIFF_TYPES] Cannot resolve "{spark_function_name}" due to data type mismatch: the left and right operands of the binary operator have incompatible types ("{snowpark_typed_args[0].typ}" and "{snowpark_typed_args[1].typ}").'
                     )
                 case _:
-                    result_type = _get_add_sub_result_type(
+                    result_type, overflow_possible = _get_add_sub_result_type(
                         snowpark_typed_args[0].typ,
                         snowpark_typed_args[1].typ,
                         spark_function_name,
                     )
-                    result_exp = snowpark_args[0] + snowpark_args[1]
-                    if isinstance(result_type, DecimalType):
-                        result_exp = _cast_helper(result_exp, result_type)
+
+                    result_exp = _arithmetic_operation(
+                        snowpark_typed_args[0],
+                        snowpark_typed_args[1],
+                        lambda x, y: x + y,
+                        overflow_possible,
+                        global_config.spark_sql_ansi_enabled,
+                        result_type,
+                    )
+
         case "-":
             spark_function_name = _get_spark_function_name(
                 snowpark_typed_args[0],
@@ -715,7 +729,7 @@ def map_unresolved_function(
                     result_type = LongType()
                     result_exp = snowpark_fn.lit(None).cast(result_type)
                 case (NullType(), _) | (_, NullType()):
-                    result_type = _get_add_sub_result_type(
+                    result_type, _ = _get_add_sub_result_type(
                         snowpark_typed_args[0].typ,
                         snowpark_typed_args[1].typ,
                         spark_function_name,
@@ -806,14 +820,20 @@ def map_unresolved_function(
                         f'[DATATYPE_MISMATCH.BINARY_OP_DIFF_TYPES] Cannot resolve "{spark_function_name}" due to data type mismatch: the left and right operands of the binary operator have incompatible types ("{snowpark_typed_args[0].typ}" and "{snowpark_typed_args[1].typ}").'
                     )
                 case _:
-                    result_type = _get_add_sub_result_type(
+                    result_type, overflow_possible = _get_add_sub_result_type(
                         snowpark_typed_args[0].typ,
                         snowpark_typed_args[1].typ,
                         spark_function_name,
                     )
-                    result_exp = snowpark_args[0] - snowpark_args[1]
-                    if isinstance(result_type, DecimalType):
-                        result_exp = _cast_helper(result_exp, result_type)
+                    result_exp = _arithmetic_operation(
+                        snowpark_typed_args[0],
+                        snowpark_typed_args[1],
+                        lambda x, y: x - y,
+                        overflow_possible,
+                        global_config.spark_sql_ansi_enabled,
+                        result_type,
+                    )
+
         case "/":
             match (snowpark_typed_args[0].typ, snowpark_typed_args[1].typ):
                 case (DecimalType() as t1, NullType()):
@@ -825,15 +845,17 @@ def map_unresolved_function(
                 ):
                     p1, s1 = _get_type_precision(snowpark_typed_args[0].typ)
                     p2, s2 = _get_type_precision(snowpark_typed_args[1].typ)
-                    result_type, overflow_detected = _get_decimal_division_result_type(
+                    result_type, overflow_possible = _get_decimal_division_result_type(
                         p1, s1, p2, s2
                     )
-                    result_exp = _get_decimal_division_result_exp(
+
+                    result_exp = _arithmetic_operation(
+                        snowpark_typed_args[0],
+                        snowpark_typed_args[1],
+                        lambda x, y: _divnull(x, y),
+                        overflow_possible,
+                        global_config.spark_sql_ansi_enabled,
                         result_type,
-                        t,
-                        overflow_detected,
-                        snowpark_args,
-                        spark_function_name,
                     )
                 case (NullType(), NullType()):
                     result_type = DoubleType()
@@ -7580,63 +7602,20 @@ def map_unresolved_function(
                     )
                     | (DecimalType(), DecimalType())
                 ):
-                    # compute new precision and scale using correct decimal division rules
-                    if isinstance(
-                        snowpark_typed_args[0].typ, DecimalType
-                    ) and isinstance(snowpark_typed_args[1].typ, DecimalType):
-                        s1, p1 = (
-                            snowpark_typed_args[0].typ.scale,
-                            snowpark_typed_args[0].typ.precision,
-                        )
-                        s2, p2 = (
-                            snowpark_typed_args[1].typ.scale,
-                            snowpark_typed_args[1].typ.precision,
-                        )
-                        # The scale and precision formula that Spark follows for DecimalType
-                        # arithmetic operations can be found in the following Spark source
-                        # code file:
-                        # https://github.com/apache/spark/blob/a584cc48ef63fefb2e035349c8684250f8b936c4/docs/sql-ref-ansi-compliance.md
-                        new_scale = max(6, s1 + p2 + 1)
-                        new_precision = p1 - s1 + s2 + new_scale
-
-                    elif isinstance(snowpark_typed_args[0].typ, DecimalType):
-                        s1, p1 = (
-                            snowpark_typed_args[0].typ.scale,
-                            snowpark_typed_args[0].typ.precision,
-                        )
-                        # INT is treated as Decimal(10, 0)
-                        new_scale = max(6, s1 + 11)
-                        new_precision = p1 - s1 + new_scale
-
-                    else:  # right is DecimalType
-                        s2, p2 = (
-                            snowpark_typed_args[1].typ.scale,
-                            snowpark_typed_args[1].typ.precision,
-                        )
-                        # INT is treated as Decimal(10, 0)
-                        new_scale = max(6, 11 + p2)
-                        new_precision = (
-                            10 - 0 + s2 + new_scale
-                        )  # INT has precision 10, scale 0
-
-                    # apply precision cap
-                    if new_precision > 38:
-                        new_scale -= new_precision - 38
-                        new_precision = 38
-                        new_scale = max(new_scale, 6)
-
-                    left_double = snowpark_fn.cast(snowpark_args[0], DoubleType())
-                    right_double = snowpark_fn.cast(snowpark_args[1], DoubleType())
-
-                    quotient = snowpark_fn.when(
-                        snowpark_args[1] == 0, snowpark_fn.lit(None)
-                    ).otherwise(left_double / right_double)
-                    quotient = snowpark_fn.cast(quotient, StringType())
+                    p1, s1 = _get_type_precision(snowpark_typed_args[0].typ)
+                    p2, s2 = _get_type_precision(snowpark_typed_args[1].typ)
+                    result_type, overflow_possible = _get_decimal_division_result_type(
+                        p1, s1, p2, s2
+                    )
 
-                    result_exp = _try_cast_helper(
-                        quotient, DecimalType(new_precision, new_scale)
+                    result_exp = _arithmetic_operation(
+                        snowpark_typed_args[0],
+                        snowpark_typed_args[1],
+                        lambda x, y: _divnull(x, y),
+                        overflow_possible,
+                        False,
+                        result_type,
                     )
-                    result_type = DecimalType(new_precision, new_scale)
                 case (_NumericType(), _NumericType()):
                     result_exp = snowpark_fn.when(
                         snowpark_args[1] == 0, snowpark_fn.lit(None)
@@ -7749,42 +7728,21 @@ def map_unresolved_function(
                     )
                     | (DecimalType(), DecimalType())
                 ):
-                    # figure out what precision to use as the overflow amount
-                    if isinstance(
-                        snowpark_typed_args[0].typ, DecimalType
-                    ) and isinstance(snowpark_typed_args[1].typ, DecimalType):
-                        new_precision = (
-                            snowpark_typed_args[0].typ.precision
-                            + snowpark_typed_args[1].typ.precision
-                            + 1
-                        )
-                        new_scale = (
-                            snowpark_typed_args[0].typ.scale
-                            + snowpark_typed_args[1].typ.scale
-                        )
-                    elif isinstance(snowpark_typed_args[0].typ, DecimalType):
-                        new_precision = snowpark_typed_args[0].typ.precision + 11
-                        new_scale = snowpark_typed_args[0].typ.scale
-                    else:
-                        new_precision = snowpark_typed_args[1].typ.precision + 11
-                        new_scale = snowpark_typed_args[1].typ.scale
-
-                    # truncating down appropriately
-                    if new_precision > 38:
-                        new_precision = 38
-                    if new_scale > new_precision:
-                        new_scale = new_precision
-
-                    left_double = snowpark_fn.cast(snowpark_args[0], DoubleType())
-                    right_double = snowpark_fn.cast(snowpark_args[1], DoubleType())
-
-                    product = left_double * right_double
-
-                    product = snowpark_fn.cast(product, StringType())
-                    result_exp = _try_cast_helper(
-                        product, DecimalType(new_precision, new_scale)
+                    p1, s1 = _get_type_precision(snowpark_typed_args[0].typ)
+                    p2, s2 = _get_type_precision(snowpark_typed_args[1].typ)
+                    (
+                        result_type,
+                        overflow_possible,
+                    ) = _get_decimal_multiplication_result_type(p1, s1, p2, s2)
+
+                    result_exp = _arithmetic_operation(
+                        snowpark_typed_args[0],
+                        snowpark_typed_args[1],
+                        lambda x, y: x * y,
+                        overflow_possible,
+                        False,
+                        result_type,
                     )
-                    result_type = DecimalType(new_precision, new_scale)
                 case (_NumericType(), _NumericType()):
                     result_exp = snowpark_args[0] * snowpark_args[1]
                     result_exp = _type_with_typer(result_exp)
@@ -8391,20 +8349,10 @@ def map_unresolved_function(
     return spark_col_names, typed_col
 
 
-def _cast_helper(column: Column, to: DataType) -> Column:
-    if global_config.spark_sql_ansi_enabled:
-        column_mediator = (
-            snowpark_fn.cast(column, StringType())
-            if isinstance(to, DecimalType)
-            else column
-        )
-        return snowpark_fn.cast(column_mediator, to)
-    else:
-        return _try_cast_helper(column, to)
-
-
 def _try_cast_helper(column: Column, to: DataType) -> Column:
     """
+    DEPRECATED because of performance issues
+
     Attempts to cast a given column to a specified data type using the same behaviour as Spark.
 
     Args:
@@ -9600,71 +9548,109 @@ def _decimal_add_sub_result_type_helper(p1, s1, p2, s2):
     return result_precision, min_scale, return_type_precision, return_type_scale
 
 
-def _get_decimal_multiplication_result_exp(
-    result_type: DecimalType | DataType,
-    other_type: DataType,
-    snowpark_args: list[Column],
-) -> Column:
-    if global_config.spark_sql_ansi_enabled:
-        result_exp = snowpark_args[0] * snowpark_args[1]
-    else:
-        if isinstance(other_type, _IntegralType):
-            result_exp = snowpark_args[0].cast(result_type) * snowpark_args[1].cast(
-                result_type
-            )
-        else:
-            result_exp = snowpark_args[0].cast(DoubleType()) * snowpark_args[1].cast(
-                DoubleType()
-            )
-        result_exp = _try_cast_helper(result_exp, result_type)
-    return result_exp
-
-
-def _get_decimal_multiplication_result_type(p1, s1, p2, s2) -> DecimalType:
+def _get_decimal_multiplication_result_type(p1, s1, p2, s2) -> tuple[DecimalType, bool]:
     result_precision = p1 + p2 + 1
     result_scale = s1 + s2
+    overflow_possible = False
     if result_precision > 38:
+        overflow_possible = True
         if result_scale > 6:
             overflow = result_precision - 38
             result_scale = max(6, result_scale - overflow)
         result_precision = 38
-    return DecimalType(result_precision, result_scale)
+    return DecimalType(result_precision, result_scale), overflow_possible
 
 
-def _get_decimal_division_result_exp(
-    result_type: DecimalType | DataType,
-    other_type: DataType,
-    overflow_detected: bool,
-    snowpark_args: list[Column],
-    spark_function_name: str,
+def _arithmetic_operation(
+    arg1: TypedColumn,
+    arg2: TypedColumn,
+    op: Callable[[Column, Column], Column],
+    overflow_possible: bool,
+    should_raise_on_overflow: bool,
+    target_type: DecimalType,
 ) -> Column:
-    if (
-        isinstance(other_type, DecimalType)
-        and overflow_detected
-        and global_config.spark_sql_ansi_enabled
-    ):
-        raise ArithmeticException(
-            f'[NUMERIC_VALUE_OUT_OF_RANGE] {spark_function_name} cannot be represented as Decimal({result_type.precision}, {result_type.scale}). If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error, and return NULL instead.'
+    def _cast_arg(tc: TypedColumn) -> Column:
+        _, s = _get_type_precision(tc.typ)
+        typ = (
+            DoubleType()
+            if s > 0
+            or (
+                isinstance(tc.typ, _FractionalType)
+                and not isinstance(tc.typ, DecimalType)
+            )
+            else LongType()
+        )
+        return tc.col.cast(typ)
+
+    op_for_overflow_check = op(arg1.col.cast(DoubleType()), arg2.col.cast(DoubleType()))
+    safe_op = op(_cast_arg(arg1), _cast_arg(arg2))
+
+    if overflow_possible:
+        return _cast_arithmetic_operation_result(
+            op_for_overflow_check, safe_op, target_type, should_raise_on_overflow
         )
     else:
-        dividend = snowpark_args[0].cast(DoubleType())
-        divisor = snowpark_args[1]
-        result_exp = _divnull(dividend, divisor)
-        result_exp = _cast_helper(result_exp, result_type)
-    return result_exp
+        return op(arg1.col, arg2.col).cast(target_type)
+
+
+def _cast_arithmetic_operation_result(
+    overflow_check_expr: Column,
+    result_expr: Column,
+    target_type: DecimalType,
+    should_raise_on_overflow: bool,
+) -> Column:
+    """
+    Casts an arithmetic operation result to the target decimal type with overflow detection.
+    This function uses a dual-expression approach for robust overflow handling:
+    Args:
+        overflow_check_expr: Arithmetic expression using DoubleType operands for overflow detection.
+                           This expression is used ONLY for boundary checking against the target
+                           decimal's min/max values. DoubleType preserves the magnitude of large
+                           intermediate results that might overflow in decimal arithmetic.
+        result_expr: Arithmetic expression using safer operand types (LongType for integers,
+                    DoubleType for fractionals) for the actual result computation.
+        target_type: Target DecimalType to cast the result to.
+        should_raise_on_overflow: If True raises ArithmeticException on overflow, if False, returns NULL on overflow.
+    """
+
+    def create_overflow_handler(min_val, max_val, type_name: str):
+        if should_raise_on_overflow:
+            raise_error = _raise_error_helper(target_type, ArithmeticException)
+            return snowpark_fn.when(
+                (overflow_check_expr < snowpark_fn.lit(min_val))
+                | (overflow_check_expr > snowpark_fn.lit(max_val)),
+                raise_error(
+                    snowpark_fn.lit(
+                        f'[NUMERIC_VALUE_OUT_OF_RANGE] Value cannot be represented as {type_name}. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error, and return NULL instead.'
+                    )
+                ),
+            ).otherwise(result_expr.cast(target_type))
+        else:
+            return snowpark_fn.when(
+                (overflow_check_expr < snowpark_fn.lit(min_val))
+                | (overflow_check_expr > snowpark_fn.lit(max_val)),
+                snowpark_fn.lit(None),
+            ).otherwise(result_expr.cast(target_type))
+
+    precision = target_type.precision
+    scale = target_type.scale
+
+    max_val = (10**precision - 1) / (10**scale)
+    min_val = -max_val
+
+    return create_overflow_handler(min_val, max_val, f"DECIMAL({precision},{scale})")
 
 
 def _get_decimal_division_result_type(p1, s1, p2, s2) -> tuple[DecimalType, bool]:
-    overflow_detected = False
+    overflow_possible = False
     result_scale = max(6, s1 + p2 + 1)
     result_precision = p1 - s1 + s2 + result_scale
     if result_precision > 38:
-        if result_precision > 40:
-            overflow_detected = True
+        overflow_possible = True
         overflow = result_precision - 38
         result_scale = max(6, result_scale - overflow)
         result_precision = 38
-    return DecimalType(result_precision, result_scale), overflow_detected
+    return DecimalType(result_precision, result_scale), overflow_possible
 
 
 def _try_arithmetic_helper(
@@ -9778,46 +9764,20 @@ def _try_arithmetic_helper(
             DecimalType(),
             DecimalType(),
         ):
-            p1, s1 = _get_type_precision(typed_args[0].typ)
-            p2, s2 = _get_type_precision(typed_args[1].typ)
-
-            if isinstance(typed_args[0].typ, _IntegralType) and isinstance(
-                typed_args[1].typ, DecimalType
-            ):
-                new_scale = s2
-                new_precision = max(p2, p1 + s2)
-            elif isinstance(typed_args[0].typ, DecimalType) and isinstance(
-                typed_args[1].typ, _IntegralType
-            ):
-                new_scale = s1
-                new_precision = max(p1, p2 + s1)
-            else:
-                # Both decimal types
-                if operation_type == 1 and s1 == s2:  # subtraction with matching scales
-                    new_scale = s1
-                    max_integral_digits = max(p1 - s1, p2 - s2)
-                    new_precision = max_integral_digits + new_scale
-                else:
-                    new_scale = max(s1, s2)
-                    max_integral_digits = max(p1 - s1, p2 - s2)
-                    new_precision = max_integral_digits + new_scale + 1
-
-            # Overflow check
-            if new_precision > 38:
-                if global_config.spark_sql_ansi_enabled:
-                    raise ArithmeticException(
-                        f'[NUMERIC_VALUE_OUT_OF_RANGE] Precision {new_precision} exceeds maximum allowed precision of 38. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error, and return NULL instead.'
-                    )
-                return snowpark_fn.lit(None)
-
-            left_operand, right_operand = snowpark_args[0], snowpark_args[1]
-
-            result = (
-                left_operand + right_operand
-                if operation_type == 0
-                else left_operand - right_operand
+            result_type, overflow_possible = _get_add_sub_result_type(
+                typed_args[0].typ,
+                typed_args[1].typ,
+                "try_add" if operation_type == 0 else "try_subtract",
+            )
+
+            return _arithmetic_operation(
+                typed_args[0],
+                typed_args[1],
+                lambda x, y: x + y if operation_type == 0 else x - y,
+                overflow_possible,
+                False,
+                result_type,
             )
-            return snowpark_fn.cast(result, DecimalType(new_precision, new_scale))
 
         # If either of the inputs is floating point, we can just let it go through to Snowflake, where overflow
         # matches Spark and goes to inf.
@@ -9863,7 +9823,8 @@ def _get_add_sub_result_type(
     type1: DataType,
     type2: DataType,
     spark_function_name: str,
-) -> DataType:
+) -> tuple[DataType, bool]:
+    overflow_possible = False
     result_type = _find_common_type([type1, type2])
     match result_type:
         case DecimalType():
@@ -9872,6 +9833,7 @@ def _get_add_sub_result_type(
             result_scale = max(s1, s2)
             result_precision = max(p1 - s1, p2 - s2) + result_scale + 1
             if result_precision > 38:
+                overflow_possible = True
                 if result_scale > 6:
                     overflow = result_precision - 38
                     result_scale = max(6, result_scale - overflow)
@@ -9900,7 +9862,7 @@ def _get_add_sub_result_type(
             raise AnalysisException(
                 f'[DATATYPE_MISMATCH.BINARY_OP_WRONG_TYPE] Cannot resolve "{spark_function_name}" due to data type mismatch: the binary operator requires the input type ("NUMERIC" or "INTERVAL DAY TO SECOND" or "INTERVAL YEAR TO MONTH" or "INTERVAL"), not "BOOLEAN".',
             )
-    return result_type
+    return result_type, overflow_possible
 
 
 def _get_spark_function_name(

snowflake/snowpark_connect/relation/io_utils.py

@@ -7,8 +7,27 @@ from urllib.parse import urlparse
 CLOUD_PREFIX_TO_CLOUD = {
     "abfss": "azure",
     "wasbs": "azure",
+    "gcs": "gcp",
+    "gs": "gcp",
 }
 
+SUPPORTED_COMPRESSION_PER_FORMAT = {
+    "csv": {"AUTO", "GZIP", "BZ2", "BROTLI", "ZSTD", "DEFLATE", "RAW_DEFLATE", "NONE"},
+    "json": {"AUTO", "GZIP", "BZ2", "BROTLI", "ZSTD", "DEFLATE", "RAW_DEFLATE", "NONE"},
+    "parquet": {"AUTO", "LZO", "SNAPPY", "NONE"},
+    "text": {"NONE"},
+}
+
+
+def supported_compressions_for_format(format: str) -> set[str]:
+    return SUPPORTED_COMPRESSION_PER_FORMAT.get(format, set())
+
+
+def is_supported_compression(format: str, compression: str | None) -> bool:
+    if compression is None:
+        return True
+    return compression in supported_compressions_for_format(format)
+
 
 def get_cloud_from_url(
     url: str,
@@ -66,7 +85,8 @@ def is_cloud_path(path: str) -> bool:
         or path.startswith("azure://")
         or path.startswith("abfss://")
         or path.startswith("wasbs://")  # Azure
-        or path.startswith("gcs://")  # GCP
+        or path.startswith("gcs://")
+        or path.startswith("gs://")  # GCP
     )