add flush for ChunkedStream

apache · May 29, 2019 · 4d58419 · 4d58419
1 parent 829ea96
commit 4d58419
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/python/pyspark/serializers.py b/python/pyspark/serializers.py
@@ -869,6 +869,14 @@ def write(self, bytes):
                 byte_pos = new_byte_pos
                 self.current_pos = 0
 
+    def flush(self):
+        if self.current_pos > 0:
+            # send buffer [0, current_pos)
+            write_int(self.current_pos, self.wrapped)
+            self.wrapped.write(self.buffer[:self.current_pos])
+            self.wrapped.flush()
+            self.current_pos = 0
+
     def close(self):
         # if there is anything left in the buffer, write it out first
         if self.current_pos > 0: