[SPARK-58456][PYTHON][SQL] Support chained comparisons in transpilation

python/pyspark/sql/tests/test_udf_transpile_hypothesis.py

@@ -206,6 +206,25 @@ def _env_opts_in(value: Optional[str]) -> bool:
         (1, -_LONG_BOUND),
     )
 
+    # Triple edges for the chained-comparison tests, pulled from the
+    # single-value edges above. A chain short-circuits, so the NULL patterns
+    # -- which position holds the None -- are the point; the boundary orderings
+    # probe sign/overflow at the type's edges.
+    _LONG_TRIPLE_EDGES = (
+        # NULL position: which link's guard fires -- or is short-circuited past
+        # -- depends on where the None sits.
+        (None, None, None),
+        (None, 0, 1),
+        (0, None, 1),
+        (0, 1, None),
+        (1, 0, None),  # first link false -> None never reached
+        # Boundary orderings.
+        (0, 1, 2),
+        (2, 1, 0),
+        (_INT32_MIN, 0, _INT32_MAX),
+        (-_LONG_BOUND, 0, _LONG_BOUND),
+    )
+
     def _seed_examples(values, key="value"):
         """Stack one ``@example`` decorator per seed value."""
 
@@ -216,21 +235,37 @@ def wrapper(method):
 
         return wrapper
 
-    def _seed_pair_examples(pairs, keys=("x", "y")):
+    def _seed_tuple_examples(tuples, keys=("x", "y", "z")):
+        """Stack one ``@example`` decorator per seed tuple, zipped onto ``keys``."""
+
         def wrapper(method):
-            for v0, v1 in reversed(pairs):
-                method = example(**{keys[0]: v0, keys[1]: v1})(method)
+            for values in reversed(tuples):
+                method = example(**dict(zip(keys, values)))(method)
             return method
 
         return wrapper
 
+    def _seed_pair_examples(pairs, keys=("x", "y")):
+        """Two-argument ``_seed_tuple_examples``."""
+        return _seed_tuple_examples(pairs, keys)
+
 
 # ---- The UDF templates we exercise --------------------------------------
 #
 # We keep these as module-level callables so that ``inspect.getsource`` works
 # (the transpiler reads source via inspection). They are deliberately written
 # the way a user would: idiomatic Python, including ``if x is not None``
 # guards and bare ``if x:`` truthiness checks.
+#
+# Where a template annotates its parameters, that is to pin the input category:
+# an untyped parameter emits both a numeric and a string variant, and these
+# tests bind Long columns, so the string one is dead weight (measured: 1 option
+# emitted instead of 2). It buys plan size, not wall clock -- measured 73s
+# untyped vs 76s typed over these five tests at 15 examples, because the cost is
+# the two Spark jobs ``_run`` executes per example against a fresh one-row
+# DataFrame. Batching rows per example is the lever that would matter. The
+# strategies still generate ``None`` -- an annotation pins the category, not
+# nullability.
 
 
 def plus_four(x):
@@ -343,6 +378,35 @@ def neq_pair(x, y):
     return x != y
 
 
+def chained_bounds(x: int) -> bool:
+    # Single-parameter chained comparison: ``ast.Compare`` with two ops, which
+    # Python evaluates as ``(0 < x) and (x < 10)`` with ``x`` evaluated once. No
+    # None guard, so a NULL input reaches the first link's raise guard, matching
+    # Python's TypeError; ``_run``'s "both raised" equivalence covers that.
+    return 0 < x < 10
+
+
+def chained_lt(x: int, y: int, z: int) -> bool:
+    # Three-column ordering chain. The interesting rows are the ones where an
+    # earlier link is False and a later operand is NULL: Python short-circuits
+    # and returns False without ever comparing against the NULL, so the
+    # transpiled form must not raise there either.
+    return x < y < z
+
+
+def chained_eq(x: int, y: int, z: int) -> bool:
+    # Equality never raises in Python, so every NULL pattern has a defined
+    # answer -- the strongest NULL probe of the chained lowering, with no "both
+    # raised" escape hatch to hide a divergence.
+    return x == y == z
+
+
+def mixed_chain(x: int, y: int, z: int) -> bool:
+    # Mixed operators in one chain (``ast.LtE`` then ``ast.Gt``), so the links
+    # disagree on direction and short-circuit on different rows again.
+    return x <= y > z
+
+
 # A lambda captured at module scope so ``inspect.getsource`` can read
 # its definition. Exercises the ``ast.Lambda`` branch in
 # ``_get_function_from_ast``.
