snowpark-connect 0.27.0__py3-none-any.whl → 0.28.1__py3-none-any.whl

snowflake/snowpark_connect/expression/map_unresolved_attribute.py

@@ -163,7 +163,6 @@ def map_unresolved_attribute(
     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
@@ -172,27 +171,13 @@ 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, has_hidden=False):
+    def get_col(snowpark_name):
         return (
             snowpark_fn.col(snowpark_name)
-            if not has_plan_id or has_hidden
+            if not has_plan_id
             else target_df.col(snowpark_name)
         )
 
@@ -277,17 +262,10 @@ 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,
-        is_qualified=has_plan_id,
-        source_qualifiers=source_qualifiers if has_plan_id else None,
+        quoted_attr_name, allow_non_exists=True
     )
-
     if snowpark_name is not None:
-        is_hidden = column_mapping.is_hidden_column_reference(
-            quoted_attr_name, source_qualifiers
-        )
-        col = get_col(snowpark_name, is_hidden)
+        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

snowflake/snowpark_connect/expression/map_unresolved_function.py

@@ -2619,9 +2619,18 @@ def map_unresolved_function(
                     result_type = input_type.element_type
                     result_exp = fn(snowpark_args[0])
                 case _:
-                    spark_col_names = ["key", "value"]
-                    result_exp = fn(snowpark_args[0])
-                    result_type = [input_type.key_type, input_type.value_type]
+                    # Check if the type has map-like attributes before accessing them
+                    if hasattr(input_type, "key_type") and hasattr(
+                        input_type, "value_type"
+                    ):
+                        spark_col_names = ["key", "value"]
+                        result_exp = fn(snowpark_args[0])
+                        result_type = [input_type.key_type, input_type.value_type]
+                    else:
+                        # Throw proper error for types without key_type/value_type attributes
+                        raise AnalysisException(
+                            f'[DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve "{function_name}({snowpark_arg_names[0]})" due to data type mismatch: Parameter 1 requires the ("ARRAY" or "MAP") type, however "{snowpark_arg_names[0]}" has the type "{str(input_type)}".'
+                        )
         case "expm1":
             spark_function_name = f"EXPM1({snowpark_arg_names[0]})"
             result_exp = snowpark_fn.exp(*snowpark_args) - 1

snowflake/snowpark_connect/expression/map_unresolved_star.py

@@ -34,7 +34,6 @@ def map_unresolved_star(
     column_mapping: ColumnNameMap,
     typer: ExpressionTyper,
 ) -> tuple[list[str], TypedColumn]:
-
     if exp.unresolved_star.HasField("unparsed_target"):
         unparsed_target = exp.unresolved_star.unparsed_target
         name_parts = split_fully_qualified_spark_name(unparsed_target)
@@ -103,7 +102,7 @@ def map_unresolved_star(
                 prefix_candidate_str = f"{prefix_candidate_str}.{name_parts[i]}"
             prefix_candidate = (
                 column_mapping.get_snowpark_column_name_from_spark_column_name(
-                    prefix_candidate_str, allow_non_exists=True, is_qualified=(i > 0)
+                    prefix_candidate_str, allow_non_exists=True
                 )
             )
             if prefix_candidate is None:
@@ -181,7 +180,7 @@ def map_unresolved_star_struct(
             prefix_candidate_str = f"{prefix_candidate_str}.{name_parts[i]}"
         prefix_candidate = (
             column_mapping.get_snowpark_column_name_from_spark_column_name(
-                prefix_candidate_str, allow_non_exists=True, is_qualified=(i > 0)
+                prefix_candidate_str, allow_non_exists=True
             )
         )
         if prefix_candidate is None:

snowflake/snowpark_connect/relation/catalogs/snowflake_catalog.py

@@ -13,7 +13,6 @@ from snowflake.core.exceptions import APIError, NotFoundError
 from snowflake.core.schema import Schema
 from snowflake.core.table import Table, TableColumn
 
-from snowflake.snowpark import functions
 from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     quote_name_without_upper_casing,
     unquote_if_quoted,
@@ -34,12 +33,19 @@ from snowflake.snowpark_connect.relation.catalogs.abstract_spark_catalog import
 )
 from snowflake.snowpark_connect.type_mapping import proto_to_snowpark_type
 from snowflake.snowpark_connect.utils.identifiers import (
+    FQN,
+    spark_to_sf_single_id_with_unquoting,
     split_fully_qualified_spark_name,
 )
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
+from snowflake.snowpark_connect.utils.temporary_view_cache import (
+    get_temp_view,
+    get_temp_view_normalized_names,
+    unregister_temp_view,
+)
 from snowflake.snowpark_connect.utils.udf_cache import cached_udf
 
 
@@ -203,6 +209,93 @@ class SnowflakeCatalog(AbstractSparkCatalog):
             exists = False
         return pandas.DataFrame({"exists": [exists]})
 
+    def _get_temp_view_prefixes(self, spark_dbName: str | None) -> list[str]:
+        if spark_dbName is None:
+            return []
+        return [
+            quote_name_without_upper_casing(part)
+            for part in split_fully_qualified_spark_name(spark_dbName)
+        ]
+
+    def _list_temp_views(
+        self,
+        spark_dbName: str | None = None,
+        pattern: str | None = None,
+    ) -> typing.Tuple[
+        list[str | None],
+        list[list[str | None]],
+        list[str],
+        list[str | None],
+        list[str | None],
+        list[bool],
+    ]:
+        catalogs: list[str | None] = list()
+        namespaces: list[list[str | None]] = list()
+        names: list[str] = list()
+        descriptions: list[str | None] = list()
+        table_types: list[str | None] = list()
+        is_temporaries: list[bool] = list()
+
+        temp_views_prefix = ".".join(self._get_temp_view_prefixes(spark_dbName))
+        normalized_spark_dbName = (
+            temp_views_prefix.lower()
+            if global_config.spark_sql_caseSensitive
+            else temp_views_prefix
+        )
+        normalized_global_temp_database_name = (
+            quote_name_without_upper_casing(
+                global_config.spark_sql_globalTempDatabase.lower()
+            )
+            if global_config.spark_sql_caseSensitive
+            else quote_name_without_upper_casing(
+                global_config.spark_sql_globalTempDatabase
+            )
+        )
+
+        temp_views = get_temp_view_normalized_names()
+        null_safe_pattern = pattern if pattern is not None else ""
+
+        for temp_view in temp_views:
+            normalized_temp_view = (
+                temp_view.lower()
+                if global_config.spark_sql_caseSensitive
+                else temp_view
+            )
+            fqn = FQN.from_string(temp_view)
+            normalized_schema = (
+                fqn.schema.lower()
+                if fqn.schema is not None and global_config.spark_sql_caseSensitive
+                else fqn.schema
+            )
+
+            is_global_view = normalized_global_temp_database_name == normalized_schema
+            is_local_temp_view = fqn.schema is None
+            # Temporary views are always shown if they match the pattern
+            matches_prefix = (
+                normalized_spark_dbName == normalized_schema or is_local_temp_view
+            )
+            if matches_prefix and bool(
+                re.match(null_safe_pattern, normalized_temp_view)
+            ):
+                names.append(unquote_if_quoted(fqn.name))
+                catalogs.append(None)
+                namespaces.append(
+                    [global_config.spark_sql_globalTempDatabase]
+                    if is_global_view
+                    else []
+                )
+                descriptions.append(None)
+                table_types.append("TEMPORARY")
+                is_temporaries.append(True)
+        return (
+            catalogs,
+            namespaces,
+            names,
+            descriptions,
+            table_types,
+            is_temporaries,
+        )
+
     def listTables(
         self,
         spark_dbName: str | None = None,
@@ -232,8 +325,7 @@ class SnowflakeCatalog(AbstractSparkCatalog):
             schema=sf_quote(sf_schema),
             pattern=_normalize_identifier(pattern),
         )
-        names: list[str] = list()
-        catalogs: list[str] = list()
+        catalogs: list[str | None] = list()
         namespaces: list[list[str | None]] = list()
         names: list[str] = list()
         descriptions: list[str | None] = list()
@@ -253,6 +345,22 @@ class SnowflakeCatalog(AbstractSparkCatalog):
             descriptions.append(o[6] if o[6] else None)
             table_types.append("PERMANENT")
             is_temporaries.append(False)
+
+        (
+            non_materialized_catalogs,
+            non_materialized_namespaces,
+            non_materialized_names,
+            non_materialized_descriptions,
+            non_materialized_table_types,
+            non_materialized_is_temporaries,
+        ) = self._list_temp_views(spark_dbName, pattern)
+        catalogs.extend(non_materialized_catalogs)
+        namespaces.extend(non_materialized_namespaces)
+        names.extend(non_materialized_names)
+        descriptions.extend(non_materialized_descriptions)
+        table_types.extend(non_materialized_table_types)
+        is_temporaries.extend(non_materialized_is_temporaries)
+
         return pandas.DataFrame(
             {
                 "name": names,
@@ -297,6 +405,36 @@ class SnowflakeCatalog(AbstractSparkCatalog):
         spark_tableName: str,
     ) -> pandas.DataFrame:
         """Listing a single table/view with provided name that's accessible in Snowflake."""
+
+        def _get_temp_view():
+            spark_table_name_parts = [
+                quote_name_without_upper_casing(part)
+                for part in split_fully_qualified_spark_name(spark_tableName)
+            ]
+            spark_view_name = ".".join(spark_table_name_parts)
+            temp_view = get_temp_view(spark_view_name)
+            if temp_view:
+                return pandas.DataFrame(
+                    {
+                        "name": [unquote_if_quoted(spark_table_name_parts[-1])],
+                        "catalog": [None],
+                        "namespace": [
+                            [unquote_if_quoted(spark_table_name_parts[-2])]
+                            if len(spark_table_name_parts) > 1
+                            else []
+                        ],
+                        "description": [None],
+                        "tableType": ["TEMPORARY"],
+                        "isTemporary": [True],
+                    }
+                )
+            return None
+
+        # Attempt to get the view from the non materialized views first
+        temp_view = _get_temp_view()
+        if temp_view is not None:
+            return temp_view
+
         sp_catalog = get_or_create_snowpark_session().catalog
         catalog, sf_database, sf_schema, table_name = _process_multi_layer_identifier(
             spark_tableName
@@ -360,12 +498,64 @@ class SnowflakeCatalog(AbstractSparkCatalog):
             exists = False
         return pandas.DataFrame({"exists": [exists]})
 
+    def _list_temp_view_columns(
+        self,
+        spark_tableName: str,
+        spark_dbName: typing.Optional[str] = None,
+    ):
+        spark_view_name_parts = [
+            quote_name_without_upper_casing(part)
+            for part in split_fully_qualified_spark_name(spark_tableName)
+        ]
+        spark_view_name_parts = (
+            self._get_temp_view_prefixes(spark_dbName) + spark_view_name_parts
+        )
+        spark_view_name = ".".join(spark_view_name_parts)
+        temp_view = get_temp_view(spark_view_name)
+
+        if not temp_view:
+            return None
+
+        names: list[str] = list()
+        descriptions: list[str | None] = list()
+        data_types: list[str] = list()
+        nullables: list[bool] = list()
+        is_partitions: list[bool] = list()
+        is_buckets: list[bool] = list()
+
+        for field, spark_column in zip(
+            temp_view.dataframe.schema.fields,
+            temp_view.column_map.get_spark_columns(),
+        ):
+            names.append(spark_column)
+            descriptions.append(None)
+            data_types.append(field.datatype.simpleString())
+            nullables.append(field.nullable)
+            is_partitions.append(False)
+            is_buckets.append(False)
+
+        return pandas.DataFrame(
+            {
+                "name": names,
+                "description": descriptions,
+                "dataType": data_types,
+                "nullable": nullables,
+                "isPartition": is_partitions,
+                "isBucket": is_buckets,
+            }
+        )
+
     def listColumns(
         self,
         spark_tableName: str,
         spark_dbName: typing.Optional[str] = None,
     ) -> pandas.DataFrame:
         """List all columns in a table/view, optionally database name filter can be provided."""
+
+        temp_view_columns = self._list_temp_view_columns(spark_tableName, spark_dbName)
+        if temp_view_columns is not None:
+            return temp_view_columns
+
         sp_catalog = get_or_create_snowpark_session().catalog
         columns: list[TableColumn] | None = None
         if spark_dbName is None:
@@ -455,17 +645,15 @@ class SnowflakeCatalog(AbstractSparkCatalog):
         spark_view_name: str,
     ) -> DataFrameContainer:
         session = get_or_create_snowpark_session()
-        schema = global_config.spark_sql_globalTempDatabase
-        result_df = session.sql(
-            "drop view if exists identifier(?)",
-            params=[f"{sf_quote(schema)}.{sf_quote(spark_view_name)}"],
-        )
-        result_df = result_df.select(
-            functions.contains('"status"', functions.lit("successfully dropped")).alias(
-                "value"
+        if not spark_view_name == "":
+            schema = global_config.spark_sql_globalTempDatabase
+            result = unregister_temp_view(
+                f"{spark_to_sf_single_id_with_unquoting(schema)}.{spark_to_sf_single_id_with_unquoting(spark_view_name)}"
             )
-        )
+        else:
+            result = False
         columns = ["value"]
+        result_df = session.createDataFrame([result], schema=columns)
         return DataFrameContainer.create_with_column_mapping(
             dataframe=result_df,
             spark_column_names=columns,
@@ -479,15 +667,14 @@ class SnowflakeCatalog(AbstractSparkCatalog):
     ) -> DataFrameContainer:
         """Drop the current temporary view."""
         session = get_or_create_snowpark_session()
-        result = session.sql(
-            "drop view if exists identifier(?)",
-            params=[sf_quote(spark_view_name)],
-        ).collect()
-        view_was_dropped = (
-            len(result) == 1 and "successfully dropped" in result[0]["status"]
-        )
-        result_df = session.createDataFrame([(view_was_dropped,)], schema=["value"])
         columns = ["value"]
+        if spark_view_name:
+            result = unregister_temp_view(
+                spark_to_sf_single_id_with_unquoting(spark_view_name)
+            )
+        else:
+            result = False
+        result_df = session.createDataFrame([result], schema=columns)
         return DataFrameContainer.create_with_column_mapping(
             dataframe=result_df,
             spark_column_names=columns,

snowflake/snowpark_connect/relation/map_extension.py

@@ -23,7 +23,6 @@ from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.typed_column import TypedColumn
 from snowflake.snowpark_connect.utils.context import (
     get_sql_aggregate_function_count,
-    not_resolving_fun_args,
     push_outer_dataframe,
     set_current_grouping_columns,
 )
@@ -336,15 +335,14 @@ def map_aggregate(
     typer = ExpressionTyper(input_df)
 
     def _map_column(exp: expression_proto.Expression) -> tuple[str, TypedColumn]:
-        with not_resolving_fun_args():
-            new_names, snowpark_column = map_expression(
-                exp, input_container.column_map, typer
+        new_names, snowpark_column = map_expression(
+            exp, input_container.column_map, typer
+        )
+        if len(new_names) != 1:
+            raise SnowparkConnectNotImplementedError(
+                "Multi-column aggregate expressions are not supported"
             )
-            if len(new_names) != 1:
-                raise SnowparkConnectNotImplementedError(
-                    "Multi-column aggregate expressions are not supported"
-                )
-            return new_names[0], snowpark_column
+        return new_names[0], snowpark_column
 
     raw_groupings: list[tuple[str, TypedColumn]] = []
     raw_aggregations: list[tuple[str, TypedColumn]] = []
@@ -431,12 +429,18 @@ def map_aggregate(
             if groupings:
                 # Normal GROUP BY with explicit grouping columns
                 result = input_df.group_by(groupings)
-            else:
+            elif not is_group_by_all:
                 # No explicit GROUP BY - this is an aggregate over the entire table
                 # Use a dummy constant that will be excluded from the final result
                 result = input_df.with_column(
                     "__dummy_group__", snowpark_fn.lit(1)
                 ).group_by("__dummy_group__")
+            else:
+                # GROUP BY ALL with only one aggregate column
+                # Snowpark doesn't support GROUP BY ALL
+                # TODO: Change in future with Snowpark Supported arguments or API for GROUP BY ALL
+                result = input_df.group_by()
+
         case snowflake_proto.Aggregate.GROUP_TYPE_ROLLUP:
             result = input_df.rollup(groupings)
         case snowflake_proto.Aggregate.GROUP_TYPE_CUBE: