snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/relation/map_map_partitions.py

@@ -0,0 +1,146 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import pyspark.sql.connect.proto.relations_pb2 as relation_proto
+from pyspark.sql.connect.proto.expressions_pb2 import CommonInlineUserDefinedFunction
+
+import snowflake.snowpark.functions as snowpark_fn
+from snowflake import snowpark
+from snowflake.snowpark.types import StructType
+from snowflake.snowpark_connect.column_name_handler import with_column_map
+from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.constants import MAP_IN_ARROW_EVAL_TYPE
+from snowflake.snowpark_connect.relation.map_relation import map_relation
+from snowflake.snowpark_connect.type_mapping import proto_to_snowpark_type
+from snowflake.snowpark_connect.utils.pandas_udtf_utils import create_pandas_udtf
+from snowflake.snowpark_connect.utils.session import get_python_udxf_import_files
+from snowflake.snowpark_connect.utils.udf_helper import (
+    SnowparkUDF,
+    process_udf_in_sproc,
+    require_creating_udf_in_sproc,
+    udf_check,
+)
+from snowflake.snowpark_connect.utils.udf_utils import (
+    ProcessCommonInlineUserDefinedFunction,
+)
+from snowflake.snowpark_connect.utils.udtf_helper import (
+    create_pandas_udtf_in_sproc,
+    require_creating_udtf_in_sproc,
+)
+
+
+def map_map_partitions(rel: relation_proto.Relation) -> snowpark.DataFrame:
+    """
+    Map a function over the partitions of the input DataFrame.
+
+    This is a simple wrapper around the `mapInPandas` method in Snowpark.
+    """
+    input_df = map_relation(rel.map_partitions.input)
+    udf_proto = rel.map_partitions.func
+    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:
+        return _map_in_arrow_with_pandas_udtf(input_df, udf_proto)
+    else:
+        return _map_partitions_with_udf(input_df, udf_proto)
+
+
+def _call_udtf(
+    udtf_name: str, input_df: snowpark.DataFrame, return_type: StructType | None = None
+) -> snowpark.DataFrame:
+    # Add a dummy column with random 1-10 values for partitioning
+    input_df_with_dummy = input_df.withColumn(
+        "_DUMMY_PARTITION_KEY",
+        (
+            snowpark_fn.uniform(
+                snowpark_fn.lit(1), snowpark_fn.lit(10), snowpark_fn.random()
+            )
+            * 10
+        ).cast("int"),
+    )
+
+    udtf_columns = input_df.columns + [snowpark_fn.col("_DUMMY_PARTITION_KEY")]
+
+    result_df_with_dummy = input_df_with_dummy.select(
+        snowpark_fn.call_table_function(udtf_name, *udtf_columns).over(
+            partition_by=[snowpark_fn.col("_DUMMY_PARTITION_KEY")]
+        )
+    )
+
+    output_cols = [field.name for field in return_type.fields]
+
+    # Only return the output columns.
+    result_df = result_df_with_dummy.select(*output_cols)
+
+    return with_column_map(
+        result_df,
+        output_cols,
+        output_cols,
+        [field.datatype for field in return_type.fields],
+    )
+
+
+def _map_in_arrow_with_pandas_udtf(
+    input_df: snowpark.DataFrame, udf_proto: CommonInlineUserDefinedFunction
+) -> snowpark.DataFrame:
+    """
+    Handle mapInArrow using pandas_udtf for partition-level Arrow processing.
+    """
+    input_schema = input_df.schema
+    spark_column_names = input_df._column_map.get_spark_columns()
+    return_type = proto_to_snowpark_type(udf_proto.python_udf.output_type)
+    if require_creating_udtf_in_sproc(udf_proto):
+        udtf_name = create_pandas_udtf_in_sproc(
+            udf_proto, spark_column_names, input_schema, return_type
+        )
+    else:
+        map_in_arrow_udtf = create_pandas_udtf(
+            udf_proto, spark_column_names, input_schema, return_type
+        )
+        udtf_name = map_in_arrow_udtf.name
+    return _call_udtf(udtf_name, input_df, return_type)
+
+
+def _map_partitions_with_udf(
+    input_df: snowpark.DataFrame, udf_proto
+) -> snowpark.DataFrame:
+    """
+    Original UDF-based approach for non-mapInArrow map_partitions cases.
+    """
+    input_column_names = input_df.columns
+    kwargs = {
+        "common_inline_user_defined_function": udf_proto,
+        "input_types": [f.datatype for f in input_df.schema.fields],
+        "called_from": "map_map_partitions",
+        "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),
+        "udf_packages": global_config.get("snowpark.connect.udf.packages", ""),
+        "udf_imports": get_python_udxf_import_files(input_df.session),
+    }
+
+    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=None,
+        )
+    udf_column_name = "UDF_OUTPUT"
+    snowpark_columns = [snowpark_fn.col(name) for name in input_df.columns]
+    result = input_df.select(snowpark_fn.call_udf(snowpark_udf.name, *snowpark_columns))
+    return with_column_map(
+        result,
+        [udf_column_name],
+        [udf_column_name],
+        [snowpark_udf.return_type],
+    )

