[SPARK-59514][PYTHON] Handle nonfinite values in assertDataFrameEqual

python/pyspark/sql/tests/test_utils.py

@@ -1884,6 +1884,48 @@ def test_assert_schema_equal_with_timestamp_nanos_types(self):
             assertSchemaEqual(s1, StructType([StructField("ts", TimestampNTZNanosType(7), True)]))
 
 
+class NonFiniteComparisonTests(unittest.TestCase):
+    def test_unequal_special_values(self):
+        nan, inf = float("nan"), float("inf")
+        for left, right in [
+            (nan, 1.0),
+            (nan, inf),
+            (nan, -inf),
+            (1.0, inf),
+            (1.0, -inf),
+            (inf, -inf),
+        ]:
+            for actual, expected in [(left, right), (right, left)]:
+                for ordered in [False, True]:
+                    for rtol, atol in [(1e-5, 1e-8), (0.0, 0.0)]:
+                        with self.subTest(actual=actual, expected=expected, ordered=ordered):
+                            with self.assertRaises(PySparkAssertionError) as error:
+                                assertDataFrameEqual(
+                                    [Row(value=actual)],
+                                    [Row(value=expected)],
+                                    checkRowOrder=ordered,
+                                    rtol=rtol,
+                                    atol=atol,
+                                )
+                            self.assertEqual(error.exception.getCondition(), "DIFFERENT_ROWS")
+
+    def test_equal_special_values_and_finite_tolerance(self):
+        for value in [float("nan"), float("inf"), float("-inf")]:
+            assertDataFrameEqual([Row(value=value)], [Row(value=value)])
+        assertDataFrameEqual([Row(value=1.001)], [Row(value=1.0)], rtol=0.01)
+        with self.assertRaises(PySparkAssertionError):
+            assertDataFrameEqual([Row(value=1.1)], [Row(value=1.0)], rtol=0.01)
+
+    def test_nested_special_values(self):
+        for wrap in [
+            lambda value: Row(nested=Row(value=value)),
+            lambda value: Row(items=[value]),
+            lambda value: Row(mapping={"key": value}),
+        ]:
+            with self.assertRaises(PySparkAssertionError):
+                assertDataFrameEqual([wrap(float("nan"))], [wrap(1.0)])
+
+
 class UtilsTests(UtilsTestsMixin, ReusedSQLTestCase):
     pass
 

python/pyspark/testing/utils.py

@@ -18,6 +18,7 @@
 import difflib
 import faulthandler
 import functools
+import math
 import os
 import signal
 import struct
@@ -1121,6 +1122,10 @@ def compare_vals(val1, val2):
                     and all(compare_vals(val1[k], val2[k]) for k in val1)
                 )
             elif isinstance(val1, float) and isinstance(val2, float):
+                if math.isnan(val1) or math.isnan(val2):
+                    return math.isnan(val1) and math.isnan(val2)
+                if math.isinf(val1) or math.isinf(val2):
+                    return val1 == val2
                 if abs(val1 - val2) > (atol + rtol * abs(val2)):
                     return False
             elif isinstance(val1, Decimal) and isinstance(val2, Decimal):