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

snowflake/snowpark_connect/expression/map_sql_expression.py

@@ -11,13 +11,18 @@ import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 import pyspark.sql.connect.proto.types_pb2 as types_proto
 from google.protobuf.any_pb2 import Any
 from pyspark.errors.exceptions.base import AnalysisException
+from pyspark.sql.connect import functions as pyspark_functions
 
 import snowflake.snowpark_connect.proto.snowflake_expression_ext_pb2 as snowflake_proto
-from snowflake import snowpark
+from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
+from snowflake.snowpark.types import DayTimeIntervalType, YearMonthIntervalType
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 from snowflake.snowpark_connect.typed_column import TypedColumn
 from snowflake.snowpark_connect.utils.context import (
+    get_jpype_jclass_lock,
     get_sql_named_arg,
     get_sql_plan,
     get_sql_pos_arg,
@@ -32,8 +37,36 @@ from .typer import ExpressionTyper
 
 DECIMAL_RE = re.compile(r"decimal\((\d+), *(\d+)\)")
 
+_INTERVAL_YEARMONTH_PATTERN_RE = re.compile(r"interval (year|month)( to (year|month))?")
+_INTERVAL_DAYTIME_PATTERN_RE = re.compile(
+    r"interval (day|hour|minute|second)( to (day|hour|minute|second))?"
+)
+
+# Interval field mappings using proper constants
+_YEAR_MONTH_FIELD_MAP = {
+    "year": YearMonthIntervalType.YEAR,
+    "month": YearMonthIntervalType.MONTH,
+}
+
+_DAY_TIME_FIELD_MAP = {
+    "day": DayTimeIntervalType.DAY,
+    "hour": DayTimeIntervalType.HOUR,
+    "minute": DayTimeIntervalType.MINUTE,
+    "second": DayTimeIntervalType.SECOND,
+}
+
 _window_specs = ContextVar[dict[str, any]]("_window_specs", default={})
 
+# Functions that can be called without parentheses in Spark SQL. Build up the list as we see more functions.
+NILARY_FUNCTIONS = frozenset(
+    [
+        "current_date",
+        "current_timestamp",
+        "current_user",
+        "user",
+    ]
+)
+
 
 def sql_parser():
     """
@@ -44,26 +77,35 @@ def sql_parser():
     """
 
     ts_type = global_config.spark_sql_timestampType
+    session_tz = global_config.spark_sql_session_timeZone
 
     if ts_type is not None:
         _get_sql_conf().get().setConfString("spark.sql.timestampType", str(ts_type))
 
+    if session_tz is not None:
+        _get_sql_conf().get().setConfString(
+            "spark.sql.session.timeZone", str(session_tz)
+        )
+
     return _get_sql_parser()
 
 
 @cache
 def _get_sql_parser():
-    return jpype.JClass("org.apache.spark.sql.execution.SparkSqlParser")()
+    with get_jpype_jclass_lock():
+        return jpype.JClass("org.apache.spark.sql.execution.SparkSqlParser")()
 
 
 @cache
 def _get_sql_conf():
-    return jpype.JClass("org.apache.spark.sql.internal.SQLConf")
+    with get_jpype_jclass_lock():
+        return jpype.JClass("org.apache.spark.sql.internal.SQLConf")
 
 
 @cache
 def _as_java_list():
-    return jpype.JClass("scala.collection.JavaConverters").seqAsJavaList
+    with get_jpype_jclass_lock():
+        return jpype.JClass("scala.collection.JavaConverters").seqAsJavaList
 
 
 def as_java_list(obj):
@@ -72,7 +114,8 @@ def as_java_list(obj):
 
 @cache
 def _as_java_map():
-    return jpype.JClass("scala.collection.JavaConverters").mapAsJavaMap
+    with get_jpype_jclass_lock():
+        return jpype.JClass("scala.collection.JavaConverters").mapAsJavaMap
 
 
 def as_java_map(obj):
@@ -89,6 +132,11 @@ def as_scala_seq(input):
     )
 
 
