# # 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 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.util 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, ) from pyspark.sql.connect.types import ( UnparsedDataType, pyspark_types_to_proto_types, proto_schema_to_pyspark_data_type, ) from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.errors.utils import current_origin from pyspark.sql.utils import is_timestamp_ntz_preferred, enum_to_value if TYPE_CHECKING: from pyspark.sql.connect.client import SparkConnectClient from pyspark.sql.connect.window import WindowSpec from pyspark.sql.connect.plan import LogicalPlan class Expression: """ Expression base class. """ def __init__(self) -> None: origin = current_origin() fragment = origin.fragment call_site = origin.call_site self.origin = None if fragment is not None and call_site is not None: self.origin = proto.Origin( python_origin=proto.PythonOrigin( fragment=origin.fragment, call_site=origin.call_site ) ) def to_plan( # type: ignore[empty-body] self, session: "SparkConnectClient" ) -> "proto.Expression": ... def __repr__(self) -> str: # type: ignore[empty-body] ... def alias(self, *alias: str, **kwargs: Any) -> "ColumnAlias": metadata = kwargs.pop("metadata", None) if len(alias) > 1 and metadata is not None: raise PySparkValueError( errorClass="ONLY_ALLOWED_FOR_SINGLE_COLUMN", messageParameters={"arg_name": "metadata"}, ) assert not kwargs, "Unexpected kwargs where passed: %s" % kwargs return ColumnAlias(self, list(alias), metadata) def name(self) -> str: # type: ignore[empty-body] ... def _create_proto_expression(self) -> proto.Expression: plan = proto.Expression() if self.origin is not None: plan.common.origin.CopyFrom(self.origin) return plan @property def children(self) -> Sequence["Expression"]: return [] def foreach(self, f: Callable[["Expression"], None]) -> None: f(self) for c in self.children: c.foreach(f) 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) @property def children(self) -> Sequence["Expression"]: children = [] for branch in self._branches: children.append(branch[0]) children.append(branch[1]) if self._else_value is not None: children.append(self._else_value) return children 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, child: Expression, alias: Sequence[str], metadata: Any): super().__init__() self._alias = alias self._metadata = metadata self._child = child def to_plan(self, session: "SparkConnectClient") -> "proto.Expression": if len(self._alias) == 1: exp = self._create_proto_expression() exp.alias.name.append(self._alias[0]) exp.alias.expr.CopyFrom(self._child.to_plan(session)) if self._metadata: exp.alias.metadata = json.dumps(self._metadata) return exp else: if self._metadata: raise PySparkValueError( errorClass="CANNOT_PROVIDE_METADATA", messageParameters={}, ) exp = self._create_proto_expression() exp.alias.name.extend(self._alias) exp.alias.expr.CopyFrom(self._child.to_plan(session)) return exp @property def children(self) -> Sequence["Expression"]: return [self._child] def __repr__(self) -> str: return f"{self._child} 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, ), ) value = enum_to_value(value) 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, np.str_)) value = str(value) 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( errorClass="UNSUPPORTED_DATA_TYPE", messageParameters={"data_type": str(dataType)}, ) self._value = value self._dataType = dataType @classmethod def _infer_type(cls, value: Any) -> DataType: value = enum_to_value(value) if value is None: return NullType() elif isinstance(value, (bytes, bytearray)): return BinaryType() elif isinstance(value, (bool, np.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( errorClass="VALUE_NOT_BETWEEN", messageParameters={ "arg_name": "value", "min": str(JVM_LONG_MIN), "max": str(JVM_SHORT_MAX), }, ) elif isinstance(value, float): return DoubleType() elif isinstance(value, (str, np.str_)): return StringType() elif isinstance(value, decimal.Decimal): return DecimalType() elif isinstance(value, datetime.datetime): return TimestampNTZType() if is_timestamp_ntz_preferred() else 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, 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> if len(value) == 0 or value[0] is None: raise PySparkTypeError( errorClass="CANNOT_INFER_ARRAY_ELEMENT_TYPE", messageParameters={}, ) return ArrayType(LiteralExpression._infer_type(value[0]), True) raise PySparkTypeError( errorClass="UNSUPPORTED_DATA_TYPE", messageParameters={"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( errorClass="UNSUPPORTED_LITERAL", messageParameters={"literal": str(literal)}, ) def to_plan(self, session: "SparkConnectClient") -> "proto.Expression": """Converts the literal expression to the literal in proto.""" expr = self._create_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( errorClass="UNSUPPORTED_DATA_TYPE", messageParameters={"data_type": str(self._dataType)}, ) return expr def __repr__(self) -> str: if self._value is None: return "NULL" elif isinstance(self._dataType, DateType): dt = DateType().fromInternal(self._value) if dt is not None and isinstance(dt, datetime.date): return dt.strftime("%Y-%m-%d") elif isinstance(self._dataType, TimestampType): ts = TimestampType().fromInternal(self._value) if ts is not None and isinstance(ts, datetime.datetime): return ts.strftime("%Y-%m-%d %H:%M:%S.%f") elif isinstance(self._dataType, TimestampNTZType): ts = TimestampNTZType().fromInternal(self._value) if ts is not None and isinstance(ts, datetime.datetime): return ts.strftime("%Y-%m-%d %H:%M:%S.%f") elif isinstance(self._dataType, DayTimeIntervalType): delta = DayTimeIntervalType().fromInternal(self._value) if delta is not None and isinstance(delta, datetime.timedelta): import pandas as pd # Note: timedelta itself does not provide isoformat method. # Both Pandas and java.time.Duration provide it, but the format # is sightly different: # java.time.Duration only applies HOURS, MINUTES, SECONDS units, # while Pandas applies all supported units. return pd.Timedelta(delta).isoformat() 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, is_metadata_column: bool = False, ) -> 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 self._is_metadata_column = is_metadata_column 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 = self._create_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 expr.unresolved_attribute.is_metadata_column = self._is_metadata_column 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], plan_id: Optional[int] = None): super().__init__() if unparsed_target is not None: assert isinstance(unparsed_target, str) and unparsed_target.endswith(".*") self._unparsed_target = unparsed_target assert plan_id is None or isinstance(plan_id, int) self._plan_id = plan_id def to_plan(self, session: "SparkConnectClient") -> "proto.Expression": expr = self._create_proto_expression() expr.unresolved_star.SetInParent() if self._unparsed_target is not None: expr.unresolved_star.unparsed_target = self._unparsed_target if self._plan_id is not None: expr.unresolved_star.plan_id = self._plan_id 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__() assert isinstance(expr, str) self._expr: str = expr def to_plan(self, session: "SparkConnectClient") -> proto.Expression: """Returns the Proto representation of the SQL expression.""" expr = self._create_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 def __repr__(self) -> str: return 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 = self._create_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 @property def children(self) -> Sequence["Expression"]: return [self._child] 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 = self._create_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 @property def children(self) -> Sequence["Expression"]: return self._args 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: Optional[Sequence[Expression]] = None, ): super().__init__() self._function_name = function_name self._deterministic = deterministic self._arguments: Sequence[Expression] = arguments or [] self._function = function def to_plan(self, session: "SparkConnectClient") -> "proto.Expression": expr = self._create_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 @property def children(self) -> Sequence["Expression"]: return self._arguments 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 = self._create_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 @property def children(self) -> Sequence["Expression"]: return [self._structExpr, self._valueExpr] def __repr__(self) -> str: return f"update_field({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 = self._create_proto_expression() expr.update_fields.struct_expression.CopyFrom(self._structExpr.to_plan(session)) expr.update_fields.field_name = self._fieldName return expr @property def children(self) -> Sequence["Expression"]: return [self._structExpr] def __repr__(self) -> str: return f"drop_field({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 = self._create_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 @property def children(self) -> Sequence["Expression"]: return [self._child, self._extraction] def __repr__(self) -> str: return f"{self._child}['{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 = self._create_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], eval_mode: Optional[str] = None, ) -> None: super().__init__() self._expr = expr assert isinstance(data_type, (DataType, str)) self._data_type = data_type if eval_mode is not None: assert isinstance(eval_mode, str) assert eval_mode in ["legacy", "ansi", "try"] self._eval_mode = eval_mode def to_plan(self, session: "SparkConnectClient") -> proto.Expression: fun = self._create_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)) if self._eval_mode is not None: if self._eval_mode == "legacy": fun.cast.eval_mode = proto.Expression.Cast.EvalMode.EVAL_MODE_LEGACY elif self._eval_mode == "ansi": fun.cast.eval_mode = proto.Expression.Cast.EvalMode.EVAL_MODE_ANSI elif self._eval_mode == "try": fun.cast.eval_mode = proto.Expression.Cast.EvalMode.EVAL_MODE_TRY return fun @property def children(self) -> Sequence["Expression"]: return [self._expr] def __repr__(self) -> str: # We cannot guarantee the string representations be exactly the same, e.g. # str(sf.col("a").cast("long")): # Column<'CAST(a AS BIGINT)'> <- Spark Classic # Column<'CAST(a AS LONG)'> <- Spark Connect if isinstance(self._data_type, DataType): str_data_type = self._data_type.simpleString().upper() else: str_data_type = str(self._data_type).upper() if self._eval_mode is not None and self._eval_mode == "try": return f"TRY_CAST({self._expr} AS {str_data_type})" else: return f"CAST({self._expr} AS {str_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 = self._create_proto_expression() expr.unresolved_named_lambda_variable.name_parts.extend(self._name_parts) return expr def __repr__(self) -> str: return ", ".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 = self._create_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 @property def children(self) -> Sequence["Expression"]: return [self._function] + self._arguments def __repr__(self) -> str: return ( f"LambdaFunction({str(self._function)}, " + f"{', '.join([str(arg) for arg in 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 = self._create_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( errorClass="VALUE_NOT_BETWEEN", messageParameters={ "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( errorClass="VALUE_NOT_BETWEEN", messageParameters={ "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 @property def children(self) -> Sequence["Expression"]: return ( [self._windowFunction] + self._windowSpec._partitionSpec + self._windowSpec._orderSpec ) 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 = self._create_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 @property def children(self) -> Sequence["Expression"]: return self._args 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}')" class NamedArgumentExpression(Expression): def __init__(self, key: str, value: Expression): super().__init__() assert isinstance(key, str) self._key = key assert isinstance(value, Expression) self._value = value def to_plan(self, session: "SparkConnectClient") -> "proto.Expression": expr = self._create_proto_expression() expr.named_argument_expression.key = self._key expr.named_argument_expression.value.CopyFrom(self._value.to_plan(session)) return expr @property def children(self) -> Sequence["Expression"]: return [self._value] def __repr__(self) -> str: return f"{self._key} => {self._value}" class SubqueryExpression(Expression): def __init__( self, plan: "LogicalPlan", subquery_type: str, partition_spec: Optional[Sequence["Expression"]] = None, order_spec: Optional[Sequence["SortOrder"]] = None, with_single_partition: Optional[bool] = None, in_subquery_values: Optional[Sequence["Expression"]] = None, ) -> None: assert isinstance(subquery_type, str) assert subquery_type in ("scalar", "exists", "table_arg", "in") super().__init__() self._plan = plan self._subquery_type = subquery_type self._partition_spec = partition_spec or [] self._order_spec = order_spec or [] self._with_single_partition = with_single_partition self._in_subquery_values = in_subquery_values or [] def to_plan(self, session: "SparkConnectClient") -> proto.Expression: expr = self._create_proto_expression() expr.subquery_expression.plan_id = self._plan._plan_id if self._subquery_type == "scalar": expr.subquery_expression.subquery_type = proto.SubqueryExpression.SUBQUERY_TYPE_SCALAR elif self._subquery_type == "exists": expr.subquery_expression.subquery_type = proto.SubqueryExpression.SUBQUERY_TYPE_EXISTS elif self._subquery_type == "table_arg": expr.subquery_expression.subquery_type = ( proto.SubqueryExpression.SUBQUERY_TYPE_TABLE_ARG ) # Populate TableArgOptions table_arg_options = expr.subquery_expression.table_arg_options if len(self._partition_spec) > 0: table_arg_options.partition_spec.extend( [p.to_plan(session) for p in self._partition_spec] ) if len(self._order_spec) > 0: table_arg_options.order_spec.extend( [o.to_plan(session).sort_order for o in self._order_spec] ) if self._with_single_partition is not None: table_arg_options.with_single_partition = self._with_single_partition elif self._subquery_type == "in": expr.subquery_expression.subquery_type = proto.SubqueryExpression.SUBQUERY_TYPE_IN expr.subquery_expression.in_subquery_values.extend( [expr.to_plan(session) for expr in self._in_subquery_values] ) return expr def __repr__(self) -> str: repr_parts = [f"plan={self._plan}", f"type={self._subquery_type}"] if self._subquery_type == "table_arg": if self._partition_spec: repr_parts.append(f"partition_spec={self._partition_spec}") if self._order_spec: repr_parts.append(f"order_spec={self._order_spec}") if self._with_single_partition is not None: repr_parts.append(f"with_single_partition={self._with_single_partition}") elif self._subquery_type == "in": repr_parts.append(f"values={self._in_subquery_values}") return f"SubqueryExpression({', '.join(repr_parts)})"