[GitHub] [beam] rohdesamuel commented on a change in pull request #12413: [BEAM-10603] Add max count to utils.to_element_list

[GitBox <gi...@apache.org>]

rohdesamuel commented on a change in pull request #12413:
URL: https://github.com/apache/beam/pull/12413#discussion_r467247182



##########
File path: sdks/python/apache_beam/runners/interactive/utils.py
##########
@@ -27,30 +27,46 @@
 import pandas as pd
 
 from apache_beam.portability.api.beam_runner_api_pb2 import TestStreamPayload
+from apache_beam.testing.test_stream import WindowedValueHolder
 
 
 def to_element_list(
     reader,  # type: Generator[Union[TestStreamPayload.Event, WindowedValueHolder]]
     coder,  # type: Coder
-    include_window_info  # type: bool
+    include_window_info,  # type: bool
+    n=None  # type: int
 ):
   # type: (...) -> List[WindowedValue]
 
   """Returns an iterator that properly decodes the elements from the reader.
   """
 
-  for e in reader:
-    if isinstance(e, TestStreamPayload.Event):
-      if (e.HasField('watermark_event') or e.HasField('processing_time_event')):
-        continue
+  # Defining a generator like this makes it easier to limit the count of
+  # elements read. Otherwise, the count limit would need to be duplicated.
+  def elements():
+    for e in reader:
+      if isinstance(e, TestStreamPayload.Event):
+        if (e.HasField('watermark_event') or
+            e.HasField('processing_time_event')):
+          continue
+        else:

Review comment:
       Done




----------------------------------------------------------------
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:
users@infra.apache.org