#

# 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 the PTransform and descendants."""

# pytype: skip-file

import collections

import operator

import os

import pickle

import random

import re

import sys

import typing

import unittest

from functools import reduce

from typing import Optional

from unittest.mock import patch

import hamcrest as hc

import numpy as np

import pytest

from parameterized import param

from parameterized import parameterized

from parameterized import parameterized_class

import apache_beam as beam

import apache_beam.transforms.combiners as combine

from apache_beam import pvalue

from apache_beam import typehints

from apache_beam.coders import coders_test_common

from apache_beam.io.iobase import Read

from apache_beam.metrics import Metrics

from apache_beam.metrics.metric import MetricsFilter

from apache_beam.options.pipeline_options import PipelineOptions

from apache_beam.options.pipeline_options import StandardOptions

from apache_beam.options.pipeline_options import StreamingOptions

from apache_beam.options.pipeline_options import TypeOptions

from apache_beam.portability import common_urns

from apache_beam.testing.test_pipeline import TestPipeline

from apache_beam.testing.test_stream import TestStream

from apache_beam.testing.util import SortLists

from apache_beam.testing.util import assert_that

from apache_beam.testing.util import equal_to

from apache_beam.transforms import WindowInto

from apache_beam.transforms import trigger

from apache_beam.transforms import window

from apache_beam.transforms.display import DisplayData

from apache_beam.transforms.display import DisplayDataItem

from apache_beam.transforms.ptransform import PTransform

from apache_beam.transforms.trigger import AccumulationMode

from apache_beam.transforms.trigger import AfterProcessingTime

from apache_beam.transforms.trigger import _AfterSynchronizedProcessingTime

from apache_beam.transforms.window import TimestampedValue

from apache_beam.typehints import with_input_types

from apache_beam.typehints import with_output_types

from apache_beam.typehints.typehints_test import TypeHintTestCase

from apache_beam.utils.timestamp import Timestamp

from apache_beam.utils.windowed_value import WindowedValue

# Disable frequent lint warning due to pipe operator for chaining transforms.

# pylint: disable=expression-not-assigned

class PTransformTest(unittest.TestCase):

def assertStartswith(self, msg, prefix):

self.assertTrue(

msg.startswith(prefix), '"%s" does not start with "%s"' % (msg, prefix))

def test_str(self):

self.assertEqual(

'<PTransform(PTransform) label=[PTransform]>', str(PTransform()))

pa = TestPipeline()

res = pa | 'ALabel' >> beam.Impulse()

self.assertEqual('AppliedPTransform(ALabel, Impulse)', str(res.producer))

pc = TestPipeline()

res = pc | beam.Impulse()

inputs_tr = res.producer.transform

inputs_tr.inputs = ('ci', )

self.assertEqual(

"<Impulse(PTransform) label=[Impulse] inputs=('ci',)>", str(inputs_tr))

pd = TestPipeline()

res = pd | beam.Impulse()

side_tr = res.producer.transform

side_tr.side_inputs = (4, )

self.assertEqual(

'<Impulse(PTransform) label=[Impulse] side_inputs=(4,)>', str(side_tr))

inputs_tr.side_inputs = ('cs', )

self.assertEqual(

"""<Impulse(PTransform) label=[Impulse] """

"""inputs=('ci',) side_inputs=('cs',)>""",

str(inputs_tr))

def test_named_annotations(self):

t = beam.Impulse()

t.annotations = lambda: {'test': 'value'}

named_t = 'Name' >> t

self.assertEqual(named_t.annotations(), {'test': 'value'})

original_annotations = named_t.annotations()

named_t.annotations = lambda: {'another': 'value', **original_annotations}

# Verify this is reflected on the original transform,

# which is what gets used in apply.

self.assertEqual(t.annotations(), {'test': 'value', 'another': 'value'})

def test_do_with_do_fn(self):

class AddNDoFn(beam.DoFn):

def process(self, element, addon):

return [element + addon]

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])

result = pcoll | 'Do' >> beam.ParDo(AddNDoFn(), 10)

assert_that(result, equal_to([11, 12, 13]))

def test_do_with_unconstructed_do_fn(self):

class MyDoFn(beam.DoFn):

def process(self):

pass

with self.assertRaises(ValueError):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])

pcoll | 'Do' >> beam.ParDo(MyDoFn) # Note the lack of ()'s

def test_do_with_callable(self):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])

result = pcoll | 'Do' >> beam.FlatMap(lambda x, addon: [x + addon], 10)

assert_that(result, equal_to([11, 12, 13]))

def test_do_with_side_input_as_arg(self):

with TestPipeline() as pipeline:

side = pipeline | 'Side' >> beam.Create([10])

pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])

result = pcoll | 'Do' >> beam.FlatMap(

lambda x, addon: [x + addon], pvalue.AsSingleton(side))

