snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/sql/connect/expressions.py

@@ -0,0 +1,1056 @@
+#
+# 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.
+#
+from pyspark.sql.connect.utils import check_dependencies
+from pyspark.sql.utils import is_timestamp_ntz_preferred
+
+check_dependencies(__name__)
+
+from typing import (
+    cast,
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Union,
+    Sequence,
+    Tuple,
+    Optional,
+)
+
+import json
+import decimal
+import datetime
+import warnings
+from threading import Lock
+
+import numpy as np
+
+from pyspark.serializers import CloudPickleSerializer
+from pyspark.sql.types import (
+    _from_numpy_type,
+    DateType,
+    ArrayType,
+    NullType,
+    BooleanType,
+    BinaryType,
+    ByteType,
+    ShortType,
+    IntegerType,
+    LongType,
+    FloatType,
+    DoubleType,
+    DecimalType,
+    StringType,
+    DataType,
+    TimestampType,
+    TimestampNTZType,
+    DayTimeIntervalType,
+)
+
+import pyspark.sql.connect.proto as proto
+from pyspark.sql.connect.types import (
+    JVM_BYTE_MIN,
+    JVM_BYTE_MAX,
+    JVM_SHORT_MIN,
+    JVM_SHORT_MAX,
+    JVM_INT_MIN,
+    JVM_INT_MAX,
+    JVM_LONG_MIN,
+    JVM_LONG_MAX,
+    UnparsedDataType,
+    pyspark_types_to_proto_types,
+    proto_schema_to_pyspark_data_type,
+)
+from pyspark.errors import PySparkTypeError, PySparkValueError
+
+if TYPE_CHECKING:
+    from pyspark.sql.connect.client import SparkConnectClient
+    from pyspark.sql.connect.window import WindowSpec
+
+
+class Expression:
+    """
+    Expression base class.
+    """
+
+    def __init__(self) -> None:
+        pass
+
+    def to_plan(self, session: "SparkConnectClient") -> "proto.Expression":
+        ...
+
+    def __repr__(self) -> str:
+        ...
+
+    def alias(self, *alias: str, **kwargs: Any) -> "ColumnAlias":
+        metadata = kwargs.pop("metadata", None)
+        assert not kwargs, "Unexpected kwargs where passed: %s" % kwargs
+        return ColumnAlias(self, list(alias), metadata)
+
+    def name(self) -> str:
+        ...
+
+
+class CaseWhen(Expression):
+    def __init__(
+        self, branches: Sequence[Tuple[Expression, Expression]], else_value: Optional[Expression]
+    ):
+        super().__init__()
+
+        assert isinstance(branches, list)
+        for branch in branches:
+            assert (
+                isinstance(branch, tuple)
+                and len(branch) == 2
+                and all(isinstance(expr, Expression) for expr in branch)
+            )
+        self._branches = branches
+
+        if else_value is not None:
+            assert isinstance(else_value, Expression)
+
+        self._else_value = else_value
+
+    def to_plan(self, session: "SparkConnectClient") -> "proto.Expression":
+        args = []
+        for condition, value in self._branches:
+            args.append(condition)
+            args.append(value)
+
+        if self._else_value is not None:
+            args.append(self._else_value)
+
+        unresolved_function = UnresolvedFunction(name="when", args=args)
+
+        return unresolved_function.to_plan(session)
+
+    def __repr__(self) -> str:
+        _cases = "".join([f" WHEN {c} THEN {v}" for c, v in self._branches])
+        _else = f" ELSE {self._else_value}" if self._else_value is not None else ""
+        return "CASE" + _cases + _else + " END"
+
+
+class ColumnAlias(Expression):
+    def __init__(self, parent: Expression, alias: Sequence[str], metadata: Any):
+        super().__init__()
+
+        self._alias = alias
+        self._metadata = metadata
+        self._parent = parent
+
+    def to_plan(self, session: "SparkConnectClient") -> "proto.Expression":
+        if len(self._alias) == 1:
+            exp = proto.Expression()
+            exp.alias.name.append(self._alias[0])
+            exp.alias.expr.CopyFrom(self._parent.to_plan(session))
+
+            if self._metadata:
+                exp.alias.metadata = json.dumps(self._metadata)
+            return exp
+        else:
+            if self._metadata:
+                raise PySparkValueError(
+                    error_class="CANNOT_PROVIDE_METADATA",
+                    message_parameters={},
+                )
+            exp = proto.Expression()
+            exp.alias.name.extend(self._alias)
+            exp.alias.expr.CopyFrom(self._parent.to_plan(session))
+            return exp
+
+    def __repr__(self) -> str:
+        return f"{self._parent} AS {','.join(self._alias)}"
+
+
+class LiteralExpression(Expression):
+    """A literal expression.
+
+    The Python types are converted best effort into the relevant proto types. On the Spark Connect
+    server side, the proto types are converted to the Catalyst equivalents."""
+
+    def __init__(self, value: Any, dataType: DataType) -> None:
+        super().__init__()
+
+        assert isinstance(
+            dataType,
+            (
+                NullType,
+                BinaryType,
+                BooleanType,
+                ByteType,
+                ShortType,
+                IntegerType,
+                LongType,
+                FloatType,
+                DoubleType,
+                DecimalType,
+                StringType,
+                DateType,
+                TimestampType,
+                TimestampNTZType,
+                DayTimeIntervalType,
+                ArrayType,
+            ),
+        )
+
+        if isinstance(dataType, NullType):
+            assert value is None
+
+        if value is not None:
+            if isinstance(dataType, BinaryType):
+                assert isinstance(value, (bytes, bytearray))
+            elif isinstance(dataType, BooleanType):
+                assert isinstance(value, (bool, np.bool_))
+                value = bool(value)
+            elif isinstance(dataType, ByteType):
+                assert isinstance(value, (int, np.int8))
+                assert JVM_BYTE_MIN <= int(value) <= JVM_BYTE_MAX
+                value = int(value)
+            elif isinstance(dataType, ShortType):
+                assert isinstance(value, (int, np.int8, np.int16))
+                assert JVM_SHORT_MIN <= int(value) <= JVM_SHORT_MAX
+                value = int(value)
+            elif isinstance(dataType, IntegerType):
+                assert isinstance(value, (int, np.int8, np.int16, np.int32))
+                assert JVM_INT_MIN <= int(value) <= JVM_INT_MAX
+                value = int(value)
+            elif isinstance(dataType, LongType):
+                assert isinstance(value, (int, np.int8, np.int16, np.int32, np.int64))
+                assert JVM_LONG_MIN <= int(value) <= JVM_LONG_MAX
+                value = int(value)
+            elif isinstance(dataType, FloatType):
+                assert isinstance(value, (float, np.float32))
+                value = float(value)
+            elif isinstance(dataType, DoubleType):
+                assert isinstance(value, (float, np.float32, np.float64))
+                value = float(value)
+            elif isinstance(dataType, DecimalType):
+                assert isinstance(value, decimal.Decimal)
+            elif isinstance(dataType, StringType):
+                assert isinstance(value, str)
+            elif isinstance(dataType, DateType):
+                assert isinstance(value, (datetime.date, datetime.datetime))
+                if isinstance(value, datetime.date):
+                    value = DateType().toInternal(value)
+                else:
+                    value = DateType().toInternal(value.date())
+            elif isinstance(dataType, TimestampType):
+                assert isinstance(value, datetime.datetime)
+                value = TimestampType().toInternal(value)
+            elif isinstance(dataType, TimestampNTZType):
+                assert isinstance(value, datetime.datetime)
+                value = TimestampNTZType().toInternal(value)
+            elif isinstance(dataType, DayTimeIntervalType):
+                assert isinstance(value, datetime.timedelta)
+                value = DayTimeIntervalType().toInternal(value)
+                assert value is not None
+            elif isinstance(dataType, ArrayType):
+                assert isinstance(value, list)
+            else:
+                raise PySparkTypeError(
+                    error_class="UNSUPPORTED_DATA_TYPE",
+                    message_parameters={"data_type": str(dataType)},
+                )
+
+        self._value = value
+        self._dataType = dataType
+
+    @classmethod
+    def _infer_type(cls, value: Any) -> DataType:
+        if value is None:
+            return NullType()
+        elif isinstance(value, (bytes, bytearray)):
+            return BinaryType()
+        elif isinstance(value, bool):
+            return BooleanType()
+        elif isinstance(value, int):
+            if JVM_INT_MIN <= value <= JVM_INT_MAX:
+                return IntegerType()
+            elif JVM_LONG_MIN <= value <= JVM_LONG_MAX:
+                return LongType()
+            else:
+                raise PySparkValueError(
+                    error_class="VALUE_NOT_BETWEEN",
+                    message_parameters={
+                        "arg_name": "value",
+                        "min": str(JVM_LONG_MIN),
+                        "max": str(JVM_SHORT_MAX),
+                    },
+                )
+        elif isinstance(value, float):
+            return DoubleType()
+        elif isinstance(value, str):
+            return StringType()
+        elif isinstance(value, decimal.Decimal):
+            return DecimalType()
+        elif isinstance(value, datetime.datetime) and is_timestamp_ntz_preferred():
+            return TimestampNTZType()
+        elif isinstance(value, datetime.datetime):
+            return TimestampType()
+        elif isinstance(value, datetime.date):
+            return DateType()
+        elif isinstance(value, datetime.timedelta):
+            return DayTimeIntervalType()
+        elif isinstance(value, np.generic):
+            dt = _from_numpy_type(value.dtype)
+            if dt is not None:
+                return dt
+            elif isinstance(value, np.bool_):
+                return BooleanType()
+        elif isinstance(value, list):
+            # follow the 'infer_array_from_first_element' strategy in 'sql.types._infer_type'
+            # right now, it's dedicated for pyspark.ml params like array<...>, array<array<...>>
+            if len(value) == 0:
+                raise PySparkValueError(
+                    error_class="CANNOT_BE_EMPTY",
+                    message_parameters={"item": "value"},
+                )
+            first = value[0]
+            if first is None:
+                raise PySparkTypeError(
+                    error_class="CANNOT_INFER_ARRAY_TYPE",
+                    message_parameters={},
+                )
+            return ArrayType(LiteralExpression._infer_type(first), True)
+
+        raise PySparkTypeError(
+            error_class="UNSUPPORTED_DATA_TYPE",
+            message_parameters={"data_type": type(value).__name__},
+        )
+
+    @classmethod
+    def _from_value(cls, value: Any) -> "LiteralExpression":
+        return LiteralExpression(value=value, dataType=LiteralExpression._infer_type(value))
+
+    @classmethod
+    def _to_value(
+        cls, literal: "proto.Expression.Literal", dataType: Optional[DataType] = None
+    ) -> Any:
+        if literal.HasField("null"):
+            return None
+        elif literal.HasField("binary"):
+            assert dataType is None or isinstance(dataType, BinaryType)
+            return literal.binary
+        elif literal.HasField("boolean"):
+            assert dataType is None or isinstance(dataType, BooleanType)
+            return literal.boolean
+        elif literal.HasField("byte"):
+            assert dataType is None or isinstance(dataType, ByteType)
+            return literal.byte
+        elif literal.HasField("short"):
+            assert dataType is None or isinstance(dataType, ShortType)
+            return literal.short
+        elif literal.HasField("integer"):
+            assert dataType is None or isinstance(dataType, IntegerType)
+            return literal.integer
+        elif literal.HasField("long"):
+            assert dataType is None or isinstance(dataType, LongType)
+            return literal.long
+        elif literal.HasField("float"):
+            assert dataType is None or isinstance(dataType, FloatType)
+            return literal.float
+        elif literal.HasField("double"):
+            assert dataType is None or isinstance(dataType, DoubleType)
+            return literal.double
+        elif literal.HasField("decimal"):
+            assert dataType is None or isinstance(dataType, DecimalType)
+            return decimal.Decimal(literal.decimal.value)
+        elif literal.HasField("string"):
+            assert dataType is None or isinstance(dataType, StringType)
+            return literal.string
+        elif literal.HasField("date"):
+            assert dataType is None or isinstance(dataType, DataType)
+            return DateType().fromInternal(literal.date)
+        elif literal.HasField("timestamp"):
+            assert dataType is None or isinstance(dataType, TimestampType)
+            return TimestampType().fromInternal(literal.timestamp)
+        elif literal.HasField("timestamp_ntz"):
+            assert dataType is None or isinstance(dataType, TimestampNTZType)
+            return TimestampNTZType().fromInternal(literal.timestamp_ntz)
+        elif literal.HasField("day_time_interval"):
+            assert dataType is None or isinstance(dataType, DayTimeIntervalType)
+            return DayTimeIntervalType().fromInternal(literal.day_time_interval)
+        elif literal.HasField("array"):
+            elementType = proto_schema_to_pyspark_data_type(literal.array.element_type)
+            if dataType is not None:
+                assert isinstance(dataType, ArrayType)
+                assert elementType == dataType.elementType
+            return [LiteralExpression._to_value(v, elementType) for v in literal.array.elements]
+
+        raise PySparkTypeError(
+            error_class="UNSUPPORTED_LITERAL",
+            message_parameters={"literal": str(literal)},
+        )
+
+    def to_plan(self, session: "SparkConnectClient") -> "proto.Expression":
+        """Converts the literal expression to the literal in proto."""
+
+        expr = proto.Expression()
+
+        if self._value is None:
+            expr.literal.null.CopyFrom(pyspark_types_to_proto_types(self._dataType))
+        elif isinstance(self._dataType, BinaryType):
+            expr.literal.binary = bytes(self._value)
+        elif isinstance(self._dataType, BooleanType):
+            expr.literal.boolean = bool(self._value)
+        elif isinstance(self._dataType, ByteType):
+            expr.literal.byte = int(self._value)
+        elif isinstance(self._dataType, ShortType):
+            expr.literal.short = int(self._value)
+        elif isinstance(self._dataType, IntegerType):
+            expr.literal.integer = int(self._value)
+        elif isinstance(self._dataType, LongType):
+            expr.literal.long = int(self._value)
+        elif isinstance(self._dataType, FloatType):
+            expr.literal.float = float(self._value)
+        elif isinstance(self._dataType, DoubleType):
+            expr.literal.double = float(self._value)
+        elif isinstance(self._dataType, DecimalType):
+            expr.literal.decimal.value = str(self._value)
+            expr.literal.decimal.precision = self._dataType.precision
+            expr.literal.decimal.scale = self._dataType.scale
+        elif isinstance(self._dataType, StringType):
+            expr.literal.string = str(self._value)
+        elif isinstance(self._dataType, DateType):
+            expr.literal.date = int(self._value)
+        elif isinstance(self._dataType, TimestampType):
+            expr.literal.timestamp = int(self._value)
+        elif isinstance(self._dataType, TimestampNTZType):
+            expr.literal.timestamp_ntz = int(self._value)
+        elif isinstance(self._dataType, DayTimeIntervalType):
+            expr.literal.day_time_interval = int(self._value)
+        elif isinstance(self._dataType, ArrayType):
+            element_type = self._dataType.elementType
+            expr.literal.array.element_type.CopyFrom(pyspark_types_to_proto_types(element_type))
+            for v in self._value:
+                expr.literal.array.elements.append(
+                    LiteralExpression(v, element_type).to_plan(session).literal
+                )
+        else:
+            raise PySparkTypeError(
+                error_class="UNSUPPORTED_DATA_TYPE",
+                message_parameters={"data_type": str(self._dataType)},
+            )
+
+        return expr
+
+    def __repr__(self) -> str:
+        return f"{self._value}"
+
+
+class ColumnReference(Expression):
+    """Represents a column reference. There is no guarantee that this column
+    actually exists. In the context of this project, we refer by its name and
+    treat it as an unresolved attribute. Attributes that have the same fully
+    qualified name are identical"""
+
+    def __init__(self, unparsed_identifier: str, plan_id: Optional[int] = None) -> None:
+        super().__init__()
+        assert isinstance(unparsed_identifier, str)
+        self._unparsed_identifier = unparsed_identifier
+
+        assert plan_id is None or isinstance(plan_id, int)
+        self._plan_id = plan_id
+
+    def name(self) -> str:
+        """Returns the qualified name of the column reference."""
+        return self._unparsed_identifier
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+        """Returns the Proto representation of the expression."""
+        expr = proto.Expression()
+        expr.unresolved_attribute.unparsed_identifier = self._unparsed_identifier
+        if self._plan_id is not None:
+            expr.unresolved_attribute.plan_id = self._plan_id
+        return expr
+
+    def __repr__(self) -> str:
+        return f"{self._unparsed_identifier}"
+
+    def __eq__(self, other: Any) -> bool:
+        return (
+            other is not None
+            and isinstance(other, ColumnReference)
+            and other._unparsed_identifier == self._unparsed_identifier
+        )
+
+
+class UnresolvedStar(Expression):
+    def __init__(self, unparsed_target: Optional[str]):
+        super().__init__()
+
+        if unparsed_target is not None:
+            assert isinstance(unparsed_target, str) and unparsed_target.endswith(".*")
+
+        self._unparsed_target = unparsed_target
+
+    def to_plan(self, session: "SparkConnectClient") -> "proto.Expression":
+        expr = proto.Expression()
+        expr.unresolved_star.SetInParent()
+        if self._unparsed_target is not None:
+            expr.unresolved_star.unparsed_target = self._unparsed_target
+        return expr
+
+    def __repr__(self) -> str:
+        if self._unparsed_target is not None:
+            return f"unresolvedstar({self._unparsed_target})"
+        else:
+            return "unresolvedstar()"
+
+    def __eq__(self, other: Any) -> bool:
+        return (
+            other is not None
+            and isinstance(other, UnresolvedStar)
+            and other._unparsed_target == self._unparsed_target
+        )
+
+
+class SQLExpression(Expression):
+    """Returns Expression which contains a string which is a SQL expression
+    and server side will parse it by Catalyst
+    """
+
+    def __init__(self, expr: str) -> None:
+        super().__init__()
+        self._expr: str = expr
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+        """Returns the Proto representation of the SQL expression."""
+        expr = proto.Expression()
+        expr.expression_string.expression = self._expr
+        return expr
+
+    def __eq__(self, other: Any) -> bool:
+        return other is not None and isinstance(other, SQLExpression) and other._expr == self._expr
+
+
+class SortOrder(Expression):
+    def __init__(self, child: Expression, ascending: bool = True, nullsFirst: bool = True) -> None:
+        super().__init__()
+        self._child = child
+        self._ascending = ascending
+        self._nullsFirst = nullsFirst
+
+    def __repr__(self) -> str:
+        return (
+            str(self._child)
+            + (" ASC" if self._ascending else " DESC")
+            + (" NULLS FIRST" if self._nullsFirst else " NULLS LAST")
+        )
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+        sort = proto.Expression()
+        sort.sort_order.child.CopyFrom(self._child.to_plan(session))
+
+        if self._ascending:
+            sort.sort_order.direction = (
+                proto.Expression.SortOrder.SortDirection.SORT_DIRECTION_ASCENDING
+            )
+        else:
+            sort.sort_order.direction = (
+                proto.Expression.SortOrder.SortDirection.SORT_DIRECTION_DESCENDING
+            )
+
+        if self._nullsFirst:
+            sort.sort_order.null_ordering = proto.Expression.SortOrder.NullOrdering.SORT_NULLS_FIRST
+        else:
+            sort.sort_order.null_ordering = proto.Expression.SortOrder.NullOrdering.SORT_NULLS_LAST
+
+        return sort
+
+
+class UnresolvedFunction(Expression):
+    def __init__(
+        self,
+        name: str,
+        args: Sequence["Expression"],
+        is_distinct: bool = False,
+    ) -> None:
+        super().__init__()
+
+        assert isinstance(name, str)
+        self._name = name
+
+        assert isinstance(args, list) and all(isinstance(arg, Expression) for arg in args)
+        self._args = args
+
+        assert isinstance(is_distinct, bool)
+        self._is_distinct = is_distinct
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+        fun = proto.Expression()
+        fun.unresolved_function.function_name = self._name
+        if len(self._args) > 0:
+            fun.unresolved_function.arguments.extend([arg.to_plan(session) for arg in self._args])
+        fun.unresolved_function.is_distinct = self._is_distinct
+        return fun
+
+    def __repr__(self) -> str:
+        # Default print handling:
+        if self._is_distinct:
+            return f"{self._name}(distinct {', '.join([str(arg) for arg in self._args])})"
+        else:
+            return f"{self._name}({', '.join([str(arg) for arg in self._args])})"
+
+
+class PythonUDF:
+    """Represents a Python user-defined function."""
+
+    def __init__(
+        self,
+        output_type: Union[DataType, str],
+        eval_type: int,
+        func: Callable[..., Any],
+        python_ver: str,
+    ) -> None:
+        self._output_type: DataType = (
+            UnparsedDataType(output_type) if isinstance(output_type, str) else output_type
+        )
+        self._eval_type = eval_type
+        self._func = func
+        self._python_ver = python_ver
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.PythonUDF:
+        if isinstance(self._output_type, UnparsedDataType):
+            parsed = session._analyze(
+                method="ddl_parse", ddl_string=self._output_type.data_type_string
+            ).parsed
+            assert isinstance(parsed, DataType)
+            output_type = parsed
+        else:
+            output_type = self._output_type
+        expr = proto.PythonUDF()
+        expr.output_type.CopyFrom(pyspark_types_to_proto_types(output_type))
+        expr.eval_type = self._eval_type
+        expr.command = CloudPickleSerializer().dumps((self._func, output_type))
+        expr.python_ver = self._python_ver
+        return expr
+
+    def __repr__(self) -> str:
+        return f"{self._output_type}, {self._eval_type}, {self._func}, f{self._python_ver}"
+
+
+class JavaUDF:
+    """Represents a Java (aggregate) user-defined function."""
+
+    def __init__(
+        self,
+        class_name: str,
+        output_type: Optional[Union[DataType, str]] = None,
+        aggregate: bool = False,
+    ) -> None:
+        self._class_name = class_name
+        self._output_type: Optional[DataType] = (
+            UnparsedDataType(output_type) if isinstance(output_type, str) else output_type
+        )
+        self._aggregate = aggregate
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.JavaUDF:
+        expr = proto.JavaUDF()
+        expr.class_name = self._class_name
+        if self._output_type is not None:
+            expr.output_type.CopyFrom(pyspark_types_to_proto_types(self._output_type))
+        expr.aggregate = self._aggregate
+        return expr
+
+    def __repr__(self) -> str:
+        return f"{self._class_name}, {self._output_type}"
+
+
+class CommonInlineUserDefinedFunction(Expression):
+    """Represents a user-defined function with an inlined defined function body of any programming
+    languages."""
+
+    def __init__(
+        self,
+        function_name: str,
+        function: Union[PythonUDF, JavaUDF],
+        deterministic: bool = False,
+        arguments: Sequence[Expression] = [],
+    ):
+        super().__init__()
+        self._function_name = function_name
+        self._deterministic = deterministic
+        self._arguments = arguments
+        self._function = function
+
+    def to_plan(self, session: "SparkConnectClient") -> "proto.Expression":
+        expr = proto.Expression()
+        expr.common_inline_user_defined_function.function_name = self._function_name
+        expr.common_inline_user_defined_function.deterministic = self._deterministic
+        if len(self._arguments) > 0:
+            expr.common_inline_user_defined_function.arguments.extend(
+                [arg.to_plan(session) for arg in self._arguments]
+            )
+        expr.common_inline_user_defined_function.python_udf.CopyFrom(
+            cast(proto.PythonUDF, self._function.to_plan(session))
+        )
+        return expr
+
+    def to_plan_udf(self, session: "SparkConnectClient") -> "proto.CommonInlineUserDefinedFunction":
+        """Compared to `to_plan`, it returns a CommonInlineUserDefinedFunction instead of an
+        Expression."""
+        expr = proto.CommonInlineUserDefinedFunction()
+        expr.function_name = self._function_name
+        expr.deterministic = self._deterministic
+        if len(self._arguments) > 0:
+            expr.arguments.extend([arg.to_plan(session) for arg in self._arguments])
+        expr.python_udf.CopyFrom(cast(proto.PythonUDF, self._function.to_plan(session)))
+        return expr
+
+    def to_plan_judf(
+        self, session: "SparkConnectClient"
+    ) -> "proto.CommonInlineUserDefinedFunction":
+        expr = proto.CommonInlineUserDefinedFunction()
+        expr.function_name = self._function_name
+        expr.java_udf.CopyFrom(cast(proto.JavaUDF, self._function.to_plan(session)))
+        return expr
+
+    def __repr__(self) -> str:
+        return f"{self._function_name}({', '.join([str(arg) for arg in self._arguments])})"
+
+
+class WithField(Expression):
+    def __init__(
+        self,
+        structExpr: Expression,
+        fieldName: str,
+        valueExpr: Expression,
+    ) -> None:
+        super().__init__()
+
+        assert isinstance(structExpr, Expression)
+        self._structExpr = structExpr
+
+        assert isinstance(fieldName, str)
+        self._fieldName = fieldName
+
+        assert isinstance(valueExpr, Expression)
+        self._valueExpr = valueExpr
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+        expr = proto.Expression()
+        expr.update_fields.struct_expression.CopyFrom(self._structExpr.to_plan(session))
+        expr.update_fields.field_name = self._fieldName
+        expr.update_fields.value_expression.CopyFrom(self._valueExpr.to_plan(session))
+        return expr
+
+    def __repr__(self) -> str:
+        return f"WithField({self._structExpr}, {self._fieldName}, {self._valueExpr})"
+
+
+class DropField(Expression):
+    def __init__(
+        self,
+        structExpr: Expression,
+        fieldName: str,
+    ) -> None:
+        super().__init__()
+
+        assert isinstance(structExpr, Expression)
+        self._structExpr = structExpr
+
+        assert isinstance(fieldName, str)
+        self._fieldName = fieldName
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+        expr = proto.Expression()
+        expr.update_fields.struct_expression.CopyFrom(self._structExpr.to_plan(session))
+        expr.update_fields.field_name = self._fieldName
+        return expr
+
+    def __repr__(self) -> str:
+        return f"DropField({self._structExpr}, {self._fieldName})"
+
+
+class UnresolvedExtractValue(Expression):
+    def __init__(
+        self,
+        child: Expression,
+        extraction: Expression,
+    ) -> None:
+        super().__init__()
+
+        assert isinstance(child, Expression)
+        self._child = child
+
+        assert isinstance(extraction, Expression)
+        self._extraction = extraction
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+        expr = proto.Expression()
+        expr.unresolved_extract_value.child.CopyFrom(self._child.to_plan(session))
+        expr.unresolved_extract_value.extraction.CopyFrom(self._extraction.to_plan(session))
+        return expr
+
+    def __repr__(self) -> str:
+        return f"UnresolvedExtractValue({str(self._child)}, {str(self._extraction)})"
+
+
+class UnresolvedRegex(Expression):
+    def __init__(self, col_name: str, plan_id: Optional[int] = None) -> None:
+        super().__init__()
+
+        assert isinstance(col_name, str)
+        self.col_name = col_name
+
+        assert plan_id is None or isinstance(plan_id, int)
+        self._plan_id = plan_id
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+        expr = proto.Expression()
+        expr.unresolved_regex.col_name = self.col_name
+        if self._plan_id is not None:
+            expr.unresolved_regex.plan_id = self._plan_id
+        return expr
+
+    def __repr__(self) -> str:
+        return f"UnresolvedRegex({self.col_name})"
+
+
+class CastExpression(Expression):
+    def __init__(
+        self,
+        expr: Expression,
+        data_type: Union[DataType, str],
+    ) -> None:
+        super().__init__()
+        self._expr = expr
+        self._data_type = data_type
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+        fun = proto.Expression()
+        fun.cast.expr.CopyFrom(self._expr.to_plan(session))
+        if isinstance(self._data_type, str):
+            fun.cast.type_str = self._data_type
+        else:
+            fun.cast.type.CopyFrom(pyspark_types_to_proto_types(self._data_type))
+        return fun
+
+    def __repr__(self) -> str:
+        return f"({self._expr} ({self._data_type}))"
+
+
+class UnresolvedNamedLambdaVariable(Expression):
+
+    _lock: Lock = Lock()
+    _nextVarNameId: int = 0
+
+    def __init__(
+        self,
+        name_parts: Sequence[str],
+    ) -> None:
+        super().__init__()
+
+        assert (
+            isinstance(name_parts, list)
+            and len(name_parts) > 0
+            and all(isinstance(p, str) for p in name_parts)
+        )
+
+        self._name_parts = name_parts
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+        expr = proto.Expression()
+        expr.unresolved_named_lambda_variable.name_parts.extend(self._name_parts)
+        return expr
+
+    def __repr__(self) -> str:
+        return f"(UnresolvedNamedLambdaVariable({', '.join(self._name_parts)})"
+
+    @staticmethod
+    def fresh_var_name(name: str) -> str:
+        assert isinstance(name, str) and str != ""
+
+        _id: Optional[int] = None
+
+        with UnresolvedNamedLambdaVariable._lock:
+            _id = UnresolvedNamedLambdaVariable._nextVarNameId
+            UnresolvedNamedLambdaVariable._nextVarNameId += 1
+
+        assert _id is not None
+
+        return f"{name}_{_id}"
+
+
+class LambdaFunction(Expression):
+    def __init__(
+        self,
+        function: Expression,
+        arguments: Sequence[UnresolvedNamedLambdaVariable],
+    ) -> None:
+        super().__init__()
+
+        assert isinstance(function, Expression)
+
+        assert (
+            isinstance(arguments, list)
+            and len(arguments) > 0
+            and all(isinstance(arg, UnresolvedNamedLambdaVariable) for arg in arguments)
+        )
+
+        self._function = function
+        self._arguments = arguments
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+        expr = proto.Expression()
+        expr.lambda_function.function.CopyFrom(self._function.to_plan(session))
+        expr.lambda_function.arguments.extend(
+            [arg.to_plan(session).unresolved_named_lambda_variable for arg in self._arguments]
+        )
+        return expr
+
+    def __repr__(self) -> str:
+        return f"(LambdaFunction({str(self._function)}, {', '.join(self._arguments)})"
+
+
+class WindowExpression(Expression):
+    def __init__(
+        self,
+        windowFunction: Expression,
+        windowSpec: "WindowSpec",
+    ) -> None:
+        super().__init__()
+
+        from pyspark.sql.connect.window import WindowSpec
+
+        assert windowFunction is not None and isinstance(windowFunction, Expression)
+
+        assert windowSpec is not None and isinstance(windowSpec, WindowSpec)
+
+        self._windowFunction = windowFunction
+
+        self._windowSpec = windowSpec
+
+    def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+        expr = proto.Expression()
+
+        expr.window.window_function.CopyFrom(self._windowFunction.to_plan(session))
+
+        if len(self._windowSpec._partitionSpec) > 0:
+            expr.window.partition_spec.extend(
+                [p.to_plan(session) for p in self._windowSpec._partitionSpec]
+            )
+        else:
+            warnings.warn(
+                "WARN WindowExpression: No Partition Defined for Window operation! "
+                "Moving all data to a single partition, this can cause serious "
+                "performance degradation."
+            )
+
+        if len(self._windowSpec._orderSpec) > 0:
+            expr.window.order_spec.extend(
+                [s.to_plan(session).sort_order for s in self._windowSpec._orderSpec]
+            )
+
+        if self._windowSpec._frame is not None:
+            if self._windowSpec._frame._isRowFrame:
+                expr.window.frame_spec.frame_type = (
+                    proto.Expression.Window.WindowFrame.FrameType.FRAME_TYPE_ROW
+                )
+
+                start = self._windowSpec._frame._start
+                if start == 0:
+                    expr.window.frame_spec.lower.current_row = True
+                elif start == JVM_LONG_MIN:
+                    expr.window.frame_spec.lower.unbounded = True
+                elif JVM_INT_MIN <= start <= JVM_INT_MAX:
+                    expr.window.frame_spec.lower.value.literal.integer = start
+                else:
+                    raise PySparkValueError(
+                        error_class="VALUE_NOT_BETWEEN",
+                        message_parameters={
+                            "arg_name": "start",
+                            "min": str(JVM_INT_MIN),
+                            "max": str(JVM_INT_MAX),
+                        },
+                    )
+
+                end = self._windowSpec._frame._end
+                if end == 0:
+                    expr.window.frame_spec.upper.current_row = True
+                elif end == JVM_LONG_MAX:
+                    expr.window.frame_spec.upper.unbounded = True
+                elif JVM_INT_MIN <= end <= JVM_INT_MAX:
+                    expr.window.frame_spec.upper.value.literal.integer = end
+                else:
+                    raise PySparkValueError(
+                        error_class="VALUE_NOT_BETWEEN",
+                        message_parameters={
+                            "arg_name": "end",
+                            "min": str(JVM_INT_MIN),
+                            "max": str(JVM_INT_MAX),
+                        },
+                    )
+
+            else:
+                expr.window.frame_spec.frame_type = (
+                    proto.Expression.Window.WindowFrame.FrameType.FRAME_TYPE_RANGE
+                )
+
+                start = self._windowSpec._frame._start
+                if start == 0:
+                    expr.window.frame_spec.lower.current_row = True
+                elif start == JVM_LONG_MIN:
+                    expr.window.frame_spec.lower.unbounded = True
+                else:
+                    expr.window.frame_spec.lower.value.literal.long = start
+
+                end = self._windowSpec._frame._end
+                if end == 0:
+                    expr.window.frame_spec.upper.current_row = True
+                elif end == JVM_LONG_MAX:
+                    expr.window.frame_spec.upper.unbounded = True
+                else:
+                    expr.window.frame_spec.upper.value.literal.long = end
+
+        return expr
+
+    def __repr__(self) -> str:
+        return f"WindowExpression({str(self._windowFunction)}, ({str(self._windowSpec)}))"
+
+
+class DistributedSequenceID(Expression):
+    def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+        unresolved_function = UnresolvedFunction(name="distributed_sequence_id", args=[])
+        return unresolved_function.to_plan(session)
+
+    def __repr__(self) -> str:
+        return "DistributedSequenceID()"
+
+
+class CallFunction(Expression):
+    def __init__(self, name: str, args: Sequence["Expression"]):
+        super().__init__()
+
+        assert isinstance(name, str)
+        self._name = name
+
+        assert isinstance(args, list) and all(isinstance(arg, Expression) for arg in args)
+        self._args = args
+
+    def to_plan(self, session: "SparkConnectClient") -> "proto.Expression":
+        expr = proto.Expression()
+        expr.call_function.function_name = self._name
+        if len(self._args) > 0:
+            expr.call_function.arguments.extend([arg.to_plan(session) for arg in self._args])
+        return expr
+
+    def __repr__(self) -> str:
+        if len(self._args) > 0:
+            return f"CallFunction('{self._name}', {', '.join([str(arg) for arg in self._args])})"
+        else:
+            return f"CallFunction('{self._name}')"