+@cache
+def _scala_some():
+    return jpype.JClass("scala.Some")
+
+
 def map_sql_expr(
     exp: expressions_proto.Expression,
     column_mapping: ColumnNameMap,
@@ -223,20 +271,52 @@ def apply_filter_clause(
 
 
 def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Expression:
-    from snowflake.snowpark_connect.expression.map_expression import (
-        map_single_column_expression,
-    )
     from snowflake.snowpark_connect.relation.map_sql import map_logical_plan_relation
 
     class_name = str(exp.getClass().getSimpleName())
     match class_name:
         case "AggregateExpression":
-            func_name = as_java_list(exp.children())[0].nodeName()
+            aggregate_func = as_java_list(exp.children())[0]
+            func_name = aggregate_func.nodeName()
             args = [
                 map_logical_plan_expression(e)
-                for e in as_java_list(as_java_list(exp.children())[0].children())
+                for e in list(as_java_list(aggregate_func.children()))
             ]
-            proto = apply_filter_clause(func_name, args, exp)
+
+            # Special handling for percentile_cont and percentile_disc
+            # These functions have a 'reverse' property that indicates sort order
+            # Pass it as a 3rd argument (sort_order expression) without modifying children
+            if func_name.lower() in ("percentile_cont", "percentiledisc"):
+                # percentile_cont/disc should always have exactly 2 children: unresolved attribute and percentile value
+                if len(args) != 2:
+                    exception = AssertionError(
+                        f"{func_name} expected 2 args but got {len(args)}"
+                    )
+                    attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+                    raise exception
+
+                reverse = bool(aggregate_func.reverse())
+
+                direction = (
+                    expressions_proto.Expression.SortOrder.SORT_DIRECTION_DESCENDING
+                    if reverse
+                    else expressions_proto.Expression.SortOrder.SORT_DIRECTION_ASCENDING
+                )
+
+                sort_order_expr = expressions_proto.Expression(
+                    sort_order=expressions_proto.Expression.SortOrder(
+                        child=args[0],
+                        direction=direction,
+                    )
+                )
+                args.append(sort_order_expr)
+                proto = apply_filter_clause(func_name, [args[0]], exp)
+                # second arg is a literal value and it doesn't make sense to apply filter on it.
+                # also skips filtering on sort_order.
+                proto.unresolved_function.arguments.append(args[1])
+                proto.unresolved_function.arguments.append(sort_order_expr)
+            else:
+                proto = apply_filter_clause(func_name, args, exp)
         case "Alias":
             proto = expressions_proto.Expression(
                 alias=expressions_proto.Expression.Alias(
@@ -253,7 +333,7 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
                     function_name="when",
                     arguments=[
                         map_logical_plan_expression(e)
-                        for e in as_java_list(exp.children())
+                        for e in list(as_java_list(exp.children()))
                     ],
                 )
             )
@@ -267,7 +347,8 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
             )
         case "Coalesce":
             arguments = [
-                map_logical_plan_expression(e) for e in as_java_list(exp.children())
+                map_logical_plan_expression(e)
+                for e in list(as_java_list(exp.children()))
             ]
 
             proto = expressions_proto.Expression(
@@ -277,6 +358,14 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
                 )
             )
         case "CreateNamedStruct":
+            # Both struct() and named_struct() Spark SQL functions produce a CreateNamedStruct
+            # logical plan. We distinguish them by checking exp.prettyName():
+            # - "named_struct" -> explicit named_struct() call, requires name-value pairs
+            # - "struct" -> struct() call, field names are inferred from column expressions
+
+            # Additionally, struct(*) with star expansion is handled as named_struct.
+            # TODO - consider refactoring the impl and handle it in "struct" impl
+
             arg_exprs = [
                 arg
                 for k_v in zip(
@@ -284,17 +373,31 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
                 )
                 for arg in k_v
             ]