assert_that(result, equal_to([11, 12, 13]))

def test_do_with_side_input_as_keyword_arg(self):

with TestPipeline() as pipeline:

side = pipeline | 'Side' >> beam.Create([10])

pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])

result = pcoll | 'Do' >> beam.FlatMap(

lambda x, addon: [x + addon], addon=pvalue.AsSingleton(side))

assert_that(result, equal_to([11, 12, 13]))

def test_callable_non_serializable_error_message(self):

class NonSerializable:

def __getstate__(self):

raise RuntimeError('nope')

bad = NonSerializable()

with self.assertRaises(RuntimeError) as context:

_ = beam.Map(lambda x: bad)

message = str(context.exception)

self.assertIn('Unable to pickle fn', message)

self.assertIn(

'User code must be serializable (picklable) for distributed execution.',

message)

self.assertIn('non-serializable objects like file handles', message)

self.assertIn(

'Try: (1) using module-level functions instead of lambdas', message)

def test_do_with_do_fn_returning_string_raises_warning(self):

ex_details = r'.*Returning a str from a ParDo or FlatMap is discouraged.'

with self.assertRaisesRegex(Exception, ex_details):

with TestPipeline() as pipeline:

pipeline._options.view_as(TypeOptions).runtime_type_check = True

pcoll = pipeline | 'Start' >> beam.Create(['2', '9', '3'])

pcoll | 'Do' >> beam.FlatMap(lambda x: x + '1')

# Since the DoFn directly returns a string we should get an

# error warning us when the pipeliene runs.

def test_do_with_do_fn_returning_dict_raises_warning(self):

ex_details = r'.*Returning a dict from a ParDo or FlatMap is discouraged.'

with self.assertRaisesRegex(Exception, ex_details):

with TestPipeline() as pipeline:

pipeline._options.view_as(TypeOptions).runtime_type_check = True

pcoll = pipeline | 'Start' >> beam.Create(['2', '9', '3'])

pcoll | 'Do' >> beam.FlatMap(lambda x: {x: '1'})

# Since the DoFn directly returns a dict we should get an error warning

# us when the pipeliene runs.

def test_do_with_multiple_outputs_maintains_unique_name(self):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])

r1 = pcoll | 'A' >> beam.FlatMap(lambda x: [x + 1]).with_outputs(main='m')

r2 = pcoll | 'B' >> beam.FlatMap(lambda x: [x + 2]).with_outputs(main='m')

assert_that(r1.m, equal_to([2, 3, 4]), label='r1')

assert_that(r2.m, equal_to([3, 4, 5]), label='r2')

@pytest.mark.it_validatesrunner

def test_impulse(self):

with TestPipeline() as pipeline:

result = pipeline | beam.Impulse() | beam.Map(lambda _: 0)

assert_that(result, equal_to([0]))

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

# Remove sickbay-streaming tag after it's resolved.

@pytest.mark.no_sickbay_streaming

@pytest.mark.it_validatesrunner

def test_read_metrics(self):

from apache_beam.io.utils import CountingSource

class CounterDoFn(beam.DoFn):

def __init__(self):

# This counter is unused.

self.received_records = Metrics.counter(

self.__class__, 'receivedRecords')

def process(self, element):

self.received_records.inc()

pipeline = TestPipeline()

(pipeline | Read(CountingSource(100)) | beam.ParDo(CounterDoFn()))

res = pipeline.run()

res.wait_until_finish()

# This counter is defined in utils.CountingSource.

metric_results = res.metrics().query(

MetricsFilter().with_name('recordsRead'))

outputs_counter = metric_results['counters'][0]

msg = outputs_counter.key.step

cont = 'SDFBoundedSourceReader'

self.assertTrue(cont in msg, '"%s" does not contain "%s"' % (msg, cont))

self.assertEqual(outputs_counter.key.metric.name, 'recordsRead')

self.assertEqual(outputs_counter.committed, 100)

@pytest.mark.it_validatesrunner

def test_par_do_with_multiple_outputs_and_using_yield(self):

class SomeDoFn(beam.DoFn):

"""A custom DoFn using yield."""

def process(self, element):

yield element

if element % 2 == 0:

yield pvalue.TaggedOutput('even', element)

else:

yield pvalue.TaggedOutput('odd', element)

with TestPipeline() as pipeline:

nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4])

results = nums | 'ClassifyNumbers' >> beam.ParDo(SomeDoFn()).with_outputs(

'odd', 'even', main='main')

assert_that(results.main, equal_to([1, 2, 3, 4]))

assert_that(results.odd, equal_to([1, 3]), label='assert:odd')

assert_that(results.even, equal_to([2, 4]), label='assert:even')

@pytest.mark.it_validatesrunner

