#

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

#

"""Unit tests for side inputs."""

# pytype: skip-file

import hashlib

import itertools

import logging

import unittest

from typing import Any

from typing import Dict

from typing import Iterable

from typing import Tuple

from typing import Union

import pytest

import apache_beam as beam

from apache_beam.testing.synthetic_pipeline import SyntheticSDFAsSource

from apache_beam.testing.test_pipeline import TestPipeline

from apache_beam.testing.test_stream import TestStream

from apache_beam.testing.util import assert_that

from apache_beam.testing.util import equal_to

from apache_beam.testing.util import equal_to_per_window

from apache_beam.transforms import Map

from apache_beam.transforms import sideinputs

from apache_beam.transforms import trigger

from apache_beam.transforms import window

from apache_beam.utils.timestamp import Timestamp

class SideInputsTest(unittest.TestCase):

def create_pipeline(self):

return TestPipeline()

def run_windowed_side_inputs(

self,

elements,

main_window_fn,

side_window_fn=None,

side_input_type=beam.pvalue.AsList,

combine_fn=None,

expected=None):

with self.create_pipeline() as p:

pcoll = p | beam.Create(elements) | beam.Map(

lambda t: window.TimestampedValue(t, t))

main = pcoll | 'WindowMain' >> beam.WindowInto(main_window_fn)

side = pcoll | 'WindowSide' >> beam.WindowInto(

side_window_fn or main_window_fn)

kw = {}

if combine_fn is not None:

side |= beam.CombineGlobally(combine_fn).without_defaults()

kw['default_value'] = 0

elif side_input_type == beam.pvalue.AsDict:

side |= beam.Map(lambda x: ('k%s' % x, 'v%s' % x))

res = main | beam.Map(lambda x, s: (x, s), side_input_type(side, **kw))

if side_input_type in (beam.pvalue.AsIter, beam.pvalue.AsList):

res |= beam.Map(lambda x_s: (x_s[0], sorted(x_s[1])))

assert_that(res, equal_to(expected))

def test_global_global_windows(self):

self.run_windowed_side_inputs([1, 2, 3],

window.GlobalWindows(),

expected=[(1, [1, 2, 3]), (2, [1, 2, 3]),

(3, [1, 2, 3])])

def test_same_fixed_windows(self):

self.run_windowed_side_inputs([1, 2, 11],

window.FixedWindows(10),

expected=[(1, [1, 2]), (2, [1, 2]),

(11, [11])])

def test_different_fixed_windows(self):

self.run_windowed_side_inputs([1, 2, 11, 21, 31],

window.FixedWindows(10),

window.FixedWindows(20),

expected=[(1, [1, 2, 11]), (2, [1, 2, 11]),

(11, [1, 2, 11]), (21, [21, 31]),

(31, [21, 31])])

def test_fixed_global_window(self):

self.run_windowed_side_inputs([1, 2, 11],

window.FixedWindows(10),

window.GlobalWindows(),

expected=[(1, [1, 2, 11]), (2, [1, 2, 11]),

(11, [1, 2, 11])])

def test_sliding_windows(self):

self.run_windowed_side_inputs(

[1, 2, 4],

window.SlidingWindows(size=6, period=2),

window.SlidingWindows(size=6, period=2),

expected=[

# Element 1 falls in three windows

(1, [1]), # [-4, 2)

(1, [1, 2]), # [-2, 4)

(1, [1, 2, 4]), # [0, 6)

# as does 2,

(2, [1, 2]), # [-2, 4)

(2, [1, 2, 4]), # [0, 6)

(2, [2, 4]), # [2, 8)

# and 4.

(4, [1, 2, 4]), # [0, 6)

(4, [2, 4]), # [2, 8)

(4, [4]), # [4, 10)

])

def test_windowed_iter(self):

self.run_windowed_side_inputs([1, 2, 11],

window.FixedWindows(10),

side_input_type=beam.pvalue.AsIter,

expected=[(1, [1, 2]), (2, [1, 2]),

(11, [11])])

def test_windowed_singleton(self):

self.run_windowed_side_inputs([1, 2, 11],

window.FixedWindows(10),

side_input_type=beam.pvalue.AsSingleton,

combine_fn=sum,

expected=[(1, 3), (2, 3), (11, 11)])

def test_windowed_dict(self):

self.run_windowed_side_inputs([1, 2, 11],

window.FixedWindows(10),

side_input_type=beam.pvalue.AsDict,

expected=[

(1, {

'k1': 'v1', 'k2': 'v2'

}),

(2, {

'k1': 'v1', 'k2': 'v2'

}),

(11, {

'k11': 'v11'

}),

])

@pytest.mark.it_validatesrunner

def test_empty_singleton_side_input(self):

pipeline = self.create_pipeline()

