Add URN registration mechanism for coders.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/ef4239ab Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/ef4239ab Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/ef4239ab Branch: refs/heads/master Commit: ef4239ab7928bfad95a4debb1517c2547473bf8f Parents: d261d6b Author: Robert Bradshaw <rober...@gmail.com> Authored: Tue Aug 22 10:01:40 2017 -0700 Committer: Robert Bradshaw <rober...@gmail.com> Committed: Thu Aug 24 15:41:46 2017 -0700 ---------------------------------------------------------------------- sdks/python/apache_beam/coders/coders.py | 66 +++++++++++++++++++++------ 1 file changed, 53 insertions(+), 13 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/beam/blob/ef4239ab/sdks/python/apache_beam/coders/coders.py ---------------------------------------------------------------------- diff --git a/sdks/python/apache_beam/coders/coders.py b/sdks/python/apache_beam/coders/coders.py index 7ced5a9..0ea5f7c 100644 --- a/sdks/python/apache_beam/coders/coders.py +++ b/sdks/python/apache_beam/coders/coders.py @@ -202,25 +202,65 @@ class Coder(object): and self._dict_without_impl() == other._dict_without_impl()) # pylint: enable=protected-access - def to_runner_api(self, context): - """For internal use only; no backwards-compatibility guarantees. + _known_urns = {} + + @classmethod + def register_urn(cls, urn, parameter_type, fn=None): + """Registeres a urn with a constructor. + + For example, if 'beam:fn:foo' had paramter type FooPayload, one could + write `RunnerApiFn.register_urn('bean:fn:foo', FooPayload, foo_from_proto)` + where foo_from_proto took as arguments a FooPayload and a PipelineContext. + This function can also be used as a decorator rather than passing the + callable in as the final parameter. + + A corresponding to_runner_api_parameter method would be expected that + returns the tuple ('beam:fn:foo', FooPayload) """ - # TODO(BEAM-115): Use specialized URNs and components. - serialized_coder = serialize_coder(self) + def register(fn): + cls._known_urns[urn] = parameter_type, fn + return staticmethod(fn) + if fn: + # Used as a statement. + register(fn) + else: + # Used as a decorator. + return register + + def to_runner_api(self, context): + from apache_beam.portability.api import beam_runner_api_pb2 + urn, typed_param, components = self.to_runner_api_parameter(context) return beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.SdkFunctionSpec( spec=beam_runner_api_pb2.FunctionSpec( - urn=urns.PICKLED_CODER, - any_param=proto_utils.pack_Any( - google.protobuf.wrappers_pb2.BytesValue( - value=serialized_coder)), - payload=serialized_coder))) + urn=urn, + any_param=proto_utils.pack_Any(typed_param), + payload=typed_param.SerializeToString() + if typed_param is not None else None)), + component_coder_ids=[context.coders.get_id(c) for c in components]) - @staticmethod - def from_runner_api(proto, context): - """For internal use only; no backwards-compatibility guarantees. + @classmethod + def from_runner_api(cls, coder_proto, context): + """Converts from an SdkFunctionSpec to a Fn object. + + Prefer registering a urn with its parameter type and constructor. """ - return deserialize_coder(proto.spec.spec.payload) + parameter_type, constructor = cls._known_urns[coder_proto.spec.spec.urn] + return constructor( + proto_utils.parse_Bytes(coder_proto.spec.spec.payload, parameter_type), + [context.coders.get_by_id(c) for c in coder_proto.component_coder_ids], + context) + + def to_runner_api_parameter(self, context): + return ( + urns.PICKLED_CODER, + google.protobuf.wrappers_pb2.BytesValue(value=serialize_coder(self)), + ()) + + +@Coder.register_urn(urns.PICKLED_CODER, google.protobuf.wrappers_pb2.BytesValue) +def _pickle_from_runner_api_parameter(payload, components, context): + return deserialize_coder(payload.value) class StrUtf8Coder(Coder):