def test_par_do_with_multiple_outputs_and_using_return(self):

def some_fn(v):

if v % 2 == 0:

return [v, pvalue.TaggedOutput('even', v)]

return [v, pvalue.TaggedOutput('odd', v)]

with TestPipeline() as pipeline:

nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4])

results = nums | 'ClassifyNumbers' >> beam.FlatMap(some_fn).with_outputs(

'odd', 'even', main='main')

assert_that(results.main, equal_to([1, 2, 3, 4]))

assert_that(results.odd, equal_to([1, 3]), label='assert:odd')

assert_that(results.even, equal_to([2, 4]), label='assert:even')

@pytest.mark.it_validatesrunner

def test_undeclared_outputs(self):

with TestPipeline() as pipeline:

nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4])

results = nums | 'ClassifyNumbers' >> beam.FlatMap(

lambda x: [

x, pvalue.TaggedOutput('even' if x % 2 == 0 else 'odd', x), pvalue

.TaggedOutput('extra', x)

]).with_outputs()

assert_that(results[None], equal_to([1, 2, 3, 4]))

assert_that(results.odd, equal_to([1, 3]), label='assert:odd')

assert_that(results.even, equal_to([2, 4]), label='assert:even')

@pytest.mark.it_validatesrunner

def test_multiple_empty_outputs(self):

with TestPipeline() as pipeline:

nums = pipeline | 'Some Numbers' >> beam.Create([1, 3, 5])

results = nums | 'ClassifyNumbers' >> beam.FlatMap(

lambda x:

[x, pvalue.TaggedOutput('even'

if x % 2 == 0 else 'odd', x)]).with_outputs()

assert_that(results[None], equal_to([1, 3, 5]))

assert_that(results.odd, equal_to([1, 3, 5]), label='assert:odd')

assert_that(results.even, equal_to([]), label='assert:even')

def test_do_requires_do_fn_returning_iterable(self):

# This function is incorrect because it returns an object that isn't an

# iterable.

def incorrect_par_do_fn(x):

return x + 5

ex_details = r'.*FlatMap and ParDo must return an iterable.'

with self.assertRaisesRegex(Exception, ex_details):

with TestPipeline() as pipeline:

pipeline._options.view_as(TypeOptions).runtime_type_check = True

pcoll = pipeline | 'Start' >> beam.Create([2, 9, 3])

pcoll | 'Do' >> beam.FlatMap(incorrect_par_do_fn)

# It's a requirement that all user-defined functions to a ParDo return

# an iterable.

def test_do_fn_with_finish(self):

class MyDoFn(beam.DoFn):

def process(self, element):

pass

def finish_bundle(self):

yield WindowedValue('finish', -1, [window.GlobalWindow()])

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])

result = pcoll | 'Do' >> beam.ParDo(MyDoFn())

# May have many bundles, but each has a start and finish.

def matcher():

def match(actual):

equal_to(['finish'])(list(set(actual)))

equal_to([1])([actual.count('finish')])

return match

assert_that(result, matcher())

def test_do_fn_with_windowing_in_finish_bundle(self):

windowfn = window.FixedWindows(2)

class MyDoFn(beam.DoFn):

def process(self, element):

yield TimestampedValue('process' + str(element), 5)

def finish_bundle(self):

yield WindowedValue('finish', 1, [windowfn])

with TestPipeline() as pipeline:

result = (

pipeline

| 'Start' >> beam.Create([1])

| beam.ParDo(MyDoFn())

| WindowInto(windowfn)

| 'create tuple' >> beam.Map(

lambda v, t=beam.DoFn.TimestampParam, w=beam.DoFn.WindowParam:

(v, t, w.start, w.end)))

expected_process = [

('process1', Timestamp(5), Timestamp(4), Timestamp(6))

]

expected_finish = [('finish', Timestamp(1), Timestamp(0), Timestamp(2))]

assert_that(result, equal_to(expected_process + expected_finish))

def test_do_fn_with_start(self):

class MyDoFn(beam.DoFn):

def __init__(self):

self.state = 'init'

def start_bundle(self):

self.state = 'started'

def process(self, element):

if self.state == 'started':

yield 'started'

self.state = 'process'

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])

result = pcoll | 'Do' >> beam.ParDo(MyDoFn())

# May have many bundles, but each has a start and finish.

def matcher():

def match(actual):

equal_to(['started'])(list(set(actual)))

equal_to([1])([actual.count('started')])

return match

assert_that(result, matcher())

def test_do_fn_with_start_error(self):

class MyDoFn(beam.DoFn):

def start_bundle(self):

return [1]

def process(self, element):

pass

with self.assertRaises(RuntimeError):

with TestPipeline() as p:

p | 'Start' >> beam.Create([1, 2, 3]) | 'Do' >> beam.ParDo(MyDoFn())

def test_map_builtin(self):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([[1, 2], [1], [1, 2, 3]])

result = pcoll | beam.Map(len)

assert_that(result, equal_to([1, 2, 3]))

def test_flatmap_builtin(self):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([

[np.array([1, 2, 3])] * 3, [np.array([5, 4, 3]), np.array([5, 6, 7])]

])

result = pcoll | beam.FlatMap(sum)

assert_that(result, equal_to([3, 6, 9, 10, 10, 10]))

def test_filter_builtin(self):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([[], [2], [], [4]])

result = pcoll | 'Filter' >> beam.Filter(len)

assert_that(result, equal_to([[2], [4]]))

def test_filter(self):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3, 4])

result = pcoll | 'Filter' >> beam.Filter(lambda x: x % 2 == 0)

assert_that(result, equal_to([2, 4]))

class _MeanCombineFn(beam.CombineFn):

def create_accumulator(self):

return (0, 0)

def add_input(self, sum_count, element):

(sum_, count) = sum_count

return sum_ + element, count + 1

def merge_accumulators(self, accumulators):

sums, counts = zip(*accumulators)

return sum(sums), sum(counts)

def extract_output(self, sum_count):

(sum_, count) = sum_count

if not count:

return float('nan')

return sum_ / float(count)

def test_combine_with_combine_fn(self):

vals = [1, 2, 3, 4, 5, 6, 7]

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create(vals)

result = pcoll | 'Mean' >> beam.CombineGlobally(self._MeanCombineFn())

assert_that(result, equal_to([sum(vals) // len(vals)]))

def test_combine_with_callable(self):

vals = [1, 2, 3, 4, 5, 6, 7]

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create(vals)

result = pcoll | beam.CombineGlobally(sum)

assert_that(result, equal_to([sum(vals)]))

def test_combine_with_side_input_as_arg(self):

values = [1, 2, 3, 4, 5, 6, 7]

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create(values)

divisor = pipeline | 'Divisor' >> beam.Create([2])

result = pcoll | 'Max' >> beam.CombineGlobally(

# Multiples of divisor only.

lambda vals, d: max(v for v in vals if v % d == 0),

pvalue.AsSingleton(divisor)).without_defaults()

filt_vals = [v for v in values if v % 2 == 0]

assert_that(result, equal_to([max(filt_vals)]))

def test_combine_per_key_with_combine_fn(self):

vals_1 = [1, 2, 3, 4, 5, 6, 7]

vals_2 = [2, 4, 6, 8, 10, 12, 14]

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create(

([('a', x) for x in vals_1] + [('b', x) for x in vals_2]))

result = pcoll | 'Mean' >> beam.CombinePerKey(self._MeanCombineFn())

assert_that(

result,

equal_to([('a', sum(vals_1) // len(vals_1)),

('b', sum(vals_2) // len(vals_2))]))

def test_combine_per_key_with_callable(self):

vals_1 = [1, 2, 3, 4, 5, 6, 7]

vals_2 = [2, 4, 6, 8, 10, 12, 14]

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create(

([('a', x) for x in vals_1] + [('b', x) for x in vals_2]))

result = pcoll | beam.CombinePerKey(sum)

assert_that(result, equal_to([('a', sum(vals_1)), ('b', sum(vals_2))]))

def test_combine_per_key_with_side_input_as_arg(self):

vals_1 = [1, 2, 3, 4, 5, 6, 7]

vals_2 = [2, 4, 6, 8, 10, 12, 14]

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create(

([('a', x) for x in vals_1] + [('b', x) for x in vals_2]))

divisor = pipeline | 'Divisor' >> beam.Create([2])

result = pcoll | beam.CombinePerKey(

lambda vals, d: max(v for v in vals if v % d == 0),

pvalue.AsSingleton(divisor)) # Multiples of divisor only.

m_1 = max(v for v in vals_1 if v % 2 == 0)

m_2 = max(v for v in vals_2 if v % 2 == 0)

assert_that(result, equal_to([('a', m_1), ('b', m_2)]))

def test_group_by_key(self):

with TestPipeline() as pipeline:

pcoll = pipeline | 'start' >> beam.Create([(1, 1), (2, 1), (3, 1), (1, 2),

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

result = pcoll | 'Group' >> beam.GroupByKey() | SortLists

assert_that(result, equal_to([(1, [1, 2, 3]), (2, [1, 2]), (3, [1])]))

def test_group_by_key_unbounded_global_default_trigger(self):

test_options = PipelineOptions()

test_options.view_as(TypeOptions).allow_unsafe_triggers = False

with self.assertRaisesRegex(

ValueError,

'GroupByKey cannot be applied to an unbounded PCollection with ' +

'global windowing and a default trigger'):

with TestPipeline(options=test_options) as pipeline:

pipeline | TestStream() | beam.GroupByKey()

def test_group_by_key_trigger(self):

options = PipelineOptions(['--allow_unsafe_triggers'])

options.view_as(StandardOptions).streaming = True

with TestPipeline(runner='BundleBasedDirectRunner',

options=options) as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([(0, 0)])

triggered = pcoll | 'Trigger' >> beam.WindowInto(

window.GlobalWindows(),

trigger=AfterProcessingTime(1),

accumulation_mode=AccumulationMode.DISCARDING)

output = triggered | 'Gbk' >> beam.GroupByKey()

self.assertTrue(

isinstance(

output.windowing.triggerfn, _AfterSynchronizedProcessingTime))

def test_group_by_key_unsafe_trigger(self):

test_options = PipelineOptions()

test_options.view_as(TypeOptions).allow_unsafe_triggers = False

with self.assertRaisesRegex(ValueError, 'Unsafe trigger'):

with TestPipeline(options=test_options) as pipeline:

_ = (

pipeline

| beam.Create([(None, None)])

| WindowInto(

window.GlobalWindows(),

trigger=trigger.AfterCount(5),

accumulation_mode=trigger.AccumulationMode.ACCUMULATING)

| beam.GroupByKey())

def test_group_by_key_allow_unsafe_triggers(self):

test_options = PipelineOptions(flags=['--allow_unsafe_triggers'])

with TestPipeline(options=test_options) as pipeline:

pcoll = (

pipeline

| beam.Create([(1, 1), (1, 2), (1, 3), (1, 4)])

| WindowInto(

window.GlobalWindows(),

trigger=trigger.AfterCount(4),

accumulation_mode=trigger.AccumulationMode.ACCUMULATING)

| beam.GroupByKey())

assert_that(pcoll, equal_to([(1, [1, 2, 3, 4])]))

def test_group_by_key_reiteration(self):

class MyDoFn(beam.DoFn):

def process(self, gbk_result):

key, value_list = gbk_result

sum_val = 0

# Iterate the GBK result for multiple times.

for _ in range(0, 17):

sum_val += sum(value_list)

return [(key, sum_val)]

with TestPipeline() as pipeline:

pcoll = pipeline | 'start' >> beam.Create([(1, 1), (1, 2), (1, 3),

(1, 4)])

result = (

pcoll | 'Group' >> beam.GroupByKey()

| 'Reiteration-Sum' >> beam.ParDo(MyDoFn()))

assert_that(result, equal_to([(1, 170)]))

def test_group_by_key_deterministic_coder(self):

# pylint: disable=global-variable-not-assigned

global MyObject # for pickling of the class instance

class MyObject:

def __init__(self, value):

self.value = value

def __eq__(self, other):

return self.value == other.value

def __hash__(self):

return hash(self.value)

class MyObjectCoder(beam.coders.Coder):

def encode(self, o):

return pickle.dumps((o.value, random.random()))

def decode(self, encoded):

return MyObject(pickle.loads(encoded)[0])

def as_deterministic_coder(self, *args, **kwargs):

return MydeterministicObjectCoder()

def to_type_hint(self):

return MyObject

class MydeterministicObjectCoder(beam.coders.Coder):

def encode(self, o):

return pickle.dumps(o.value)

def decode(self, encoded):

return MyObject(pickle.loads(encoded))

def is_deterministic(self):

return True

beam.coders.registry.register_coder(MyObject, MyObjectCoder)

with TestPipeline() as pipeline:

pcoll = pipeline | beam.Create([(MyObject(k % 2), k) for k in range(10)])

grouped = pcoll | beam.GroupByKey() | beam.MapTuple(

lambda k, vs: (k.value, sorted(vs)))

combined = pcoll | beam.CombinePerKey(sum) | beam.MapTuple(

lambda k, v: (k.value, v))

assert_that(

grouped,

equal_to([(0, [0, 2, 4, 6, 8]), (1, [1, 3, 5, 7, 9])]),

'CheckGrouped')

assert_that(combined, equal_to([(0, 20), (1, 25)]), 'CheckCombined')

def test_group_by_key_non_deterministic_coder(self):

with self.assertRaisesRegex(Exception, r'deterministic'):

with TestPipeline() as pipeline:

_ = (

pipeline

| beam.Create([(PickledObject(10), None)])

| beam.GroupByKey()

| beam.MapTuple(lambda k, v: list(v)))

def test_group_by_key_allow_non_deterministic_coder(self):

with TestPipeline() as pipeline:

# The GroupByKey below would fail without this option.

pipeline._options.view_as(

TypeOptions).allow_non_deterministic_key_coders = True

grouped = (

pipeline

| beam.Create([(PickledObject(10), None)])

| beam.GroupByKey()

| beam.MapTuple(lambda k, v: list(v)))

assert_that(grouped, equal_to([[None]]))

def test_group_by_key_fake_deterministic_coder(self):

fresh_registry = beam.coders.typecoders.CoderRegistry()

with patch.object(

beam.coders, 'registry', fresh_registry), patch.object(

beam.coders.typecoders, 'registry', fresh_registry):

with TestPipeline() as pipeline:

# The GroupByKey below would fail without this registration.

beam.coders.registry.register_fallback_coder(

beam.coders.coders.FakeDeterministicFastPrimitivesCoder())

grouped = (

pipeline

| beam.Create([(PickledObject(10), None)])

| beam.GroupByKey()

| beam.MapTuple(lambda k, v: list(v)))

assert_that(grouped, equal_to([[None]]))

def test_partition_with_partition_fn(self):

class SomePartitionFn(beam.PartitionFn):

def partition_for(self, element, num_partitions, offset):

return (element % 3) + offset

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8])

# Attempt nominal partition operation.

partitions = pcoll | 'Part 1' >> beam.Partition(SomePartitionFn(), 4, 1)

assert_that(partitions[0], equal_to([]))

assert_that(partitions[1], equal_to([0, 3, 6]), label='p1')

assert_that(partitions[2], equal_to([1, 4, 7]), label='p2')

assert_that(partitions[3], equal_to([2, 5, 8]), label='p3')

# Check that a bad partition label will yield an error. For the

# DirectRunner, this error manifests as an exception.

with self.assertRaises(Exception):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8])

partitions = pcoll | beam.Partition(SomePartitionFn(), 4, 10000)

def test_partition_with_callable(self):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8])

partitions = (

pcoll |

'part' >> beam.Partition(lambda e, n, offset: (e % 3) + offset, 4, 1))

assert_that(partitions[0], equal_to([]))

assert_that(partitions[1], equal_to([0, 3, 6]), label='p1')

assert_that(partitions[2], equal_to([1, 4, 7]), label='p2')

assert_that(partitions[3], equal_to([2, 5, 8]), label='p3')

def test_partition_with_callable_and_side_input(self):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8])

side_input = pipeline | 'Side Input' >> beam.Create([100, 1000])

partitions = (

pcoll | 'part' >> beam.Partition(

lambda e, n, offset, si_list: ((e + len(si_list)) % 3) + offset,

4,

1,

pvalue.AsList(side_input)))

assert_that(partitions[0], equal_to([]))

assert_that(partitions[1], equal_to([1, 4, 7]), label='p1')

assert_that(partitions[2], equal_to([2, 5, 8]), label='p2')

assert_that(partitions[3], equal_to([0, 3, 6]), label='p3')

def test_partition_followed_by_flatten_and_groupbykey(self):

"""Regression test for an issue with how partitions are handled."""

with TestPipeline() as pipeline:

contents = [('aa', 1), ('bb', 2), ('aa', 2)]

created = pipeline | 'A' >> beam.Create(contents)

partitioned = created | 'B' >> beam.Partition(lambda x, n: len(x) % n, 3)

flattened = partitioned | 'C' >> beam.Flatten()

grouped = flattened | 'D' >> beam.GroupByKey() | SortLists

assert_that(grouped, equal_to([('aa', [1, 2]), ('bb', [2])]))

@pytest.mark.it_validatesrunner

def test_flatten_pcollections(self):

with TestPipeline() as pipeline:

pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3])

pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7])

result = (pcoll_1, pcoll_2) | 'Flatten' >> beam.Flatten()

assert_that(result, equal_to([0, 1, 2, 3, 4, 5, 6, 7]))

def test_flatten_no_pcollections(self):

with TestPipeline() as pipeline:

with self.assertRaises(ValueError):

() | 'PipelineArgMissing' >> beam.Flatten()

result = () | 'Empty' >> beam.Flatten(pipeline=pipeline)

assert_that(result, equal_to([]))

@pytest.mark.it_validatesrunner

def test_flatten_one_single_pcollection(self):

with TestPipeline() as pipeline:

input = [0, 1, 2, 3]

pcoll = pipeline | 'Input' >> beam.Create(input)

result = (pcoll, ) | 'Single Flatten' >> beam.Flatten()

assert_that(result, equal_to(input))

@parameterized.expand([

param(compat_version=None),

param(compat_version="2.66.0"),

])

@pytest.mark.it_validatesrunner

@pytest.mark.uses_dill

def test_group_by_key_importable_special_types(self, compat_version):

def generate(_):

for _ in range(100):

yield (coders_test_common.MyTypedNamedTuple(1, 'a'), 1)

pipeline = TestPipeline(is_integration_test=True)

if compat_version:

pytest.importorskip("dill")

