robertwb commented on a change in pull request #11314: [BEAM-9562] Send Timers
over Data Channel as Elements
URL: https://github.com/apache/beam/pull/11314#discussion_r406003474
##########
File path: sdks/python/apache_beam/runners/worker/data_plane.py
##########
@@ -408,27 +470,67 @@ def close_callback(data):
return ClosableOutputStream.create(
close_callback, add_to_send_queue, self._data_buffer_time_limit_ms)
+ def output_timer_stream(self, instruction_id, transform_id, timer_family_id):
+ def add_to_send_queue(timer):
+ if timer:
+ self._to_send.put(
+ beam_fn_api_pb2.Elements.Timer(
+ instruction_id=instruction_id,
+ transform_id=transform_id,
+ timer_family_id=timer_family_id,
+ timers=timer,
+ is_last=False))
+
+ def close_callback(timer):
+ add_to_send_queue(timer)
+ self._to_send.put(
+ beam_fn_api_pb2.Elements.Timer(
+ instruction_id=instruction_id,
+ transform_id=transform_id,
+ timer_family_id=timer_family_id,
+ timers=b'',
+ is_last=True))
+
+ return ClosableOutputStream.create(
+ close_callback, add_to_send_queue, self._data_buffer_time_limit_ms)
+
def _write_outputs(self):
# type: () -> Iterator[beam_fn_api_pb2.Elements]
- 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)
+ stream_done = False
+ while not stream_done:
+ streams = None
+ if 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 and streams[-1] is self._WRITES_FINISHED:
+ stream_done = True
+ streams.pop()
+ if streams:
+ elements = beam_fn_api_pb2.Elements()
+ data_stream = []
+ timer_stream = []
+ for stream in streams:
+ if isinstance(stream, beam_fn_api_pb2.Elements.Timer):
+ timer_stream.append(stream)
+ if isinstance(stream, beam_fn_api_pb2.Elements.Data):
Review comment:
else
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
