| # | |
| # 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 windowing classes.""" | |
| # pytype: skip-file | |
| import unittest | |
| import apache_beam as beam | |
| from apache_beam.coders import coders | |
| from apache_beam.runners import pipeline_context | |
| from apache_beam.testing.test_pipeline import TestPipeline | |
| from apache_beam.testing.util import assert_that | |
| from apache_beam.testing.util import equal_to | |
| from apache_beam.transforms import CombinePerKey | |
| from apache_beam.transforms import Create | |
| from apache_beam.transforms import FlatMapTuple | |
| from apache_beam.transforms import GroupByKey | |
| from apache_beam.transforms import Map | |
| from apache_beam.transforms import MapTuple | |
| from apache_beam.transforms import WindowInto | |
| from apache_beam.transforms import combiners | |
| from apache_beam.transforms import core | |
| from apache_beam.transforms.core import Windowing | |
| from apache_beam.transforms.trigger import AccumulationMode | |
| from apache_beam.transforms.trigger import AfterCount | |
| from apache_beam.transforms.window import FixedWindows | |
| from apache_beam.transforms.window import GlobalWindow | |
| from apache_beam.transforms.window import GlobalWindows | |
| from apache_beam.transforms.window import IntervalWindow | |
| from apache_beam.transforms.window import NonMergingWindowFn | |
| from apache_beam.transforms.window import Sessions | |
| from apache_beam.transforms.window import SlidingWindows | |
| from apache_beam.transforms.window import TimestampCombiner | |
| from apache_beam.transforms.window import TimestampedValue | |
| from apache_beam.transforms.window import WindowedValue | |
| from apache_beam.transforms.window import WindowFn | |
| from apache_beam.utils.timestamp import MAX_TIMESTAMP | |
| from apache_beam.utils.timestamp import MIN_TIMESTAMP | |
| def context(element, timestamp): | |
| return WindowFn.AssignContext(timestamp, element) | |
| class ReifyWindowsFn(core.DoFn): | |
| def process(self, element, window=core.DoFn.WindowParam): | |
| key, values = element | |
| yield "%s @ %s" % (key, window), values | |
| class TestCustomWindows(NonMergingWindowFn): | |
| """A custom non merging window fn which assigns elements into interval windows | |
| [0, 3), [3, 5) and [5, element timestamp) based on the element timestamps. | |
| """ | |
| def assign(self, context): | |
| timestamp = context.timestamp | |
| if timestamp < 3: | |
| return [IntervalWindow(0, 3)] | |
| elif timestamp < 5: | |
| return [IntervalWindow(3, 5)] | |
| else: | |
| return [IntervalWindow(5, timestamp)] | |
| def get_window_coder(self): | |
| return coders.IntervalWindowCoder() | |
| class WindowTest(unittest.TestCase): | |
| def test_timestamped_value_cmp(self): | |
| self.assertEqual(TimestampedValue('a', 2), TimestampedValue('a', 2)) | |
| self.assertEqual(TimestampedValue('a', 2), TimestampedValue('a', 2.0)) | |
| self.assertNotEqual(TimestampedValue('a', 2), TimestampedValue('a', 2.1)) | |
| self.assertNotEqual(TimestampedValue('a', 2), TimestampedValue('b', 2)) | |
| def test_global_window(self): | |
| self.assertEqual(GlobalWindow(), GlobalWindow()) | |
| self.assertNotEqual( | |
| GlobalWindow(), IntervalWindow(MIN_TIMESTAMP, MAX_TIMESTAMP)) | |
| self.assertNotEqual( | |
| IntervalWindow(MIN_TIMESTAMP, MAX_TIMESTAMP), GlobalWindow()) | |
| self.assertTrue(GlobalWindow().max_timestamp() < MAX_TIMESTAMP) | |
| def test_fixed_windows(self): | |
| # Test windows with offset: 2, 7, 12, 17, ... | |
| windowfn = FixedWindows(size=5, offset=2) | |
| self.assertEqual([IntervalWindow(7, 12)], windowfn.assign(context('v', 7))) | |
| self.assertEqual([IntervalWindow(7, 12)], windowfn.assign(context('v', 11))) | |
| self.assertEqual([IntervalWindow(12, 17)], | |
| windowfn.assign(context('v', 12))) | |
| # Test windows without offset: 0, 5, 10, 15, ... | |
| windowfn = FixedWindows(size=5) | |
| self.assertEqual([IntervalWindow(5, 10)], windowfn.assign(context('v', 5))) | |
| self.assertEqual([IntervalWindow(5, 10)], windowfn.assign(context('v', 9))) | |
| self.assertEqual([IntervalWindow(10, 15)], | |
| windowfn.assign(context('v', 10))) | |
| # Test windows with offset out of range. | |
| windowfn = FixedWindows(size=5, offset=12) | |
| self.assertEqual([IntervalWindow(7, 12)], windowfn.assign(context('v', 11))) | |
| def test_sliding_windows_assignment(self): | |
| windowfn = SlidingWindows(size=15, period=5, offset=2) | |
| expected = [ | |
| IntervalWindow(7, 22), IntervalWindow(2, 17), IntervalWindow(-3, 12) | |
| ] | |
| self.assertEqual(expected, windowfn.assign(context('v', 7))) | |
| self.assertEqual(expected, windowfn.assign(context('v', 8))) | |
| self.assertEqual(expected, windowfn.assign(context('v', 11))) | |
| def test_sliding_windows_assignment_fraction(self): | |
| windowfn = SlidingWindows(size=3.5, period=2.5, offset=1.5) | |
| self.assertEqual([IntervalWindow(1.5, 5.0), IntervalWindow(-1.0, 2.5)], | |
| windowfn.assign(context('v', 1.7))) | |
| self.assertEqual([IntervalWindow(1.5, 5.0)], | |
| windowfn.assign(context('v', 3))) | |
| def test_sliding_windows_assignment_fraction_large_offset(self): | |
| windowfn = SlidingWindows(size=3.5, period=2.5, offset=4.0) | |
| self.assertEqual([IntervalWindow(1.5, 5.0), IntervalWindow(-1.0, 2.5)], | |
| windowfn.assign(context('v', 1.7))) | |
| self.assertEqual([IntervalWindow(4.0, 7.5), IntervalWindow(1.5, 5.0)], | |
| windowfn.assign(context('v', 4.5))) | |
| def test_sessions_merging(self): | |
| windowfn = Sessions(10) | |
| def merge(*timestamps): | |
| windows = [windowfn.assign(context(None, t)) for t in timestamps] | |
| running = set() | |
| class TestMergeContext(WindowFn.MergeContext): | |
| def __init__(self): | |
| super().__init__(running) | |
| def merge(self, to_be_merged, merge_result): | |
| for w in to_be_merged: | |
| if w in running: | |
| running.remove(w) | |
| running.add(merge_result) | |
| for ws in windows: | |
| running.update(ws) | |
| windowfn.merge(TestMergeContext()) | |
| windowfn.merge(TestMergeContext()) | |
| return sorted(running) | |
| self.assertEqual([IntervalWindow(2, 12)], merge(2)) | |
| self.assertEqual([IntervalWindow(2, 12), IntervalWindow(19, 29)], | |
| merge(2, 19)) | |
| self.assertEqual([IntervalWindow(2, 19)], merge(2, 9)) | |
| self.assertEqual([IntervalWindow(2, 19)], merge(9, 2)) | |
| self.assertEqual([IntervalWindow(2, 19), IntervalWindow(19, 29)], | |
| merge(2, 9, 19)) | |
| self.assertEqual([IntervalWindow(2, 19), IntervalWindow(19, 29)], | |
| merge(19, 9, 2)) | |
| self.assertEqual([IntervalWindow(2, 25)], merge(2, 15, 10)) | |
| def timestamped_key_values(self, pipeline, key, *timestamps): | |
| return ( | |
| pipeline | 'start' >> Create(timestamps) | |
| | Map(lambda x: WindowedValue((key, x), x, [GlobalWindow()]))) | |
| def test_sliding_windows(self): | |
| with TestPipeline() as p: | |
| pcoll = self.timestamped_key_values(p, 'key', 1, 2, 3) | |
| result = ( | |
| pcoll | |
| | 'w' >> WindowInto(SlidingWindows(period=2, size=4)) | |
| | GroupByKey() | |
| | beam.MapTuple(lambda k, vs: (k, sorted(vs))) | |
| | beam.ParDo(ReifyWindowsFn())) | |
| expected = [('key @ [-2.0, 2.0)', [1]), ('key @ [0.0, 4.0)', [1, 2, 3]), | |
| ('key @ [2.0, 6.0)', [2, 3])] | |
| assert_that(result, equal_to(expected)) | |
| def test_sliding_windows_period_longer_than_size(self): | |
| with TestPipeline() as p: | |
| pcoll = self.timestamped_key_values(p, 'key', 1, 2, 3, 4, 5, 6, 7, 8) | |
| result = ( | |
| pcoll | |
| | 'w' >> WindowInto(SlidingWindows(period=4, size=2)) | |
| | GroupByKey() | |
| | beam.MapTuple(lambda k, vs: (k, sorted(vs))) | |
| | beam.ParDo(ReifyWindowsFn())) | |
| expected = [('key @ [0.0, 2.0)', [1]), ('key @ [4.0, 6.0)', [4, 5]), | |
| ('key @ [8.0, 10.0)', [8])] | |
| assert_that(result, equal_to(expected)) | |
| def test_sessions(self): | |
| with TestPipeline() as p: | |
| pcoll = self.timestamped_key_values(p, 'key', 1, 2, 3, 20, 35, 27) | |
| sort_values = Map(lambda k_vs: (k_vs[0], sorted(k_vs[1]))) | |
| result = ( | |
| pcoll | |
| | 'w' >> WindowInto(Sessions(10)) | |
| | GroupByKey() | |
| | sort_values | |
| | beam.ParDo(ReifyWindowsFn())) | |
| expected = [('key @ [1.0, 13.0)', [1, 2, 3]), | |
| ('key @ [20.0, 45.0)', [20, 27, 35])] | |
| assert_that(result, equal_to(expected)) | |
| def test_timestamped_value(self): | |
| with TestPipeline() as p: | |
| result = ( | |
| p | |
| | 'start' >> Create([(k, k) for k in range(10)]) | |
| | Map(lambda x_t: TimestampedValue(x_t[0], x_t[1])) | |
| | 'w' >> WindowInto(FixedWindows(5)) | |
| | Map(lambda v: ('key', v)) | |
| | GroupByKey() | |
| | beam.MapTuple(lambda k, vs: (k, sorted(vs)))) | |
| assert_that( | |
| result, | |
| equal_to([('key', [0, 1, 2, 3, 4]), ('key', [5, 6, 7, 8, 9])])) | |
| def test_rewindow(self): | |
| with TestPipeline() as p: | |
| result = ( | |
| p | |
| | Create([(k, k) for k in range(10)]) | |
| | Map(lambda x_t1: TimestampedValue(x_t1[0], x_t1[1])) | |
| | 'window' >> WindowInto(SlidingWindows(period=2, size=6)) | |
| # Per the model, each element is now duplicated across | |
| # three windows. Rewindowing must preserve this duplication. | |
| | 'rewindow' >> WindowInto(FixedWindows(5)) | |
| | 'rewindow2' >> WindowInto(FixedWindows(5)) | |
| | Map(lambda v: ('key', v)) | |
| | GroupByKey() | |
| | beam.MapTuple(lambda k, vs: (k, sorted(vs)))) | |
| assert_that( | |
| result, | |
| equal_to([('key', sorted([0, 1, 2, 3, 4] * 3)), | |
| ('key', sorted([5, 6, 7, 8, 9] * 3))])) | |
| def test_rewindow_regroup(self): | |
| with TestPipeline() as p: | |
| grouped = ( | |
| p | |
| | Create(range(5)) | |
| | Map(lambda t: TimestampedValue(('key', t), t)) | |
| | 'window' >> WindowInto(FixedWindows(5, offset=3)) | |
| | GroupByKey() | |
| | MapTuple(lambda k, vs: (k, sorted(vs)))) | |
| # Both of these group-and-ungroup sequences should be idempotent. | |
| regrouped1 = ( | |
| grouped | |
| | 'w1' >> WindowInto(FixedWindows(5, offset=3)) | |
| | 'g1' >> GroupByKey() | |
| | FlatMapTuple(lambda k, vs: [(k, v) for v in vs])) | |
| regrouped2 = ( | |
| grouped | |
| | FlatMapTuple(lambda k, vs: [(k, v) for v in vs]) | |
| | 'w2' >> WindowInto(FixedWindows(5, offset=3)) | |
| | 'g2' >> GroupByKey() | |
| | MapTuple(lambda k, vs: (k, sorted(vs)))) | |
| with_windows = Map(lambda e, w=beam.DoFn.WindowParam: (e, w)) | |
| expected = [(('key', [0, 1, 2]), IntervalWindow(-2, 3)), | |
| (('key', [3, 4]), IntervalWindow(3, 8))] | |
| assert_that(grouped | 'ww' >> with_windows, equal_to(expected)) | |
| assert_that( | |
| regrouped1 | 'ww1' >> with_windows, equal_to(expected), label='r1') | |
| assert_that( | |
| regrouped2 | 'ww2' >> with_windows, equal_to(expected), label='r2') | |
| def test_timestamped_with_combiners(self): | |
| with TestPipeline() as p: | |
| result = ( | |
| p | |
| # Create some initial test values. | |
| | 'start' >> Create([(k, k) for k in range(10)]) | |
| # The purpose of the WindowInto transform is to establish a | |
| # FixedWindows windowing function for the PCollection. | |
| # It does not bucket elements into windows since the timestamps | |
| # from Create are not spaced 5 ms apart and very likely they all | |
| # fall into the same window. | |
| | 'w' >> WindowInto(FixedWindows(5)) | |
| # Generate timestamped values using the values as timestamps. | |
| # Now there are values 5 ms apart and since Map propagates the | |
| # windowing function from input to output the output PCollection | |
| # will have elements falling into different 5ms windows. | |
| | Map(lambda x_t2: TimestampedValue(x_t2[0], x_t2[1])) | |
| # We add a 'key' to each value representing the index of the | |
| # window. This is important since there is no guarantee of | |
| # order for the elements of a PCollection. | |
| | Map(lambda v: (v // 5, v))) | |
| # Sum all elements associated with a key and window. Although it | |
| # is called CombinePerKey it is really CombinePerKeyAndWindow the | |
| # same way GroupByKey is really GroupByKeyAndWindow. | |
| sum_per_window = result | CombinePerKey(sum) | |
| # Compute mean per key and window. | |
| mean_per_window = result | combiners.Mean.PerKey() | |
| assert_that( | |
| sum_per_window, equal_to([(0, 10), (1, 35)]), label='assert:sum') | |
| assert_that( | |
| mean_per_window, equal_to([(0, 2.0), (1, 7.0)]), label='assert:mean') | |
| def test_custom_windows(self): | |
| with TestPipeline() as p: | |
| pcoll = self.timestamped_key_values(p, 'key', 0, 1, 2, 3, 4, 5, 6) | |
| # pylint: disable=abstract-class-instantiated | |
| result = ( | |
| pcoll | |
| | 'custom window' >> WindowInto(TestCustomWindows()) | |
| | GroupByKey() | |
| | 'sort values' >> MapTuple(lambda k, vs: (k, sorted(vs)))) | |
| assert_that( | |
| result, | |
| equal_to([('key', [0, 1, 2]), ('key', [3, 4]), ('key', [5]), | |
| ('key', [6])])) | |
| def test_window_assignment_idempotency(self): | |
| with TestPipeline() as p: | |
| pcoll = self.timestamped_key_values(p, 'key', 0, 2, 4) | |
| result = ( | |
| pcoll | |
| | 'window' >> WindowInto(FixedWindows(2)) | |
| | 'same window' >> WindowInto(FixedWindows(2)) | |
| | 'same window again' >> WindowInto(FixedWindows(2)) | |
| | GroupByKey()) | |
| assert_that(result, equal_to([('key', [0]), ('key', [2]), ('key', [4])])) | |
| def test_window_assignment_through_multiple_gbk_idempotency(self): | |
| with TestPipeline() as p: | |
| pcoll = self.timestamped_key_values(p, 'key', 0, 2, 4) | |
| result = ( | |
| pcoll | |
| | 'window' >> WindowInto(FixedWindows(2)) | |
| | 'gbk' >> GroupByKey() | |
| | 'same window' >> WindowInto(FixedWindows(2)) | |
| | 'another gbk' >> GroupByKey() | |
| | 'same window again' >> WindowInto(FixedWindows(2)) | |
| | 'gbk again' >> GroupByKey()) | |
| assert_that( | |
| result, | |
| equal_to([('key', [[[0]]]), ('key', [[[2]]]), ('key', [[[4]]])])) | |
| class RunnerApiTest(unittest.TestCase): | |
| def test_windowfn_encoding(self): | |
| for window_fn in (GlobalWindows(), | |
| FixedWindows(37), | |
| SlidingWindows(2, 389), | |
| Sessions(5077)): | |
| context = pipeline_context.PipelineContext() | |
| self.assertEqual( | |
| window_fn, | |
| WindowFn.from_runner_api(window_fn.to_runner_api(context), context)) | |
| def test_windowing_encoding(self): | |
| for windowing in (Windowing(GlobalWindows()), | |
| Windowing( | |
| FixedWindows(1, 3), | |
| AfterCount(6), | |
| accumulation_mode=AccumulationMode.ACCUMULATING), | |
| Windowing( | |
| SlidingWindows(10, 15, 21), | |
| AfterCount(28), | |
| timestamp_combiner=TimestampCombiner.OUTPUT_AT_LATEST, | |
| accumulation_mode=AccumulationMode.DISCARDING)): | |
| context = pipeline_context.PipelineContext() | |
| self.assertEqual( | |
| windowing, | |
| Windowing.from_runner_api(windowing.to_runner_api(context), context)) | |
| if __name__ == '__main__': | |
| unittest.main() |