| # | |
| # 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. | |
| # | |
| import logging | |
| import unittest | |
| import apache_beam as beam | |
| from apache_beam.testing.util import assert_that | |
| from apache_beam.testing.util import equal_to | |
| from apache_beam.yaml.readme_test import create_test_method | |
| from apache_beam.yaml.yaml_transform import YamlTransform | |
| DATA = [ | |
| beam.Row(animal='cat', weight=1), | |
| beam.Row(animal='cat', weight=5), | |
| beam.Row(animal='dog', weight=10), | |
| ] | |
| class ProgrammingGuideTest(unittest.TestCase): | |
| test_pipelines_constructing_reading = create_test_method( | |
| 'RUN', | |
| 'test_pipelines_constructing_reading', | |
| ''' | |
| # [START pipelines_constructing_reading] | |
| pipeline: | |
| source: | |
| type: ReadFromText | |
| config: | |
| path: ... | |
| # [END pipelines_constructing_reading] | |
| transforms: [] | |
| ''') | |
| test_create_pcollection = create_test_method( | |
| 'RUN', | |
| 'test_create_pcollection', | |
| ''' | |
| # [START create_pcollection] | |
| pipeline: | |
| transforms: | |
| - type: Create | |
| config: | |
| elements: | |
| - A | |
| - B | |
| - ... | |
| # [END create_pcollection] | |
| transforms: [] | |
| ''') | |
| def test_group_by(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| elements = p | beam.Create(DATA) | |
| result = elements | YamlTransform( | |
| ''' | |
| # [START group_by] | |
| type: Combine | |
| config: | |
| group_by: animal | |
| combine: | |
| weight: group | |
| # [END group_by] | |
| ''') | |
| assert_that( | |
| result | beam.Map(lambda x: beam.Row(**x._asdict())), | |
| equal_to([ | |
| beam.Row(animal='cat', weight=[1, 5]), | |
| beam.Row(animal='dog', weight=[10]), | |
| ])) | |
| def test_co_group_by(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| result = p | YamlTransform( | |
| ''' | |
| type: composite | |
| transforms: | |
| # [START cogroupbykey_inputs] | |
| - type: Create | |
| name: CreateEmails | |
| config: | |
| elements: | |
| - { name: "amy", email: "amy@example.com" } | |
| - { name: "carl", email: "carl@example.com" } | |
| - { name: "julia", email: "julia@example.com" } | |
| - { name: "carl", email: "carl@email.com" } | |
| - type: Create | |
| name: CreatePhones | |
| config: | |
| elements: | |
| - { name: "amy", phone: "111-222-3333" } | |
| - { name: "james", phone: "222-333-4444" } | |
| - { name: "amy", phone: "333-444-5555" } | |
| - { name: "carl", phone: "444-555-6666" } | |
| # [END cogroupbykey_inputs] | |
| # [START cogroupbykey] | |
| - type: MapToFields | |
| name: PrepareEmails | |
| input: CreateEmails | |
| config: | |
| language: python | |
| fields: | |
| name: name | |
| email: "[email]" | |
| phone: "[]" | |
| - type: MapToFields | |
| name: PreparePhones | |
| input: CreatePhones | |
| config: | |
| language: python | |
| fields: | |
| name: name | |
| email: "[]" | |
| phone: "[phone]" | |
| - type: Combine | |
| name: CoGropuBy | |
| input: [PrepareEmails, PreparePhones] | |
| config: | |
| group_by: [name] | |
| combine: | |
| email: concat | |
| phone: concat | |
| - type: MapToFields | |
| name: FormatResults | |
| input: CoGropuBy | |
| config: | |
| language: python | |
| fields: | |
| formatted: | |
| "'%s; %s; %s' % (name, sorted(email), sorted(phone))" | |
| # [END cogroupbykey] | |
| output: FormatResults | |
| ''') | |
| assert_that( | |
| result | beam.Map(lambda x: x.formatted), | |
| equal_to([ | |
| # [START cogroupbykey_formatted_outputs] | |
| "amy; ['amy@example.com']; ['111-222-3333', '333-444-5555']", | |
| "carl; ['carl@email.com', 'carl@example.com']; ['444-555-6666']", | |
| "james; []; ['222-333-4444']", | |
| "julia; ['julia@example.com']; []", | |
| # [END cogroupbykey_formatted_outputs] | |
| ])) | |
| def test_combine_ref(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| elements = p | beam.Create(DATA) | |
| result = elements | YamlTransform( | |
| ''' | |
| # [START combine_simple_py] | |
| type: Combine | |
| config: | |
| language: python | |
| group_by: animal | |
| combine: | |
| biggest: | |
| fn: | |
| type: 'apache_beam.transforms.combiners.TopCombineFn' | |
| config: | |
| n: 2 | |
| value: weight | |
| # [END combine_simple_py] | |
| ''') | |
| assert_that( | |
| result | beam.Map(lambda x: beam.Row(**x._asdict())), | |
| equal_to([ | |
| beam.Row(animal='cat', biggest=[5, 1]), | |
| beam.Row(animal='dog', biggest=[10]), | |
| ])) | |
| def test_combine_globally(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| elements = p | beam.Create(DATA) | |
| result = elements | YamlTransform( | |
| ''' | |
| # [START combine_global_sum] | |
| type: Combine | |
| config: | |
| group_by: [] | |
| combine: | |
| weight: sum | |
| # [END combine_global_sum] | |
| ''') | |
| assert_that( | |
| result | beam.Map(lambda x: beam.Row(**x._asdict())), | |
| equal_to([ | |
| beam.Row(weight=16), | |
| ])) | |
| def test_combine_per_key(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| elements = p | beam.Create(DATA) | |
| result = elements | YamlTransform( | |
| ''' | |
| # [START combine_per_key] | |
| type: Combine | |
| config: | |
| group_by: [animal] | |
| combine: | |
| total_weight: | |
| fn: sum | |
| value: weight | |
| average_weight: | |
| fn: mean | |
| value: weight | |
| # [END combine_per_key] | |
| ''') | |
| assert_that( | |
| result | beam.Map(lambda x: beam.Row(**x._asdict())), | |
| equal_to([ | |
| beam.Row(animal='cat', total_weight=6, average_weight=3), | |
| beam.Row(animal='dog', total_weight=10, average_weight=10), | |
| ])) | |
| def test_flatten(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| result = p | YamlTransform( | |
| ''' | |
| type: composite | |
| transforms: | |
| - type: Create | |
| name: SomeProducingTransform | |
| config: | |
| elements: [a, b, c] | |
| - type: Create | |
| name: AnotherProducingTransform | |
| config: | |
| elements: [d, e, f] | |
| # [START model_multiple_pcollections_flatten] | |
| - type: Flatten | |
| input: [SomeProducingTransform, AnotherProducingTransform] | |
| # [END model_multiple_pcollections_flatten] | |
| output: Flatten | |
| ''') | |
| assert_that( | |
| result | beam.Map(lambda x: x.element), | |
| equal_to(['a', 'b', 'c', 'd', 'e', 'f'])) | |
| def test_expode(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| result = p | YamlTransform( | |
| ''' | |
| type: composite | |
| transforms: | |
| - type: Create | |
| name: SomeProducingTransform | |
| config: | |
| elements: | |
| - {line: "a aa"} | |
| - {line: "b"} | |
| # [START model_multiple_output_dofn] | |
| - type: MapToFields | |
| input: SomeProducingTransform | |
| config: | |
| language: python | |
| fields: | |
| word: "line.split()" | |
| - type: Explode | |
| input: MapToFields | |
| config: | |
| fields: word | |
| # [END model_multiple_output_dofn] | |
| output: Explode | |
| ''') | |
| assert_that( | |
| result | beam.Map(lambda x: x.word), equal_to(['a', 'aa', 'b'])) | |
| def test_schema_output_type(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| elements = p | beam.Create(DATA) | |
| result = elements | YamlTransform( | |
| ''' | |
| # [START schema_output_type] | |
| type: MapToFields | |
| config: | |
| language: python | |
| fields: | |
| new_field: | |
| expression: "hex(weight)" | |
| output_type: { "type": "string" } | |
| # [END schema_output_type] | |
| ''') | |
| assert_that( | |
| result | beam.Map(lambda x: x.new_field), | |
| equal_to(['0x1', '0x5', '0xa'])) | |
| def test_fixed_windows(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| elements = p | beam.Create(DATA) | |
| _ = elements | YamlTransform( | |
| ''' | |
| # [START setting_fixed_windows] | |
| type: WindowInto | |
| windowing: | |
| type: fixed | |
| size: 60s | |
| # [END setting_fixed_windows] | |
| ''') | |
| def test_sliding_windows(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| elements = p | beam.Create(DATA) | |
| _ = elements | YamlTransform( | |
| ''' | |
| # [START setting_sliding_windows] | |
| type: WindowInto | |
| windowing: | |
| type: sliding | |
| size: 5m | |
| period: 30s | |
| # [END setting_sliding_windows] | |
| ''') | |
| def test_session_windows(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| elements = p | beam.Create(DATA) | |
| _ = elements | YamlTransform( | |
| ''' | |
| # [START setting_session_windows] | |
| type: WindowInto | |
| windowing: | |
| type: sessions | |
| gap: 60s | |
| # [END setting_session_windows] | |
| ''') | |
| def test_global_windows(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| elements = p | beam.Create(DATA) | |
| _ = elements | YamlTransform( | |
| ''' | |
| # [START setting_global_window] | |
| type: WindowInto | |
| windowing: | |
| type: global | |
| # [END setting_global_window] | |
| ''') | |
| def test_assign_timestamps(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| elements = p | beam.Create([ | |
| beam.Row(external_timestamp_field='2024-01-01', value=1), | |
| beam.Row(external_timestamp_field='2024-01-01', value=2), | |
| beam.Row(external_timestamp_field='2024-01-02', value=10), | |
| ]) | |
| _ = elements | YamlTransform( | |
| ''' | |
| # [START setting_timestamp] | |
| type: AssignTimestamps | |
| config: | |
| language: python | |
| timestamp: | |
| callable: | | |
| import datetime | |
| def extract_timestamp(x): | |
| raw = datetime.datetime.strptime( | |
| x.external_timestamp_field, "%Y-%m-%d") | |
| return raw.astimezone(datetime.timezone.utc) | |
| # [END setting_timestamp] | |
| ''') | |
| def test_partition(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| elements = p | beam.Create([ | |
| beam.Row(percentile=1), | |
| beam.Row(percentile=20), | |
| beam.Row(percentile=90), | |
| ]) | |
| _ = elements | YamlTransform( | |
| ''' | |
| # [START model_multiple_pcollections_partition] | |
| type: Partition | |
| config: | |
| by: str(percentile // 10) | |
| language: python | |
| outputs: ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10"] | |
| # [END model_multiple_pcollections_partition] | |
| ''') | |
| if __name__ == '__main__': | |
| logging.getLogger().setLevel(logging.INFO) | |
| unittest.main() |