pipeline.get_pipeline_options().view_as(

StreamingOptions).update_compatibility_version = compat_version

with pipeline as p:

result = (

p

| 'Create' >> beam.Create([i for i in range(100)])

| 'Generate' >> beam.ParDo(generate)

| 'Reshuffle' >> beam.Reshuffle()

| 'GBK' >> beam.GroupByKey())

assert_that(

result,

equal_to([(

coders_test_common.MyTypedNamedTuple(1, 'a'),

[1 for i in range(10000)])]))

@pytest.mark.it_validatesrunner

def test_group_by_key_dynamic_special_types(self):

def create_dynamic_named_tuple():

return collections.namedtuple('DynamicNamedTuple', ['x', 'y'])

dynamic_named_tuple = create_dynamic_named_tuple()

# Standard FastPrimitivesCoder falls back to python PickleCoder which

# cannot serialize dynamic types or types defined in __main__. Use

# CloudPickleCoder as fallback coder for non-deterministic steps.

class FastPrimitivesCoderV2(beam.coders.FastPrimitivesCoder):

def __init__(self):

super().__init__(fallback_coder=beam.coders.CloudpickleCoder())

beam.coders.typecoders.registry.register_coder(

dynamic_named_tuple, FastPrimitivesCoderV2)

def generate(_):

for _ in range(100):

yield (dynamic_named_tuple(1, 'a'), 1)

pipeline = TestPipeline(is_integration_test=True)

with pipeline as p:

result = (

p

| 'Create' >> beam.Create([i for i in range(100)])

| 'Reshuffle' >> beam.Reshuffle()

| 'Generate' >> beam.ParDo(generate).with_output_types(

tuple[dynamic_named_tuple, int])

| 'GBK' >> beam.GroupByKey()

| 'Count Elements' >> beam.Map(lambda x: len(x[1])))

assert_that(result, equal_to([10000]))

# TODO(https://github.com/apache/beam/issues/20067): Does not work in

# streaming mode on Dataflow.

@pytest.mark.no_sickbay_streaming

@pytest.mark.it_validatesrunner

def test_flatten_same_pcollections(self):

with TestPipeline() as pipeline:

pc = pipeline | beam.Create(['a', 'b'])

assert_that((pc, pc, pc) | beam.Flatten(), equal_to(['a', 'b'] * 3))

def test_flatten_pcollections_in_iterable(self):

with TestPipeline() as pipeline:

pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3])

pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7])

result = [pcoll for pcoll in (pcoll_1, pcoll_2)] | beam.Flatten()

assert_that(result, equal_to([0, 1, 2, 3, 4, 5, 6, 7]))

@pytest.mark.it_validatesrunner

def test_flatten_a_flattened_pcollection(self):

with TestPipeline() as pipeline:

pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3])

pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7])

pcoll_3 = pipeline | 'Start 3' >> beam.Create([8, 9])

pcoll_12 = (pcoll_1, pcoll_2) | 'Flatten' >> beam.Flatten()

pcoll_123 = (pcoll_12, pcoll_3) | 'Flatten again' >> beam.Flatten()

assert_that(pcoll_123, equal_to([x for x in range(10)]))

def test_flatten_input_type_must_be_iterable(self):

# Inputs to flatten *must* be an iterable.

with self.assertRaises(ValueError):

4 | beam.Flatten()

def test_flatten_input_type_must_be_iterable_of_pcolls(self):

# Inputs to flatten *must* be an iterable of PCollections.

with self.assertRaises(TypeError):

{'l': 'test'} | beam.Flatten()

with self.assertRaises(TypeError):

set([1, 2, 3]) | beam.Flatten()

@pytest.mark.it_validatesrunner

def test_flatten_multiple_pcollections_having_multiple_consumers(self):

with TestPipeline() as pipeline:

input = pipeline | 'Start' >> beam.Create(['AA', 'BBB', 'CC'])

def split_even_odd(element):

tag = 'even_length' if len(element) % 2 == 0 else 'odd_length'

return pvalue.TaggedOutput(tag, element)

even_length, odd_length = (input | beam.Map(split_even_odd)

.with_outputs('even_length', 'odd_length'))

merged = (even_length, odd_length) | 'Flatten' >> beam.Flatten()

assert_that(merged, equal_to(['AA', 'BBB', 'CC']))

assert_that(even_length, equal_to(['AA', 'CC']), label='assert:even')

assert_that(odd_length, equal_to(['BBB']), label='assert:odd')

def test_flatten_with(self):

with TestPipeline() as pipeline:

input = pipeline | 'Start' >> beam.Create(['AA', 'BBB', 'CC'])

result = (

input

| 'WithPCollection' >> beam.FlattenWith(input | beam.Map(str.lower))

| 'WithPTransform' >> beam.FlattenWith(beam.Create(['x', 'y'])))