-            struct_args = [
-                map_logical_plan_expression(e)
-                for e in arg_exprs
-                if e.prettyName() != "NamePlaceholder"
-            ]
-            proto = expressions_proto.Expression(
-                unresolved_function=expressions_proto.Expression.UnresolvedFunction(
-                    function_name="named_struct",
-                    arguments=struct_args,
+
+            if (
+                "unresolvedstar" in [e.prettyName() for e in arg_exprs]
+                or exp.prettyName() == "named_struct"
+            ):
+                struct_args = [
+                    map_logical_plan_expression(e)
+                    for e in arg_exprs
+                    if e.prettyName() != "NamePlaceholder"
+                ]
+                proto = expressions_proto.Expression(
+                    unresolved_function=expressions_proto.Expression.UnresolvedFunction(
+                        function_name="named_struct",
+                        arguments=struct_args,
+                    )
+                )
+            else:
+                arg_exprs = [arg for arg in as_java_list(exp.valExprs())]
+                struct_args = [map_logical_plan_expression(e) for e in arg_exprs]
+                proto = expressions_proto.Expression(
+                    unresolved_function=expressions_proto.Expression.UnresolvedFunction(
+                        function_name="struct",
+                        arguments=struct_args,
+                    )
                 )
-            )
         case "Exists":
             rel_proto = map_logical_plan_relation(exp.plan())
             any_proto = Any()
@@ -308,22 +411,23 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
             )
             proto = expressions_proto.Expression(extension=any_proto)
         case "ExpressionWithUnresolvedIdentifier":
-            plan_id = None
-            identifierExpr = map_logical_plan_expression(exp.identifierExpr())
-            session = snowpark.Session.get_active_session()
-            m = ColumnNameMap([], [], None)
-            expr = map_single_column_expression(
-                identifierExpr, m, ExpressionTyper.dummy_typer(session)
+            from snowflake.snowpark_connect.relation.map_sql import (
+                get_relation_identifier_name,
             )
-            value = session.range(1).select(expr[1].col).collect()[0][0]
 
-            proto = expressions_proto.Expression(
-                unresolved_attribute=expressions_proto.Expression.UnresolvedAttribute(
-                    unparsed_identifier=str(value),
-                    plan_id=plan_id,
-                ),
-            )
-            # TODO: support identifier referencing unresolved function
+            value = unquote_if_quoted(get_relation_identifier_name(exp))
+            if getattr(pyspark_functions, value.lower(), None) is not None:
+                unresolved_function = exp.exprBuilder().apply(
+                    _scala_some()(value).toList()
+                )
+                proto = map_logical_plan_expression(unresolved_function)
+            else:
+                proto = expressions_proto.Expression(
+                    unresolved_attribute=expressions_proto.Expression.UnresolvedAttribute(
+                        unparsed_identifier=str(value),
+                        plan_id=None,
+                    ),
+                )
         case "InSubquery":
             rel_proto = map_logical_plan_relation(exp.query().plan())
             any_proto = Any()
@@ -334,7 +438,7 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
                         subquery_type=snowflake_proto.SubqueryExpression.SUBQUERY_TYPE_IN,
                         in_subquery_values=[
                             map_logical_plan_expression(value)
-                            for value in as_java_list(exp.values())
+                            for value in list(as_java_list(exp.values()))
                         ],
                     )
                 )
@@ -343,7 +447,7 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
         case "LambdaFunction":
             arguments = [
                 map_logical_plan_expression(arg).unresolved_named_lambda_variable
-                for arg in as_java_list(exp.arguments())
+                for arg in list(as_java_list(exp.arguments()))
             ]
             proto = expressions_proto.Expression(
                 lambda_function=expressions_proto.Expression.LambdaFunction(
@@ -352,19 +456,28 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
                 )
             )
         case "Like" | "ILike" | "RLike":
