Source code for apache_beam.testing.test_stream

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Provides TestStream for verifying streaming runner semantics.

For internal use only; no backwards-compatibility guarantees.
"""
# pytype: skip-file

from abc import ABCMeta
from abc import abstractmethod
from enum import Enum
from functools import total_ordering

import apache_beam as beam
from apache_beam import coders
from apache_beam import pvalue
from apache_beam.portability import common_urns
from apache_beam.portability.api import beam_interactive_api_pb2
from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.portability.api import endpoints_pb2
from apache_beam.transforms import PTransform
from apache_beam.transforms import core
from apache_beam.transforms import window
from apache_beam.transforms.timeutil import TimeDomain
from apache_beam.transforms.userstate import TimerSpec
from apache_beam.transforms.userstate import on_timer
from apache_beam.transforms.window import TimestampedValue
from apache_beam.utils import timestamp
from apache_beam.utils.timestamp import MIN_TIMESTAMP
from apache_beam.utils.timestamp import Duration
from apache_beam.utils.timestamp import Timestamp
from apache_beam.utils.windowed_value import WindowedValue

__all__ = [
    'Event',
    'ElementEvent',
    'WatermarkEvent',
    'ProcessingTimeEvent',
    'TestStream',
]


[docs]@total_ordering class Event(metaclass=ABCMeta): # type: ignore[misc] """Test stream event to be emitted during execution of a TestStream.""" @abstractmethod def __eq__(self, other): raise NotImplementedError @abstractmethod def __hash__(self): raise NotImplementedError @abstractmethod def __lt__(self, other): raise NotImplementedError
[docs] @abstractmethod def to_runner_api(self, element_coder): raise NotImplementedError
[docs] @staticmethod def from_runner_api(proto, element_coder): if proto.HasField('element_event'): event = proto.element_event tag = None if event.tag == 'None' else event.tag return ElementEvent([ TimestampedValue( element_coder.decode(tv.encoded_element), Timestamp(micros=1000 * tv.timestamp)) for tv in proto.element_event.elements ], tag=tag) # yapf: disable elif proto.HasField('watermark_event'): event = proto.watermark_event tag = None if event.tag == 'None' else event.tag return WatermarkEvent( Timestamp(micros=1000 * proto.watermark_event.new_watermark), tag=tag) elif proto.HasField('processing_time_event'): return ProcessingTimeEvent( timestamp.Duration( micros=1000 * proto.processing_time_event.advance_duration)) else: raise ValueError( 'Unknown TestStream Event type: %s' % proto.WhichOneof('event'))
[docs]class ElementEvent(Event): """Element-producing test stream event.""" def __init__(self, timestamped_values, tag=None): self.timestamped_values = timestamped_values self.tag = tag def __eq__(self, other): if not isinstance(other, ElementEvent): return False return ( self.timestamped_values == other.timestamped_values and self.tag == other.tag) def __hash__(self): return hash(self.timestamped_values) def __lt__(self, other): if not isinstance(other, ElementEvent): raise TypeError return self.timestamped_values < other.timestamped_values
[docs] def to_runner_api(self, element_coder): tag = 'None' if self.tag is None else self.tag return beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event.AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=element_coder.encode(tv.value), timestamp=tv.timestamp.micros // 1000) for tv in self.timestamped_values ], tag=tag))
def __repr__(self): return 'ElementEvent: <{}, {}>'.format([(e.value, e.timestamp) for e in self.timestamped_values], self.tag)
[docs]class WatermarkEvent(Event): """Watermark-advancing test stream event.""" def __init__(self, new_watermark, tag=None): self.new_watermark = Timestamp.of(new_watermark) self.tag = tag def __eq__(self, other): if not isinstance(other, WatermarkEvent): return False return self.new_watermark == other.new_watermark and self.tag == other.tag def __hash__(self): return hash(str(self.new_watermark) + str(self.tag)) def __lt__(self, other): if not isinstance(other, WatermarkEvent): raise TypeError return self.new_watermark < other.new_watermark
[docs] def to_runner_api(self, unused_element_coder): tag = 'None' if self.tag is None else self.tag # Assert that no precision is lost. assert self.new_watermark.micros % 1000 == 0 return beam_runner_api_pb2.TestStreamPayload.Event( watermark_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceWatermark( new_watermark=self.new_watermark.micros // 1000, tag=tag))
def __repr__(self): return 'WatermarkEvent: <{}, {}>'.format(self.new_watermark, self.tag)
[docs]class ProcessingTimeEvent(Event): """Processing time-advancing test stream event.""" def __init__(self, advance_by): self.advance_by = Duration.of(advance_by) def __eq__(self, other): if not isinstance(other, ProcessingTimeEvent): return False return self.advance_by == other.advance_by def __hash__(self): return hash(self.advance_by) def __lt__(self, other): if not isinstance(other, ProcessingTimeEvent): raise TypeError return self.advance_by < other.advance_by
[docs] def to_runner_api(self, unused_element_coder): assert self.advance_by.micros % 1000 == 0 return beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=self.advance_by.micros // 1000))
def __repr__(self): return 'ProcessingTimeEvent: <{}>'.format(self.advance_by)
class WindowedValueHolderMeta(type): """A metaclass that overrides the isinstance check for WindowedValueHolder. Python does a quick test for exact match. If an instance is exactly of type WindowedValueHolder, the overridden isinstance check is omitted. The override is needed because WindowedValueHolder elements encoded then decoded become Row elements. """ def __instancecheck__(cls, other): """Checks if a beam.Row typed instance is a WindowedValueHolder. """ return ( isinstance(other, beam.Row) and hasattr(other, 'windowed_value') and hasattr(other, 'urn') and isinstance(other.windowed_value, WindowedValue) and other.urn == common_urns.coders.ROW.urn) class WindowedValueHolder(beam.Row, metaclass=WindowedValueHolderMeta): """A class that holds a WindowedValue. This is a special class that can be used by the runner that implements the TestStream as a signal that the underlying value should be unreified to the specified window. """ # Register WindowedValueHolder to always use RowCoder. coders.registry.register_coder(WindowedValueHolderMeta, coders.RowCoder) def __init__(self, windowed_value): assert isinstance(windowed_value, WindowedValue), ( 'WindowedValueHolder can only hold %s type. Instead, %s is given.') % ( WindowedValue, windowed_value) super().__init__( **{ 'windowed_value': windowed_value, 'urn': common_urns.coders.ROW.urn }) @classmethod def from_row(cls, row): """Converts a beam.Row typed instance to WindowedValueHolder. """ if isinstance(row, WindowedValueHolder): return WindowedValueHolder(row.windowed_value) assert isinstance(row, beam.Row), 'The given row %s must be a %s type' % ( row, beam.Row) assert hasattr(row, 'windowed_value'), ( 'The given %s must have a windowed_value attribute.') % row assert isinstance(row.windowed_value, WindowedValue), ( 'The windowed_value attribute of %s must be a %s type') % ( row, WindowedValue)
[docs]class TestStream(PTransform): """Test stream that generates events on an unbounded PCollection of elements. Each event emits elements, advances the watermark or advances the processing time. After all of the specified elements are emitted, ceases to produce output. Applying the PTransform will return a single PCollection if only the default output or only one output tag has been used. Otherwise a dictionary of output names to PCollections will be returned. """ def __init__( self, coder=coders.FastPrimitivesCoder(), events=None, output_tags=None, endpoint=None): """ Args: coder: (apache_beam.Coder) the coder to encode/decode elements. events: (List[Event]) a list of instructions for the TestStream to execute. If specified, the events tags must exist in the output_tags. output_tags: (List[str]) Initial set of outputs. If no event references an output tag, no output will be produced for that tag. endpoint: (str) a URL locating a TestStreamService. """ super().__init__() assert coder is not None self.coder = coder self.watermarks = {None: timestamp.MIN_TIMESTAMP} self.output_tags = set(output_tags) if output_tags else set() self._events = [] if events is None else list(events) self._endpoint = endpoint event_tags = set( e.tag for e in self._events if isinstance(e, (WatermarkEvent, ElementEvent))) assert event_tags.issubset(self.output_tags), \ '{} is not a subset of {}'.format(event_tags, output_tags) assert not (self._events and self._endpoint), \ 'Only either events or an endpoint can be given at once.'
[docs] def get_windowing(self, unused_inputs): return core.Windowing(window.GlobalWindows())
def _infer_output_coder(self, input_type=None, input_coder=None): return self.coder
[docs] def expand(self, pbegin): assert isinstance(pbegin, pvalue.PBegin) self.pipeline = pbegin.pipeline if not self.output_tags: self.output_tags = {None} # For backwards compatibility return a single PCollection. if self.output_tags == {None}: return pvalue.PCollection( self.pipeline, is_bounded=False, tag=list(self.output_tags)[0]) return { tag: pvalue.PCollection(self.pipeline, is_bounded=False, tag=tag) for tag in self.output_tags }
def _add(self, event): if isinstance(event, ElementEvent): for tv in event.timestamped_values: assert tv.timestamp < timestamp.MAX_TIMESTAMP, ( 'Element timestamp must be before timestamp.MAX_TIMESTAMP.') elif isinstance(event, WatermarkEvent): if event.tag not in self.watermarks: self.watermarks[event.tag] = timestamp.MIN_TIMESTAMP assert event.new_watermark > self.watermarks[event.tag], ( 'Watermark must strictly-monotonically advance.') self.watermarks[event.tag] = event.new_watermark elif isinstance(event, ProcessingTimeEvent): assert event.advance_by > 0, ( 'Must advance processing time by positive amount.') else: raise ValueError('Unknown event: %s' % event) self._events.append(event)
[docs] def add_elements(self, elements, tag=None, event_timestamp=None): """Add elements to the TestStream. Elements added to the TestStream will be produced during pipeline execution. These elements can be TimestampedValue, WindowedValue or raw unwrapped elements that are serializable using the TestStream's specified Coder. When a TimestampedValue or a WindowedValue element is used, the timestamp of the TimestampedValue or WindowedValue will be the timestamp of the produced element; otherwise, the current watermark timestamp will be used for that element. The windows of a given WindowedValue are ignored by the TestStream. """ self.output_tags.add(tag) timestamped_values = [] if tag not in self.watermarks: self.watermarks[tag] = timestamp.MIN_TIMESTAMP for element in elements: if isinstance(element, TimestampedValue): timestamped_values.append(element) elif isinstance(element, WindowedValue): # Drop windows for elements in test stream. timestamped_values.append( TimestampedValue(element.value, element.timestamp)) else: # Add elements with timestamp equal to current watermark. if event_timestamp is None: event_timestamp = self.watermarks[tag] timestamped_values.append(TimestampedValue(element, event_timestamp)) self._add(ElementEvent(timestamped_values, tag)) return self
[docs] def advance_watermark_to(self, new_watermark, tag=None): """Advance the watermark to a given Unix timestamp. The Unix timestamp value used must be later than the previous watermark value and should be given as an int, float or utils.timestamp.Timestamp object. """ self.output_tags.add(tag) self._add(WatermarkEvent(new_watermark, tag)) return self
[docs] def advance_watermark_to_infinity(self, tag=None): """Advance the watermark to the end of time, completing this TestStream.""" self.advance_watermark_to(timestamp.MAX_TIMESTAMP, tag) return self
[docs] def advance_processing_time(self, advance_by): """Advance the current processing time by a given duration in seconds. The duration must be a positive second duration and should be given as an int, float or utils.timestamp.Duration object. """ self._add(ProcessingTimeEvent(advance_by)) return self
[docs] def to_runner_api_parameter(self, context): # Sort the output tags so that the order is deterministic and we are able # to test equality on a roundtrip through the to/from proto apis. return ( common_urns.primitives.TEST_STREAM.urn, beam_runner_api_pb2.TestStreamPayload( coder_id=context.coders.get_id(self.coder), events=[e.to_runner_api(self.coder) for e in self._events], endpoint=endpoints_pb2.ApiServiceDescriptor(url=self._endpoint)))
[docs] @staticmethod @PTransform.register_urn( common_urns.primitives.TEST_STREAM.urn, beam_runner_api_pb2.TestStreamPayload) def from_runner_api_parameter(ptransform, payload, context): coder = context.coders.get_by_id(payload.coder_id) output_tags = set( None if k == 'None' else k for k in ptransform.outputs.keys()) return TestStream( coder=coder, events=[Event.from_runner_api(e, coder) for e in payload.events], output_tags=output_tags, endpoint=payload.endpoint.url)
class TimingInfo(object): def __init__(self, processing_time, watermark): self._processing_time = Timestamp.of(processing_time) self._watermark = Timestamp.of(watermark) @property def processing_time(self): return self._processing_time @property def watermark(self): return self._watermark def __repr__(self): return '({}, {})'.format(self.processing_time, self.watermark) class PairWithTiming(PTransform): """Pairs the input element with timing information. Input: element; output: KV(element, timing information) Where timing information := (processing time, watermark) This is used in the ReverseTestStream implementation to replay watermark advancements. """ URN = "beam:transform:pair_with_timing:v1" def expand(self, pcoll): return pvalue.PCollection.from_(pcoll) class OutputFormat(Enum): TEST_STREAM_EVENTS = 1 TEST_STREAM_FILE_RECORDS = 2 SERIALIZED_TEST_STREAM_FILE_RECORDS = 3 class ReverseTestStream(PTransform): """A Transform that can create TestStream events from a stream of elements. This currently assumes that this the pipeline being run on a single machine and elements come in order and are outputted in the same order that they came in. """ def __init__( self, sample_resolution_sec, output_tag, coder=None, output_format=None): self._sample_resolution_sec = sample_resolution_sec self._output_tag = output_tag self._output_format = output_format if output_format \ else OutputFormat.TEST_STREAM_EVENTS self._coder = coder if coder else beam.coders.FastPrimitivesCoder() def expand(self, pcoll): ret = ( pcoll | beam.WindowInto(beam.window.GlobalWindows()) # First get the initial timing information. This will be used to start # the periodic timers which will generate processing time and watermark # advancements every `sample_resolution_sec`. | 'initial timing' >> PairWithTiming() # Next, map every element to the same key so that only a single timer is # started for this given ReverseTestStream. | 'first key' >> beam.Map(lambda x: (0, x)) # Next, pass-through each element which will be paired with its timing # info in the next step. Also, start the periodic timers. We use timers # in this situation to capture watermark advancements that occur when # there are no elements being produced upstream. | beam.ParDo( _TimingEventGenerator( output_tag=self._output_tag, sample_resolution_sec=self._sample_resolution_sec)) # Next, retrieve the timing information for watermark events that were # generated in the previous step. This is because elements generated # through the timers don't have their timing information yet. | 'timing info for watermarks' >> PairWithTiming() # Re-key to the same key to keep global state. | 'second key' >> beam.Map(lambda x: (0, x)) # Format the events properly. | beam.ParDo(_TestStreamFormatter(self._coder, self._output_format))) if self._output_format == OutputFormat.SERIALIZED_TEST_STREAM_FILE_RECORDS: def serializer(e): return e.SerializeToString() ret = ret | 'serializer' >> beam.Map(serializer) return ret class _TimingEventGenerator(beam.DoFn): """Generates ProcessingTimeEvents and WatermarkEvents at a regular cadence. The runner keeps the state of the clock (which may be faked) and the watermarks, which are inaccessible to SDKs. This DoFn generates ProcessingTimeEvents and WatermarkEvents at a specified sampling rate to capture any clock or watermark advancements between elements. """ # Used to return the initial timing information. EXECUTE_ONCE_STATE = beam.transforms.userstate.BagStateSpec( name='execute_once_state', coder=beam.coders.FastPrimitivesCoder()) # A processing time timer in an infinite loop that generates the events that # will be paired with the TimingInfo from the runner. TIMING_SAMPLER = TimerSpec('timing_sampler', TimeDomain.REAL_TIME) def __init__(self, output_tag, sample_resolution_sec=0.1): self._output_tag = output_tag self._sample_resolution_sec = sample_resolution_sec @on_timer(TIMING_SAMPLER) def on_timing_sampler( self, timestamp=beam.DoFn.TimestampParam, window=beam.DoFn.WindowParam, timing_sampler=beam.DoFn.TimerParam(TIMING_SAMPLER)): """Yields an unbounded stream of ProcessingTimeEvents and WatermarkEvents. The returned events will be paired with the TimingInfo. This loop's only purpose is to generate these events even when there are no elements. """ next_sample_time = (timestamp.micros * 1e-6) + self._sample_resolution_sec timing_sampler.set(next_sample_time) # Generate two events, the delta since the last sample and a place-holder # WatermarkEvent. This is a placeholder because we can't otherwise add the # watermark from the runner to the event. yield ProcessingTimeEvent(self._sample_resolution_sec) yield WatermarkEvent(MIN_TIMESTAMP) def process( self, e, timestamp=beam.DoFn.TimestampParam, window=beam.DoFn.WindowParam, timing_sampler=beam.DoFn.TimerParam(TIMING_SAMPLER), execute_once_state=beam.DoFn.StateParam(EXECUTE_ONCE_STATE)): _, (element, timing_info) = e # Only set the timers once and only send the header once. first_time = next(execute_once_state.read(), True) if first_time: # Generate the initial timing events. execute_once_state.add(False) now_sec = timing_info.processing_time.micros * 1e-6 timing_sampler.set(now_sec + self._sample_resolution_sec) # Here we capture the initial time offset and initial watermark. This is # where we emit the TestStreamFileHeader. yield beam_interactive_api_pb2.TestStreamFileHeader(tag=self._output_tag) yield ProcessingTimeEvent( Duration(micros=timing_info.processing_time.micros)) yield WatermarkEvent(MIN_TIMESTAMP) yield element class _TestStreamFormatter(beam.DoFn): """Formats the events to the specified output format. """ # In order to generate the processing time deltas, we need to keep track of # the previous clock time we got from the runner. PREV_SAMPLE_TIME_STATE = beam.transforms.userstate.BagStateSpec( name='prev_sample_time_state', coder=beam.coders.FastPrimitivesCoder()) def __init__(self, coder, output_format): self._coder = coder self._output_format = output_format def start_bundle(self): self.elements = [] self.timing_events = [] self.header = None def finish_bundle(self): """Outputs all the buffered elements. """ if self._output_format == OutputFormat.TEST_STREAM_EVENTS: return self._output_as_events() return self._output_as_records() def process( self, e, timestamp=beam.DoFn.TimestampParam, prev_sample_time_state=beam.DoFn.StateParam(PREV_SAMPLE_TIME_STATE)): """Buffers elements until the end of the bundle. This buffers elements instead of emitting them immediately to keep elements that come in the same bundle to be outputted in the same bundle. """ _, (element, timing_info) = e if isinstance(element, beam_interactive_api_pb2.TestStreamFileHeader): self.header = element elif isinstance(element, WatermarkEvent): # WatermarkEvents come in with a watermark of MIN_TIMESTAMP. Fill in the # correct watermark from the runner here. element.new_watermark = timing_info.watermark.micros if element not in self.timing_events: self.timing_events.append(element) elif isinstance(element, ProcessingTimeEvent): # Because the runner holds the clock, calculate the processing time delta # here. The TestStream may have faked out the clock, and thus the # delta calculated in the SDK with time.time() will be wrong. prev_sample = next(prev_sample_time_state.read(), Timestamp()) prev_sample_time_state.clear() prev_sample_time_state.add(timing_info.processing_time) advance_by = timing_info.processing_time - prev_sample element.advance_by = advance_by self.timing_events.append(element) else: self.elements.append(TimestampedValue(element, timestamp)) def _output_as_events(self): """Outputs buffered elements as TestStream events. """ if self.timing_events: yield WindowedValue( self.timing_events, timestamp=0, windows=[beam.window.GlobalWindow()]) if self.elements: yield WindowedValue([ElementEvent(self.elements)], timestamp=0, windows=[beam.window.GlobalWindow()]) def _output_as_records(self): """Outputs buffered elements as TestStreamFileRecords. """ if self.header: yield WindowedValue( self.header, timestamp=0, windows=[beam.window.GlobalWindow()]) if self.timing_events: timing_events = self._timing_events_to_records(self.timing_events) for r in timing_events: yield WindowedValue( r, timestamp=0, windows=[beam.window.GlobalWindow()]) if self.elements: elements = self._elements_to_record(self.elements) yield WindowedValue( elements, timestamp=0, windows=[beam.window.GlobalWindow()]) def _timing_events_to_records(self, timing_events): """Returns given timing_events as TestStreamFileRecords. """ records = [] for e in self.timing_events: if isinstance(e, ProcessingTimeEvent): processing_time_event = beam_runner_api_pb2.\ TestStreamPayload.Event.AdvanceProcessingTime( advance_duration=e.advance_by.micros) records.append( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=processing_time_event))) elif isinstance(e, WatermarkEvent): watermark_event = beam_runner_api_pb2.\ TestStreamPayload.Event.AdvanceWatermark( new_watermark=int(e.new_watermark)) records.append( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload.Event( watermark_event=watermark_event))) return records def _elements_to_record(self, elements): """Returns elements as TestStreamFileRecords. """ elements = [] for tv in self.elements: element_timestamp = tv.timestamp.micros element = beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=self._coder.encode(tv.value), timestamp=element_timestamp) elements.append(element) element_event = beam_runner_api_pb2.TestStreamPayload.Event.AddElements( elements=elements) return beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload.Event( element_event=element_event))