[SPARK-21985][PySpark] PairDeserializer is broken for double-zipped RDDs

python/pyspark/serializers.py

@@ -333,6 +333,9 @@ class PairDeserializer(Serializer):
     Deserializes the JavaRDD zip() of two PythonRDDs.
     Due to pyspark batching we cannot simply use the result of the Java RDD zip,
     we additionally need to do the zip within each pair of batches.
+
+    It is the responsibility of the user of this class to ensure the batch sizes of the key and
+    value serializer are the same size. If they are not this will give incorrect results.
     """
 
     def __init__(self, key_ser, val_ser):
@@ -343,9 +346,6 @@ def _load_stream_without_unbatching(self, stream):
         key_batch_stream = self.key_ser._load_stream_without_unbatching(stream)
         val_batch_stream = self.val_ser._load_stream_without_unbatching(stream)
         for (key_batch, val_batch) in zip(key_batch_stream, val_batch_stream):
-            if len(key_batch) != len(val_batch):
-                raise ValueError("Can not deserialize PairRDD with different number of items"
-                                 " in batches: (%d, %d)" % (len(key_batch), len(val_batch)))
             # for correctness with repeated cartesian/zip this must be returned as one batch
             yield zip(key_batch, val_batch)
 

python/pyspark/tests.py

@@ -644,6 +644,18 @@ def test_cartesian_chaining(self):
             set([(x, (y, y)) for x in range(10) for y in range(10)])
         )
 
+    def test_zip_chaining(self):
+        # Tests for SPARK-21985
+        rdd = self.sc.parallelize(range(10), 2)
+        self.assertSetEqual(
+            set(rdd.zip(rdd).zip(rdd).collect()),
+            set([((x, x), x) for x in range(10)])
+        )
+        self.assertSetEqual(
+            set(rdd.zip(rdd.zip(rdd)).collect()),
+            set([(x, (x, x)) for x in range(10)])
+        )
+
     def test_deleting_input_files(self):
         # Regression test for SPARK-1025
         tempFile = tempfile.NamedTemporaryFile(delete=False)