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

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,13 +86,21 @@ 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:
     udf_check(udf_proto)
-    processed_return_type, original_return_type = process_udf_return_type(
-        udf_proto.python_udf.output_type
-    )
+    match udf_proto.WhichOneof("function"):
+        case "python_udf":
+            output_type = udf_proto.python_udf.output_type
+        case "scalar_scala_udf":
+            output_type = udf_proto.scalar_scala_udf.outputType
+        case _:
+            raise ValueError(
+                f"Unsupported UDF type: {udf_proto.WhichOneof('function')}"
+            )
+    processed_return_type, original_return_type = process_udf_return_type(output_type)
     session = get_or_create_snowpark_session()
     kwargs = {
         "common_inline_user_defined_function": udf_proto,
@@ -77,9 +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
         return udf
 
 
@@ -94,31 +140,44 @@ def map_common_inline_user_defined_udf(
         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,
-        )
+    match udf_proto.WhichOneof("function"):
+        case "python_udf":
+            processed_return_type, original_return_type = process_udf_return_type(
+                udf_proto.python_udf.output_type
+            )
+        case "scalar_scala_udf":
+            processed_return_type, original_return_type = process_udf_return_type(
+                udf_proto.scalar_scala_udf.outputType
+            )
+
+    @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,

snowflake/snowpark_connect/expression/map_unresolved_attribute.py

@@ -12,12 +12,13 @@ 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.types import ArrayType, LongType, 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.context import (
+    get_current_grouping_columns,
     get_is_evaluating_sql,
     get_outer_dataframes,
     get_plan_id_map,
@@ -30,6 +31,41 @@ from snowflake.snowpark_connect.utils.identifiers import (
 SPARK_QUOTED = re.compile("^(`.*`)$", re.DOTALL)
 
 
+def _get_catalog_database_from_column_map(
+    column_name: str, column_map: ColumnNameMap
+) -> dict[str, str]:
+    """
+    Get catalog/database info from the column map for a given column name.
+    This replaces the previous extraction logic by looking up stored metadata.
+
+    Returns:
+        dict: catalog_database_info containing catalog/database metadata if found
+    """
+    catalog_database_info = {}
+
+    # Look up in the column map using case-sensitive or case-insensitive matching
+    matching_columns = []
+    if hasattr(column_map, "is_case_sensitive") and column_map.is_case_sensitive():
+        matching_columns = column_map.spark_to_col.get(column_name, [])
+    elif hasattr(column_map, "uppercase_spark_to_col"):
+        matching_columns = column_map.uppercase_spark_to_col.get(
+            column_name.upper(), []
+        )
+    elif hasattr(column_map, "spark_to_col"):
+        matching_columns = column_map.spark_to_col.get(column_name, [])
+
+    # If we found a matching column with catalog/database info, use it
+    for col_names in matching_columns:
+        if col_names.catalog_info and col_names.database_info:
+            catalog_database_info = {
+                "catalog": col_names.catalog_info,
+                "database": col_names.database_info,
+            }
+            break
+
+    return catalog_database_info
+
+
 def map_unresolved_attribute(
     exp: expressions_proto.Expression,
     column_mapping: ColumnNameMap,
@@ -40,9 +76,93 @@ def map_unresolved_attribute(
 
     assert len(name_parts) > 0, f"Unable to parse input attribute: {original_attr_name}"
 
+    # Special handling for Spark's automatic grouping__id column
+    # In Spark SQL, when using GROUP BY CUBE/ROLLUP/GROUPING SETS, an automatic
+    # virtual column called 'grouping__id' (with double underscores) is available.
+    # In Snowflake, we need to convert this to a GROUPING_ID() function call.
+    if len(name_parts) == 1 and name_parts[0].lower() == "grouping__id":
+        grouping_spark_columns = get_current_grouping_columns()
+        if not grouping_spark_columns:
+            # grouping__id can only be used with GROUP BY CUBE/ROLLUP/GROUPING SETS
+            raise AnalysisException(
+                "[MISSING_GROUP_BY] grouping__id can only be used with GROUP BY (CUBE | ROLLUP | GROUPING SETS)"
+            )
+        # Convert to GROUPING_ID() function call with the grouping columns
+        # Map Spark column names to Snowpark column names
+        snowpark_cols = []
+        for spark_col_name in grouping_spark_columns:
+            # Get the Snowpark column name from the mapping
+            snowpark_name = (
+                column_mapping.get_snowpark_column_name_from_spark_column_name(
+                    spark_col_name
+                )
+            )
+            if not snowpark_name:
+                raise AnalysisException(
+                    f"[INTERNAL_ERROR] Cannot find Snowpark column mapping for grouping column '{spark_col_name}'"
+                )
+            snowpark_cols.append(snowpark_fn.col(snowpark_name))
+
+        # Call GROUPING_ID with all grouping columns using Snowpark names
+        result_col = snowpark_fn.grouping_id(*snowpark_cols)
+
+        # TypedColumn expects a callable that returns a list of types
+        # GROUPING_ID returns a BIGINT (LongType) in both Spark and Snowflake
+        # representing the bit vector of grouping indicators
+        typed_col = TypedColumn(result_col, lambda: [LongType()])
+        return ("grouping__id", typed_col)
+
+    # Validate that DataFrame API doesn't allow catalog.database.column patterns
+    # These patterns should only work in SQL, not DataFrame API
+    if len(name_parts) >= 4:
+        # For 4+ parts, check if this looks like catalog.database.column.field
+        # (as opposed to a valid table.column.field pattern)
+
+        # Heuristic: if the pattern looks like catalog.database.column.field,
+        # reject it in DataFrame API context (but allow in SQL)
+
+        # Check if first part looks like a catalog name (not a column)
+        first_part = name_parts[0]
+        first_part_snowpark = (
+            column_mapping.get_snowpark_column_name_from_spark_column_name(
+                first_part, allow_non_exists=True
+            )
+        )
+
+        # If first part is not a column and we have 4+ parts, check if it's a catalog reference
+        if first_part_snowpark is None and len(name_parts) >= 4:
+            # Import here to avoid circular import issues
+            from snowflake.snowpark_connect.relation.catalogs import CATALOGS
+
+            # Check if the first part is a registered catalog name OR looks like a catalog pattern
+            is_registered_catalog = first_part.lower() in CATALOGS
+            is_catalog_like = (
+                # Contains "catalog" in the name
+                "catalog" in first_part.lower()
+                # Follows catalog naming patterns (no numbers, shorter descriptive names)
+                or (
+                    len(first_part) < 20
+                    and not any(char.isdigit() for char in first_part)
+                    and not first_part.startswith("mydb")
+                    and not first_part.endswith(  # Skip test-generated database names
+                        "_dbmsu"
+                    )
+                )  # Skip test-generated database names
+            )
+
+            is_catalog = is_registered_catalog or is_catalog_like
+
+            if is_catalog:
+                # This looks like a catalog.database.column.field pattern
+                raise AnalysisException(
+                    f"[UNRESOLVED_COLUMN.WITH_SUGGESTION] A column or function parameter with name `{original_attr_name}` cannot be resolved. "
+                    f"Cross-catalog column references are not supported in DataFrame API."
+                )
+
     attr_name = ".".join(name_parts)
 
     has_plan_id = exp.unresolved_attribute.HasField("plan_id")
+    source_qualifiers = None
 
     if has_plan_id:
         plan_id = exp.unresolved_attribute.plan_id
@@ -51,13 +171,27 @@ def map_unresolved_attribute(
         assert (
             target_df is not None
         ), f"resolving an attribute of a unresolved dataframe {plan_id}"
+
+        # Get the qualifiers for this column from the target DataFrame
+        source_qualifiers = (
+            target_df_container.column_map.get_qualifier_for_spark_column(
+                name_parts[-1]
+            )
+        )
+
+        if hasattr(column_mapping, "hidden_columns"):
+            hidden = column_mapping.hidden_columns
+        else:
+            hidden = None
+
         column_mapping = target_df_container.column_map
+        column_mapping.hidden_columns = hidden
         typer = ExpressionTyper(target_df)
 
-    def get_col(snowpark_name):
+    def get_col(snowpark_name, has_hidden=False):
         return (
             snowpark_fn.col(snowpark_name)
-            if not has_plan_id
+            if not has_plan_id or has_hidden
             else target_df.col(snowpark_name)
         )
 
@@ -88,8 +222,14 @@ def map_unresolved_attribute(
             raise AnalysisException(f"Invalid regex pattern '{regex_pattern}': {e}")
 
         if not matched_columns:
+            # Keep the improved error message for SQL regex patterns
+            # This is only hit for SQL queries like SELECT `(e|f)` FROM table
+            # when spark.sql.parser.quotedRegexColumnNames is enabled
             raise AnalysisException(
-                f"No columns match the regex pattern '{regex_pattern}'"
+                f"No columns match the regex pattern '{regex_pattern}'. "
+                f"Snowflake SQL does not support SELECT statements with no columns. "
+                f"Please ensure your regex pattern matches at least one column. "
+                f"Available columns: {', '.join(available_columns[:10])}{'...' if len(available_columns) > 10 else ''}"
             )
 
         # When multiple columns match, we need to signal that this should expand to multiple columns
@@ -136,24 +276,55 @@ def map_unresolved_attribute(
         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
+        quoted_attr_name,
+        allow_non_exists=True,
+        is_qualified=has_plan_id,
+        source_qualifiers=source_qualifiers if has_plan_id else None,
     )
+
     if snowpark_name is not None:
-        col = get_col(snowpark_name)
+        is_hidden = column_mapping.is_hidden_column_reference(
+            quoted_attr_name, source_qualifiers
+        )
+        col = get_col(snowpark_name, is_hidden)
         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
+        snowpark_name: str | None = None
+        column_part_index: int = 0
+
+        # Get catalog/database info from column map if available
+        catalog_database_info = _get_catalog_database_from_column_map(
+            original_attr_name, column_mapping
         )
-        if snowpark_name is None:
-            for outer_df_container in get_outer_dataframes():
-                snowpark_name = outer_df_container.column_map.get_snowpark_column_name_from_spark_column_name(
-                    name_parts[0], allow_non_exists=True
+
+        # Try to find the column name in different parts of the name_parts array
+        # For qualified names like "table.column.field", we need to find the column part
+        for i in range(len(name_parts)):
+            candidate_column = name_parts[i]
+            snowpark_name = (
+                column_mapping.get_snowpark_column_name_from_spark_column_name(
+                    candidate_column, allow_non_exists=True
+                )
+            )
+            if snowpark_name is not None:
+                column_part_index = i
+                break
+
+            # Also try in outer dataframes
+            for outer_df in get_outer_dataframes():
+                snowpark_name = (
+                    outer_df.column_map.get_snowpark_column_name_from_spark_column_name(
+                        candidate_column, allow_non_exists=True
+                    )
                 )
                 if snowpark_name is not None:
+                    column_part_index = i
                     break
 
+            if snowpark_name is not None:
+                break
+
         if snowpark_name is None:
             # Attempt LCA fallback.
             alias_tc = resolve_lca_alias(attr_name)
@@ -164,7 +335,8 @@ def map_unresolved_attribute(
 
             # 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:
+            remaining_parts = name_parts
+            if len(remaining_parts) > 1:
                 unqualified_name = name_parts[-1]
                 snowpark_name = (
                     column_mapping.get_snowpark_column_name_from_spark_column_name(
@@ -178,6 +350,9 @@ def map_unresolved_attribute(
                     )
                     typed_col = TypedColumn(col, lambda: typer.type(col))
                     typed_col.set_qualifiers(qualifiers)
+                    # Store catalog/database info if found in column map
+                    if catalog_database_info:
+                        typed_col.set_catalog_database_info(catalog_database_info)
                     return (unqualified_name, typed_col)
 
         if snowpark_name is None:
@@ -193,7 +368,7 @@ def map_unresolved_attribute(
         try:
             col_type = typer.type(col)[0]
         except SnowparkSQLException as e:
-            if "invalid identifier" in e.raw_message:
+            if e.raw_message is not None and "invalid identifier" in e.raw_message:
                 raise AnalysisException(
                     f'[COLUMN_NOT_FOUND] The column "{attr_name}" does not exist in the target dataframe.'
                 )
@@ -204,7 +379,8 @@ def map_unresolved_attribute(
         if is_struct:
             attr_name = name_parts[-1]
 
-        path = name_parts[1:]
+        # Calculate the field path correctly based on where we found the column
+        path = name_parts[column_part_index + 1 :]
         if is_struct and not global_config.spark_sql_caseSensitive:
             path = _match_path_to_struct(path, col_type)
 
@@ -215,6 +391,14 @@ def map_unresolved_attribute(
 
     typed_col = TypedColumn(col, lambda: typer.type(col))
     typed_col.set_qualifiers(qualifiers)
+
+    # Store catalog/database info if available from column map
+    final_catalog_database_info = _get_catalog_database_from_column_map(
+        original_attr_name, column_mapping
+    )
+    if final_catalog_database_info:
+        typed_col.set_catalog_database_info(final_catalog_database_info)
+
     return (name_parts[-1], typed_col)
 
 

snowflake/snowpark_connect/expression/map_unresolved_extract_value.py

@@ -6,7 +6,13 @@ import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake.snowpark._internal.analyzer.expression import Literal
-from snowflake.snowpark.types import ArrayType, MapType, StructType, _IntegralType
+from snowflake.snowpark.types import (
+    ArrayType,
+    MapType,
+    NullType,
+    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
@@ -22,7 +28,7 @@ def _check_if_array_type(
         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)
+        and isinstance(container_type[0], (_IntegralType, NullType))
     )
 
 
@@ -56,29 +62,51 @@ def map_unresolved_extract_value(
         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)
     is_array = _check_if_array_type(extract_typed_column, child_typed_column)
     if is_array:
-        # 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)
-        )
+        if isinstance(extract_typed_column.typ, NullType):
+            result_exp = snowpark_fn.lit(None)
+        else:
+            if (
+                isinstance(extract_typed_column.col._expression, Literal)
+                and extract_typed_column.col._expression.value is not None
+            ):
+                # Using NULL in NVL triggers Snowflake Optimiser to be much more efficient comparing to using a number.
+                # This unfortunately has a side effect of throwing and error when attempting to get the item from array.
+                # That's why we need to have a separate branch for fetching Nullable literals and non-literal expressions.
+                extracted_index = snowpark_fn.nvl(
+                    extract_typed_column.col, snowpark_fn.lit(None)
+                )
+            else:
+                extracted_index = snowpark_fn.nvl(
+                    extract_typed_column.col, snowpark_fn.lit(0)
+                )
+
+            result_exp = snowpark_fn.when(
+                snowpark_fn.nvl(
+                    (extract_typed_column.col < 0)
+                    | (extract_typed_column.col > 2_147_483_647),
+                    snowpark_fn.lit(True),
+                ),
+                snowpark_fn.lit(None),
+            ).otherwise(
+                snowpark_fn.get(
+                    child_typed_column.col,
+                    extracted_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 is_array:
+        invalid_array_index = (
+            snowpark_fn.array_size(child_typed_column.col) <= extract_typed_column.col
+        ) | (extract_typed_column.col < 0)
         result_exp = snowpark_fn.when(
-            index.isNull(),
+            invalid_array_index,
             child_typed_column.col.getItem("[snowpark_connect::INVALID_ARRAY_INDEX]"),
         ).otherwise(result_exp)