pcol = pipeline | 'start' >> beam.Create([1, 2])

side = pipeline | 'side' >> beam.Create([]) # Empty side input.

def my_fn(k, s):

# TODO(robertwb): Should this be an error as in Java?

v = ('empty' if isinstance(s, beam.pvalue.EmptySideInput) else 'full')

return [(k, v)]

result = pcol | 'compute' >> beam.FlatMap(

my_fn, beam.pvalue.AsSingleton(side))

assert_that(result, equal_to([(1, 'empty'), (2, 'empty')]))

pipeline.run()

# TODO(https://github.com/apache/beam/issues/19012): Disable this test in

# streaming temporarily. Remove sickbay-streaming tag after it's fixed.

@pytest.mark.no_sickbay_streaming

@pytest.mark.it_validatesrunner

def test_multi_valued_singleton_side_input(self):

pipeline = self.create_pipeline()

pcol = pipeline | 'start' >> beam.Create([1, 2])

side = pipeline | 'side' >> beam.Create([3, 4]) # 2 values in side input.

pcol | 'compute' >> beam.FlatMap( # pylint: disable=expression-not-assigned

lambda x, s: [x * s], beam.pvalue.AsSingleton(side))

with self.assertRaises(Exception):

pipeline.run()

@pytest.mark.it_validatesrunner

def test_default_value_singleton_side_input(self):

pipeline = self.create_pipeline()

pcol = pipeline | 'start' >> beam.Create([1, 2])

side = pipeline | 'side' >> beam.Create([]) # 0 values in side input.

result = pcol | beam.FlatMap(

lambda x, s: [x * s], beam.pvalue.AsSingleton(side, 10))

assert_that(result, equal_to([10, 20]))

pipeline.run()

@pytest.mark.it_validatesrunner

def test_iterable_side_input(self):

pipeline = self.create_pipeline()

pcol = pipeline | 'start' >> beam.Create([1, 2])

side = pipeline | 'side' >> beam.Create([3, 4]) # 2 values in side input.

result = pcol | 'compute' >> beam.FlatMap(

lambda x, s: [x * y for y in s], beam.pvalue.AsIter(side))

assert_that(result, equal_to([3, 4, 6, 8]))

pipeline.run()

@pytest.mark.it_validatesrunner

def test_reiterable_side_input(self):

expected_side = frozenset(range(100))

def check_reiteration(main, side):

assert expected_side == set(side), side

# Iterate a second time.

assert expected_side == set(side), side

# Iterate over two copies of the input at the same time.

both = zip(side, side)

first, second = zip(*both)

assert expected_side == set(first), first

assert expected_side == set(second), second

# This will iterate over two copies of the side input, but offset.

