#

# 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 coders that must be consistent across all Beam SDKs.

"""

# pytype: skip-file

import json

import logging

import math

import os.path

import sys

import unittest

from copy import deepcopy

from typing import Dict

from typing import Tuple

import numpy as np

import yaml

from numpy.testing import assert_array_equal

from apache_beam.coders import coder_impl

from apache_beam.portability.api import beam_runner_api_pb2

from apache_beam.portability.api import schema_pb2

from apache_beam.runners import pipeline_context

from apache_beam.transforms import userstate

from apache_beam.transforms import window

from apache_beam.transforms.window import IntervalWindow

from apache_beam.typehints import schemas

from apache_beam.utils import windowed_value

from apache_beam.utils.sharded_key import ShardedKey

from apache_beam.utils.timestamp import Timestamp

from apache_beam.utils.windowed_value import PaneInfo

from apache_beam.utils.windowed_value import PaneInfoTiming

STANDARD_CODERS_YAML = os.path.normpath(

os.path.join(

os.path.dirname(__file__), '../portability/api/standard_coders.yaml'))

def _load_test_cases(test_yaml):

"""Load test data from yaml file and return an iterable of test cases.

See ``standard_coders.yaml`` for more details.

"""

if not os.path.exists(test_yaml):

raise ValueError('Could not find the test spec: %s' % test_yaml)

with open(test_yaml, 'rb') as coder_spec:

for ix, spec in enumerate(

yaml.load_all(coder_spec, Loader=yaml.SafeLoader)):

spec['index'] = ix

name = spec.get('name', spec['coder']['urn'].split(':')[-2])

yield [name, spec]

def parse_float(s):

x = float(s)

if math.isnan(x):

# In Windows, float('NaN') has opposite sign from other platforms.

# For the purpose of this test, we just need consistency.

x = abs(x)

return x

def value_parser_from_schema(schema):

def attribute_parser_from_type(type_):

parser = nonnull_attribute_parser_from_type(type_)

if type_.nullable:

return lambda x: None if x is None else parser(x)

else:

return parser

def nonnull_attribute_parser_from_type(type_):

# TODO: This should be exhaustive

type_info = type_.WhichOneof("type_info")

if type_info == "atomic_type":

if type_.atomic_type == schema_pb2.BYTES:

return lambda x: x.encode("utf-8")

else:

return schemas.ATOMIC_TYPE_TO_PRIMITIVE[type_.atomic_type]

elif type_info == "array_type":

element_parser = attribute_parser_from_type(type_.array_type.element_type)

return lambda x: list(map(element_parser, x))

elif type_info == "map_type":

key_parser = attribute_parser_from_type(type_.map_type.key_type)

value_parser = attribute_parser_from_type(type_.map_type.value_type)

return lambda x: dict(

(key_parser(k), value_parser(v)) for k, v in x.items())

elif type_info == "row_type":

return value_parser_from_schema(type_.row_type.schema)

elif type_info == "logical_type":

# In YAML logical types are represented with their representation types.

to_language_type = schemas.LogicalType.from_runner_api(

type_.logical_type).to_language_type

parse_representation = attribute_parser_from_type(

type_.logical_type.representation)

return lambda x: to_language_type(parse_representation(x))

parsers = [(field.name, attribute_parser_from_type(field.type))

for field in schema.fields]

constructor = schemas.named_tuple_from_schema(schema)

def value_parser(x):

result = []

x = deepcopy(x)

for name, parser in parsers:

value = x.pop(name)

result.append(None if value is None else parser(value))

if len(x):

raise ValueError(

"Test data contains attributes that don't exist in the schema: {}".

format(', '.join(x.keys())))

return constructor(*result)

return value_parser

class StandardCodersTest(unittest.TestCase):

_urn_to_json_value_parser = {

'beam:coder:bytes:v1': lambda x: x.encode('utf-8'),

'beam:coder:bool:v1': lambda x: x,

'beam:coder:string_utf8:v1': lambda x: x,

'beam:coder:varint:v1': lambda x: x,

'beam:coder:kv:v1': lambda x, key_parser, value_parser:

(key_parser(x['key']), value_parser(x['value'])),

'beam:coder:interval_window:v1': lambda x: IntervalWindow(

start=Timestamp(micros=(x['end'] - x['span']) * 1000), end=Timestamp(

micros=x['end'] * 1000)),

'beam:coder:iterable:v1': lambda x, parser: list(map(parser, x)),

'beam:coder:state_backed_iterable:v1': lambda x, parser: list(

map(parser, x)),

'beam:coder:global_window:v1': lambda x: window.GlobalWindow(),

'beam:coder:windowed_value:v1': lambda x, value_parser, window_parser:

windowed_value.create(

value_parser(x['value']), x['timestamp'] * 1000, tuple(

window_parser(w) for w in x['windows'])),

'beam:coder:param_windowed_value:v1': lambda x, value_parser,

window_parser: windowed_value.create(

value_parser(x['value']), x['timestamp'] * 1000, tuple(

window_parser(w) for w in x['windows']), PaneInfo(

x['pane']['is_first'], x['pane']['is_last'], PaneInfoTiming.

from_string(x['pane']['timing']), x['pane']['index'], x[

'pane']['on_time_index'])),

'beam:coder:timer:v1': lambda x, value_parser, window_parser: userstate.

Timer(

user_key=value_parser(x['userKey']), dynamic_timer_tag=x[

'dynamicTimerTag'], clear_bit=x['clearBit'], windows=tuple(

window_parser(w) for w in x['windows']), fire_timestamp=None,

hold_timestamp=None, paneinfo=None)

if x['clearBit'] else userstate.Timer(

user_key=value_parser(x['userKey']), dynamic_timer_tag=x[

'dynamicTimerTag'], clear_bit=x['clearBit'], fire_timestamp=

Timestamp(micros=x['fireTimestamp'] * 1000), hold_timestamp=Timestamp(

micros=x['holdTimestamp'] * 1000), windows=tuple(

window_parser(w) for w in x['windows']), paneinfo=PaneInfo(

x['pane']['is_first'], x['pane']['is_last'],

PaneInfoTiming.from_string(x['pane']['timing']), x[

'pane']['index'], x['pane']['on_time_index'])),

'beam:coder:double:v1': parse_float,

'beam:coder:sharded_key:v1': lambda x, value_parser: ShardedKey(

key=value_parser(x['key']), shard_id=x['shardId'].encode('utf-8')),

'beam:coder:custom_window:v1': lambda x, window_parser: window_parser(

x['window']),

'beam:coder:nullable:v1': lambda x, value_parser: x.encode('utf-8')

if x else None

}

def test_standard_coders(self):

for name, spec in _load_test_cases(STANDARD_CODERS_YAML):

logging.info('Executing %s test.', name)

self._run_standard_coder(name, spec)

def _run_standard_coder(self, name, spec):

def assert_equal(actual, expected):

"""Handle nan values which self.assertEqual fails on."""

if (isinstance(actual, float) and isinstance(expected, float) and

math.isnan(actual) and math.isnan(expected)):

return

self.assertEqual(actual, expected)

coder = self.parse_coder(spec['coder'])

parse_value = self.json_value_parser(spec['coder'])

nested_list = [spec['nested']] if 'nested' in spec else [True, False]

for nested in nested_list:

for expected_encoded, json_value in spec['examples'].items():

value = parse_value(json_value)

expected_encoded = expected_encoded.encode('latin1')

if not spec['coder'].get('non_deterministic', False):

actual_encoded = encode_nested(coder, value, nested)

if self.fix and actual_encoded != expected_encoded:

self.to_fix[spec['index'], expected_encoded] = actual_encoded

else:

self.assertEqual(expected_encoded, actual_encoded)

decoded = decode_nested(coder, expected_encoded, nested)

assert_equal(decoded, value)

else:

# Only verify decoding for a non-deterministic coder

self.assertEqual(

decode_nested(coder, expected_encoded, nested), value)

if spec['coder']['urn'] == 'beam:coder:row:v1':

# Test batch encoding/decoding as well.

values = [

parse_value(json_value) for json_value in spec['examples'].values()

]

columnar = {

field.name: np.array([getattr(value, field.name) for value in values])

for field in coder.schema.fields

}

dest = {

field: np.empty_like(values)

for field, values in columnar.items()

}

for column in dest.values():

column[:] = 0 if 'int' in column.dtype.name else None

expected_encoded = ''.join(spec['examples'].keys()).encode('latin1')

actual_encoded = encode_batch(coder, columnar)

assert_equal(expected_encoded, actual_encoded)

decoded_count = decode_batch(coder, expected_encoded, dest)

assert_equal(len(spec['examples']), decoded_count)

for field, values in dest.items():

assert_array_equal(columnar[field], dest[field])

def parse_coder(self, spec):

context = pipeline_context.PipelineContext()

coder_id = str(hash(str(spec)))

component_ids = [

context.coders.get_id(self.parse_coder(c))

for c in spec.get('components', ())

]

if spec.get('state'):

def iterable_state_read(state_token, elem_coder):

state = spec.get('state').get(state_token.decode('latin1'))

if state is None:

state = ''

input_stream = coder_impl.create_InputStream(state.encode('latin1'))

while input_stream.size() > 0:

yield elem_coder.decode_from_stream(input_stream, True)

context.iterable_state_read = iterable_state_read

context.coders.put_proto(

coder_id,

beam_runner_api_pb2.Coder(

spec=beam_runner_api_pb2.FunctionSpec(

urn=spec['urn'],

payload=spec.get('payload', '').encode('latin1')),

component_coder_ids=component_ids))

return context.coders.get_by_id(coder_id)

def json_value_parser(self, coder_spec):

# TODO: integrate this with the logic for the other parsers

if coder_spec['urn'] == 'beam:coder:row:v1':

schema = schema_pb2.Schema.FromString(

coder_spec['payload'].encode('latin1'))

return value_parser_from_schema(schema)

component_parsers = [

self.json_value_parser(c) for c in coder_spec.get('components', ())

]

return lambda x: self._urn_to_json_value_parser[coder_spec['urn']](

x, *component_parsers)

# Used when --fix is passed.

fix = False

to_fix: Dict[Tuple[int, bytes], bytes] = {}

@classmethod

def tearDownClass(cls):

if cls.fix and cls.to_fix:

print("FIXING", len(cls.to_fix), "TESTS")

doc_sep = '\n---\n'

docs = open(STANDARD_CODERS_YAML).read().split(doc_sep)

def quote(s):

return json.dumps(s.decode('latin1')).replace(r'\u0000', r'\0')

for (doc_ix, expected_encoded), actual_encoded in cls.to_fix.items():

print(quote(expected_encoded), "->", quote(actual_encoded))

docs[doc_ix] = docs[doc_ix].replace(

quote(expected_encoded) + ':', quote(actual_encoded) + ':')

open(STANDARD_CODERS_YAML, 'w').write(doc_sep.join(docs))

def encode_nested(coder, value, nested=True):

out = coder_impl.create_OutputStream()

coder.get_impl().encode_to_stream(value, out, nested)

return out.get()

def decode_nested(coder, encoded, nested=True):

return coder.get_impl().decode_from_stream(

coder_impl.create_InputStream(encoded), nested)

def encode_batch(row_coder, values):

out = coder_impl.create_OutputStream()

row_coder.get_impl().encode_batch_to_stream(values, out)

return out.get()

def decode_batch(row_coder, encoded, dest):

return row_coder.get_impl().decode_batch_from_stream(

dest, coder_impl.create_InputStream(encoded))

if __name__ == '__main__':

if '--fix' in sys.argv:

StandardCodersTest.fix = True

sys.argv.remove('--fix')

unittest.main()