ibzib commented on a change in pull request #11136: [BEAM-7505] Add side input load test to Python SDK URL: https://github.com/apache/beam/pull/11136#discussion_r395229525
########## File path: sdks/python/apache_beam/testing/load_tests/sideinput_test.py ########## @@ -47,65 +49,154 @@ or: ./gradlew -PloadTest.args=" - --publish_to_big_query=true - --project=... - --metrics_dataset=python_load_tests - --metrics_table=side_input + --side_input_type=iter --input_options='{ - \"num_records\": 1, - \"key_size\": 1, - \"value_size\": 1}' - --runner=DirectRunner" \ + \"num_records\": 300, + \"key_size\": 5, + \"value_size\": 15}'" \ -PloadTest.mainClass=apache_beam.testing.load_tests.sideinput_test \ -Prunner=DirectRunner :sdks:python:apache_beam:testing:load_tests:run """ # pytype: skip-file from __future__ import absolute_import +from __future__ import division import logging import apache_beam as beam -from apache_beam.pvalue import AsIter from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.synthetic_pipeline import SyntheticSource class SideInputTest(LoadTest): + SIDE_INPUT_TYPES = { + 'iter': beam.pvalue.AsIter, + 'list': beam.pvalue.AsList, + 'dict': beam.pvalue.AsDict, + } + def __init__(self): super(SideInputTest, self).__init__() - self.iterations = self.get_option_or_default( - 'number_of_counter_operations', 1) + self.windows = self.get_option_or_default('window_count', default=0) + self.access_percentage = self.get_option_or_default( + 'access_percentage', default=100) + if self.access_percentage < 0 or self.access_percentage > 100: + raise ValueError( + 'access_percentage: Invalid value. Should be in range ' + 'from 0 to 100, got {} instead'.format(self.access_percentage)) + + self.side_input_size = self.get_option_or_default( + 'side_input_size', default=0) + if self.side_input_size == 0: + self.side_input_size = self.input_options.get('num_records') + + self.side_input_type = self.pipeline.get_option('side_input_type') + if self.side_input_type is None: + raise ValueError('side_input_type is required') + + def materialize_as(self): + try: + return self.SIDE_INPUT_TYPES[self.side_input_type] + except KeyError: + raise ValueError( + 'Unknown side input type. You have to provide one of ' + 'these: {}'.format(list(self.SIDE_INPUT_TYPES.keys()))) def test(self): - def join_fn(element, side_input, iterations): - result = [] - for i in range(iterations): - for key, value in side_input: - if i == iterations - 1: - result.append({key: element[1] + value}) - yield result - - main_input = ( + class SequenceSideInputTestDoFn(beam.DoFn): + """Iterate over first n side_input elements. Iterate over all + elements if `first_n` is :data:`None`.""" + def __init__(self, first_n=None): + self._first_n = first_n + + def process(self, unused_element, side_input): + for i, _ in enumerate(side_input): + if self._first_n and i >= self._first_n: + return + # No-op. We only make sure that the element is accessed. + + class MappingSideInputTestDoFn(beam.DoFn): + """Take a sequence of keys as an additional side input and for each + key in the sequence checks the value for key in the dictionary.""" + def process(self, unused_element, dict_side_input, keys_to_check): + for key in keys_to_check: + # No-op. We only make sure that the element is accessed. + dict_side_input[key] + + class GetRandomKeys(beam.DoFn): + def __init__(self, n): + self._n = n + + def process(self, unused_element, dict_side_input): + import random + n = min(self._n, len(dict_side_input)) + return random.sample(dict_side_input.keys(), n) + + class AddEventTimestamps(beam.DoFn): + """Assign timestamp to each element of PCollection, starting from the + current Unix seconds-since-epoch timestamp.""" + def setup(self): + import time + self.current_time = int(time.time()) Review comment: Is there any particular reason we need to use the real time? If not, it would be easier to read timestamps if we start at 0. ---------------------------------------------------------------- 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: us...@infra.apache.org With regards, Apache Git Services