nehsyc commented on a change in pull request #13144:
URL: https://github.com/apache/beam/pull/13144#discussion_r510588684
##########
File path: sdks/python/apache_beam/transforms/util.py
##########
@@ -777,33 +795,56 @@ def process(
window=DoFn.WindowParam,
element_state=DoFn.StateParam(ELEMENT_STATE),
count_state=DoFn.StateParam(COUNT_STATE),
- expiry_timer=DoFn.TimerParam(EXPIRY_TIMER)):
+ window_timer=DoFn.TimerParam(WINDOW_TIMER),
+ buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)):
# Allowed lateness not supported in Python SDK
#
https://beam.apache.org/documentation/programming-guide/#watermarks-and-late-data
- expiry_timer.set(window.end)
+ window_timer.set(window.end)
element_state.add(element)
count_state.add(1)
count = count_state.read()
+ if count == 1 and max_buffering_duration_secs is not None:
+ # This is the first element in batch. Start counting buffering time if
a
+ # limit was set.
+ buffering_timer.set(clock() + max_buffering_duration_secs)
if count >= batch_size:
- batch = [element for element in element_state.read()]
- key, _ = batch[0]
- batch_values = [v for (k, v) in batch]
- yield (key, batch_values)
- element_state.clear()
- count_state.clear()
-
- @on_timer(EXPIRY_TIMER)
- def expiry(
+ key, batch_values = self.flush_batch(
+ element_state, count_state, buffering_timer)
+ if key is not None:
Review comment:
Done
##########
File path: sdks/python/apache_beam/transforms/util.py
##########
@@ -777,33 +795,56 @@ def process(
window=DoFn.WindowParam,
element_state=DoFn.StateParam(ELEMENT_STATE),
count_state=DoFn.StateParam(COUNT_STATE),
- expiry_timer=DoFn.TimerParam(EXPIRY_TIMER)):
+ window_timer=DoFn.TimerParam(WINDOW_TIMER),
+ buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)):
# Allowed lateness not supported in Python SDK
#
https://beam.apache.org/documentation/programming-guide/#watermarks-and-late-data
- expiry_timer.set(window.end)
+ window_timer.set(window.end)
element_state.add(element)
count_state.add(1)
count = count_state.read()
+ if count == 1 and max_buffering_duration_secs is not None:
+ # This is the first element in batch. Start counting buffering time if
a
+ # limit was set.
+ buffering_timer.set(clock() + max_buffering_duration_secs)
if count >= batch_size:
- batch = [element for element in element_state.read()]
- key, _ = batch[0]
- batch_values = [v for (k, v) in batch]
- yield (key, batch_values)
- element_state.clear()
- count_state.clear()
-
- @on_timer(EXPIRY_TIMER)
- def expiry(
+ key, batch_values = self.flush_batch(
+ element_state, count_state, buffering_timer)
+ if key is not None:
+ yield key, batch_values
+
+ @on_timer(WINDOW_TIMER)
+ def on_window_timer(
self,
element_state=DoFn.StateParam(ELEMENT_STATE),
- count_state=DoFn.StateParam(COUNT_STATE)):
+ count_state=DoFn.StateParam(COUNT_STATE),
+ buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)):
+ key, batch_values = self.flush_batch(
+ element_state, count_state, buffering_timer)
+ if key is not None:
+ yield key, batch_values
+
+ @on_timer(BUFFERING_TIMER)
+ def on_buffering_timer(
+ self,
+ element_state=DoFn.StateParam(ELEMENT_STATE),
+ count_state=DoFn.StateParam(COUNT_STATE),
+ buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)):
+ key, batch_values = self.flush_batch(
+ element_state, count_state, buffering_timer)
+ if key is not None:
+ yield key, batch_values
+
+ def flush_batch(self, element_state, count_state, buffering_timer):
batch = [element for element in element_state.read()]
- if batch:
- key, _ = batch[0]
- batch_values = [v for (k, v) in batch]
- yield (key, batch_values)
- element_state.clear()
- count_state.clear()
+ if not batch:
+ return None, None
Review comment:
Gotcha. I missed the return case previously so it didn't work correctly.
Thanks for the suggestion!
----------------------------------------------------------------
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]