+            arguments = [
+                map_logical_plan_expression(e)
+                for e in list(as_java_list(exp.children()))
+            ]
+            # exp.escapeChar() returns a JPype JChar - convert to string and create a literal
+            if getattr(exp, "escapeChar", None) is not None:
+                escape_char_str = str(exp.escapeChar())
+                escape_literal = expressions_proto.Expression(
+                    literal=expressions_proto.Expression.Literal(string=escape_char_str)
+                )
+                arguments.append(escape_literal)
             proto = expressions_proto.Expression(
                 unresolved_function=expressions_proto.Expression.UnresolvedFunction(
                     function_name=class_name.lower(),
-                    arguments=[
-                        map_logical_plan_expression(e)
-                        for e in as_java_list(exp.children())
-                    ],
+                    arguments=arguments,
                 )
             )
         case "LikeAny" | "NotLikeAny" | "LikeAll" | "NotLikeAll":
-            patterns = as_java_list(exp.patterns())
+            patterns = list(as_java_list(exp.patterns()))
             arguments = [
-                map_logical_plan_expression(e) for e in as_java_list(exp.children())
+                map_logical_plan_expression(e)
+                for e in list(as_java_list(exp.children()))
             ]
             arguments += [map_logical_plan_expression(e) for e in patterns]
             proto = expressions_proto.Expression(
@@ -384,8 +497,106 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
                 type_value = types_proto.DataType()
             elif type_name == "binary":
                 type_value = bytes(type_value)
-            elif type_name.startswith("interval "):
-                type_name = "day_time_interval"  # TODO
+            elif year_month_match := _INTERVAL_YEARMONTH_PATTERN_RE.match(type_name):
+                # Extract start and end fields for year-month intervals
+                start_field_name = year_month_match.group(1)  # 'year' or 'month'
+                end_field_name = (
+                    year_month_match.group(3)
+                    if year_month_match.group(3)
+                    else start_field_name
+                )
+
+                # Validate field names exist in mapping
+                start_field = _YEAR_MONTH_FIELD_MAP.get(start_field_name)
+                end_field = _YEAR_MONTH_FIELD_MAP.get(end_field_name)
+
+                if start_field is None:
+                    exception = AnalysisException(
+                        f"Invalid year-month interval start field: '{start_field_name}'. Expected 'year' or 'month'."
+                    )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
+                if end_field is None:
+                    exception = AnalysisException(
+                        f"Invalid year-month interval end field: '{end_field_name}'. Expected 'year' or 'month'."
+                    )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
+
+                # Validate field ordering (start_field should be <= end_field)
+                if start_field > end_field:
+                    exception = AnalysisException(
+                        f"Invalid year-month interval: start field '{start_field_name}' must come before or equal to end field '{end_field_name}'."
+                    )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
+
+                # Use extension for year-month intervals to preserve start/end field info
+                literal = expressions_proto.Expression.Literal(
+                    year_month_interval=type_value
+                )
+                any_proto = Any()
+                any_proto.Pack(
+                    snowflake_proto.ExpExtension(
+                        interval_literal=snowflake_proto.IntervalLiteralExpression(
+                            literal=literal,
+                            start_field=start_field,
+                            end_field=end_field,
+                        )
+                    )
+                )
+                return expressions_proto.Expression(extension=any_proto)
+            elif day_time_match := _INTERVAL_DAYTIME_PATTERN_RE.match(type_name):
+                # Extract start and end fields for day-time intervals
+                start_field_name = day_time_match.group(
+                    1
+                )  # 'day', 'hour', 'minute', 'second'
+                end_field_name = (
+                    day_time_match.group(3)
+                    if day_time_match.group(3)
+                    else start_field_name
+                )
+
+                # Validate field names exist in mapping
+                start_field = _DAY_TIME_FIELD_MAP.get(start_field_name)
+                end_field = _DAY_TIME_FIELD_MAP.get(end_field_name)
+
+                if start_field is None:
+                    exception = AnalysisException(
+                        f"Invalid day-time interval start field: '{start_field_name}'. Expected 'day', 'hour', 'minute', or 'second'."
+                    )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
+                if end_field is None:
+                    exception = AnalysisException(
+                        f"Invalid day-time interval end field: '{end_field_name}'. Expected 'day', 'hour', 'minute', or 'second'."
+                    )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
+
+                # Validate field ordering (start_field should be <= end_field)
+                if start_field > end_field:
+                    exception = AnalysisException(
+                        f"Invalid day-time interval: start field '{start_field_name}' must come before or equal to end field '{end_field_name}'."
+                    )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
+
+                # Use extension for day-time intervals to preserve start/end field info
+                literal = expressions_proto.Expression.Literal(
+                    day_time_interval=type_value
+                )
+                any_proto = Any()
+                any_proto.Pack(
+                    snowflake_proto.ExpExtension(
+                        interval_literal=snowflake_proto.IntervalLiteralExpression(
+                            literal=literal,
+                            start_field=start_field,
+                            end_field=end_field,
+                        )
+                    )
+                )
+                return expressions_proto.Expression(extension=any_proto)
             elif m := DECIMAL_RE.fullmatch(type_name):
                 type_name = "decimal"
                 type_value = expressions_proto.Expression.Literal.Decimal(
@@ -425,19 +636,27 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
             name = str(exp.name())
             value = get_sql_named_arg(name)
             if not value.HasField("literal_type"):
-                raise AnalysisException(f"Found an unbound parameter {name!r}")
+                exception = AnalysisException(f"Found an unbound parameter {name!r}")
+                attach_custom_error_code(exception, ErrorCodes.INVALID_SQL_SYNTAX)
+                raise exception
             proto = expressions_proto.Expression(literal=value)
         case "NamePlaceholder$":
             # This is a placeholder for an expression name to be resolved later.
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 "NamePlaceholder is not supported in SQL expressions."
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
         case "PosParameter":
             pos = exp.pos()
             try:
                 value = get_sql_pos_arg(pos)
             except KeyError:
-                raise AnalysisException(f"Found an unbound parameter at position {pos}")
+                exception = AnalysisException(
+                    f"Found an unbound parameter at position {pos}"
+                )
+                attach_custom_error_code(exception, ErrorCodes.INVALID_SQL_SYNTAX)
+                raise exception
             proto = expressions_proto.Expression(literal=value)
         case "ScalarSubquery":
             rel_proto = map_logical_plan_relation(exp.plan())
@@ -507,7 +726,7 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
                     ]
                     + [
                         map_logical_plan_expression(e)
-                        for e in as_java_list(exp.children())
+                        for e in list(as_java_list(exp.children()))
                     ],
                 )
             )
