snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/relation/map_aggregate.py

@@ -0,0 +1,322 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import re
+from dataclasses import dataclass
+
+import pyspark.sql.connect.proto.relations_pb2 as relation_proto
+
+from snowflake import snowpark
+from snowflake.snowpark.types import DataType
+from snowflake.snowpark_connect.column_name_handler import (
+    make_column_names_snowpark_compatible,
+    with_column_map,
+)
+from snowflake.snowpark_connect.expression.literal import get_literal_field_and_name
+from snowflake.snowpark_connect.expression.map_expression import (
+    map_single_column_expression,
+)
+from snowflake.snowpark_connect.expression.typer import ExpressionTyper
+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 temporary_pivot_expression
+
+
+def map_group_by_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
+    """
+    Groups the DataFrame using the specified columns.
+
+    Aggregations come in as expressions, which are mapped to `snowpark.Column`
+    objects.
+    """
+    input_df, columns = map_aggregate_helper(rel)
+    if len(columns.grouping_expressions()) == 0:
+        result = input_df.agg(*columns.aggregation_expressions())
+    else:
+        result = input_df.group_by(*columns.grouping_expressions()).agg(
+            *columns.aggregation_expressions()
+        )
+    return with_column_map(
+        result,
+        columns.spark_names(),
+        columns.snowpark_names(),
+        columns.data_types(),
+        column_qualifiers=columns.get_qualifiers(),
+        parent_column_name_map=input_df._column_map,
+    )
+
+
+def map_rollup_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
+    """
+    Create a multidimensional rollup for the current DataFrame using the specified columns.
+
+    Aggregations come in as expressions, which are mapped to `snowpark.Column`
+    objects.
+    """
+    input_df, columns = map_aggregate_helper(rel)
+    if len(columns.grouping_expressions()) == 0:
+        result = input_df.agg(*columns.aggregation_expressions())
+    else:
+        result = input_df.rollup(*columns.grouping_expressions()).agg(
+            *columns.aggregation_expressions()
+        )
+    return with_column_map(
+        result,
+        columns.spark_names(),
+        columns.snowpark_names(),
+        columns.data_types(),
+        column_qualifiers=columns.get_qualifiers(),
+        parent_column_name_map=input_df._column_map,
+    )
+
+
+def map_cube_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
+    """
+    Create a multidimensional cube for the current DataFrame using the specified columns.
+
+    Aggregations come in as expressions, which are mapped to `snowpark.Column`
+    objects.
+    """
+    input_df, columns = map_aggregate_helper(rel)
+    if len(columns.grouping_expressions()) == 0:
+        result = input_df.agg(*columns.aggregation_expressions())
+    else:
+        result = input_df.cube(*columns.grouping_expressions()).agg(
+            *columns.aggregation_expressions()
+        )
+    return with_column_map(
+        result,
+        columns.spark_names(),
+        columns.snowpark_names(),
+        columns.data_types(),
+        column_qualifiers=columns.get_qualifiers(),
+        parent_column_name_map=input_df._column_map,
+    )
+
+
+def map_pivot_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
+    """
+    Pivots a column of the current DataFrame and performs the specified aggregation.
+
+    There are 2 versions of the pivot function: one that requires the caller to specify the list of the distinct values
+    to pivot on and one that does not.
+    """
+    input_df, columns = map_aggregate_helper(rel, pivot=True, skip_alias=True)
+    pivot_column = map_single_column_expression(
+        rel.aggregate.pivot.col, input_df._column_map, ExpressionTyper(input_df)
+    )
+    pivot_values = [
+        get_literal_field_and_name(lit)[0] for lit in rel.aggregate.pivot.values
+    ]
+
+    if len(columns.grouping_expressions()) == 0:
+        result = input_df.pivot(
+            pivot_column[1].col, pivot_values if pivot_values else None
+        ).agg(*columns.aggregation_expressions())
+    else:
+        result = (
+            input_df.group_by(*columns.grouping_expressions())
+            .pivot(pivot_column[1].col, pivot_values if pivot_values else None)
+            .agg(*columns.aggregation_expressions())
+        )
+
+    spark_columns = []
+    for col in [string_parser(s) for s in result.columns]:
+        spark_col = (
+            input_df._column_map.get_spark_column_name_from_snowpark_column_name(
+                col, allow_non_exists=True
+            )
+        )
+
+        if spark_col is not None:
+            spark_columns.append(spark_col)
+        else:
+            spark_columns.append(col)
+
+    agg_name_list = [c.spark_name for c in columns.grouping_columns]
+    return with_column_map(
+        result,
+        agg_name_list + spark_columns[len(agg_name_list) :],
+        result.columns,
+        column_qualifiers=(
+            columns.get_qualifiers()[: len(agg_name_list)]
+            + [[]] * (len(spark_columns) - len(agg_name_list))
+        ),
+        parent_column_name_map=input_df._column_map,
+    )
+
+
+def string_parser(s):
+    # For values that are used as pivoted columns, the input and output are in the following format (outermost double quotes are part of the input):
+
+    # 1. "'Java'" -> Java
+    # 2. "'""C++""'" -> "C++"
+    # 3. "'""""''Scala''""""'" -> ""'Scala'""
+
+    # As we can see:
+    # 1. the whole content is always nested in a double quote followed by a single quote ("'<content>'").
+    # 2. the string content is nested in single quotes ('<string_content>')
+    # 3. double quote is escased by another double quote, this is snowflake behavior
+    # 4. if there is a single quote followed by a single quote, the first single quote needs to be preserved in the output
+
+    try:
+        # handling values that are used as pivoted columns
+        match = re.match(r'^"\'(.*)\'"$', s)
+        # extract the content between the outermost double quote followed by a single quote "'
+        content = match.group(1)
+        # convert the escaped double quote to the actual double quote
+        content = content.replace('""', '"')
+        escape_single_quote_placeholder = "__SAS_PLACEHOLDER_ESCAPE_SINGLE_QUOTE__"
+        # replace two consecutive single quote in the content with a placeholder, the first single quote needs to be preserved
+        content = re.sub(r"''", escape_single_quote_placeholder, content)
+        # remove the solo single quote, they are not part of the string content
+        content = re.sub(r"'", "", content)
+        # replace the placeholder with the single quote which we want to preserve
+        result = content.replace(escape_single_quote_placeholder, "'")
+        return result
+    except Exception:
+        # fallback to the original logic, handling aliased column names
+        double_quote_list = re.findall(r'"(.*?)"', s)
+        spark_string = ""
+        for entry in list(filter(None, double_quote_list)):
+            if "'" in entry:
+                entry = entry.replace("'", "")
+                if len(entry) > 0:
+                    spark_string += entry
+            elif entry.isdigit() or re.compile(r"^\d+?\.\d+?$").match(entry):
+                # skip quoting digits or decimal numbers as column names.
+                spark_string += entry
+            else:
+                spark_string += '"' + entry + '"'
+        return s if spark_string == "" else spark_string
+
+
+@dataclass(frozen=True)
+class _ColumnMetadata:
+    expression: snowpark.Column
+    spark_name: str
+    snowpark_name: str
+    data_type: DataType
+    qualifiers: list[str]
+
+
+@dataclass(frozen=True)
+class _Columns:
+    grouping_columns: list[_ColumnMetadata]
+    aggregation_columns: list[_ColumnMetadata]
+    can_infer_schema: bool
+
+    def grouping_expressions(self) -> list[snowpark.Column]:
+        return [col.expression for col in self.grouping_columns]
+
+    def aggregation_expressions(self) -> list[snowpark.Column]:
+        return [col.expression for col in self.aggregation_columns]
+
+    def expressions(self) -> list[snowpark.Column]:
+        return self.grouping_expressions() + self.aggregation_expressions()
+
+    def snowpark_names(self) -> list[str]:
+        return [
+            col.snowpark_name
+            for col in self.grouping_columns + self.aggregation_columns
+            if col.snowpark_name is not None
+        ]
+
+    def spark_names(self) -> list[str]:
+        return [
+            col.spark_name for col in self.grouping_columns + self.aggregation_columns
+        ]
+
+    def get_qualifiers(self) -> list[list[str]]:
+        return [
+            col.qualifiers for col in self.grouping_columns + self.aggregation_columns
+        ]
+
+    def data_types(self) -> list[DataType] | None:
+        if not self.can_infer_schema:
+            return None
+        return [
+            col.data_type
+            for col in self.grouping_columns + self.aggregation_columns
+            if col.data_type is not None
+        ]
+
+
+def map_aggregate_helper(
+    rel: relation_proto.Relation, pivot: bool = False, skip_alias: bool = False
+):
+    input_df = map_relation(rel.aggregate.input)
+    grouping_expressions = rel.aggregate.grouping_expressions
+    expressions = rel.aggregate.aggregate_expressions
+    groupings: list[_ColumnMetadata] = []
+    aggregations: list[_ColumnMetadata] = []
+
+    typer = ExpressionTyper(input_df)
+    schema_inferrable = True
+
+    with temporary_pivot_expression(pivot):
+        for exp in grouping_expressions:
+            new_name, snowpark_column = map_single_column_expression(
+                exp, input_df._column_map, typer
+            )
+            alias = make_column_names_snowpark_compatible(
+                [new_name], rel.common.plan_id, len(groupings)
+            )[0]
+            groupings.append(
+                _ColumnMetadata(
+                    snowpark_column.col
+                    if skip_alias
+                    else snowpark_column.col.alias(alias),
+                    new_name,
+                    None if skip_alias else alias,
+                    None if pivot else snowpark_column.typ,
+                    snowpark_column.get_qualifiers(),
+                )
+            )
+
+        for exp in expressions:
+            new_name, snowpark_column = map_single_column_expression(
+                exp, input_df._column_map, typer
+            )
+            alias = make_column_names_snowpark_compatible(
+                [new_name], rel.common.plan_id, len(groupings) + len(aggregations)
+            )[0]
+
+            def type_agg_expr(
+                agg_exp: TypedColumn, schema_inferrable: bool
+            ) -> DataType | None:
+                if pivot or not schema_inferrable:
+                    return None
+                try:
+                    return agg_exp.typ
+                except Exception:
+                    # This type used for schema inference optimization purposes.
+                    # typer may not be able to infer the type of some expressions
+                    # in that case we return None, and the optimization will not be applied.
+                    return None
+
+            agg_col_typ = type_agg_expr(snowpark_column, schema_inferrable)
+            if agg_col_typ is None:
+                schema_inferrable = False
+
+            aggregations.append(
+                _ColumnMetadata(
+                    snowpark_column.col
+                    if skip_alias
+                    else snowpark_column.col.alias(alias),
+                    new_name,
+                    None if skip_alias else alias,
+                    agg_col_typ,
+                    [],
+                )
+            )
+
+        return (
+            input_df,
+            _Columns(
+                grouping_columns=groupings,
+                aggregation_columns=aggregations,
+                can_infer_schema=schema_inferrable,
+            ),
+        )

snowflake/snowpark_connect/relation/map_catalog.py

@@ -0,0 +1,151 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import re
+
+import pandas
+import pyspark.sql.connect.proto.catalog_pb2 as catalog_proto
+
+from snowflake import snowpark
+from snowflake.snowpark_connect.relation.catalogs import CATALOGS
+from snowflake.snowpark_connect.relation.catalogs.utils import (
+    CURRENT_CATALOG_NAME,
+    get_current_catalog,
+    set_current_catalog,
+)
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
+from snowflake.snowpark_connect.utils.telemetry import (
+    SnowparkConnectNotImplementedError,
+)
+
+
+def map_catalog(
+    rel: catalog_proto.Catalog,
+) -> pandas.DataFrame | snowpark.DataFrame:
+    match rel.WhichOneof("cat_type"):
+        # Database related APIs
+        case "current_database":
+            return get_current_catalog().currentDatabase()
+        case "database_exists":
+            return get_current_catalog().databaseExists(
+                rel.database_exists.db_name,
+            )
+        case "get_database":
+            return get_current_catalog().getDatabase(
+                rel.get_database.db_name,
+            )
+        case "list_databases":
+            return get_current_catalog().listDatabases(
+                (
+                    rel.list_databases.pattern
+                    if rel.list_databases.HasField("pattern")
+                    else None
+                ),
+            )
+        case "set_current_database":
+            return get_current_catalog().setCurrentDatabase(
+                rel.set_current_database.db_name,
+            )
+        # Table related APIs
+        case "get_table":
+            return get_current_catalog().getTable(rel.get_table.table_name)
+        case "list_tables":
+            return get_current_catalog().listTables(
+                (
+                    rel.list_tables.db_name
+                    if rel.list_tables.HasField("db_name")
+                    else None
+                ),
+                (
+                    rel.list_tables.pattern
+                    if rel.list_tables.HasField("pattern")
+                    else None
+                ),
+            )
+        case "list_catalogs":
+            if (
+                rel.list_catalogs.HasField("pattern")
+                and rel.list_catalogs.pattern == ""
+            ):
+                return pandas.DataFrame([])
+            # Handle None pattern by treating it as matching all catalogs
+            pattern_str = (
+                rel.list_catalogs.pattern
+                if rel.list_catalogs.pattern is not None
+                else ".*"
+            )
+            pattern = re.compile(pattern_str)
+            return pandas.DataFrame(
+                [
+                    (name, cat.description)
+                    for name, cat in CATALOGS.items()
+                    if pattern.match(name) is not None
+                ]
+            )
+        case "current_catalog":
+            return pandas.DataFrame({"current_catalog": [CURRENT_CATALOG_NAME]})
+        case "set_current_catalog":
+            set_current_catalog(rel.set_current_catalog.catalog_name)
+            return pandas.DataFrame()
+        case "table_exists":
+            return get_current_catalog().tableExists(
+                rel.table_exists.table_name,
+                (
+                    rel.table_exists.db_name
+                    if rel.table_exists.HasField("db_name")
+                    else None
+                ),
+            )
+        # Column related APIs
+        case "list_columns":
+            return get_current_catalog().listColumns(
+                rel.list_columns.table_name,
+                (
+                    rel.list_columns.db_name
+                    if rel.list_columns.HasField("db_name")
+                    else None
+                ),
+            )
+        # View related APIs
+        case "drop_global_temp_view":
+            return get_current_catalog().dropGlobalTempView(
+                rel.drop_global_temp_view.view_name
+            )
+        case "drop_temp_view":
+            return get_current_catalog().dropTempView(rel.drop_temp_view.view_name)
+        case "recover_partitions":
+            # 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.
+            logger.warning("recover_partitions is triggered with no-op")
+            return pandas.DataFrame()
+        # Table related APIs
+        case "create_table":
+            return get_current_catalog().createTable(
+                rel.create_table.table_name,
+                rel.create_table.path,
+                rel.create_table.source,
+                rel.create_table.schema,
+                rel.create_table.description,
+                **rel.create_table.options,
+            )
+        # Caching related APIs
+        case "cache_table":
+            return get_current_catalog().cacheTable(
+                rel.cache_table.table_name,
+                rel.cache_table.storage_level,
+            )
+        case "clear_cache":
+            return get_current_catalog().clearCache()
+        case "is_cached":
+            return get_current_catalog().isCached(rel.is_cached.table_name)
+        case "refresh_by_path":
+            return get_current_catalog().refreshByPath(rel.refresh_by_path.path)
+        case "refresh_table":
+            return get_current_catalog().refreshTable(rel.refresh_table.table_name)
+        case "uncache_table":
+            return get_current_catalog().uncacheTable(rel.uncache_table.table_name)
+        case other:
+            # TODO: list_function implementation is blocked on SNOW-1787268
+            raise SnowparkConnectNotImplementedError(f"Other Relation {other}")