def test_unsupported_combinefn_detection(self):
class CombinerWithNonDefaultSetupTeardown(combiners.CountCombineFn):
def setup(self, *args, **kwargs):
pass
def teardown(self, *args, **kwargs):
pass
runner = DataflowRunner()
with self.assertRaisesRegex(
ValueError, 'CombineFn.setup and CombineFn.'
'teardown are not supported'):
with beam.Pipeline(runner=runner,
options=PipelineOptions(
self.default_properties)) as p:
_ = (p | beam.Create([1])
| beam.CombineGlobally(
CombinerWithNonDefaultSetupTeardown()))
try:
with beam.Pipeline(runner=runner,
options=PipelineOptions(
self.default_properties)) as p:
_ = (p | beam.Create([1])
| beam.CombineGlobally(
combiners.SingleInputTupleCombineFn(
combiners.CountCombineFn(),
combiners.CountCombineFn())))
except ValueError:
self.fail('ValueError raised unexpectedly')
def run_combine(pipeline, input_elements=5, lift_combiners=True):
# Calculate the expected result, which is the sum of an arithmetic sequence.
# By default, this is equal to: 0 + 1 + 2 + 3 + 4 = 10
expected_result = input_elements * (input_elements - 1) / 2
# Enable runtime type checking in order to cover TypeCheckCombineFn by
# the test.
pipeline.get_pipeline_options().view_as(
TypeOptions).runtime_type_check = True
pipeline.get_pipeline_options().view_as(
TypeOptions).allow_unsafe_triggers = True
with pipeline as p:
pcoll = p | 'Start' >> beam.Create(range(input_elements))
# Certain triggers, such as AfterCount, are incompatible with combiner
# lifting. We can use that fact to prevent combiners from being lifted.
if not lift_combiners:
pcoll |= beam.WindowInto(
window.GlobalWindows(),
trigger=trigger.AfterCount(input_elements),
accumulation_mode=trigger.AccumulationMode.DISCARDING)
# Pass an additional 'None' in order to cover _CurriedFn by the test.
pcoll |= 'Do' >> beam.CombineGlobally(
combiners.SingleInputTupleCombineFn(
CallSequenceEnforcingCombineFn(),
CallSequenceEnforcingCombineFn()), None).with_fanout(fanout=1)
assert_that(pcoll, equal_to([(expected_result, expected_result)]))
