devin-petersohn commented on code in PR #56327: URL: https://github.com/apache/spark/pull/56327#discussion_r3686763772
########## python/pyspark/sql/transpile.py: ########## @@ -0,0 +1,1098 @@ +# +# 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. +# +""" +Experimental tools for transpiling UDFS. + +Transpilation is only attempted when both +``spark.sql.experimental.optimizer.transpilePyUDFs=true`` and +``spark.sql.ansi.enabled=true``. The generated Catalyst expressions +target ANSI-mode SQL semantics (overflow raises, divide-by-zero raises, +etc.); running them under non-ANSI mode would silently diverge from the +Python interpretation in ways we don't currently track. If you flip +transpilation on with ANSI off the UDF will fall back to interpreted +Python execution and a warning is logged at UDF construction time. + +Python's ``+`` and ``*`` are overloaded for text (concat / repeat), so an +untyped parameter is transpiled into one option per input-type category +(numeric and string) and the JVM picks the one matching the bound column +types -- falling back to interpreted Python when none fit. Annotating the +UDF's parameters (e.g. ``def f(a: int, b: str)``) pins each category and +keeps the option matrix small; prefer doing so. To bound plan growth, +functions with more than three untyped parameters only emit the +all-numeric and all-string variants. +""" + +import ast +from typing import Any, Callable, List, Optional, Tuple, TYPE_CHECKING +import inspect +import itertools +import textwrap +from pyspark.errors import UnsupportedOperationException +from pyspark.sql.column import Column +from pyspark.sql.types import ( + BinaryType, + BooleanType, + DataType, + DecimalType, + NumericType, + StringType, +) +from pyspark.sql.functions import ( + abs as _abs, + coalesce, + col, + concat, + lit, + pmod, + raise_error, + repeat, + when, +) + + +if TYPE_CHECKING: + from pyspark.sql import SparkSession + from pyspark.sql._typing import DataTypeOrString + + +class AbstractTranspiler(object): + """Base class for transpilers. All experimental.""" + + varieties: dict[str, type["AbstractTranspiler"]] = {} + # Specify the "friendly" name a user can add to spark.sql.experimental.optimizer.pyTranspilers + # to enable this transpiler. + variety: str = "" + + @classmethod + def register(cls) -> None: + AbstractTranspiler.varieties[cls.variety] = cls + + def _transpile_from_ast( + self, + src: Optional[str], + ast_info: ast.AST, + function_ast: ast.FunctionDef, + params: List[str], + returnType: "DataTypeOrString", + param_categories: Optional[dict] = None, + ) -> Optional[Column]: + pass + + +def _is_definitely_basic_type(node: ast.AST) -> bool: + """ + Return True when ``node`` is statically guaranteed to produce a Python + basic/builtin type (int, float, str, bool, None, lists, etc.). + All ast.Name's are treated as basic types for now this will need to be updated + if/when we add free variables / closures to transpilation. + """ + match node: + case ast.Constant(): + return True + case ast.BinOp(left=left, right=right): + return _is_definitely_basic_type(left) and _is_definitely_basic_type(right) + case ast.UnaryOp(operand=operand): + return _is_definitely_basic_type(operand) + case ast.Name(): + return True + case _: + return False + + +def _is_definitely_boolean(node: ast.AST) -> bool: + """Return True when ``node`` is statically guaranteed to produce a Python + ``bool`` (or ``None``, which round-trips through ``coalesce``). + + Used to gate ``if``/ternary lowering: we only allow the test expression + into Catalyst's ``when(coalesce(test, false), ...)`` form when it provably + produces a boolean. Everything else (bare Name, arithmetic, function calls, + subscript, ...) must force a fallback to interpreted Python instead of + silently diverging. + """ + match node: + case ast.Constant(value=v): + return v is None or isinstance(v, bool) + case ast.Compare(left=left, comparators=comparators): + # All comparison operators of simple types bool + return all(_is_definitely_basic_type(v) for v in comparators + [left]) + case ast.BoolOp(values=values): + return all(_is_definitely_boolean(v) for v in values) + case ast.UnaryOp(op=ast.Not()): + # `not x` always produces bool. + return True + case ast.IfExp(body=body, orelse=orelse): + # Ternary is boolean only if both branches are. + return _is_definitely_boolean(body) and _is_definitely_boolean(orelse) + case _: + return False + + +class CatalystTranspiler(AbstractTranspiler): + """Transpiler that attempts to convert a Python UDF into native Spark SQL expressions.""" + + variety = "catalyst" + + # TODO (SPARK-55218): handle implicit-None return bodies like + # ``def f(x): x + x`` -- no return statement means return None; + # we should lower to lit(None) and optionally warn since it's + # likely a mistake. + def _convert_branch(self, params: List[str], statements: List[ast.stmt], slot: str) -> Column: + """Lower a single-statement if-body / if-else block. + + ``slot`` is just used to disambiguate the multi-statement error + message between the body and the else arm. + """ + if len(statements) > 1: + raise UnsupportedOperationException( + f"if statements with more than one expression in the {slot} " + "are not currently supported by the transpiler" + ) + if len(statements) == 0: + return lit(None) + return self._convert_chunk(params, statements[0]) + + def _safe_category(self, params: List[str], node: Optional[ast.AST]) -> Optional[str]: + """Best-effort input-type category for an if/else branch, or ``None`` when + it can't be pinned down statically. + + Used only to compare the two branches of an if/ternary. A ``None`` result + means "treat as compatible" (don't force a fallback): the node is absent, + is a bare ``None`` literal (which unifies with any branch type via + ``coalesce``/``Cast``), or its category can't be determined. + """ + if node is None: + return None + # If-statement branches arrive as ``Return`` statements; classify the + # returned value, not the statement wrapper (``_is_definitely_boolean`` + # has no ``Return`` case, so without this a boolean-returning branch + # would fall through to ``_category``'s numeric catch-all). + if isinstance(node, ast.Return): + return self._safe_category(params, node.value) + # An if-statement's category is its branches' common category (the + # ``_category`` catch-all would mislabel every ``ast.If`` "numeric"). + # Mismatched branches return None ("can't be pinned down"); the + # branch-compatibility check in ``_convert_if_like`` raises for them. + if isinstance(node, ast.If): + body_c = self._safe_category(params, node.body[0]) if node.body else None + else_c = self._safe_category(params, node.orelse[0]) if node.orelse else None + if body_c is not None and else_c is not None and body_c != else_c: + return None + return body_c if body_c is not None else else_c + if isinstance(node, ast.Constant) and node.value is None: + return None + # Comparisons / ``not`` / boolean ops produce a boolean column; classify + # them as "bool" (``_category``'s catch-all would mislabel them numeric). + if _is_definitely_boolean(node): + return "bool" + try: + return self._category(params, node) + except UnsupportedOperationException: + return None + + def _convert_if_like( + self, + params: List[str], + test_col: Column, + body_col: Column, + else_col: Column, + test_node: ast.AST, + body_node: Optional[ast.AST], + else_node: Optional[ast.AST], + ) -> Column: + # We cannot soundly lower a generic Python truthiness test here. + # Python truthiness depends on the runtime input type and value: + # for example, 0, 0.0, "", empty collections, and None are all + # falsy, while most other values are truthy. The transpiler does + # not have enough input type information at this point to decide + # whether ``test_col`` is a boolean expression or a bare value + # whose truthiness would need Python-specific handling. Emitting + # ``when(coalesce(test_col, false), ...)`` is therefore unsound: + # it can either fail Spark analysis for non-boolean columns or + # silently diverge from Python semantics. Fail closed so the UDF + # falls back to interpreted Python execution instead. + if not _is_definitely_boolean(test_node): + raise UnsupportedOperationException( + f"bare truthiness tests ({ast.dump(test_node)}) in if-expressions are " + " not currently supported by the transpiler" + ) + # When the two branches resolve to concrete but different categories + # (e.g. numeric vs string), the lowered ``when(...).otherwise(...)`` is a + # CASE WHEN whose branch values share no common type under ANSI. That node + # is carried as a child of the TranspiledPythonUDF and is type-checked by + # CheckAnalysis *before* ConvertToCatalyst can drop it, so it would fail + # the whole query rather than fall back. Refuse here so the UDF runs as + # interpreted Python instead. Branches whose category we can't pin down + # (e.g. a bare ``None``) are treated as compatible and don't force this. + body_cat = self._safe_category(params, body_node) + else_cat = self._safe_category(params, else_node) + if body_cat is not None and else_cat is not None and body_cat != else_cat: + raise UnsupportedOperationException( + f"if/else branches have incompatible categories ({body_cat} vs " + f"{else_cat}); the lowered CASE WHEN has no common type under ANSI, " + "so the transpiler falls back to interpreted Python" + ) + safe_test = coalesce(test_col, lit(False)) + return when(safe_test, body_col).otherwise(else_col) + + def _lower_eq( + self, + params: List[str], + left_node: ast.AST, + right_node: ast.AST, + equal: bool, + ) -> Column: + """Lower ``==`` / ``!=`` with Python's None-equality semantics. + + Unlike ordering operators, Python doesn't raise on ``None == x`` / + ``None != x``: ``None == None`` is True, ``None == 0`` is False, + and ``!=`` is the negation. Spark's ``==`` returns NULL on NULL + operands (three-valued logic), which would round-trip through + the UDF as ``None`` rather than the bool Python would have + produced. Hand-roll the four cases via ``when`` branches. + + When the two operands resolve to concrete but DIFFERENT categories + (e.g. ``x == True`` on a numeric column, or ``x == "5"`` under the + numeric variant), the lowered ``=`` either fails analysis under ANSI + (bool vs bigint) -- which would break a working UDF since the option + is type-checked before ConvertToCatalyst can drop it -- or coerces + where Python's ``==`` is simply False. Refuse those so the UDF falls + back to interpreted Python. A ``None`` literal operand stays allowed + (the four-branch NULL handling above reproduces Python exactly). + + One value-level difference remains (needs runtime values, so it is + documented, not guarded): Spark treats ``NaN = NaN`` as true, while + Python's ``nan == nan`` is False. + """ + lc = self._safe_category(params, left_node) + rc = self._safe_category(params, right_node) + if lc is not None and rc is not None and lc != rc: + raise UnsupportedOperationException( + f"`==`/`!=` operands have incompatible categories ({lc} vs {rc}); " + "Python compares across types as unequal while Spark would coerce " + "or fail analysis, so the transpiler falls back to interpreted Python" + ) + left_col = self._convert_chunk(params, left_node) + right_col = self._convert_chunk(params, right_node) + left_null = left_col.isNull() + right_null = right_col.isNull() + if equal: + both_null_val: Column = lit(True) + one_null_val: Column = lit(False) + value_cmp = left_col == right_col + else: + both_null_val = lit(False) + one_null_val = lit(True) + value_cmp = left_col != right_col + return ( + when(left_null & right_null, both_null_val) + .when(left_null | right_null, one_null_val) + .otherwise(value_cmp) + ) + + def _lower_value_compare( + self, + params: List[str], + left_node: ast.AST, + right_node: ast.AST, + op: Callable[[Column, Column], Column], + op_repr: str, + ) -> Column: + """Lower a value comparison (``<``, ``<=``, ``>``, ``>=``). + + Python raises ``TypeError`` when an operand of these operators is + ``None`` (e.g. ``None > 0``), whereas Spark's three-valued logic + returns ``NULL``. To stay faithful to the source UDF we guard the + comparison: if either operand is ``NULL`` we raise via + ``raise_error``, otherwise we evaluate ``left op right`` as usual. + Callers that have already proven the operand non-null (``if x is + not None: x > 0``) take the otherwise branch, so they never trip + the raise. + + Python also forbids ordering across types (``1 < "a"`` -> TypeError), + whereas Spark would coerce the operands and return a (wrong) boolean. + We therefore only lower when both operands share a category; a + mismatch raises so this variant is dropped and the UDF falls back to + interpreted Python rather than silently diverging. + + One value-level difference from Python remains (it needs runtime + value info, so it is documented, not guarded): Spark orders ``NaN`` + as greater than every value, whereas Python's ``NaN`` comparisons + are all ``False``. + """ + lc = self._category(params, left_node) + rc = self._category(params, right_node) + if lc != rc: + raise UnsupportedOperationException( + f"`{op_repr}` compares operands of different categories " + f"({lc} vs {rc}); Python would raise TypeError, so the " + "transpiler falls back to interpreted Python" + ) + left_col = self._convert_chunk(params, left_node) + right_col = self._convert_chunk(params, right_node) + null_guard = left_col.isNull() | right_col.isNull() + err = lit( + "Python UDF transpiler: cannot compare NULL with operator " + f"`{op_repr}`; Python would raise TypeError here. Add an " + "`is not None` guard or filter NULLs upstream." + ) + return when(null_guard, raise_error(err)).otherwise(op(left_col, right_col)) + + def _category(self, params: List[str], node: ast.AST) -> str: + """Infer ``"numeric"`` or ``"string"`` for ``node`` under the current + ``self._param_categories`` assumption (set per input-type variant). + + Drives operator selection (``+`` -> add vs concat, ``*`` -> multiply vs + repeat) and raises ``UnsupportedOperationException`` when an operator's + operands are type-incompatible, so the caller drops that variant and the + JVM picks another option / falls back to the Python UDF. + """ + match node: + case ast.Constant(value=v): + # bool subclasses int, so classify it first: int/float -> numeric, + # str -> string, bool -> bool, bytes -> binary. None/complex/ + # Ellipsis have no usable Spark column type, so raise to drop this + # variant and fall back rather than emit an option that fails + # CheckAnalysis or silently diverges (e.g. `x + None` -> NULL where + # Python raises TypeError). + if isinstance(v, bool): + return "bool" + if isinstance(v, bytes): + return "binary" + if isinstance(v, (int, float)): + return "numeric" + if isinstance(v, str): + return "string" + raise UnsupportedOperationException( + f"constant {v!r} ({type(v).__name__}) has no usable column " + "category; falling back to interpreted Python" + ) + case ast.Name(id=name) if name in params: + index = params.index(name) + if params and params[0] == "self": + index -= 1 + return self._param_categories.get(index, "numeric") + case ast.BinOp(left=left, op=op, right=right): + lc = self._category(params, left) + rc = self._category(params, right) + if isinstance(op, ast.Add) and lc == rc: + return lc # str + str -> str, num + num -> num + if isinstance(op, ast.Mult): + if lc == "numeric" and rc == "numeric": Review Comment: ```suggestion if {lc, rc} == {"numeric", "numeric"}: ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
