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

snowflake/snowpark_connect/expression/map_unresolved_extract_value.py

@@ -5,6 +5,7 @@
 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_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.config import global_config
@@ -57,7 +58,8 @@ def map_unresolved_extract_value(
         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):
+    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.
@@ -74,12 +76,37 @@ def map_unresolved_extract_value(
 
     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
-    ):
+    if spark_sql_ansi_enabled and is_array:
         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))
+    def _get_extracted_value_type():
+        if is_array:
+            return [child_typed_column.typ.element_type]
+        elif isinstance(child_typed_column.typ, MapType):
+            return [child_typed_column.typ.value_type]
+        elif (
+            isinstance(child_typed_column.typ, StructType)
+            and isinstance(extract_typed_column.col._expr1, Literal)
+            and isinstance(extract_typed_column.col._expr1.value, str)
+        ):
+            struct = dict(
+                {
+                    (
+                        f.name
+                        if global_config.spark_sql_caseSensitive
+                        else f.name.lower(),
+                        f.datatype,
+                    )
+                    for f in child_typed_column.typ.fields
+                }
+            )
+            key = extract_typed_column.col._expr1.value
+            key = key if global_config.spark_sql_caseSensitive else key.lower()
+
+            return [struct[key]] if key in struct else typer.type(result_exp)
+        return typer.type(result_exp)
+
+    return spark_function_name, TypedColumn(result_exp, _get_extracted_value_type)

snowflake/snowpark_connect/expression/map_unresolved_function.py

@@ -80,6 +80,9 @@ from snowflake.snowpark_connect.constants import (
     SPARK_TZ_ABBREVIATIONS_OVERRIDES,
     STRUCTURED_TYPES_ENABLED,
 )