snowflake/snowpark_connect/relation/map_relation.py

@@ -0,0 +1,253 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import copy
+
+import pandas
+import pyspark.sql.connect.proto.relations_pb2 as relation_proto
+
+from snowflake import snowpark
+from snowflake.snowpark_connect.column_name_handler import set_schema_getter
+from snowflake.snowpark_connect.utils.cache import (
+    df_cache_map_get,
+    df_cache_map_put_if_absent,
+)
+from snowflake.snowpark_connect.utils.context import (
+    get_plan_id_map,
+    get_session_id,
+    push_operation_scope,
+    set_is_aggregate_function,
+    set_plan_id_map,
+)
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
+from snowflake.snowpark_connect.utils.telemetry import (
+    SnowparkConnectNotImplementedError,
+)
+
+# Spark doesn't have proto for NaturalJoin as Spark DF doesn't have API for it.
+# BUT spark supports NATURAL JOIN via SQL
+NATURAL_JOIN_TYPE_BASE = 22
+
+
+def map_relation(
+    rel: relation_proto.Relation, reuse_parsed_plan: bool = True
+) -> snowpark.DataFrame | pandas.DataFrame:
+    """Map a Spark Protobuf Relation message to a Snowpark DataFrame or pandas DataFrame.
+
+    NOTE: We return a pandas DataFrame object when the return value of the operation is a
+    scalar value. The client expects these as an Arrow buffer with this return value packed
+    into a single cell. Rather than do what Spark does (which is put the data back into
+    Spark and then convert it to an Arrow buffer), we leave the value we have already extracted
+    from Snowflake here in the TCM as a pandas DataFrame and convert it to an Arrow Buffer.
+
+    Args:
+        rel (relation_proto.Relation): The Spark Protobuf Relation message to map.
+        reuse_parsed_plan (bool, optional): If True, reuses previously parsed df from cache
+            to avoid redundant operations.
+
+    Returns:
+        snowpark.DataFrame | pandas.DataFrame: The Snowpark DataFrame or pandas DataFrame
+        that corresponds to the input Spark Protobuf Relation message.
+    """
+    # TODO: from snowflake_connect_server.relation import map_extension
+    from snowflake.snowpark_connect.relation import (
+        map_aggregate,
+        map_catalog,
+        map_column_ops,
+        map_crosstab,
+        map_extension,
+        map_join,
+        map_local_relation,
+        map_map_partitions,
+        map_row_ops,
+        map_sample_by,
+        map_show_string,
+        map_sql,
+        map_stats,
+        map_subquery_alias,
+        map_udtf,
+        read,
+    )
+
+    if reuse_parsed_plan and rel.HasField("common") and rel.common.HasField("plan_id"):
+        # TODO: remove get_session_id() when we host SAS in Snowflake server
+        # Check for cached relation
+        cache_entry = df_cache_map_get((get_session_id(), rel.common.plan_id))
+        if cache_entry is not None:
+            if isinstance(cache_entry, snowpark.DataFrame):
+                set_plan_id_map(rel.common.plan_id, cache_entry)
+            return cache_entry
+
+        # If df is not cached, check if we have parsed the plan
+        cached_df = get_plan_id_map(rel.common.plan_id)
+        if cached_df is not None:
+            result = copy.copy(cached_df)
+            result._column_map = copy.deepcopy(cached_df._column_map)
+            result._table_name = copy.deepcopy(cached_df._table_name)
+            set_schema_getter(result, lambda: cached_df.schema)
+            # 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.
+            result._output = cached_df._output
+
+            return result
+
+    if rel.WhichOneof("rel_type") is not None:
+        logger.info(rel.WhichOneof("rel_type").upper())
+    else:
+        # This happens when the relation is empty, usually because the incoming message
+        # type was incorrectly routed here.
+        raise SnowparkConnectNotImplementedError("No Relation Type")
+
+    result: snowpark.DataFrame | pandas.DataFrame
+    operation = rel.WhichOneof("rel_type")
+    with push_operation_scope(operation):
+        match operation:
+            case "aggregate":
+                set_is_aggregate_function(("default", True))
+                match rel.aggregate.group_type:
+                    case relation_proto.Aggregate.GroupType.GROUP_TYPE_GROUPBY:
+                        result = map_aggregate.map_group_by_aggregate(rel)
+                    case relation_proto.Aggregate.GroupType.GROUP_TYPE_ROLLUP:
+                        result = map_aggregate.map_rollup_aggregate(rel)
+                    case relation_proto.Aggregate.GroupType.GROUP_TYPE_CUBE:
+                        result = map_aggregate.map_cube_aggregate(rel)
+                    case relation_proto.Aggregate.GroupType.GROUP_TYPE_PIVOT:
+                        result = map_aggregate.map_pivot_aggregate(rel)
+                    case other:
+                        raise SnowparkConnectNotImplementedError(f"AGGREGATE {other}")
+            case "approx_quantile":
+                result = map_stats.map_approx_quantile(rel)
+            case "as_of_join":
+                raise SnowparkConnectNotImplementedError("AS_OF_JOIN")
+            case "catalog":  # TODO: order these alphabetically
+                result = map_catalog.map_catalog(rel.catalog)
+            case "collect_metrics":
+                # no-op, Snowflake doesn't support metrics
+                result = map_relation(rel.collect_metrics.input)
+            case "common_inline_user_defined_table_function":
+                result = map_udtf.map_common_inline_user_defined_table_function(
+                    rel.common_inline_user_defined_table_function
+                )
+            case "corr":
+                result = map_stats.map_corr(rel)
+            case "cov":
+                result = map_stats.map_cov(rel)
+            case "crosstab":
+                result = map_crosstab.map_crosstab(rel)
+            case "deduplicate":
+                result = map_row_ops.map_deduplicate(rel)
+            case "describe":
+                result = map_stats.map_describe(rel)
+            case "drop":
+                result = map_column_ops.map_drop(rel)
+            case "drop_na":
+                result = map_row_ops.map_dropna(rel)
+            case "extension":
+                result = map_extension.map_extension(rel)
+            case "fill_na":
+                result = map_row_ops.map_fillna(rel)
+            case "filter":
+                result = map_row_ops.map_filter(rel)
+            case "freq_items":
+                result = map_stats.map_freq_items(rel)
+            case "hint":
+                # no-op, Snowflake doesn't support hints
+                result = map_relation(rel.hint.input)
+            case "html_string":
+                result = map_show_string.map_repr_html(rel)
+            case "join":
+                result = map_join.map_join(rel)
+            case "limit":
+                result = map_row_ops.map_limit(rel)
+            case "local_relation":
+                result = map_local_relation.map_local_relation(rel)
+                df_cache_map_put_if_absent(
+                    (get_session_id(), rel.common.plan_id), lambda: result, False
+                )
+            case "cached_local_relation":
+                cached_df = df_cache_map_get(
+                    (get_session_id(), rel.cached_local_relation.hash)
+                )
+                if cached_df is None:
+                    raise ValueError(
+                        f"Local relation with hash {rel.cached_local_relation.hash} not found in cache."
+                    )
+                return cached_df
+            case "map_partitions":
+                result = map_map_partitions.map_map_partitions(rel)
+            case "offset":
+                result = map_row_ops.map_offset(rel)
+            case "project":
+                result = map_column_ops.map_project(rel)
+            case "range":
+                result = map_local_relation.map_range(rel)
+            case "read":
+                result = read.map_read(rel)
+            case "repartition":
+                # TODO: Snowpark df identity transform with annotation
+                result = map_relation(rel.repartition.input)
+            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)
+            case "replace":
+                result = map_row_ops.map_replace(rel)
+            case "sample":
+                sampled_df_not_evaluated = map_row_ops.map_sample(rel)
+                df_cache_map_put_if_absent(
+                    (get_session_id(), rel.common.plan_id),
+                    lambda: sampled_df_not_evaluated,
+                    True,
+                )
+
+                # We will retrieve from cache and return that, because insertion to cache
+                # triggers evaluation.
+                result = df_cache_map_get((get_session_id(), rel.common.plan_id))
+            case "sample_by":
+                result = map_sample_by.map_sample_by(rel)
+            case "set_op":
+                match rel.set_op.set_op_type:
+                    case relation_proto.SetOperation.SetOpType.SET_OP_TYPE_INTERSECT:
+                        result = map_row_ops.map_intersect(rel)
+                    case relation_proto.SetOperation.SetOpType.SET_OP_TYPE_UNION:
+                        result = map_row_ops.map_union(rel)
+                    case relation_proto.SetOperation.SetOpType.SET_OP_TYPE_EXCEPT:
+                        result = map_row_ops.map_except(rel)
+                    case other:
+                        raise SnowparkConnectNotImplementedError(f"SET_OP {other}")
+            case "show_string":
+                result = map_show_string.map_show_string(rel)
+            case "sort":
+                result = map_column_ops.map_sort(
+                    rel.sort
+                )  # TODO: follow this pattern elsewhere.
+            case "sql":
+                result = map_sql.map_sql(rel)
+            case "subquery_alias":
+                result = map_subquery_alias.map_alias(rel)
+            case "summary":
+                result = map_stats.map_summary(rel)
+            case "tail":
+                result = map_row_ops.map_tail(rel)
+            case "to_df":
+                result = map_column_ops.map_to_df(rel)
+            case "to_schema":
+                result = map_column_ops.map_to_schema(rel)
+            case "unpivot":
+                result = map_column_ops.map_unpivot(rel)
+            case "with_columns":
+                result = map_column_ops.map_with_columns(rel)
+            case "with_columns_renamed":
+                result = map_column_ops.map_with_columns_renamed(rel)
+            case "with_relations":
+                raise SnowparkConnectNotImplementedError("WITH_RELATIONS")
+            case "group_map":
+                result = map_column_ops.map_group_map(rel)
+            case other:
+                raise SnowparkConnectNotImplementedError(f"Other Relation {other}")
+
+        if isinstance(result, snowpark.DataFrame):
+            set_plan_id_map(rel.common.plan_id, result)
+        return result