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

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":

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,

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,10 +27,14 @@ 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
+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]]:
@@ -125,7 +157,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:
@@ -163,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
 
@@ -211,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 = []
 

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/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