assert_that(

result, equal_to(['AA', 'BBB', 'CC', 'aa', 'bbb', 'cc', 'x', 'y']))

def test_group_by_key_input_must_be_kv_pairs(self):

with self.assertRaises(typehints.TypeCheckError) as e:

with TestPipeline() as pipeline:

pcolls = pipeline | 'A' >> beam.Create([1, 2, 3, 4, 5])

pcolls | 'D' >> beam.GroupByKey()

self.assertStartswith(

e.exception.args[0],

'Input type hint violation at D: expected '

'Tuple[TypeVariable[K], TypeVariable[V]]')

def test_group_by_key_only_input_must_be_kv_pairs(self):

with self.assertRaises(typehints.TypeCheckError) as cm:

with TestPipeline() as pipeline:

pcolls = pipeline | 'A' >> beam.Create(['a', 'b', 'f'])

pcolls | 'D' >> beam.GroupByKey()

expected_error_prefix = (

'Input type hint violation at D: expected '

'Tuple[TypeVariable[K], TypeVariable[V]]')

self.assertStartswith(cm.exception.args[0], expected_error_prefix)

def test_keys_and_values(self):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([(3, 1), (2, 1), (1, 1), (3, 2),

(2, 2), (3, 3)])

keys = pcoll.apply(beam.Keys('keys'))

vals = pcoll.apply(beam.Values('vals'))

assert_that(keys, equal_to([1, 2, 2, 3, 3, 3]), label='assert:keys')

assert_that(vals, equal_to([1, 1, 1, 2, 2, 3]), label='assert:vals')

def test_kv_swap(self):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create([(6, 3), (1, 2), (7, 1), (5, 2),

(3, 2)])

result = pcoll.apply(beam.KvSwap(), label='swap')

assert_that(result, equal_to([(1, 7), (2, 1), (2, 3), (2, 5), (3, 6)]))

def test_distinct(self):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Start' >> beam.Create(

[6, 3, 1, 1, 9, 'pleat', 'pleat', 'kazoo', 'navel'])

result = pcoll.apply(beam.Distinct())

assert_that(result, equal_to([1, 3, 6, 9, 'pleat', 'kazoo', 'navel']))

def test_chained_ptransforms(self):

with TestPipeline() as pipeline:

t = (

beam.Map(lambda x: (x, 1))

| beam.GroupByKey()

| beam.Map(lambda x_ones: (x_ones[0], sum(x_ones[1]))))

result = pipeline | 'Start' >> beam.Create(['a', 'a', 'b']) | t

assert_that(result, equal_to([('a', 2), ('b', 1)]))

def test_apply_to_list(self):

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

[0, 1, 2] | 'AddOne' >> beam.Map(lambda x: x + 1))

self.assertCountEqual([1],

[0, 1, 2] | 'Odd' >> beam.Filter(lambda x: x % 2))

self.assertCountEqual([1, 2, 100, 3], ([1, 2, 3], [100]) | beam.Flatten())

join_input = ([('k', 'a')], [('k', 'b'), ('k', 'c')])

self.assertCountEqual([('k', (['a'], ['b', 'c']))],

join_input | beam.CoGroupByKey() | SortLists)

def test_multi_input_ptransform(self):

class DisjointUnion(PTransform):

def expand(self, pcollections):

return (

pcollections

| beam.Flatten()

| beam.Map(lambda x: (x, None))

| beam.GroupByKey()

| beam.Map(lambda kv: kv[0]))

self.assertEqual([1, 2, 3], sorted(([1, 2], [2, 3]) | DisjointUnion()))

def test_apply_to_crazy_pvaluish(self):

class NestedFlatten(PTransform):

"""A PTransform taking and returning nested PValueish.

Takes as input a list of dicts, and returns a dict with the corresponding

values flattened.

"""

def _extract_input_pvalues(self, pvalueish):

pvalueish = list(pvalueish)

return pvalueish, sum([list(p.values()) for p in pvalueish], [])

def expand(self, pcoll_dicts):

keys = reduce(operator.or_, [set(p.keys()) for p in pcoll_dicts])

res = {}

for k in keys:

res[k] = [p[k] for p in pcoll_dicts if k in p] | k >> beam.Flatten()

return res

res = [{

'a': [1, 2, 3]

}, {

'a': [4, 5, 6], 'b': ['x', 'y', 'z']

}, {

'a': [7, 8], 'b': ['x', 'y'], 'c': []

}] | NestedFlatten()

self.assertEqual(3, len(res))

self.assertEqual([1, 2, 3, 4, 5, 6, 7, 8], sorted(res['a']))

self.assertEqual(['x', 'x', 'y', 'y', 'z'], sorted(res['b']))

self.assertEqual([], sorted(res['c']))