@@ -515,32 +734,42 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
             proto = map_logical_plan_expression(exp.child())
         case "UnresolvedAttribute":
             *parents, name = as_java_list(exp.nameParts())
-            plan_id = None
-            is_fully_qualified_name = False
-            if parents:
-                parent_name = ".".join(str(p) for p in parents)
-                plan_id = get_sql_plan(parent_name)
-                # If no plan_id exists, treat the column name as fully qualified by its parent.
-                if plan_id is None:
+            if not parents and name.lower() in NILARY_FUNCTIONS:
+                # this is very likely a function call without parentheses disguised as an attribute, e.g. `CURRENT_TIMESTAMP` instead of `CURRENT_TIMESTAMP()`.
+                # Note limitation: this only works when these names are not real table column names, which should be very rare and is bad practice. E.g., mytable(CURRENT_TIMESTAMP, col2, col3).
+                proto = expressions_proto.Expression(
+                    unresolved_function=expressions_proto.Expression.UnresolvedFunction(
+                        function_name=name.lower(),
+                        is_distinct=False,
+                    )
+                )
+            else:
+                plan_id = None
+                is_fully_qualified_name = False
+                if parents:
+                    parent_name = ".".join(str(p) for p in parents)
+                    plan_id = get_sql_plan(parent_name)
+                    # If no plan_id exists, treat the column name as fully qualified by its parent.
+                    if plan_id is None:
+                        # There's a difference in how Spark sql and dataframe operation passes backticks in column names.
+                        # Spark sql un-escapes the backticks instead of passing the raw string. This
+                        # logic is to escape backticks again to make it consistent with regular spark functions.
+                        parent_chain = ".".join(escape_spark_quoted(p) for p in parents)
+                        name = f"{parent_chain}.{escape_spark_quoted(name)}"
+                        is_fully_qualified_name = True
+
+                if not is_fully_qualified_name:
                     # There's a difference in how Spark sql and dataframe operation passes backticks in column names.
                     # Spark sql un-escapes the backticks instead of passing the raw string. This
                     # logic is to escape backticks again to make it consistent with regular spark functions.
