def expand(self, pbegin):
if not isinstance(pbegin, pvalue.PBegin):
raise Exception("GenerateSequence must be a root transform")
coder = VarIntCoder()
coder_urn = ['beam:coder:varint:v1']
args = {
'start':
ConfigValue(coder_urn=coder_urn, payload=coder.encode(self.start))
}
if self.stop:
args['stop'] = ConfigValue(coder_urn=coder_urn,
payload=coder.encode(self.stop))
if self.elements_per_period:
args['elements_per_period'] = ConfigValue(
coder_urn=coder_urn,
payload=coder.encode(self.elements_per_period))
if self.max_read_time:
args['max_read_time'] = ConfigValue(coder_urn=coder_urn,
payload=coder.encode(
self.max_read_time))
payload = ExternalConfigurationPayload(configuration=args)
return pbegin.apply(
ExternalTransform(self._urn, payload.SerializeToString(),
self.expansion_service))
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)
def __init__(self,
start,
stop=None,
elements_per_period=None,
max_read_time=None,
expansion_service=None):
coder = VarIntCoder()
coder_urn = 'beam:coder:varint:v1'
args = {
'start': ConfigValue(coder_urn=coder_urn,
payload=coder.encode(start))
}
if stop:
args['stop'] = ConfigValue(coder_urn=coder_urn,
payload=coder.encode(stop))
if elements_per_period:
args['elements_per_period'] = ConfigValue(
coder_urn=coder_urn, payload=coder.encode(elements_per_period))
if max_read_time:
args['max_read_time'] = ConfigValue(
coder_urn=coder_urn, payload=coder.encode(max_read_time))
payload = ExternalConfigurationPayload(configuration=args)
super(GenerateSequence,
self).__init__('beam:external:java:generate_sequence:v1',
payload.SerializeToString(), expansion_service)
def _config_value(cls, obj, typehint):
"""
Helper to create a ConfigValue with an encoded value.
"""
coder = registry.get_coder(typehint)
urns = list(iter_urns(coder))
if 'beam:coder:pickled_python:v1' in urns:
raise RuntimeError("Found non-portable coder for %s" % (typehint, ))
return ConfigValue(
coder_urn=urns, payload=coder.get_impl().encode_nested(obj))
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))
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_str(str_obj):
encoded_str = str_obj.encode('utf-8')
coder = LengthPrefixCoder(BytesCoder())
coder_urns = ['beam:coder:bytes:v1']
return ConfigValue(coder_urn=coder_urns, payload=coder.encode(encoded_str))
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))