Commit 4c02f009 authored by Yuxin Wu's avatar Yuxin Wu

Fix None mapping in parallel MapData

parent ba855cdb
...@@ -341,12 +341,16 @@ class _ParallelMapData(ProxyDataFlow): ...@@ -341,12 +341,16 @@ class _ParallelMapData(ProxyDataFlow):
def get_data_non_strict(self): def get_data_non_strict(self):
for dp in self._iter: for dp in self._iter:
self._send(dp) self._send(dp)
yield self._recv() ret = self._recv()
if ret is not None:
yield ret
self._iter = self.ds.get_data() # refresh self._iter = self.ds.get_data() # refresh
for _ in range(self._buffer_size): for _ in range(self._buffer_size):
self._send(next(self._iter)) self._send(next(self._iter))
yield self._recv() ret = self._recv()
if ret is not None:
yield ret
def get_data_strict(self): def get_data_strict(self):
for dp in self._iter: for dp in self._iter:
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment