snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/expression/map_udf.py

@@ -0,0 +1,142 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
+import pyspark.sql.connect.proto.types_pb2 as types_proto
+
+import snowflake.snowpark.functions as snowpark_fn
+from snowflake import snowpark
+from snowflake.snowpark.types import MapType, StructType, VariantType
+from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+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.session import (
+    get_or_create_snowpark_session,
+    get_python_udxf_import_files,
+)
+from snowflake.snowpark_connect.utils.udf_helper import (
+    SnowparkUDF,
+    gen_input_types,
+    infer_snowpark_arguments,
+    process_udf_in_sproc,
+    require_creating_udf_in_sproc,
+    udf_check,
+)
+from snowflake.snowpark_connect.utils.udf_utils import (
+    ProcessCommonInlineUserDefinedFunction,
+)
+
+
+def process_udf_return_type(
+    return_type: types_proto.DataType,
+) -> tuple[snowpark.types.DataType, snowpark.types.DataType]:
+    """Process UDF return type, handling DDL strings if present.
+
+    Returns a tuple of (processed_type, original_type) where:
+    - processed_type: The type to use for UDF registration ((MapType, StructType) -> VariantType)
+    - original_type: The original type for result processing
+    """
+    original_snowpark_type = proto_to_snowpark_type(return_type)
+
+    # Snowflake UDF does not support MapType or StructType, so we convert them to VariantType.
+    # We return both the converted type and original type for proper result processing.
+    if isinstance(original_snowpark_type, (MapType, StructType)):
+        return VariantType(), original_snowpark_type
+
+    return original_snowpark_type, original_snowpark_type
+
+
+def register_udf(
+    udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
+) -> SnowparkUDF:
+    udf_check(udf_proto)
+    processed_return_type, original_return_type = process_udf_return_type(
+        udf_proto.python_udf.output_type
+    )
+    session = get_or_create_snowpark_session()
+    kwargs = {
+        "common_inline_user_defined_function": udf_proto,
+        "called_from": "register_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):
+        return process_udf_in_sproc(**kwargs)
+    else:
+        udf_processor = ProcessCommonInlineUserDefinedFunction(**kwargs)
+        udf = udf_processor.create_udf()
+        udf = SnowparkUDF(
+            name=udf.name,
+            input_types=udf._input_types,
+            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
+        return udf
+
+
+def map_common_inline_user_defined_udf(
+    exp: expressions_proto.Expression,
+    column_mapping: ColumnNameMap,
+    typer: ExpressionTyper,
+) -> tuple[str, TypedColumn]:
+    udf_proto = exp.common_inline_user_defined_function
+    udf_check(udf_proto)
+    snowpark_udf_arg_names, snowpark_udf_args = infer_snowpark_arguments(
+        udf_proto, column_mapping, typer
+    )
+    input_types = gen_input_types(snowpark_udf_args, typer)
+    processed_return_type, original_return_type = process_udf_return_type(
+        udf_proto.python_udf.output_type
+    )
+    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,
+        )
+
+    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,
+    # we need to parse the JSON result back to the original type
+    if isinstance(original_return_type, (MapType, StructType)) and isinstance(
+        processed_return_type, VariantType
+    ):
+        # Parse JSON and cast back to original type
+        result_expr = snowpark_fn.parse_json(udf_call_expr).cast(original_return_type)
+        result_type = original_return_type
+    else:
+        result_expr = udf_call_expr
+        result_type = snowpark_udf.return_type
+
+    return (
+        f"{udf_proto.function_name}({', '.join(snowpark_udf_arg_names)})",
+        TypedColumn(
+            result_expr,
+            lambda: [result_type],
+        ),
+    )

snowflake/snowpark_connect/expression/map_unresolved_attribute.py

@@ -0,0 +1,241 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import re
+
+import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
+from pyspark.errors.exceptions.connect import AnalysisException
+
+import snowflake.snowpark.functions as snowpark_fn
+from snowflake.snowpark._internal.analyzer.analyzer_utils import (
+    quote_name_without_upper_casing,
+)
+from snowflake.snowpark.exceptions import SnowparkSQLException
+from snowflake.snowpark.types import ArrayType, MapType, StructType
+from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.expression.typer import ExpressionTyper
+from snowflake.snowpark_connect.typed_column import TypedColumn
+from snowflake.snowpark_connect.utils.attribute_handling import (
+    split_fully_qualified_spark_name,
+)
+from snowflake.snowpark_connect.utils.context import (
+    get_is_evaluating_sql,
+    get_outer_dataframes,
+    get_plan_id_map,
+    resolve_lca_alias,
+)
+
+SPARK_QUOTED = re.compile("^(`.*`)$", re.DOTALL)
+
+
+def map_unresolved_attribute(
+    exp: expressions_proto.Expression,
+    column_mapping: ColumnNameMap,
+    typer: ExpressionTyper,
+) -> tuple[str, TypedColumn]:
+    original_attr_name = exp.unresolved_attribute.unparsed_identifier
+    name_parts = split_fully_qualified_spark_name(original_attr_name)
+
+    assert len(name_parts) > 0, f"Unable to parse input attribute: {original_attr_name}"
+
+    attr_name = ".".join(name_parts)
+
+    has_plan_id = exp.unresolved_attribute.HasField("plan_id")
+
+    if has_plan_id:
+        plan_id = exp.unresolved_attribute.plan_id
+        target_df = get_plan_id_map(plan_id)
+        assert (
+            target_df is not None
+        ), f"resolving an attribute of a unresolved dataframe {plan_id}"
+        column_mapping = target_df._column_map
+        typer = ExpressionTyper(target_df)
+
+    def get_col(snowpark_name):
+        return (
+            snowpark_fn.col(snowpark_name)
+            if not has_plan_id
+            else target_df.col(snowpark_name)
+        )
+
+    # Check if regex column names are enabled and this is a quoted identifier
+    # We need to check the original attribute name before split_fully_qualified_spark_name processes it
+    if (
+        get_is_evaluating_sql()
+        and global_config.spark_sql_parser_quotedRegexColumnNames
+        and SPARK_QUOTED.match(original_attr_name)
+    ):
+        # Extract regex pattern by removing backticks
+        regex_pattern = original_attr_name[1:-1]  # Remove first and last backtick
+
+        # Get all available column names from the column mapping
+        available_columns = column_mapping.get_spark_columns()
+
+        # Match the regex pattern against available columns
+        matched_columns = []
+        try:
+            compiled_regex = re.compile(
+                regex_pattern,
+                re.IGNORECASE if not global_config.spark_sql_caseSensitive else 0,
+            )
+            for col_name in available_columns:
+                if compiled_regex.fullmatch(col_name):
+                    matched_columns.append(col_name)
+        except re.error as e:
+            raise AnalysisException(f"Invalid regex pattern '{regex_pattern}': {e}")
+
+        if not matched_columns:
+            raise AnalysisException(
+                f"No columns match the regex pattern '{regex_pattern}'"
+            )
+
+        # When multiple columns match, we need to signal that this should expand to multiple columns
+        # Since map_unresolved_attribute can only return one column, we'll use a special marker
+        # to indicate that this is a multi-column regex expansion
+        if len(matched_columns) > 1:
+            # Create a special column name that indicates multi-column expansion
+            # The higher-level logic will need to handle this
+            multi_col_name = "__REGEX_MULTI_COL__"
+            # For now, return the first column but mark it specially
+            quoted_col_name = matched_columns[0]
+            snowpark_name = (
+                column_mapping.get_snowpark_column_name_from_spark_column_name(
+                    quoted_col_name
+                )
+            )
+            col = get_col(snowpark_name)
+            qualifiers = column_mapping.get_qualifier_for_spark_column(quoted_col_name)
+            typed_col = TypedColumn(col, lambda: typer.type(col))
+            typed_col.set_qualifiers(qualifiers)
+            # Store matched columns info for later use
+            typed_col._regex_matched_columns = matched_columns
+            return (multi_col_name, typed_col)
+        else:
+            # Single column match - return that column
+            quoted_col_name = matched_columns[0]
+            snowpark_name = (
+                column_mapping.get_snowpark_column_name_from_spark_column_name(
+                    quoted_col_name
+                )
+            )
+            col = get_col(snowpark_name)
+            qualifiers = column_mapping.get_qualifier_for_spark_column(quoted_col_name)
+            typed_col = TypedColumn(col, lambda: typer.type(col))
+            typed_col.set_qualifiers(qualifiers)
+            return (matched_columns[0], typed_col)
+
+    quoted_attr_name = ".".join(
+        quote_name_without_upper_casing(x) for x in name_parts[:-1]
+    )
+    if len(name_parts) > 1:
+        quoted_attr_name = f"{quoted_attr_name}.{name_parts[-1]}"
+    else:
+        quoted_attr_name = name_parts[0]
+
+    snowpark_name = column_mapping.get_snowpark_column_name_from_spark_column_name(
+        quoted_attr_name, allow_non_exists=True
+    )
+    if snowpark_name is not None:
+        col = get_col(snowpark_name)
+        qualifiers = column_mapping.get_qualifier_for_spark_column(quoted_attr_name)
+    else:
+        # this means it has to be a struct column with a field name
+        snowpark_name = column_mapping.get_snowpark_column_name_from_spark_column_name(
+            name_parts[0], allow_non_exists=True
+        )
+        if snowpark_name is None:
+            for outer_df in get_outer_dataframes():
+                snowpark_name = outer_df._column_map.get_snowpark_column_name_from_spark_column_name(
+                    name_parts[0], allow_non_exists=True
+                )
+                if snowpark_name is not None:
+                    break
+
+        if snowpark_name is None:
+            # Attempt LCA fallback.
+            alias_tc = resolve_lca_alias(attr_name)
+
+            if alias_tc is not None:
+                # Return the TypedColumn that represents the alias.
+                return (attr_name, alias_tc)
+
+            # If qualified name not found, try to resolve as unqualified column name
+            # This handles cases like "d.name" where we need to find "name" after a JOIN
+            if len(name_parts) > 1:
+                unqualified_name = name_parts[-1]
+                snowpark_name = (
+                    column_mapping.get_snowpark_column_name_from_spark_column_name(
+                        unqualified_name, allow_non_exists=True
+                    )
+                )
+                if snowpark_name is not None:
+                    col = get_col(snowpark_name)
+                    qualifiers = column_mapping.get_qualifier_for_spark_column(
+                        unqualified_name
+                    )
+                    typed_col = TypedColumn(col, lambda: typer.type(col))
+                    typed_col.set_qualifiers(qualifiers)
+                    return (unqualified_name, typed_col)
+
+        if snowpark_name is None:
+            if has_plan_id:
+                raise AnalysisException(
+                    f'[RESOLVED_REFERENCE_COLUMN_NOT_FOUND] The column "{attr_name}" does not exist in the target dataframe.'
+                )
+            else:
+                # Column does not exist. Pass in dummy column name for lazy error throwing as it could be a built-in function
+                snowpark_name = attr_name
+
+        col = get_col(snowpark_name)
+        try:
+            col_type = typer.type(col)[0]
+        except SnowparkSQLException as e:
+            if "invalid identifier" in e.raw_message:
+                raise AnalysisException(
+                    f'[COLUMN_NOT_FOUND] The column "{attr_name}" does not exist in the target dataframe.'
+                )
+            else:
+                raise
+        is_struct = isinstance(col_type, StructType)
+        # for struct columns when accessed, spark use just the leaf field name rather than fully attributed one
+        if is_struct:
+            attr_name = name_parts[-1]
+
+        path = name_parts[1:]
+        if is_struct and not global_config.spark_sql_caseSensitive:
+            path = _match_path_to_struct(path, col_type)
+
+        for field_name in path:
+            col = col.getItem(field_name)
+
+        qualifiers = []
+
+    typed_col = TypedColumn(col, lambda: typer.type(col))
+    typed_col.set_qualifiers(qualifiers)
+    return (name_parts[-1], typed_col)
+
+
+def _match_path_to_struct(path: list[str], col_type: StructType) -> list[str]:
+    """Takes a path of names and adjusts them to strictly match the field names in a StructType."""
+    adjusted_path = []
+    typ = col_type
+    for i, name in enumerate(path):
+        if isinstance(typ, StructType):
+            lowercase_name = name.lower()
+            for field in typ.fields:
+                if field.name.lower() == lowercase_name:
+                    adjusted_path.append(field.name)
+                    typ = field.datatype
+                    break
+        elif isinstance(typ, MapType) or isinstance(typ, ArrayType):
+            # For MapType and ArrayType, we can use the name as is.
+            adjusted_path.append(name)
+            typ = typ.value_type if isinstance(typ, MapType) else typ.element_type
+        else:
+            # If the type is not a struct, map, or array, we cannot access the field.
+            raise AnalysisException(
+                f"[INVALID_EXTRACT_BASE_FIELD_TYPE] Can't extract a value from \"{'.'.join(path[:i])}\". Need a complex type [STRUCT, ARRAY, MAP] but got \"{typ}\"."
+            )
+    return adjusted_path

snowflake/snowpark_connect/expression/map_unresolved_extract_value.py

@@ -0,0 +1,85 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
+
+import snowflake.snowpark.functions as snowpark_fn
+from snowflake.snowpark.types import ArrayType, MapType, StructType, _IntegralType
+from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.expression.typer import ExpressionTyper
+from snowflake.snowpark_connect.typed_column import TypedColumn
+
+
+def _check_if_array_type(
+    child_typed_column: TypedColumn, extract_typed_column: TypedColumn
+):
+    extract_typed_column_type = extract_typed_column.types
+    container_type = child_typed_column.types
+    return (
+        len(extract_typed_column_type) == 1
+        and isinstance(extract_typed_column_type[0], ArrayType)
+        and len(container_type) == 1
+        and isinstance(container_type[0], _IntegralType)
+    )
+
+
+def map_unresolved_extract_value(
+    exp: expressions_proto.Expression,
+    column_mapping: ColumnNameMap,
+    typer: ExpressionTyper,
+) -> tuple[str, TypedColumn]:
+    from snowflake.snowpark_connect.expression.map_expression import (
+        map_single_column_expression,
+    )
+
+    child_name, child_typed_column = map_single_column_expression(
+        exp.unresolved_extract_value.child, column_mapping, typer
+    )
+    extract_name, extract_typed_column = map_single_column_expression(
+        exp.unresolved_extract_value.extraction,
+        column_mapping,
+        typer,
+    )
+    spark_function_name = (
+        f"{child_name}.{extract_name}"
+        if isinstance(child_typed_column.typ, StructType)
+        else f"{child_name}[{extract_name}]"
+    )
+    # Spark respects "spark.sql.caseSensitive" for struct fields
+    # map keys are compared as-is
+    if global_config.spark_sql_caseSensitive or isinstance(
+        child_typed_column.typ, MapType
+    ):
+        extract_fn = snowpark_fn.get
+    else:
+        extract_fn = snowpark_fn.get_ignore_case
+    # Set index to a dummy value before we use it later in the ansi mode check.
+    index = snowpark_fn.lit(1)
+    if _check_if_array_type(extract_typed_column, child_typed_column):
+        # Set all non-valid array indices to NULL.
+        # This is done because both conditions of a CASE WHEN statement are executed regardless of if the condition is true or not.
+        # Getting a negative index in Snowflake throws an error; thus, we convert all non-valid array indices to NULL before getting the index.
+        index = snowpark_fn.when(
+            (snowpark_fn.array_size(child_typed_column.col) > extract_typed_column.col)
+            & (extract_typed_column.col >= 0),
+            extract_typed_column.col,
+        ).otherwise(snowpark_fn.lit(None))
+        result_exp = snowpark_fn.when(index.isNull(), snowpark_fn.lit(None)).otherwise(
+            snowpark_fn.get(child_typed_column.col, index)
+        )
+    else:
+        result_exp = extract_fn(child_typed_column.col, extract_typed_column.col)
+
+    spark_sql_ansi_enabled = global_config.spark_sql_ansi_enabled
+
+    if spark_sql_ansi_enabled and _check_if_array_type(
+        extract_typed_column, child_typed_column
+    ):
+        result_exp = snowpark_fn.when(
+            index.isNull(),
+            child_typed_column.col.getItem("[snowpark_connect::INVALID_ARRAY_INDEX]"),
+        ).otherwise(result_exp)
+
+    return spark_function_name, TypedColumn(result_exp, lambda: typer.type(result_exp))