+from snowflake.snowpark_connect.expression.function_defaults import (
+    inject_function_defaults,
+)
 from snowflake.snowpark_connect.expression.literal import get_literal_field_and_name
 from snowflake.snowpark_connect.expression.map_cast import (
     CAST_FUNCTIONS,
@@ -299,6 +302,9 @@ def map_unresolved_function(
     function_name = exp.unresolved_function.function_name.lower()
     is_udtf_call = function_name in session._udtfs
 
+    # Inject default parameters for functions that need them (especially for Scala clients)
+    inject_function_defaults(exp.unresolved_function)
+
     def _resolve_args_expressions(exp: expressions_proto.Expression):
         def _resolve_fn_arg(exp):
             with resolving_fun_args():
@@ -1107,7 +1113,7 @@ def map_unresolved_function(
             result_exp = TypedColumn(
                 result_exp, lambda: [ArrayType(snowpark_typed_args[0].typ)]
             )
-        case "array_size" | "cardinality":
+        case "array_size":
             array_type = snowpark_typed_args[0].typ
             if not isinstance(array_type, ArrayType):
                 raise AnalysisException(
@@ -1116,6 +1122,16 @@ def map_unresolved_function(
             result_exp = TypedColumn(
                 snowpark_fn.array_size(*snowpark_args), lambda: [LongType()]
             )
+        case "cardinality":
+            arg_type = snowpark_typed_args[0].typ
+            if isinstance(arg_type, (ArrayType, MapType)):
+                result_exp = TypedColumn(
+                    snowpark_fn.size(*snowpark_args), lambda: [LongType()]
+                )
+            else:
+                raise AnalysisException(
+                    f"Expected argument '{snowpark_arg_names[0]}' to have an ArrayType or MapType, but got {arg_type.simpleString()}."
+                )
         case "array_sort":
             result_exp = TypedColumn(
                 snowpark_fn.array_sort(*snowpark_args),
@@ -1295,10 +1311,35 @@ def map_unresolved_function(
             )
             result_exp = TypedColumn(result_exp, lambda: [LongType()])
         case "bit_get" | "getbit":
-            bit_get_function = snowpark_fn.function("GETBIT")
-            result_exp = TypedColumn(
-                bit_get_function(*snowpark_args), lambda: [LongType()]
+            snowflake_compat = get_boolean_session_config_param(
+                "enable_snowflake_extension_behavior"
             )
+            col, pos = snowpark_args
+            if snowflake_compat:
+                bit_get_function = snowpark_fn.function("GETBIT")(col, pos)
+            else:
+                raise_error = _raise_error_helper(LongType())
+                bit_get_function = snowpark_fn.when(
+                    (snowpark_fn.lit(0) <= pos) & (pos <= snowpark_fn.lit(63))
+                    | snowpark_fn.is_null(pos),
+                    snowpark_fn.function("GETBIT")(col, pos),
+                ).otherwise(
+                    raise_error(
+                        snowpark_fn.concat(
+                            snowpark_fn.lit(
+                                "Invalid bit position: ",
+                            ),
+                            snowpark_fn.cast(
+                                pos,
+                                StringType(),
+                            ),
+                            snowpark_fn.lit(
+                                " exceeds the bit upper limit",
+                            ),
+                        )
+                    )
+                )
+            result_exp = TypedColumn(bit_get_function, lambda: [LongType()])
         case "bit_length":
             bit_length_function = snowpark_fn.function("bit_length")
             result_exp = TypedColumn(
@@ -3726,7 +3767,7 @@ def map_unresolved_function(
                     snowpark_args[1] <= 0, snowpark_fn.lit("")
                 ).otherwise(snowpark_fn.left(*snowpark_args))
             result_type = StringType()
-        case "length" | "char_length" | "character_length":
+        case "length" | "char_length" | "character_length" | "len":
             if exp.unresolved_function.arguments[0].HasField("literal"):
                 # Only update the name if it has the literal field.
                 # If it doesn't, it means it's binary data.
@@ -4636,7 +4677,7 @@ def map_unresolved_function(
                     snowpark_args[0],
                 )
             result_type = DateType()
-        case "not":
+        case "not" | "!":
             spark_function_name = f"(NOT {snowpark_arg_names[0]})"
             result_exp = ~snowpark_args[0]
             result_type = BooleanType()
@@ -5212,9 +5253,8 @@ def map_unresolved_function(
             # TODO: Seems like more validation of the arguments is appropriate.
             args = exp.unresolved_function.arguments
             if len(args) > 0:
-                if not (
-                    isinstance(snowpark_typed_args[0].typ, IntegerType)
-                    or isinstance(snowpark_typed_args[0].typ, NullType)
+                if not isinstance(
+                    snowpark_typed_args[0].typ, (IntegerType, LongType, NullType)
                 ):
                     raise AnalysisException(
                         f"""[DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve "{spark_function_name}" due to data type mismatch: Parameter 1 requires the ("INT" or "BIGINT") type, however {snowpark_arg_names[0]} has the type "{snowpark_typed_args[0].typ}"""
@@ -5496,9 +5536,27 @@ def map_unresolved_function(
             ):
                 result_exp = snowpark_fn.lit(None)
             else:
+                right_expr = snowpark_fn.right(*snowpark_args)
+                if isinstance(snowpark_typed_args[0].typ, TimestampType):
+                    # Spark format is always displayed as YYY-MM-DD HH:mm:ss.FF6
+                    # When microseconds are equal to 0 .FF6 part is removed
+                    # When microseconds are equal to 0 at the end, they are removed i.e. .123000 -> .123 when displayed
+
+                    formated_timestamp = snowpark_fn.to_varchar(
+                        snowpark_args[0], "YYYY-MM-DD HH:MI:SS.FF6"
+                    )
+                    right_expr = snowpark_fn.right(
+                        snowpark_fn.regexp_replace(
+                            snowpark_fn.regexp_replace(formated_timestamp, "0+$", ""),
+                            "\\.$",
+                            "",
+                        ),
+                        snowpark_args[1],
+                    )
+
                 result_exp = snowpark_fn.when(
                     snowpark_args[1] <= 0, snowpark_fn.lit("")
-                ).otherwise(snowpark_fn.right(*snowpark_args))
+                ).otherwise(right_expr)
             result_type = StringType()
         case "rint":
             result_exp = snowpark_fn.cast(
@@ -6729,6 +6787,7 @@ def map_unresolved_function(
                         if value == "" or any(
                             c in value for c in [",", "\n", "\r", '"', "'"]
                         ):
+                            value = value.replace("\\", "\\\\").replace('"', '\\"')
                             result.append(f'"{value}"')
                         else:
                             result.append(value)

snowflake/snowpark_connect/includes/python/pyspark/pandas/spark/__init__.py

@@ -0,0 +1,16 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#