review comments

python/pyspark/serializers.py

@@ -356,33 +356,6 @@ def __repr__(self):
         return "ArrowStreamPandasSerializer"
 
 
-class InterleavedArrowReader(object):
-
-    def __init__(self, stream):
-        self._stream = stream
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        dataframes_in_group = read_int(self._stream)
-        if dataframes_in_group == 2:
-            return self._read_df(), self._read_df()
-        elif dataframes_in_group == 0:
-            raise StopIteration
-        else:
-            raise ValueError(
-                'Received Invalid number of dataframes in group {0}'.format(dataframes_in_group))
-
-    def next(self):
-        return self.__next__()
-
-    def _read_df(self):
-        import pyarrow as pa
-        reader = pa.ipc.open_stream(self._stream)
-        return [b for b in reader]
-
-
 class ArrowStreamPandasUDFSerializer(ArrowStreamPandasSerializer):
     """
     Serializer used by Python worker to evaluate Pandas UDFs
@@ -428,21 +401,31 @@ def __repr__(self):
         return "ArrowStreamPandasUDFSerializer"
 
 
-class PandasCogroupSerializer(ArrowStreamPandasUDFSerializer):
+class CogroupUDFSerializer(ArrowStreamPandasUDFSerializer):
 
     def __init__(self, timezone, safecheck, assign_cols_by_name):
-        super(PandasCogroupSerializer, self).__init__(timezone, safecheck, assign_cols_by_name)
+        super(CogroupUDFSerializer, self).__init__(timezone, safecheck, assign_cols_by_name)
 
     def load_stream(self, stream):
         """
         Deserialize Cogrouped ArrowRecordBatches to a tuple of Arrow tables and return as a two
         lists of pandas.Series.
         """
-        reader = InterleavedArrowReader(stream)
-        for batch1, batch2 in reader:
-            import pyarrow as pa
-            yield ([self.arrow_to_pandas(c) for c in pa.Table.from_batches(batch1).itercolumns()],
-                   [self.arrow_to_pandas(c) for c in pa.Table.from_batches(batch2).itercolumns()])
+        import pyarrow as pa
+        dataframes_in_group = None
+
+        while dataframes_in_group is None or dataframes_in_group > 0:
+            dataframes_in_group = read_int(stream)
+
+            if dataframes_in_group == 2:
+                batch1 = [batch for batch in ArrowStreamSerializer.load_stream(self, stream)]
+                batch2 = [batch for batch in ArrowStreamSerializer.load_stream(self, stream)]
+                yield ([self.arrow_to_pandas(c) for c in pa.Table.from_batches(batch1).itercolumns()],
+                       [self.arrow_to_pandas(c) for c in pa.Table.from_batches(batch2).itercolumns()])
+
+            elif dataframes_in_group != 0:
+                raise ValueError(
+                    'Received Invalid number of dataframes in group {0}'.format(dataframes_in_group))
 
 
 class BatchedSerializer(Serializer):

python/pyspark/sql/cogroup.py

@@ -44,7 +44,7 @@ def apply(self, udf):
         as a `DataFrame`.
 
         The user-defined function should take two `pandas.DataFrame` and return another
-        `pandas.DataFrame`. For each side of the cogroup, all columns are passed together
+        ``pandas.DataFrame``. For each side of the cogroup, all columns are passed together
         as a `pandas.DataFrame` to the user-function and the returned `pandas.DataFrame`
         are combined as a :class:`DataFrame`.
 

python/pyspark/sql/tests/test_pandas_udf_cogrouped_map.py

@@ -96,20 +96,20 @@ def test_complex_group_by(self):
             'v2': [90, 100, 110]
         })
 
-        left_df = self.spark\
+        left_gdf = self.spark\
             .createDataFrame(left)\
             .groupby(col('id') % 2 == 0)
 
-        right_df = self.spark \
+        right_gdf = self.spark \
             .createDataFrame(right) \
             .groupby(col('id') % 2 == 0)
 
         @pandas_udf('k long, v long, v2 long', PandasUDFType.COGROUPED_MAP)
         def merge_pandas(l, r):
             return pd.merge(l[['k', 'v']], r[['k', 'v2']], on=['k'])
 
-        result = left_df \
-            .cogroup(right_df) \
+        result = left_gdf \
+            .cogroup(right_gdf) \
             .apply(merge_pandas) \
             .sort(['k']) \
             .toPandas()

python/pyspark/worker.py

@@ -39,7 +39,7 @@
 from pyspark.rdd import PythonEvalType
 from pyspark.serializers import write_with_length, write_int, read_long, read_bool, \
     write_long, read_int, SpecialLengths, UTF8Deserializer, PickleSerializer, \
-    BatchedSerializer, ArrowStreamPandasUDFSerializer, PandasCogroupSerializer
+    BatchedSerializer, ArrowStreamPandasUDFSerializer, CogroupUDFSerializer
 from pyspark.sql.types import to_arrow_type, StructType
 from pyspark.util import _get_argspec, fail_on_stopiteration
 from pyspark import shuffle
@@ -314,7 +314,7 @@ def read_udfs(pickleSer, infile, eval_type):
         # Scalar Pandas UDF handles struct type arguments as pandas DataFrames instead of
         # pandas Series. See SPARK-27240.
         if eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
-            ser = PandasCogroupSerializer(timezone, safecheck, assign_cols_by_name)
+            ser = CogroupUDFSerializer(timezone, safecheck, assign_cols_by_name)
         else:
             df_for_struct = (eval_type == PythonEvalType.SQL_SCALAR_PANDAS_UDF or
                              eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF or
@@ -418,8 +418,8 @@ def extract_key_value_indexes():
             pickleSer, infile, eval_type, runner_conf, udf_index=0)
         udfs['f'] = udf
         parsed_offsets = extract_key_value_indexes()
-        keys = ["a[%d]" % o for o in parsed_offsets[0][0]]
-        vals = ["a[%d]" % o for o in parsed_offsets[0][1]]
+        keys = ["a[%d]" % (o,) for o in parsed_offsets[0][0]]
+        vals = ["a[%d]" % (o, ) for o in parsed_offsets[0][1]]
         mapper_str = "lambda a: f([%s], [%s])" % (", ".join(keys), ", ".join(vals))
     elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
         # We assume there is only one UDF here because cogrouped map doesn't
@@ -429,10 +429,10 @@ def extract_key_value_indexes():
             pickleSer, infile, eval_type, runner_conf, udf_index=0)
         udfs['f'] = udf
         parsed_offsets = extract_key_value_indexes()
-        df1_keys = ["a[0][%d]" % o for o in parsed_offsets[0][0]]
-        df1_vals = ["a[0][%d]" % o for o in parsed_offsets[0][1]]
-        df2_keys = ["a[1][%d]" % o for o in parsed_offsets[1][0]]
-        df2_vals = ["a[1][%d]" % o for o in parsed_offsets[1][1]]
+        df1_keys = ["a[0][%d]" % (o, ) for o in parsed_offsets[0][0]]
+        df1_vals = ["a[0][%d]" % (o, )for o in parsed_offsets[0][1]]
+        df2_keys = ["a[1][%d]" % (o, ) for o in parsed_offsets[1][0]]
+        df2_vals = ["a[1][%d]" % (o, ) for o in parsed_offsets[1][1]]
         mapper_str = "lambda a: f([%s], [%s], [%s], [%s])" % (
             ", ".join(df1_keys), ", ".join(df1_vals), ", ".join(df2_keys), ", ".join(df2_vals))
     else: