snowpark-connect 0.25.0__tar.gz → 0.26.0__tar.gz

{snowpark_connect-0.25.0/src/snowpark_connect.egg-info → snowpark_connect-0.26.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-connect
-Version: 0.25.0
+Version: 0.26.0
 Summary: Snowpark Connect for Spark
 Author: Snowflake, Inc
 License: Apache License, Version 2.0

{snowpark_connect-0.25.0 → snowpark_connect-0.26.0}/src/snowflake/snowpark_connect/config.py

@@ -22,6 +22,9 @@ from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.types import TimestampTimeZone, TimestampType
 from snowflake.snowpark_connect.utils.concurrent import SynchronizedDict
 from snowflake.snowpark_connect.utils.context import get_session_id
+from snowflake.snowpark_connect.utils.external_udxf_cache import (
+    clear_external_udxf_cache,
+)
 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 (
@@ -136,6 +139,9 @@ class GlobalConfig:
         "spark.sql.parser.quotedRegexColumnNames": "false",
         # custom configs
         "snowpark.connect.version": ".".join(map(str, sas_version)),
+        # Control whether repartition(n) on a DataFrame forces splitting into n files during writes
+        # This matches spark behavior more closely, but introduces overhead.
+        "snowflake.repartition.for.writes": "false",
     }
 
     boolean_config_list = [
@@ -148,6 +154,7 @@ class GlobalConfig:
         "spark.sql.legacy.allowHashOnMapType",
         "spark.Catalog.databaseFilterInformationSchema",
         "spark.sql.parser.quotedRegexColumnNames",
+        "snowflake.repartition.for.writes",
     ]
 
     int_config_list = [
@@ -592,6 +599,9 @@ def parse_imports(session: snowpark.Session, imports: str | None) -> None:
     if not imports:
         return
 
+    # UDF needs to be recreated to include new imports
+    clear_external_udxf_cache(session)
+
     for udf_import in imports.strip("[] ").split(","):
         session.add_import(udf_import)
 

{snowpark_connect-0.25.0 → snowpark_connect-0.26.0}/src/snowflake/snowpark_connect/dataframe_container.py

@@ -29,6 +29,7 @@ class DataFrameContainer:
         table_name: str | None = None,
         alias: str | None = None,
         cached_schema_getter: Callable[[], StructType] | None = None,
+        partition_hint: int | None = None,
     ) -> None:
         """
         Initialize a new DataFrameContainer.
@@ -39,11 +40,13 @@ class DataFrameContainer:
             table_name: Optional table name for the DataFrame
             alias: Optional alias for the DataFrame
             cached_schema_getter: Optional function to get cached schema
+            partition_hint: Optional partition count from repartition() operations
         """
         self._dataframe = dataframe
         self._column_map = self._create_default_column_map(column_map)
         self._table_name = table_name
         self._alias = alias
+        self._partition_hint = partition_hint
 
         if cached_schema_getter is not None:
             self._apply_cached_schema_getter(cached_schema_getter)
@@ -62,6 +65,7 @@ class DataFrameContainer:
         table_name: str | None = None,
         alias: str | None = None,
         cached_schema_getter: Callable[[], StructType] | None = None,
+        partition_hint: int | None = None,
     ) -> DataFrameContainer:
         """
         Create a new container with complete column mapping configuration.
@@ -78,6 +82,7 @@ class DataFrameContainer:
             table_name: Optional table name
             alias: Optional alias
             cached_schema_getter: Optional function to get cached schema
+            partition_hint: Optional partition count from repartition() operations
 
         Returns:
             A new DataFrameContainer instance
@@ -123,6 +128,7 @@ class DataFrameContainer:
             table_name=table_name,
             alias=alias,
             cached_schema_getter=final_schema_getter,
+            partition_hint=partition_hint,
         )
 
     @property
@@ -163,6 +169,16 @@ class DataFrameContainer:
         """Set the alias name."""
         self._alias = value
 
+    @property
+    def partition_hint(self) -> int | None:
+        """Get the partition hint count."""
+        return self._partition_hint
+
+    @partition_hint.setter
+    def partition_hint(self, value: int | None) -> None:
+        """Set the partition hint count."""
+        self._partition_hint = value
+
     def _create_default_column_map(
         self, column_map: ColumnNameMap | None
     ) -> ColumnNameMap:

{snowpark_connect-0.25.0 → snowpark_connect-0.26.0}/src/snowflake/snowpark_connect/expression/map_udf.py

@@ -13,6 +13,10 @@ from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.type_mapping import proto_to_snowpark_type
 from snowflake.snowpark_connect.typed_column import TypedColumn
+from snowflake.snowpark_connect.utils.external_udxf_cache import (
+    cache_external_udf,
+    get_external_udf_from_cache,
+)
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.udf_helper import (
     SnowparkUDF,
@@ -30,6 +34,39 @@ from snowflake.snowpark_connect.utils.udxf_import_utils import (
 )
 
 
+def cache_external_udf_wrapper(from_register_udf: bool):
+    def outer_wrapper(wrapper_func):
+        def wrapper(
+            udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
+        ) -> SnowparkUDF | None:
+            udf_hash = hash(str(udf_proto))
+            cached_udf = get_external_udf_from_cache(udf_hash)
+
+            if cached_udf:
+                session = get_or_create_snowpark_session()
+                function_type = udf_proto.WhichOneof("function")
+                # TODO: Align this with SNOW-2316798 after merge
+                match function_type:
+                    case "scalar_scala_udf":
+                        session._udfs[cached_udf.name] = cached_udf
+                    case "python_udf" if from_register_udf:
+                        session._udfs[udf_proto.function_name.lower()] = cached_udf
+                    case "python_udf":
+                        pass
+                    case _:
+                        raise ValueError(f"Unsupported UDF type: {function_type}")
+
+                return cached_udf
+
+            snowpark_udf = wrapper_func(udf_proto)
+            cache_external_udf(udf_hash, snowpark_udf)
+            return snowpark_udf
+
+        return wrapper
+
+    return outer_wrapper
+
+
 def process_udf_return_type(
     return_type: types_proto.DataType,
 ) -> tuple[snowpark.types.DataType, snowpark.types.DataType]:
@@ -49,6 +86,7 @@ def process_udf_return_type(
     return original_snowpark_type, original_snowpark_type
 
 
+@cache_external_udf_wrapper(from_register_udf=True)
 def register_udf(
     udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
 ) -> SnowparkUDF:
@@ -84,12 +122,10 @@ def register_udf(
             return_type=udf._return_type,
             original_return_type=original_return_type,
         )
-        # the create udf does register the udf but this seems to be for the client side check
-        # TODO: check if this is needed
+        session._udfs[udf_proto.function_name.lower()] = udf
+        # scala udfs can be also accessed using `udf.name`
         if udf_processor._function_type == "scalar_scala_udf":
             session._udfs[udf.name] = udf
-        else:
-            session._udfs[udf_proto.function_name.lower()] = udf
         return udf
 
 
@@ -114,29 +150,34 @@ def map_common_inline_user_defined_udf(
                 udf_proto.scalar_scala_udf.outputType
             )
 
-    session = get_or_create_snowpark_session()
-    kwargs = {
-        "common_inline_user_defined_function": udf_proto,
-        "input_types": input_types,
-        "called_from": "map_common_inline_user_defined_udf",
-        "return_type": processed_return_type,
-        "udf_packages": global_config.get("snowpark.connect.udf.packages", ""),
-        "udf_imports": get_python_udxf_import_files(session),
-        "original_return_type": original_return_type,
-    }
-    if require_creating_udf_in_sproc(udf_proto):
-        snowpark_udf = process_udf_in_sproc(**kwargs)
-    else:
-        udf_processor = ProcessCommonInlineUserDefinedFunction(**kwargs)
-        udf = udf_processor.create_udf()
-        snowpark_udf = SnowparkUDF(
-            name=udf.name,
-            input_types=udf._input_types,
-            return_type=udf._return_type,
-            original_return_type=original_return_type,
-        )
-        if udf_processor._function_type == "scalar_scala_udf":
-            session._udfs[udf.name] = snowpark_udf
+    @cache_external_udf_wrapper(from_register_udf=False)
+    def get_snowpark_udf(
+        udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
+    ) -> SnowparkUDF:
+        session = get_or_create_snowpark_session()
+        kwargs = {
+            "common_inline_user_defined_function": udf_proto,
+            "input_types": input_types,
+            "called_from": "map_common_inline_user_defined_udf",
+            "return_type": processed_return_type,
+            "udf_packages": global_config.get("snowpark.connect.udf.packages", ""),
+            "udf_imports": get_python_udxf_import_files(session),
+            "original_return_type": original_return_type,
+        }
+        if require_creating_udf_in_sproc(udf_proto):
+            snowpark_udf = process_udf_in_sproc(**kwargs)
+        else:
+            udf_processor = ProcessCommonInlineUserDefinedFunction(**kwargs)
+            udf = udf_processor.create_udf()
+            snowpark_udf = SnowparkUDF(
+                name=udf.name,
+                input_types=udf._input_types,
+                return_type=udf._return_type,
+                original_return_type=original_return_type,
+            )
+        return snowpark_udf
+
+    snowpark_udf = get_snowpark_udf(udf_proto)
     udf_call_expr = snowpark_fn.call_udf(snowpark_udf.name, *snowpark_udf_args)
 
     # If the original return type was MapType or StructType but we converted it to VariantType,

{snowpark_connect-0.25.0 → snowpark_connect-0.26.0}/src/snowflake/snowpark_connect/expression/map_unresolved_function.py

@@ -476,11 +476,8 @@ def map_unresolved_function(
             return TypedColumn(result, lambda: expected_types)
 
     match function_name:
-        case func_name if (
-            get_is_evaluating_sql() and func_name.lower() in session._udfs
-        ):
-            # TODO: In Spark, UDFs can override built-in functions in SQL,
-            # but not in DataFrame ops.
+        case func_name if func_name.lower() in session._udfs:
+            # In Spark, UDFs can override built-in functions
             udf = session._udfs[func_name.lower()]
             result_exp = snowpark_fn.call_udf(
                 udf.name,
@@ -6479,6 +6476,18 @@ def map_unresolved_function(
                     if pattern_value is None:
                         return snowpark_fn.lit(None)
 
+                    # Optimization: treat escaped regex that resolves to a pure literal delimiter
+                    # - Single char: "\\."
+                    # - Multi char: e.g., "\\.505\\."
+                    if re.fullmatch(r"(?:\\.)+", pattern_value):
+                        literal_delim = re.sub(r"\\(.)", r"\1", pattern_value)
+                        return snowpark_fn.when(
+                            limit <= 0,
+                            snowpark_fn.split(
+                                str_, snowpark_fn.lit(literal_delim)
+                            ).cast(result_type),
+                        ).otherwise(native_split)
+
                     is_regexp = re.match(
                         ".*[\\[\\.\\]\\*\\?\\+\\^\\$\\{\\}\\|\\(\\)\\\\].*",
                         pattern_value,
@@ -8285,15 +8294,6 @@ def map_unresolved_function(
                 ),
             )
             result_type = BinaryType()
-        case udf_name if udf_name.lower() in session._udfs:
-            # TODO: In Spark, UDFs can override built-in functions in SQL,
-            # but not in DataFrame ops.
-            udf = session._udfs[udf_name.lower()]
-            result_exp = snowpark_fn.call_udf(
-                udf.name,
-                *(snowpark_fn.cast(arg, VariantType()) for arg in snowpark_args),
-            )
-            result_type = udf.return_type
         case udtf_name if udtf_name.lower() in session._udtfs:
             udtf, spark_col_names = session._udtfs[udtf_name.lower()]
             result_exp = snowpark_fn.call_table_function(
@@ -9623,13 +9623,14 @@ def _get_decimal_division_result_exp(
     snowpark_args: list[Column],
     spark_function_name: str,
 ) -> Column:
-    if isinstance(other_type, DecimalType) and overflow_detected:
-        if 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.'
-            )
-        else:
-            result_exp = snowpark_fn.lit(None).cast(result_type)
+    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.'
+        )
     else:
         dividend = snowpark_args[0].cast(DoubleType())
         divisor = snowpark_args[1]

{snowpark_connect-0.25.0 → snowpark_connect-0.26.0}/src/snowflake/snowpark_connect/relation/map_map_partitions.py

@@ -46,9 +46,10 @@ def map_map_partitions(
     udf_check(udf_proto)
 
     # Check if this is mapInArrow (eval_type == 207)
-    eval_type = udf_proto.python_udf.eval_type
-
-    if eval_type == MAP_IN_ARROW_EVAL_TYPE:
+    if (
+        udf_proto.WhichOneof("function") == "python_udf"
+        and udf_proto.python_udf.eval_type == MAP_IN_ARROW_EVAL_TYPE
+    ):
         return _map_in_arrow_with_pandas_udtf(input_container, udf_proto)
     else:
         return _map_partitions_with_udf(input_df, udf_proto)
@@ -126,7 +127,11 @@ def _map_partitions_with_udf(
         "udf_name": "spark_map_partitions_udf",
         "input_column_names": input_column_names,
         "replace": True,
-        "return_type": proto_to_snowpark_type(udf_proto.python_udf.output_type),
+        "return_type": proto_to_snowpark_type(
+            udf_proto.python_udf.output_type
+            if udf_proto.WhichOneof("function") == "python_udf"
+            else udf_proto.scalar_scala_udf.outputType
+        ),
         "udf_packages": global_config.get("snowpark.connect.udf.packages", ""),
         "udf_imports": get_python_udxf_import_files(input_df.session),
     }

{snowpark_connect-0.25.0 → snowpark_connect-0.26.0}/src/snowflake/snowpark_connect/relation/map_relation.py

@@ -90,6 +90,7 @@ def map_relation(
                 table_name=copy.deepcopy(cached_container.table_name),
                 alias=cached_container.alias,
                 cached_schema_getter=lambda: cached_df.schema,
+                partition_hint=cached_container.partition_hint,
             )
             # If we don't make a copy of the df._output, the expression IDs for attributes in Snowpark DataFrames will differ from those stored in the cache,
             # leading to errors during query execution.
@@ -189,13 +190,23 @@ def map_relation(
             case "read":
                 result = read.map_read(rel)
             case "repartition":
-                # TODO: Snowpark df identity transform with annotation
+                # Preserve partition hint for file output control
+                # This handles both repartition(n) with shuffle=True and coalesce(n) with shuffle=False
                 result = map_relation(rel.repartition.input)
+                if rel.repartition.num_partitions > 0:
+                    result.partition_hint = rel.repartition.num_partitions
             case "repartition_by_expression":
                 # This is a no-op operation in SAS as Snowpark doesn't have the concept of partitions.
                 # All the data in the dataframe will be treated as a single partition, and this will not
                 # have any side effects.
                 result = map_relation(rel.repartition_by_expression.input)
+                # Only preserve partition hint if num_partitions is explicitly specified and > 0
+                # Column-based repartitioning without count should clear any existing partition hints
+                if rel.repartition_by_expression.num_partitions > 0:
+                    result.partition_hint = rel.repartition_by_expression.num_partitions
+                else:
+                    # Column-based repartitioning clears partition hint (resets to default behavior)
+                    result.partition_hint = None
             case "replace":
                 result = map_row_ops.map_replace(rel)
             case "sample":

{snowpark_connect-0.25.0 → snowpark_connect-0.26.0}/src/snowflake/snowpark_connect/relation/map_row_ops.py

@@ -553,7 +553,14 @@ def map_filter(
         rel.filter.condition, input_container.column_map, typer
     )
 
-    result = input_df.filter(condition.col)
+    if rel.filter.input.WhichOneof("rel_type") == "subquery_alias":
+        # map_subquery_alias does not actually wrap the DataFrame in an alias or subquery.
+        # Apparently, there are cases (e.g., TpcdsQ53) where this is required, without it, we get
+        # SQL compilation error.
+        # To mitigate it, we are doing .select("*"), .alias() introduces additional describe queries
+        result = input_df.select("*").filter(condition.col)
+    else:
+        result = input_df.filter(condition.col)
 
     return DataFrameContainer(
         result,

{snowpark_connect-0.25.0 → snowpark_connect-0.26.0}/src/snowflake/snowpark_connect/relation/map_udtf.py

@@ -31,6 +31,10 @@ from snowflake.snowpark_connect.type_mapping import (
     proto_to_snowpark_type,
 )
 from snowflake.snowpark_connect.utils.context import push_udtf_context
+from snowflake.snowpark_connect.utils.external_udxf_cache import (
+    cache_external_udtf,
+    get_external_udtf_from_cache,
+)
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.udtf_helper import (
     SnowparkUDTF,
@@ -44,6 +48,34 @@ from snowflake.snowpark_connect.utils.udxf_import_utils import (
 )
 
 
+def cache_external_udtf_wrapper(from_register_udtf: bool):
+    def outer_wrapper(wrapper_func):
+        def wrapper(
+            udtf_proto: relation_proto.CommonInlineUserDefinedTableFunction,
+            spark_column_names,
+        ) -> SnowparkUDTF | None:
+            udf_hash = hash(str(udtf_proto))
+            cached_udtf = get_external_udtf_from_cache(udf_hash)
+
+            if cached_udtf:
+                if from_register_udtf:
+                    session = get_or_create_snowpark_session()
+                    session._udtfs[udtf_proto.function_name.lower()] = (
+                        cached_udtf,
+                        spark_column_names,
+                    )
+
+                return cached_udtf
+
+            snowpark_udf = wrapper_func(udtf_proto, spark_column_names)
+            cache_external_udtf(udf_hash, snowpark_udf)
+            return snowpark_udf
+
+        return wrapper
+
+    return outer_wrapper
+
+
 def build_expected_types_from_parsed(
     parsed_return: types_proto.DataType,
 ) -> List[Tuple[str, Any]]:
@@ -165,26 +197,37 @@ def register_udtf(
     ) = process_return_type(python_udft.return_type)
     function_name = udtf_proto.function_name
 
-    kwargs = {
-        "session": session,
-        "udtf_proto": udtf_proto,
-        "expected_types": expected_types,
-        "output_schema": output_schema,
-        "packages": global_config.get("snowpark.connect.udf.packages", ""),
-        "imports": get_python_udxf_import_files(session),
-        "called_from": "register_udtf",
-        "is_arrow_enabled": is_arrow_enabled_in_udtf(),
-        "is_spark_compatible_udtf_mode_enabled": is_spark_compatible_udtf_mode_enabled(),
-    }
-
-    if require_creating_udtf_in_sproc(udtf_proto):
-        snowpark_udtf = create_udtf_in_sproc(**kwargs)
-    else:
-        udtf = create_udtf(**kwargs)
-        snowpark_udtf = SnowparkUDTF(
-            name=udtf.name, input_types=udtf._input_types, output_schema=output_schema
-        )
+    @cache_external_udtf_wrapper(from_register_udtf=True)
+    def _register_udtf(
+        udtf_proto: relation_proto.CommonInlineUserDefinedTableFunction,
+        spark_column_names,
+    ):
+        kwargs = {
+            "session": session,
+            "udtf_proto": udtf_proto,
+            "expected_types": expected_types,
+            "output_schema": output_schema,
+            "packages": global_config.get("snowpark.connect.udf.packages", ""),
+            "imports": get_python_udxf_import_files(session),
+            "called_from": "register_udtf",
+            "is_arrow_enabled": is_arrow_enabled_in_udtf(),
+            "is_spark_compatible_udtf_mode_enabled": is_spark_compatible_udtf_mode_enabled(),
+        }
+
+        if require_creating_udtf_in_sproc(udtf_proto):
+            snowpark_udtf = create_udtf_in_sproc(**kwargs)
+        else:
+            udtf = create_udtf(**kwargs)
+            snowpark_udtf = SnowparkUDTF(
+                name=udtf.name,
+                input_types=udtf._input_types,
+                output_schema=output_schema,
+            )
+
+        return snowpark_udtf
 
+    snowpark_udtf = _register_udtf(udtf_proto, spark_column_names)
+    # We have to update cached _udtfs here, because function could have been cached in map_common_inline_user_defined_table_function
     session._udtfs[function_name.lower()] = (snowpark_udtf, spark_column_names)
     return snowpark_udtf
 
@@ -213,32 +256,41 @@ def map_common_inline_user_defined_table_function(
         spark_column_names,
     ) = process_return_type(python_udft.return_type)
 
-    kwargs = {
-        "session": session,
-        "udtf_proto": rel,
-        "expected_types": expected_types,
-        "output_schema": output_schema,
-        "packages": global_config.get("snowpark.connect.udf.packages", ""),
-        "imports": get_python_udxf_import_files(session),
-        "called_from": "map_common_inline_user_defined_table_function",
-        "is_arrow_enabled": is_arrow_enabled_in_udtf(),
-        "is_spark_compatible_udtf_mode_enabled": is_spark_compatible_udtf_mode_enabled(),
-    }
-
-    if require_creating_udtf_in_sproc(rel):
-        snowpark_udtf_or_error = create_udtf_in_sproc(**kwargs)
-        if isinstance(snowpark_udtf_or_error, str):
-            raise PythonException(snowpark_udtf_or_error)
-        snowpark_udtf = snowpark_udtf_or_error
-    else:
-        udtf_or_error = create_udtf(**kwargs)
-        if isinstance(udtf_or_error, str):
-            raise PythonException(udtf_or_error)
-        udtf = udtf_or_error
-        snowpark_udtf = SnowparkUDTF(
-            name=udtf.name, input_types=udtf._input_types, output_schema=output_schema
-        )
+    @cache_external_udtf_wrapper(from_register_udtf=False)
+    def _get_udtf(
+        udtf_proto: relation_proto.CommonInlineUserDefinedTableFunction,
+        spark_column_names,
+    ):
+        kwargs = {
+            "session": session,
+            "udtf_proto": udtf_proto,
+            "expected_types": expected_types,
+            "output_schema": output_schema,
+            "packages": global_config.get("snowpark.connect.udf.packages", ""),
+            "imports": get_python_udxf_import_files(session),
+            "called_from": "map_common_inline_user_defined_table_function",
+            "is_arrow_enabled": is_arrow_enabled_in_udtf(),
+            "is_spark_compatible_udtf_mode_enabled": is_spark_compatible_udtf_mode_enabled(),
+        }
+
+        if require_creating_udtf_in_sproc(udtf_proto):
+            snowpark_udtf_or_error = create_udtf_in_sproc(**kwargs)
+            if isinstance(snowpark_udtf_or_error, str):
+                raise PythonException(snowpark_udtf_or_error)
+            snowpark_udtf = snowpark_udtf_or_error
+        else:
+            udtf_or_error = create_udtf(**kwargs)
+            if isinstance(udtf_or_error, str):
+                raise PythonException(udtf_or_error)
+            udtf = udtf_or_error
+            snowpark_udtf = SnowparkUDTF(
+                name=udtf.name,
+                input_types=udtf._input_types,
+                output_schema=output_schema,
+            )
+        return snowpark_udtf
 
+    snowpark_udtf = _get_udtf(rel, spark_column_names)
     column_map = ColumnNameMap([], [])
     snowpark_udtf_args = []
 

{snowpark_connect-0.25.0 → snowpark_connect-0.26.0}/src/snowflake/snowpark_connect/relation/utils.py

@@ -6,6 +6,7 @@ import random
 import re
 import string
 import time
+import uuid
 from typing import Sequence
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
@@ -153,6 +154,49 @@ def random_string(
     return "".join([prefix, random_part, suffix])
 
 
+def generate_spark_compatible_filename(
+    task_id: int = 0,
+    attempt_number: int = 0,
+    compression: str = None,
+    format_ext: str = "parquet",
+) -> str:
+    """Generate a Spark-compatible filename following the convention:
+    part-<task-id>-<uuid>-c<attempt-number>.<compression>.<format>
+
+    Args:
+        task_id: Task ID (usually 0 for single partition)
+        attempt_number: Attempt number (usually 0)
+        compression: Compression type (e.g., 'snappy', 'gzip', 'none')
+        format_ext: File format extension (e.g., 'parquet', 'csv', 'json')
+
+    Returns:
+        A filename string following Spark's naming convention
+    """
+    # Generate a UUID for uniqueness
+    file_uuid = str(uuid.uuid4())
+
+    # Format task ID with leading zeros (5 digits)
+    formatted_task_id = f"{task_id:05d}"
+
+    # Format attempt number with leading zeros (3 digits)
+    formatted_attempt = f"{attempt_number:03d}"
+
+    # Build the base filename
+    base_name = f"part-{formatted_task_id}-{file_uuid}-c{formatted_attempt}"
+
+    # Add compression if specified and not 'none'
+    if compression and compression.lower() not in ("none", "uncompressed"):
+        compression_part = f".{compression.lower()}"
+    else:
+        compression_part = ""
+
+    # Add format extension if specified
+    if format_ext:
+        return f"{base_name}{compression_part}.{format_ext}"
+    else:
+        return f"{base_name}{compression_part}"
+
+
 def _normalize_query_for_semantic_hash(query_str: str) -> str:
     """
     Normalize a query string for semantic comparison by extracting original names from