@@ -440,9 +504,9 @@ def _run(self, func, return_type, df, *udf_arg_columns, kwargs=None):
                 interpreted_value = _SENTINEL_RAISED
                 interpreted_error = e
 
-        # If the transpiled path raises an exception we also need the interpreted path to raise one,
-        # however if the Python code (that in the interpreted path) raises an exception, the transpiled
-        # path may return a valid value.
+        # Asymmetric on purpose. If the transpiled path raises, the interpreted path
+        # must raise too -- that is the dangerous direction. The reverse is allowed:
+        # interpreted Python may raise where the transpiled form returns a value.
         if transpiled_error is not None:
             self.assertIsNotNone(
                 interpreted_error,
@@ -466,6 +530,16 @@ def _two_long_arg_df(self, x, y):
         )
         return self.spark.createDataFrame([Row(a=x, b=y)], schema=schema)
 
+    def _three_long_arg_df(self, x, y, z):
+        schema = StructType(
+            [
+                StructField("a", LongType(), nullable=True),
+                StructField("b", LongType(), nullable=True),
+                StructField("c", LongType(), nullable=True),
+            ]
+        )
+        return self.spark.createDataFrame([Row(a=x, b=y, c=z)], schema=schema)
+
     if _have_hypothesis:
 
         @_hyp_settings
@@ -664,6 +738,52 @@ def test_lambda_plus_four_matches_python(self, value):
             transpiled, interpreted = self._run(lambda_plus_four, LongType(), df, "a")
             self.assertEqual(transpiled, interpreted, f"lambda_plus_four mismatch on {value!r}")
 
+        # ---- Chained comparisons ----------------------------------------
+        #
+        # Python's ``a OP b OP c`` short-circuits, so a NULL operand the second
+        # link would have touched must not raise. ``_run`` catches the dangerous
+        # direction: the transpiled form raising where interpreted Python does
+        # not. It tolerates the reverse, which test_udf_transpile_unit.py pins
+        # explicitly instead.
+
+        @_hyp_settings
+        @given(value=_long_strategy)
+        @_seed_examples(_LONG_EDGES)
+        def test_chained_bounds_matches_python(self, value):
+            df = self._single_arg_df(value, LongType())
+            transpiled, interpreted = self._run(chained_bounds, BooleanType(), df, "a")
+            self.assertEqual(transpiled, interpreted, f"chained_bounds mismatch on {value!r}")
+
+        @_hyp_settings
+        @given(x=_long_strategy, y=_long_strategy, z=_long_strategy)
+        @_seed_tuple_examples(_LONG_TRIPLE_EDGES)
+        def test_chained_lt_matches_python(self, x, y, z):
+            df = self._three_long_arg_df(x, y, z)
+            transpiled, interpreted = self._run(chained_lt, BooleanType(), df, "a", "b", "c")
+            self.assertEqual(
+                transpiled, interpreted, f"chained_lt mismatch on ({x!r}, {y!r}, {z!r})"
+            )
+
+        @_hyp_settings
+        @given(x=_long_strategy, y=_long_strategy, z=_long_strategy)
+        @_seed_tuple_examples(_LONG_TRIPLE_EDGES)
+        def test_chained_eq_matches_python(self, x, y, z):
+            df = self._three_long_arg_df(x, y, z)
+            transpiled, interpreted = self._run(chained_eq, BooleanType(), df, "a", "b", "c")
+            self.assertEqual(
+                transpiled, interpreted, f"chained_eq mismatch on ({x!r}, {y!r}, {z!r})"
+            )
+
+        @_hyp_settings
+        @given(x=_long_strategy, y=_long_strategy, z=_long_strategy)
+        @_seed_tuple_examples(_LONG_TRIPLE_EDGES)
+        def test_mixed_chain_matches_python(self, x, y, z):
+            df = self._three_long_arg_df(x, y, z)
+            transpiled, interpreted = self._run(mixed_chain, BooleanType(), df, "a", "b", "c")
+            self.assertEqual(
+                transpiled, interpreted, f"mixed_chain mismatch on ({x!r}, {y!r}, {z!r})"
+            )
+
 
 class UDFTranspileHypothesisGatingTests(unittest.TestCase):
     """Smoke tests that always run, regardless of the env gate.