Add big collection test for GBK and BagState

Author: baeminbo

big-collection/python/bigbag.py

@@ -0,0 +1,72 @@
+import logging
+from typing import Tuple
+
+from apache_beam import Create
+from apache_beam import DoFn
+from apache_beam import FlatMap
+from apache_beam import ParDo
+from apache_beam import Pipeline
+from apache_beam import TimeDomain
+from apache_beam import WithKeys
+from apache_beam import typehints
+from apache_beam.coders import BytesCoder
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.options.pipeline_options import SetupOptions
+from apache_beam.transforms.userstate import BagStateSpec
+from apache_beam.transforms.userstate import TimerSpec
+from apache_beam.transforms.userstate import on_timer
+
+# The total bytes processed is NUM_SHARDS * NUM_ELEMENTS_PER_SHARD * ELEMENT_BYTES ~= 3 GiB
+NUM_SHARDS = 100
+NUM_ELEMENTS_PER_SHARD = 10
+ELEMENT_BYTES = 3 * 1024 * 1024  # 3 MiB
+
+
+@typehints.with_input_types(Tuple[str, bytes])
+@typehints.with_output_types(None)
+class BigBagDoFn(DoFn):
+  VALUES_STATE = BagStateSpec('values', BytesCoder())
+  END_OF_WINDOW_TIMER = TimerSpec('end_of_window', TimeDomain.WATERMARK)
+
+  def process(self, element: Tuple[str, bytes], window=DoFn.WindowParam,
+      values_state=DoFn.StateParam(VALUES_STATE),
+      end_of_window_timer=DoFn.TimerParam(END_OF_WINDOW_TIMER)):
+    logging.info('start process.')
+    key, value = element
+    end_of_window_timer.set(window.end)
+    values_state.add(value)
+    logging.info('end process.')
+
+  @on_timer(END_OF_WINDOW_TIMER)
+  def end_of_window(self, values_state=DoFn.StateParam(VALUES_STATE)):
+    logging.info('start end_of_window.')
+
+    read_count = 0
+    read_bytes = 0
+    values = values_state.read()
+    for value in values:
+      read_count += 1
+      read_bytes += len(value)
+
+    logging.info('read_count: %s, read_bytes: %s', read_count, read_bytes)
+    logging.info('end end_of_window.')
+
+
+def main():
+  options = PipelineOptions()
+  options.view_as(SetupOptions).save_main_session = True
+
+  p = Pipeline(options=options)
+  (p
+   | Create(list(range(NUM_SHARDS)))
+   | FlatMap(lambda _:
+             (bytes(ELEMENT_BYTES) for _ in range(NUM_ELEMENTS_PER_SHARD)))
+   | WithKeys('')
+   | ParDo(BigBagDoFn()))
+
+  p.run()
+
+
+if __name__ == '__main__':
+  logging.getLogger().setLevel(logging.INFO)
+  main()

big-collection/python/biggbk.py

@@ -0,0 +1,52 @@
+import logging
+from typing import Iterable
+from typing import Tuple
+
+from apache_beam import Create
+from apache_beam import FlatMap
+from apache_beam import GroupByKey
+from apache_beam import ParDo
+from apache_beam import Pipeline
+from apache_beam import WithKeys
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.options.pipeline_options import SetupOptions
+
+# The total bytes processed is NUM_SHARDS * NUM_ELEMENTS_PER_SHARD * ELEMENT_BYTES ~= 3 GiB
+
+NUM_SHARDS = 100
+NUM_ELEMENTS_PER_SHARD = 10
+ELEMENT_BYTES = 3 * 1024 * 1024  # 3 MiB
+
+def print_bytes(element: Tuple[str, Iterable[bytes]]) -> None:
+  logging.info('start print_bytes.')
+  key, values = element
+
+  read_count = 0
+  read_bytes = 0
+  for value in values:
+    read_count += 1
+    read_bytes += len(value)
+
+  logging.info('read_count: %s, read_bytes: %s', read_count, read_bytes)
+  logging.info('end print_bytes.')
+
+
+def main():
+  options = PipelineOptions()
+  options.view_as(SetupOptions).save_main_session = True
+
+  p = Pipeline(options=options)
+  (p
+   | Create(list(range(NUM_SHARDS)))
+   | FlatMap(lambda _:
+             (bytes(ELEMENT_BYTES) for _ in range(NUM_ELEMENTS_PER_SHARD)))
+   | WithKeys('')
+   | GroupByKey()
+   | ParDo(print_bytes))
+
+  p.run()
+
+
+if __name__ == '__main__':
+  logging.getLogger().setLevel(logging.INFO)
+  main()

big-collection/python/requirements.txt

@@ -0,0 +1 @@
+apache-beam[gcp]==2.35.0

big-collection/python/run_bigbag.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+PROJECT=$(gcloud config get-value project)
+PIPELINE=bigbag.py
+
+# NOTE: Job has no progress if --number_of_worker_harness_threads=1
+python $PIPELINE --runner=DataflowRunner \
+  --project=$PROJECT \
+  --region=us-central1 \
+  --streaming \
+  --worker_machine_type custom-1-102400-ext \
+  --experiments=use_runner_v2 \
+  --experiments=no_use_multiple_sdk_containers \
+  --max_num_workers=1 \
+  --number_of_worker_harness_threads=2

big-collection/python/run_biggbk.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+PROJECT=$(gcloud config get-value project)
+PIPELINE=biggbk.py
+
+# NOTE: Job has no progress if --number_of_worker_harness_threads=1
+python $PIPELINE --runner=DataflowRunner \
+  --project=$PROJECT \
+  --region=us-central1 \
+  --streaming \
+  --worker_machine_type custom-1-102400-ext \
+  --experiments=use_runner_v2 \
+  --experiments=no_use_multiple_sdk_containers \
+  --max_num_workers=1 \
+  --number_of_worker_harness_threads=2