Adjust to cleaned up pyarrow FileReader API, support multiple record …

…batches in a stream closes apache#22
icexelloss · Feb 23, 2017 · f44e6d7 · f44e6d7
1 parent fbe3b7c
commit f44e6d7
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/python/pyspark/serializers.py b/python/pyspark/serializers.py
@@ -192,9 +192,9 @@ def dumps(self, obj):
         raise NotImplementedError
 
     def loads(self, obj):
-        from pyarrow.ipc import ArrowFileReader
-        reader = ArrowFileReader(obj)
-        return reader.get_record_batch(0)
+        from pyarrow import FileReader, BufferReader
+        reader = FileReader(BufferReader(obj))
+        return reader.read_all()
 
     def __repr__(self):
         return "ArrowSerializer"