snowpark-connect 0.32.0__py3-none-any.whl → 0.33.0__py3-none-any.whl

snowflake/snowpark_connect/relation/map_sql.py

@@ -61,6 +61,9 @@ from snowflake.snowpark_connect.relation.map_relation import (
     NATURAL_JOIN_TYPE_BASE,
     map_relation,
 )
+
+# Import from utils for consistency
+from snowflake.snowpark_connect.relation.utils import is_aggregate_function
 from snowflake.snowpark_connect.type_mapping import map_snowpark_to_pyspark_types
 from snowflake.snowpark_connect.utils.context import (
     _accessing_temp_object,
@@ -588,25 +591,28 @@ def map_sql_to_pandas_df(
                     f"CREATE TABLE {if_not_exists}{name} LIKE {source}"
                 ).collect()
             case "CreateTempViewUsing":
+                parsed_sql = sqlglot.parse_one(sql_string, dialect="spark")
+
+                spark_view_name = next(parsed_sql.find_all(sqlglot.exp.Table)).name
+
+                num_columns = len(list(parsed_sql.find_all(sqlglot.exp.ColumnDef)))
+                null_list = (
+                    ", ".join(["NULL"] * num_columns) if num_columns > 0 else "*"
+                )
                 empty_select = (
-                    " AS SELECT * WHERE 1 = 0"
+                    f" AS SELECT {null_list} WHERE 1 = 0"
                     if logical_plan.options().isEmpty()
                     and logical_plan.children().isEmpty()
                     else ""
                 )
-                parsed_sql = (
-                    sqlglot.parse_one(sql_string, dialect="spark")
-                    .transform(_normalize_identifiers)
+
+                transformed_sql = (
+                    parsed_sql.transform(_normalize_identifiers)
                     .transform(_remove_column_data_type)
                     .transform(_remove_file_format_property)
                 )
-                snowflake_sql = parsed_sql.sql(dialect="snowflake")
+                snowflake_sql = transformed_sql.sql(dialect="snowflake")
                 session.sql(f"{snowflake_sql}{empty_select}").collect()
-                spark_view_name = next(
-                    sqlglot.parse_one(sql_string, dialect="spark").find_all(
-                        sqlglot.exp.Table
-                    )
-                ).name
                 snowflake_view_name = spark_to_sf_single_id_with_unquoting(
                     spark_view_name
                 )
@@ -877,16 +883,71 @@ def map_sql_to_pandas_df(
                 overwrite_str = "OVERWRITE" if logical_plan.overwrite() else ""
                 cols_str = "(" + ", ".join(user_columns) + ")" if user_columns else ""
 
-                try:
-                    target_table = session.table(name)
-                    target_schema = target_table.schema
+                # Extract partition spec if any
+                partition_spec = logical_plan.partitionSpec()
+                partition_map = as_java_map(partition_spec)
+
+                partition_columns = {}
+                for entry in partition_map.entrySet():
+                    col_name = str(entry.getKey())
+                    value_option = entry.getValue()
+                    if value_option.isDefined():
+                        partition_columns[col_name] = value_option.get()
+
+                # Add partition columns to the dataframe
+                if partition_columns:
+                    """
+                    Spark sends them in the partition spec and the values won't be present in the values array.
+                    As snowflake does not support static partitions in INSERT INTO statements,
+                    we need to add the partition columns to the dataframe as literal columns.
+
+                    ex: INSERT INTO TABLE test_table PARTITION (ds='2021-01-01', hr=10) VALUES ('k1', 100), ('k2', 200), ('k3', 300)
+
+                    Spark sends: VALUES ('k1', 100), ('k2', 200), ('k3', 300) with partition spec (ds='2021-01-01', hr=10)
+                    Snowflake expects: VALUES ('k1', 100, '2021-01-01', 10), ('k2', 200, '2021-01-01', 10), ('k3', 300, '2021-01-01', 10)
+
+                    We need to add the partition columns to the dataframe as literal columns.
+
+                    ex: df = df.withColumn('ds', snowpark_fn.lit('2021-01-01'))
+                        df = df.withColumn('hr', snowpark_fn.lit(10))
+
+                    Then the final query will be:
+                    INSERT INTO TABLE test_table VALUES ('k1', 100, '2021-01-01', 10), ('k2', 200, '2021-01-01', 10), ('k3', 300, '2021-01-01', 10)
+                    """
+                    for partition_col, partition_value in partition_columns.items():
+                        df = df.withColumn(
+                            partition_col, snowpark_fn.lit(partition_value)
+                        )
+
+                target_table = session.table(name)
+                target_schema = target_table.schema
 
-                    # Modify df with NaN → NULL conversion for DECIMAL columns
+                expected_number_of_columns = (
+                    len(user_columns) if user_columns else len(target_schema.fields)
+                )
+                if expected_number_of_columns != len(df.schema.fields):
+                    reason = (
+                        "too many data columns"
+                        if len(df.schema.fields) > expected_number_of_columns
+                        else "not enough data columns"
+                    )
+                    exception = AnalysisException(
+                        f'[INSERT_COLUMN_ARITY_MISMATCH.{reason.replace(" ", "_").upper()}] Cannot write to {name}, the reason is {reason}:\n'
+                        f'Table columns: {", ".join(target_schema.names)}.\n'
+                        f'Data columns: {", ".join(df.schema.names)}.'
+                    )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
+
+                try:
+                    # Modify df with type conversions and struct field name mapping
                     modified_columns = []
                     for source_field, target_field in zip(
                         df.schema.fields, target_schema.fields
                     ):
                         col_name = source_field.name
+
+                        # Handle different type conversions
                         if isinstance(
                             target_field.datatype, snowpark.types.DecimalType
                         ) and isinstance(
@@ -904,12 +965,25 @@ def map_sql_to_pandas_df(
                                 .alias(col_name)
                             )
                             modified_columns.append(modified_col)
+                        elif (
+                            isinstance(target_field.datatype, snowpark.types.StructType)
+                            and source_field.datatype != target_field.datatype
+                        ):
+                            # Cast struct with field name mapping (e.g., col1,col2 -> i1,i2)
+                            # This fixes INSERT INTO table with struct literals like (2, 3)
+                            modified_col = (
+                                snowpark_fn.col(col_name)
+                                .cast(target_field.datatype, rename_fields=True)
+                                .alias(col_name)
+                            )
+                            modified_columns.append(modified_col)
                         else:
                             modified_columns.append(snowpark_fn.col(col_name))
 
                     df = df.select(modified_columns)
                 except Exception:
                     pass
+
                 queries = df.queries["queries"]
                 final_query = queries[0]
                 session.sql(
@@ -1345,7 +1419,7 @@ def map_sql_to_pandas_df(
 
                 return pandas.DataFrame({"": [""]}), ""
             case "RepairTable":
-                # No-Op. Snowflake doesn't have explicit partitions to repair.
+                # No-Op: Snowflake doesn't have explicit partitions to repair.
                 table_relation = logical_plan.child()
                 db_and_table_name = as_java_list(table_relation.multipartIdentifier())
                 multi_part_len = len(db_and_table_name)
@@ -1598,7 +1672,19 @@ def map_logical_plan_relation(
                     attr_parts = as_java_list(expr.nameParts())
                     if len(attr_parts) == 1:
                         attr_name = str(attr_parts[0])
-                        return alias_map.get(attr_name, expr)
+                        if attr_name in alias_map:
+                            # Check if the alias references an aggregate function
+                            # If so, don't substitute because you can't GROUP BY an aggregate
+                            aliased_expr = alias_map[attr_name]
+                            aliased_expr_class = str(
+                                aliased_expr.getClass().getSimpleName()
+                            )
+                            if aliased_expr_class == "UnresolvedFunction":
+                                func_name = str(aliased_expr.nameParts().head())
+                                if is_aggregate_function(func_name):
+                                    return expr
+                            return aliased_expr
+                        return expr
 
                     return expr
 

snowflake/snowpark_connect/relation/read/map_read_json.py

@@ -117,6 +117,10 @@ def map_read_json(
                     if unquote_if_quoted(sf.name) in columns_with_valid_contents
                 ]
 
+        new_schema, fields_changed = validate_and_update_schema(schema)
+        if fields_changed:
+            schema = new_schema
+
         df = construct_dataframe_by_schema(
             schema, df.to_local_iterator(), session, snowpark_options, batch_size
         )
@@ -134,6 +138,84 @@ def map_read_json(
         )
 
 
+def should_drop_field(field: StructField) -> bool:
+    if isinstance(field.datatype, StructType):
+        # "a" : {} => drop the field
+        if len(field.datatype.fields) == 0:
+            return True
+    elif (
+        isinstance(field.datatype, ArrayType)
+        and field.datatype.element_type is not None
+        and isinstance(field.datatype.element_type, StructType)
+    ):
+        if len(field.datatype.element_type.fields) == 0:
+            # "a" : [{}] => drop the field
+            return True
+    return False
+
+
+# Validate the schema to ensure it is valid for Snowflake
+# Handles these cases:
+#   1. Drops StructField([])
+#   2. Drops ArrayType(StructType([]))
+#   3. ArrayType() -> ArrayType(StringType())
+def validate_and_update_schema(schema: StructType | None) -> (StructType | None, bool):
+    if not isinstance(schema, StructType):
+        return schema, False
+    new_fields = []
+    fields_changed = False
+    for sf in schema.fields:
+        if should_drop_field(sf):
+            fields_changed = True
+            continue
+        if isinstance(sf.datatype, StructType):
+            # If the schema is a struct, validate the child schema
+            if len(sf.datatype.fields) == 0:
+                # No fields in the struct, drop the field
+                fields_changed = True
+                continue
+            child_field = StructField(sf.name, sf.datatype, sf.nullable)
+            # Recursively validate the child schema
+            child_field.datatype, child_field_changes = validate_and_update_schema(
+                sf.datatype
+            )
+            if should_drop_field(child_field):
+                fields_changed = True
+                continue
+            new_fields.append(child_field)
+            fields_changed = fields_changed or child_field_changes
+        elif isinstance(sf.datatype, ArrayType):
+            # If the schema is an array, validate the element schema
+            if sf.datatype.element_type is not None and isinstance(
+                sf.datatype.element_type, StructType
+            ):
+                # If the element schema is a struct, validate the element schema
+                if len(sf.datatype.element_type.fields) == 0:
+                    # No fields in the struct, drop the field
+                    fields_changed = True
+                    continue
+                else:
+                    # Recursively validate the element schema
+                    element_schema, element_field_changes = validate_and_update_schema(
+                        sf.datatype.element_type
+                    )
+                    if element_field_changes:
+                        sf.datatype.element_type = element_schema
+                        fields_changed = True
+                    if should_drop_field(sf):
+                        fields_changed = True
+                        continue
+            elif sf.datatype.element_type is None:
+                fields_changed = True
+                sf.datatype.element_type = StringType()
+            new_fields.append(sf)
+        else:
+            new_fields.append(sf)
+    if fields_changed:
+        schema.fields = new_fields
+    return schema, fields_changed
+
+
 def merge_json_schema(
     content: typing.Any,
     schema: StructType | None,
@@ -378,8 +460,11 @@ def construct_row_by_schema(
         inner_schema = schema.element_type
         if isinstance(content, str):
             content = json.loads(content)
-        for ele in content:
-            result.append(construct_row_by_schema(ele, inner_schema, snowpark_options))
+        if inner_schema is not None:
+            for ele in content:
+                result.append(
+                    construct_row_by_schema(ele, inner_schema, snowpark_options)
+                )
         return result
     elif isinstance(schema, DateType):
         return cast_to_match_snowpark_type(

snowflake/snowpark_connect/relation/utils.py

@@ -284,3 +284,49 @@ def snowpark_functions_col(name: str, column_map: ColumnNameMap) -> snowpark.Col
     """
     is_qualified_name = name not in column_map.get_snowpark_columns()
     return snowpark_fn.col(name, _is_qualified_name=is_qualified_name)
+
+
+def is_aggregate_function(func_name: str) -> bool:
+    """
+    Check if a function name is an aggregate function.
+
+    Uses a hybrid approach:
+    1. First checks PySpark's docstring convention (docstrings starting with "Aggregate function:")
+    2. Falls back to a hardcoded list for functions with missing/incorrect docstrings
+
+    This ensures comprehensive coverage while automatically supporting new PySpark aggregate functions.
+
+    Args:
+        func_name: The function name to check (case-insensitive)
+
+    Returns:
+        True if the function is an aggregate function, False otherwise
+    """
+    try:
+        import pyspark.sql.functions as pyspark_functions
+
+        # TODO:
+        """
+        Check we can leverage scala classes to determine agg functions:
+        https://github.com/apache/spark/blob/master/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/interfaces.scala#L207
+        """
+
+        # Try PySpark docstring approach first (covers most aggregate functions)
+        pyspark_func = getattr(pyspark_functions, func_name.lower(), None)
+        if pyspark_func and pyspark_func.__doc__:
+            if pyspark_func.__doc__.lstrip().startswith("Aggregate function:"):
+                return True
+
+        # Fallback list for aggregate functions with missing/incorrect docstrings
+        # These are known aggregate functions that don't have proper docstring markers
+        fallback_aggregates = {
+            "percentile_cont",
+            "percentile_disc",
+            "any_value",
+            "grouping",
+            "grouping_id",
+        }
+        return func_name.lower() in fallback_aggregates
+
+    except Exception:
+        return False