def _encode_map(dict_obj):
kv_list = [(key.encode('utf-8'), val.encode('utf-8'))
for key, val in dict_obj.items()]
coder = IterableCoder(
TupleCoder(
[LengthPrefixCoder(BytesCoder()),
LengthPrefixCoder(BytesCoder())]))
coder_urns = [
'beam:coder:iterable:v1', 'beam:coder:kv:v1', 'beam:coder:bytes:v1',
'beam:coder:bytes:v1'
]
return ConfigValue(coder_urn=coder_urns, payload=coder.encode(kv_list))
def test_implicit_payload_builder_with_bytes(self):
values = PayloadBase.bytes_values
builder = ImplicitSchemaPayloadBuilder(values)
result = builder.build()
if sys.version_info[0] < 3:
# in python 2.x bytes coder will be inferred
args = {
'integer_example': ConfigValue(
coder_urn=['beam:coder:varint:v1'],
payload=VarIntCoder()
.get_impl().encode_nested(values['integer_example'])),
'string_example': ConfigValue(
coder_urn=['beam:coder:bytes:v1'],
payload=StrUtf8Coder()
.get_impl().encode_nested(values['string_example'])),
'list_of_strings': ConfigValue(
coder_urn=['beam:coder:iterable:v1',
'beam:coder:bytes:v1'],
payload=IterableCoder(StrUtf8Coder())
.get_impl().encode_nested(values['list_of_strings'])),
'optional_kv': ConfigValue(
coder_urn=['beam:coder:kv:v1',
'beam:coder:bytes:v1',
'beam:coder:double:v1'],
payload=TupleCoder([StrUtf8Coder(), FloatCoder()])
.get_impl().encode_nested(values['optional_kv'])),
}
expected = get_payload(args)
self.assertEqual(result, expected)
else:
expected = get_payload(PayloadBase.args)
self.assertEqual(result, expected)
class SimpleTestStatefulDoFn(DoFn):
BUFFER_STATE = CombiningValueStateSpec(
'buffer',
IterableCoder(VarIntCoder()), ToListCombineFn())
EXPIRY_TIMER = TimerSpec('expiry1', TimeDomain.WATERMARK)
def process(self, element, buffer=DoFn.StateParam(BUFFER_STATE),
timer1=DoFn.TimerParam(EXPIRY_TIMER)):
unused_key, value = element
buffer.add(value)
timer1.set(20)
@on_timer(EXPIRY_TIMER)
def expiry_callback(self, buffer=DoFn.StateParam(BUFFER_STATE),
timer=DoFn.TimerParam(EXPIRY_TIMER)):
yield ''.join(str(x) for x in sorted(buffer.read()))
class PayloadBase(object):
values = {
'integer_example': 1,
'boolean': True,
'string_example': u'thing',
'list_of_strings': [u'foo', u'bar'],
'optional_kv': (u'key', 1.1),
'optional_integer': None,
}
bytes_values = {
'integer_example': 1,
'boolean': True,
'string_example': 'thing',
'list_of_strings': ['foo', 'bar'],
'optional_kv': ('key', 1.1),
'optional_integer': None,
}
args = {
'integer_example':
ConfigValue(coder_urn=['beam:coder:varint:v1'],
payload=VarIntCoder().get_impl().encode_nested(
values['integer_example'])),
'boolean':
ConfigValue(coder_urn=['beam:coder:bool:v1'],
payload=BooleanCoder().get_impl().encode_nested(
values['boolean'])),
'string_example':
ConfigValue(coder_urn=['beam:coder:string_utf8:v1'],
payload=StrUtf8Coder().get_impl().encode_nested(
values['string_example'])),
'list_of_strings':
ConfigValue(
coder_urn=['beam:coder:iterable:v1', 'beam:coder:string_utf8:v1'],
payload=IterableCoder(StrUtf8Coder()).get_impl().encode_nested(
values['list_of_strings'])),
'optional_kv':
ConfigValue(coder_urn=[
'beam:coder:kv:v1', 'beam:coder:string_utf8:v1',
'beam:coder:double:v1'
],
payload=TupleCoder([
StrUtf8Coder(), FloatCoder()
]).get_impl().encode_nested(values['optional_kv'])),
}
def get_payload_from_typing_hints(self, values):
"""Return ExternalConfigurationPayload based on python typing hints"""
raise NotImplementedError
def get_payload_from_beam_typehints(self, values):
"""Return ExternalConfigurationPayload based on beam typehints"""
raise NotImplementedError
def test_typing_payload_builder(self):
result = self.get_payload_from_typing_hints(self.values)
expected = get_payload(self.args)
self.assertEqual(result, expected)
def test_typing_payload_builder_with_bytes(self):
"""
string_utf8 coder will be used even if values are not unicode in python 2.x
"""
result = self.get_payload_from_typing_hints(self.bytes_values)
expected = get_payload(self.args)
self.assertEqual(result, expected)
def test_typehints_payload_builder(self):
result = self.get_payload_from_beam_typehints(self.values)
expected = get_payload(self.args)
self.assertEqual(result, expected)
def test_typehints_payload_builder_with_bytes(self):
"""
string_utf8 coder will be used even if values are not unicode in python 2.x
"""
result = self.get_payload_from_beam_typehints(self.bytes_values)
expected = get_payload(self.args)
self.assertEqual(result, expected)
def test_optional_error(self):
"""
value can only be None if typehint is Optional
"""
with self.assertRaises(RuntimeError):
self.get_payload_from_typing_hints({k: None for k in self.values})
def _encode_list(list_obj):
encoded_list = [val.encode('utf-8') for val in list_obj]
coder = IterableCoder(LengthPrefixCoder(BytesCoder()))
coder_urns = ['beam:coder:iterable:v1', 'beam:coder:bytes:v1']
return ConfigValue(coder_urn=coder_urns,
payload=coder.encode(encoded_list))
