| # |
| # 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() |
