holdenk commented on code in PR #59216:
URL: https://github.com/apache/spark/pull/59216#discussion_r4342683841
##########
python/pyspark/sql/transpile.py:
##########
@@ -673,60 +890,49 @@ def _convert_chunk(self, params: List[str], body: ast.AST
| None) -> Column:
orelse[0] if orelse else None,
)
case ast.Compare(left, ops, comps):
- if len(ops) != 1 or len(comps) != 1:
+ # `in` / `not in` are container membership, not a chainable
value
+ # comparison, and need their own NULL analysis: Python's `in`
is
+ # equality-based and never raises on a `None` probe. Not
lowered
+ # yet, so refuse them and fall back.
+ if any(isinstance(o, (ast.In, ast.NotIn)) for o in ops):
raise UnsupportedOperationException(
- "chained comparisons (e.g. `a < b < c`) are not
supported by the transpiler"
+ "`in` / `not in` is not supported by the transpiler;
the "
+ "UDF falls back to interpreted Python"
)
- comp = comps[0]
- match ops[0]:
- case ast.Is() | ast.IsNot():
- # Only lower `x is None` / `None is x` (and their
- # `is not` variants) to isNull/isNotNull. For any
- # other comparator (e.g. `x is 0`, `x is y`) Python
- # performs an object-identity check that has no SQL
- # equivalent, so we must fall back to interpreted
- # Python rather than silently emitting a null check.
- is_none_left = isinstance(left, ast.Constant) and
left.value is None
- is_none_right = isinstance(comp, ast.Constant) and
comp.value is None
- if not (is_none_left or is_none_right):
- raise UnsupportedOperationException(
- "`is`/`is not` is only supported when one "
- "operand is the literal None; other identity "
- "checks (e.g. `x is 0`, `x is y`) cannot be "
- "lowered to SQL and the UDF falls back to "
- "interpreted Python"
- )
- subject_node = comp if is_none_left else left
- subject_col = self._convert_chunk(params, subject_node)
- if isinstance(ops[0], ast.Is):
- return subject_col.isNull()
- else:
- return subject_col.isNotNull()
- case ast.Eq():
- return self._lower_eq(params, left, comp, equal=True)
- case ast.NotEq():
- return self._lower_eq(params, left, comp, equal=False)
- case ast.Lt():
- return self._lower_value_compare(
- params, left, comp, lambda l, r: l < r, "<"
- )
- case ast.LtE():
- return self._lower_value_compare(
- params, left, comp, lambda l, r: l <= r, "<="
- )
- case ast.Gt():
- return self._lower_value_compare(
- params, left, comp, lambda l, r: l > r, ">"
- )
- case ast.GtE():
- return self._lower_value_compare(
- params, left, comp, lambda l, r: l >= r, ">="
- )
- case _:
- raise UnsupportedOperationException(
- f"comparison operator {type(ops[0]).__name__} "
- "is not supported by the transpiler"
- )
+ if len(ops) > _MAX_COMPARISON_OPS:
+ raise UnsupportedOperationException(
+ f"chained comparison has {len(ops)} operators; the "
+ f"transpiler lowers at most {_MAX_COMPARISON_OPS}, so
the "
+ "UDF falls back to interpreted Python"
+ )
+ # Python evaluates `a OP1 b OP2 c` as `(a OP1 b) and (b OP2 c)`
+ # with `b` evaluated once, so the fold below re-lowers each
+ # interior operand once per adjacent pair. Safe for the operand
+ # grammar we accept, all of it deterministic and side-effect
free.
+ # It does mean the body READS a parameter once per adjacent
pair,
+ # which SPARK-58626 is what makes safe: `_udf_param_N` is a
+ # reference, and `ConvertToCatalyst` gives an argument read
more
+ # than once one column per row -- or declines to lower the call
+ # where no column fits -- so a non-deterministic argument is
drawn
+ # once and every read agrees. While each read still carried
its own
+ # copy of the argument, `50 < x < 60` over `rand(7) * 100` was
~30%
+ # true where Python is ~10%.
+ #
+ # Folding into nested CASE WHENs instead -- atomic to predicate
+ # splitting -- was rejected: it costs pushdown for every
+ # `and` / `or` UDF, including those over non-nullable columns
that
+ # need no guard at all.
+ #
+ # See the module docstring for the short-circuit and NULL-guard
+ # semantics this fold relies on.
+ nodes = [left, *comps]
+ return functools.reduce(
Review Comment:
(Also this should probably be higher than a Nit comment)
--
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]
