davidyan74 commented on a change in pull request #12411:
URL: https://github.com/apache/beam/pull/12411#discussion_r462561306
##########
File path:
sdks/python/apache_beam/runners/interactive/options/capture_limiters.py
##########
@@ -71,3 +87,62 @@ def _trigger(self):
def is_triggered(self):
return self._triggered
+
+
+class CountLimiter(ElementLimiter):
+ """Limits by counting the number of elements seen."""
+ def __init__(self, max_count):
+ self._max_count = max_count
+ self._count = 0
+
+ def update(self, e):
+ # A TestStreamFileRecord can contain many elements at once. If e is a file
+ # record, then count the number of elements in the bundle.
+ if isinstance(e, TestStreamFileRecord):
+ if not e.recorded_event.element_event:
+ return
+ self._count += len(e.recorded_event.element_event.elements)
+
+ # Otherwise, count everything else but the header of the file since it is
+ # not an element.
+ elif not isinstance(e, TestStreamFileHeader):
+ self._count += 1
+
+ def is_triggered(self):
+ return self._count >= self._max_count
+
+
+class ProcessingTimeLimiter(ElementLimiter):
+ """Limits by how long the ProcessingTime passed in the element stream.
+
+ This measures the duration from either a constructor argument or from the
+ first element in the stream. Each subsequent element has a delta
+ "advance_duration" that moves the internal clock forward. This triggers when
+ the duration from the internal clock and the start exceeds the given
duration.
+ """
+ def __init__(self, max_duration_secs, start_secs=None):
Review comment:
nit: would it be simpler if we default this to 0 rather than None? It
would let us eliminate the None checks in the methods.
----------------------------------------------------------------
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]
