def test_top_shorthands(self): with TestPipeline() as pipeline: pcoll = pipeline | 'start' >> Create( [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]) result_top = pcoll | 'top' >> beam.CombineGlobally( combine.Largest(5)) result_bot = pcoll | 'bot' >> beam.CombineGlobally( combine.Smallest(4)) assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='assert:top') assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='assert:bot') pcoll = pipeline | 'start-perkey' >> Create( [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]]) result_ktop = pcoll | 'top-perkey' >> beam.CombinePerKey( combine.Largest(5)) result_kbot = pcoll | 'bot-perkey' >> beam.CombinePerKey( combine.Smallest(4)) assert_that(result_ktop, equal_to([('a', [9, 6, 6, 5, 3])]), label='ktop') assert_that(result_kbot, equal_to([('a', [0, 1, 1, 1])]), label='kbot')
def test_combine_globally_display_data(self): transform = beam.CombineGlobally(combine.Smallest(5)) dd = DisplayData.create_from(transform) expected_items = [ DisplayDataItemMatcher('combine_fn', combine.Smallest), DisplayDataItemMatcher('n', 5), DisplayDataItemMatcher('compare', 'gt')] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
def test_combine_per_key_top_display_data(self): def individual_test_per_key_dd(combineFn): transform = beam.CombinePerKey(combineFn) dd = DisplayData.create_from(transform) expected_items = [ DisplayDataItemMatcher('combine_fn', combineFn.__class__), DisplayDataItemMatcher('n', combineFn._n), DisplayDataItemMatcher('compare', combineFn._compare.__name__)] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) individual_test_per_key_dd(combine.Largest(5)) individual_test_per_key_dd(combine.Smallest(3)) individual_test_per_key_dd(combine.TopCombineFn(8)) individual_test_per_key_dd(combine.Largest(5))