-                    parent_chain = ".".join(escape_spark_quoted(p) for p in parents)
-                    name = f"{parent_chain}.{escape_spark_quoted(name)}"
-                    is_fully_qualified_name = True
-
-            if not is_fully_qualified_name:
-                # There's a difference in how Spark sql and dataframe operation passes backticks in column names.
-                # Spark sql un-escapes the backticks instead of passing the raw string. This
-                # logic is to escape backticks again to make it consistent with regular spark functions.
-                name = escape_spark_quoted(name)
+                    name = escape_spark_quoted(name)
 
-            proto = expressions_proto.Expression(
-                unresolved_attribute=expressions_proto.Expression.UnresolvedAttribute(
-                    unparsed_identifier=str(name),
-                    plan_id=plan_id,
-                ),
-            )
+                proto = expressions_proto.Expression(
+                    unresolved_attribute=expressions_proto.Expression.UnresolvedAttribute(
+                        unparsed_identifier=str(name),
+                        plan_id=plan_id,
+                    ),
+                )
         case "UnresolvedExtractValue":
             proto = expressions_proto.Expression(
                 unresolved_extract_value=expressions_proto.Expression.UnresolvedExtractValue(
@@ -550,18 +779,20 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
             )
         case "UnresolvedFunction":
             func_name = ".".join(
-                str(part) for part in as_java_list(exp.nameParts())
+                str(part) for part in list(as_java_list(exp.nameParts()))
             ).lower()
             args = [
                 map_logical_plan_expression(arg)
-                for arg in as_java_list(exp.arguments())
+                for arg in list(as_java_list(exp.arguments()))
             ]
 
             proto = apply_filter_clause(func_name, args, exp, exp.isDistinct())
         case "UnresolvedNamedLambdaVariable":
             proto = expressions_proto.Expression(
                 unresolved_named_lambda_variable=expressions_proto.Expression.UnresolvedNamedLambdaVariable(
-                    name_parts=[str(part) for part in as_java_list(exp.nameParts())],
+                    name_parts=[
+                        str(part) for part in list(as_java_list(exp.nameParts()))
+                    ],
                 )
             )
         case "UnresolvedStar":
@@ -582,9 +813,11 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
                 # Build Window expression
                 proto = get_window_expression_proto(window_spec, exp.child())
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Window specification not found {window_spec_reference!r}"
                 )
+                attach_custom_error_code(exception, ErrorCodes.INSUFFICIENT_INPUT)
+                raise exception
         case "UTF8String":
             proto = expressions_proto.Expression(
                 literal=expressions_proto.Expression.Literal(
@@ -614,13 +847,15 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
                     function_name=proto_func,
                     arguments=[
                         map_logical_plan_expression(arg)
-                        for arg in as_java_list(exp.children())
+                        for arg in list(as_java_list(exp.children()))
                     ],
                 )
             )
 
         case other:
-            raise SnowparkConnectNotImplementedError(f"Not implemented: {other}")
+            exception = SnowparkConnectNotImplementedError(f"Not implemented: {other}")
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
     return proto
 
@@ -643,11 +878,11 @@ def get_window_expression_proto(
             window_function=map_logical_plan_expression(window_function),
             partition_spec=[
                 map_logical_plan_expression(e)
-                for e in as_java_list(window_spec.partitionSpec())
+                for e in list(as_java_list(window_spec.partitionSpec()))
             ],
             order_spec=[
                 map_logical_plan_expression(e).sort_order
-                for e in as_java_list(window_spec.orderSpec())
+                for e in list(as_java_list(window_spec.orderSpec()))
             ],
             frame_spec=frame_spec_proto,
         )