#

# 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',

]

@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

@abstractmethod

def to_runner_api(self, element_coder):

raise NotImplementedError

@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'))

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

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)

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

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)

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

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)

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.'

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

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)

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

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

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

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

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)))

@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))