def _write_outputs(self): # type: () -> Iterator[beam_fn_api_pb2.Elements] stream_done = False while not stream_done: streams = [self._to_send.get()] try: # Coalesce up to 100 other items. for _ in range(100): streams.append(self._to_send.get_nowait()) except queue.Empty: pass if streams[-1] is self._WRITES_FINISHED: stream_done = True streams.pop() if streams: data_stream = [] timer_stream = [] for stream in streams: if isinstance(stream, beam_fn_api_pb2.Elements.Timers): timer_stream.append(stream) elif isinstance(stream, beam_fn_api_pb2.Elements.Data): data_stream.append(stream) else: raise ValueError('Unexpected output element type %s' % type(stream)) yield beam_fn_api_pb2.Elements(data=data_stream, timers=timer_stream)
def _write_outputs(self): done = False while not done: data = [self._to_send.get()] try: # Coalesce up to 100 other items. for _ in range(100): data.append(self._to_send.get_nowait()) except queue.Empty: pass if data[-1] is self._WRITES_FINISHED: done = True data.pop() if data: yield beam_fn_api_pb2.Elements(data=data)