#

# 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: "[email protected]" }

- { name: "carl", email: "[email protected]" }

- { name: "julia", email: "[email protected]" }

- { name: "carl", email: "[email protected]" }

- 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; ['[email protected]']; ['111-222-3333', '333-444-5555']",

"carl; ['[email protected]', '[email protected]']; ['444-555-6666']",

"james; []; ['222-333-4444']",

"julia; ['[email protected]']; []",

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