[PYTHON] Compare Row lengths in assertDataFrameEqual

Author: laserninja

python/pyspark/sql/tests/test_utils.py

@@ -1884,6 +1884,32 @@ def test_assert_schema_equal_with_timestamp_nanos_types(self):
             assertSchemaEqual(s1, StructType([StructField("ts", TimestampNTZNanosType(7), True)]))
 
 
+class RowComparisonTests(unittest.TestCase):
+    def test_different_row_lengths(self):
+        pairs = [(Row(), Row(x=1)), (Row(x=1), Row(x=1, y=2))]
+        for left, right in pairs:
+            for actual, expected in [(left, right), (right, left)]:
+                for wrap in [
+                    lambda row: row,
+                    lambda row: Row(nested=row),
+                    lambda row: Row(items=[row]),
+                    lambda row: Row(mapping={"key": row}),
+                ]:
+                    for ordered in [False, True]:
+                        with self.subTest(actual=actual, expected=expected, ordered=ordered):
+                            with self.assertRaises(PySparkAssertionError) as error:
+                                assertDataFrameEqual(
+                                    [wrap(actual)], [wrap(expected)], checkRowOrder=ordered
+                                )
+                            self.assertEqual(error.exception.getCondition(), "DIFFERENT_ROWS")
+
+    def test_equal_row_lengths(self):
+        for row in [Row(), Row(x=1), Row(nested=Row(x=1, y=2))]:
+            assertDataFrameEqual([row], [row])
+        with self.assertRaises(PySparkAssertionError):
+            assertDataFrameEqual([Row(x=1)], [Row(x=2)])
+
+
 class UtilsTests(UtilsTestsMixin, ReusedSQLTestCase):
     pass
 

python/pyspark/testing/utils.py

@@ -1113,7 +1113,9 @@ def compare_vals(val1, val2):
                     compare_vals(x, y) for x, y in zip(val1, val2)
                 )
             elif isinstance(val1, Row) and isinstance(val2, Row):
-                return all(compare_vals(x, y) for x, y in zip(val1, val2))
+                return len(val1) == len(val2) and all(
+                    compare_vals(x, y) for x, y in zip(val1, val2)
+                )
             elif isinstance(val1, dict) and isinstance(val2, dict):
                 return (
                     len(val1) == len(val2)