offset = [None] * (len(expected_side) // 2)

both = zip(itertools.chain(side, offset), itertools.chain(offset, side))

first, second = zip(*both)

expected_and_none = frozenset.union(expected_side, [None])

assert expected_and_none == set(first), first

assert expected_and_none == set(second), second

pipeline = self.create_pipeline()

pcol = pipeline | 'start' >> beam.Create(['A', 'B'])

side = pipeline | 'side' >> beam.Create(expected_side)

_ = pcol | 'check' >> beam.Map(check_reiteration, beam.pvalue.AsIter(side))

pipeline.run()

@pytest.mark.it_validatesrunner

def test_as_list_and_as_dict_side_inputs(self):

a_list = [5, 1, 3, 2, 9]

some_pairs = [('crouton', 17), ('supreme', None)]

pipeline = self.create_pipeline()

main_input = pipeline | 'main input' >> beam.Create([1])

side_list = pipeline | 'side list' >> beam.Create(a_list)

side_pairs = pipeline | 'side pairs' >> beam.Create(some_pairs)

results = main_input | 'concatenate' >> beam.Map(

lambda x, the_list, the_dict: [x, the_list, the_dict],

beam.pvalue.AsList(side_list),

beam.pvalue.AsDict(side_pairs))

def matcher(expected_elem, expected_list, expected_pairs):

def match(actual):

[[actual_elem, actual_list, actual_dict]] = actual

equal_to([expected_elem])([actual_elem])

equal_to(expected_list)(actual_list)

equal_to(expected_pairs)(actual_dict.items())

return match

assert_that(results, matcher(1, a_list, some_pairs))

pipeline.run()

@pytest.mark.it_validatesrunner

def test_as_singleton_without_unique_labels(self):

# This should succeed as calling beam.pvalue.AsSingleton on the same

# PCollection twice with the same defaults will return the same

# view.

a_list = [2]

pipeline = self.create_pipeline()

main_input = pipeline | 'main input' >> beam.Create([1])

side_list = pipeline | 'side list' >> beam.Create(a_list)

results = main_input | beam.Map(

lambda x, s1, s2: [x, s1, s2],

beam.pvalue.AsSingleton(side_list),

beam.pvalue.AsSingleton(side_list))

def matcher(expected_elem, expected_singleton):

def match(actual):

[[actual_elem, actual_singleton1, actual_singleton2]] = actual

equal_to([expected_elem])([actual_elem])

equal_to([expected_singleton])([actual_singleton1])

equal_to([expected_singleton])([actual_singleton2])

return match

assert_that(results, matcher(1, 2))

pipeline.run()

@pytest.mark.it_validatesrunner

def test_as_singleton_with_different_defaults(self):

a_list = []

pipeline = self.create_pipeline()

main_input = pipeline | 'main input' >> beam.Create([1])

side_list = pipeline | 'side list' >> beam.Create(a_list)

results = main_input | beam.Map(

lambda x, s1, s2: [x, s1, s2],

beam.pvalue.AsSingleton(side_list, default_value=2),

beam.pvalue.AsSingleton(side_list, default_value=3))

def matcher(expected_elem, expected_singleton1, expected_singleton2):

def match(actual):

[[actual_elem, actual_singleton1, actual_singleton2]] = actual

equal_to([expected_elem])([actual_elem])

equal_to([expected_singleton1])([actual_singleton1])

equal_to([expected_singleton2])([actual_singleton2])

return match

assert_that(results, matcher(1, 2, 3))

pipeline.run()

@pytest.mark.it_validatesrunner

def test_as_list_twice(self):

# This should succeed as calling beam.pvalue.AsList on the same

# PCollection twice will return the same view.

a_list = [1, 2, 3]

pipeline = self.create_pipeline()

main_input = pipeline | 'main input' >> beam.Create([1])

side_list = pipeline | 'side list' >> beam.Create(a_list)

results = main_input | beam.Map(

lambda x, ls1, ls2: [x, ls1, ls2],

beam.pvalue.AsList(side_list),

beam.pvalue.AsList(side_list))

def matcher(expected_elem, expected_list):

def match(actual):

[[actual_elem, actual_list1, actual_list2]] = actual

equal_to([expected_elem])([actual_elem])

equal_to(expected_list)(actual_list1)

equal_to(expected_list)(actual_list2)

return match

assert_that(results, matcher(1, [1, 2, 3]))

pipeline.run()

@pytest.mark.it_validatesrunner

def test_as_dict_twice(self):

some_kvs = [('a', 1), ('b', 2)]

pipeline = self.create_pipeline()

main_input = pipeline | 'main input' >> beam.Create([1])

side_kvs = pipeline | 'side kvs' >> beam.Create(some_kvs)

results = main_input | beam.Map(

lambda x, dct1, dct2: [x, dct1, dct2],

beam.pvalue.AsDict(side_kvs),

beam.pvalue.AsDict(side_kvs))

def matcher(expected_elem, expected_kvs):

def match(actual):

[[actual_elem, actual_dict1, actual_dict2]] = actual

equal_to([expected_elem])([actual_elem])

equal_to(expected_kvs)(actual_dict1.items())

equal_to(expected_kvs)(actual_dict2.items())

return match

assert_that(results, matcher(1, some_kvs))

pipeline.run()

@pytest.mark.it_validatesrunner

def test_flattened_side_input(self):

pipeline = self.create_pipeline()

main_input = pipeline | 'main input' >> beam.Create([None])

side_input = (

pipeline | 'side1' >> beam.Create(['a']),

pipeline | 'side2' >> beam.Create(['b'])) | beam.Flatten()

results = main_input | beam.FlatMap(

lambda _, ab: ab, beam.pvalue.AsList(side_input))

assert_that(results, equal_to(['a', 'b']))

pipeline.run()

# TODO(BEAM-9499): Disable this test in streaming temporarily.

@pytest.mark.no_sickbay_batch

@pytest.mark.no_sickbay_streaming

@pytest.mark.it_validatesrunner

def test_multi_triggered_gbk_side_input(self):

"""Test a GBK sideinput, with multiple triggering."""

# TODO(https://github.com/apache/beam/issues/20065): Remove use of this

# experiment. This flag is only necessary when using the multi-output

# TestStream b/c it relies on using the PCollection output tags as the

# PCollection output ids.

with TestPipeline() as p:

test_stream = (

p

| 'Mixed TestStream' >> TestStream().advance_watermark_to(

3,

tag='main').add_elements(['a1'], tag='main').advance_watermark_to(

8, tag='main').add_elements(['a2'], tag='main').add_elements(

[window.TimestampedValue(('k', 100), 2)], tag='side').

add_elements([window.TimestampedValue(

('k', 400), 7)], tag='side').advance_watermark_to_infinity(

tag='main').advance_watermark_to_infinity(tag='side'))

main_data = (

test_stream['main']

| 'Main windowInto' >> beam.WindowInto(

window.FixedWindows(5),

accumulation_mode=trigger.AccumulationMode.DISCARDING))

side_data = (

test_stream['side']

| 'Side windowInto' >> beam.WindowInto(

window.FixedWindows(5),

trigger=trigger.AfterWatermark(early=trigger.AfterCount(1)),

accumulation_mode=trigger.AccumulationMode.DISCARDING)

| beam.CombinePerKey(sum)

| 'Values' >> Map(lambda k_vs: k_vs[1]))

class RecordFn(beam.DoFn):

def process(

self,

elm=beam.DoFn.ElementParam,

ts=beam.DoFn.TimestampParam,

side=beam.DoFn.SideInputParam):

yield (elm, ts, side)

records = (

main_data

| beam.ParDo(RecordFn(), beam.pvalue.AsList(side_data)))

expected_window_to_elements = {

window.IntervalWindow(0, 5): [

('a1', Timestamp(3), [100, 0]),

],

window.IntervalWindow(5, 10): [('a2', Timestamp(8), [400, 0])],

}

assert_that(

records,

equal_to_per_window(expected_window_to_elements),

use_global_window=False,

label='assert per window')

@pytest.mark.it_validatesrunner

def test_side_input_with_sdf(self):

"""Test a side input with SDF.

This test verifies consisency of side input when it is split due to

SDF (Splittable DoFns). The consistency is verified by checking the size

and fingerprint of the side input.

This test needs to run with at least 2 workers (--num_workers=2) and

autoscaling disabled (--autoscaling_algorithm=NONE). Otherwise it might

provide false positives (i.e. not fail on bad state).

"""

initial_elements = 1000

num_records = 10000

key_size = 10

value_size = 100

expected_fingerprint = '00f7eeac8514721e2683d14a504b33d1'

class GetSyntheticSDFOptions(beam.DoFn):

"""A DoFn that emits elements for genenrating SDF."""

def process(self, element: Any) -> Iterable[Dict[str, Union[int, str]]]:

yield {

'num_records': num_records // initial_elements,

'key_size': key_size,

'value_size': value_size,

'initial_splitting_num_bundles': 0,

'initial_splitting_desired_bundle_size': 0,

'sleep_per_input_record_sec': 0,

'initial_splitting': 'const',

}

class SideInputTrackingDoFn(beam.DoFn):

"""A DoFn that emits the size and fingerprint of the side input.

In this context, the size is the number of elements and the fingerprint

is the hash of the sorted serialized elements.

"""

def process(

self, element: Any,

side_input: Iterable[Tuple[bytes,

bytes]]) -> Iterable[Tuple[int, str]]:

# Sort for consistent hashing.

sorted_side_input = sorted(side_input)

size = len(sorted_side_input)

m = hashlib.md5()

for key, value in sorted_side_input:

m.update(key)

m.update(value)

yield (size, m.hexdigest())

pipeline = self.create_pipeline()

main_input = pipeline | 'Main input: Create' >> beam.Create([0])

side_input = pipeline | 'Side input: Create' >> beam.Create(

range(initial_elements))

side_input |= 'Side input: Get synthetic SDF options' >> beam.ParDo(

GetSyntheticSDFOptions())

side_input |= 'Side input: Process and split' >> beam.ParDo(

SyntheticSDFAsSource())

results = main_input | 'Emit side input' >> beam.ParDo(

SideInputTrackingDoFn(), beam.pvalue.AsIter(side_input))

assert_that(results, equal_to([(num_records, expected_fingerprint)]))

pipeline.run()

def test_default_window_mapping_fn_source_window(self):

"""Test that the default window mapping function will propagate the

source window when attempting to assign context.

"""

class StringIDWindow(window.BoundedWindow):

"""A window defined by an arbitrary string ID."""

def __init__(self, window_id: str):

super().__init__(self._getTimestampFromProto())

self.id = window_id

@staticmethod

def _getTimestampFromProto() -> Timestamp:

return Timestamp(micros=0)

class StringIDWindows(window.NonMergingWindowFn):

""" A windowing function that assigns each element a window with ID."""

def assign(

self, assign_context: window.WindowFn.AssignContext

) -> Iterable[window.BoundedWindow]:

if assign_context.element is None:

assert assign_context.window is not None

return [assign_context.window]

return [StringIDWindow(str(assign_context.element))]

def get_window_coder(self):

return None

mapping_fn = sideinputs.default_window_mapping_fn(StringIDWindows())

source_window = StringIDWindows().assign(

window.WindowFn.AssignContext(Timestamp(10), element='element'))[0]

bounded_window = mapping_fn(source_window)

assert bounded_window is not None

assert bounded_window.id == 'element'

if __name__ == '__main__':

logging.getLogger().setLevel(logging.DEBUG)

unittest.main()