#

# 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.

#

"""Type coders registration.

This module contains functionality to define and use coders for custom classes.

Let's say we have a class Xyz and we are processing a PCollection with elements

of type Xyz. If we do not register a coder for Xyz, a default pickle-based

fallback coder will be used. This can be undesirable for two reasons. First, we

may want a faster coder or a more space efficient one. Second, the pickle-based

coder is not deterministic in the sense that objects like dictionaries or sets

are not guaranteed to be encoded in the same way every time (elements are not

really ordered).

Two (sometimes three) steps are needed to define and use a custom coder:

- define the coder class

- associate the code with the class (a.k.a. coder registration)

- typehint DoFns or transforms with the new class or composite types using

the class.

A coder class is defined by subclassing from CoderBase and defining the

encode_to_bytes and decode_from_bytes methods. The framework uses duck-typing

for coders so it is not strictly required to subclass from CoderBase as long as

the encode/decode methods are defined.

Registering a coder class is made with a register_coder() call::

from apache_beam import coders

...

coders.registry.register_coder(Xyz, XyzCoder)

Additionally, DoFns and PTransforms may need type hints. This is not always

necessary since there is functionality to infer the return types of DoFns by

analyzing the code. For instance, for the function below the return type of

'Xyz' will be inferred::

def MakeXyzs(v):

return Xyz(v)

If Xyz is inferred then its coder will be used whenever the framework needs to

serialize data (e.g., writing to the shuffler subsystem responsible for group by

key operations). If a typehint is needed it can be specified by decorating the

DoFns or using with_input_types/with_output_types methods on PTransforms. For

example, the above function can be decorated::

@with_output_types(Xyz)

def MakeXyzs(v):

return complex_operation_returning_Xyz(v)

See apache_beam.typehints.decorators module for more details.

"""

# pytype: skip-file

from typing import Any

from typing import Dict

from typing import Iterable

from typing import List

from typing import Type

from apache_beam.coders import coders

from apache_beam.typehints import typehints

from apache_beam.utils import windowed_value

__all__ = ['registry']

class CoderRegistry(object):

"""A coder registry for typehint/coder associations."""

def __init__(self, fallback_coder=None):

self._coders: Dict[Any, Type[coders.Coder]] = {}

self.custom_types: List[Any] = []

self.register_standard_coders(fallback_coder)

def register_standard_coders(self, fallback_coder):

"""Register coders for all basic and composite types."""

# Coders without subclasses.

self._register_coder_internal(int, coders.VarIntCoder)

self._register_coder_internal(float, coders.FloatCoder)

self._register_coder_internal(bytes, coders.BytesCoder)

self._register_coder_internal(bool, coders.BooleanCoder)

self._register_coder_internal(str, coders.StrUtf8Coder)

self._register_coder_internal(windowed_value.PaneInfo, coders.PaneInfoCoder)

self._register_coder_internal(typehints.TupleConstraint, coders.TupleCoder)

self._register_coder_internal(typehints.DictConstraint, coders.MapCoder)

self._register_coder_internal(

typehints.WindowedTypeConstraint, coders.WindowedValueCoder)

# Default fallback coders applied in that order until the first matching

# coder found.

default_fallback_coders = [

coders.ProtoCoder, coders.ProtoPlusCoder, coders.FastPrimitivesCoder

]

self._fallback_coder = fallback_coder or FirstOf(default_fallback_coders)

def register_fallback_coder(self, fallback_coder):

self._fallback_coder = FirstOf([fallback_coder, self._fallback_coder])

def _register_coder_internal(

self, typehint_type: Any,

typehint_coder_class: Type[coders.Coder]) -> None:

self._coders[typehint_type] = typehint_coder_class

@staticmethod

def _normalize_typehint_type(typehint_type):

if typehint_type.__module__ == '__main__':

# See https://github.com/apache/beam/issues/21541

# TODO(robertwb): Remove once all runners are portable.

return getattr(typehint_type, '__name__', str(typehint_type))

return typehint_type

def register_coder(

self, typehint_type: Any,

typehint_coder_class: Type[coders.Coder]) -> None:

if not isinstance(typehint_coder_class, type):

raise TypeError(

'Coder registration requires a coder class object. '

'Received %r instead.' % typehint_coder_class)

if typehint_type not in self.custom_types:

self.custom_types.append(typehint_type)

self._register_coder_internal(

self._normalize_typehint_type(typehint_type), typehint_coder_class)

def get_coder(self, typehint: Any) -> coders.Coder:

if typehint and typehint.__module__ == '__main__':

# See https://github.com/apache/beam/issues/21541

# TODO(robertwb): Remove once all runners are portable.

typehint = getattr(typehint, '__name__', str(typehint))

coder = self._coders.get(

typehint.__class__

if isinstance(typehint, typehints.TypeConstraint) else typehint,

None)

if isinstance(typehint, typehints.TypeConstraint) and coder is not None:

return coder.from_type_hint(typehint, self)

if coder is None:

# We use the fallback coder when there is no coder registered for a

# typehint. For example a user defined class with no coder specified.

if not hasattr(self, '_fallback_coder'):

raise RuntimeError(

'Coder registry has no fallback coder. This can happen if the '

'fast_coders module could not be imported.')

if isinstance(typehint, typehints.IterableTypeConstraint):

return coders.IterableCoder.from_type_hint(typehint, self)

elif isinstance(typehint, typehints.ListConstraint):

return coders.ListCoder.from_type_hint(typehint, self)

elif typehints.is_nullable(typehint):

return coders.NullableCoder.from_type_hint(typehint, self)

elif typehint is None:

# In some old code, None is used for Any.

# TODO(robertwb): Clean this up.

pass

elif typehint is object or typehint == typehints.Any:

# We explicitly want the fallback coder.

pass

elif isinstance(typehint, typehints.TypeVariable):

# TODO(robertwb): Clean this up when type inference is fully enabled.

pass

else:

# TODO(robertwb): Re-enable this warning when it's actionable.

# warnings.warn('Using fallback coder for typehint: %r.' % typehint)

pass

coder = self._fallback_coder

return coder.from_type_hint(typehint, self)

def get_custom_type_coder_tuples(self, types=None):

"""Returns type/coder tuples for all custom types passed in."""

return [(t, self._coders[self._normalize_typehint_type(t)])

for t in self.custom_types if (types is None or t in types)]

def load_custom_type_coder_tuples(self, type_coder):

"""Load type/coder tuples into coder registry."""

for t, c in type_coder:

self.register_coder(t, c)

def verify_deterministic(self, key_coder, op_name, silent=True):

if not key_coder.is_deterministic():

error_msg = (

'The key coder "%s" for %s '

'is not deterministic. This may result in incorrect '

'pipeline output. This can be fixed by adding a type '

'hint to the operation preceding the GroupByKey step, '

'and for custom key classes, by writing a '

'deterministic custom Coder. Please see the '

'documentation for more details.' % (key_coder, op_name))

return key_coder.as_deterministic_coder(op_name, error_msg)

else:

return key_coder

class FirstOf(object):

"""For internal use only; no backwards-compatibility guarantees.

A class used to get the first matching coder from a list of coders."""

def __init__(self, coders: Iterable[Type[coders.Coder]]) -> None:

self._coders = coders

def from_type_hint(self, typehint, registry):

messages = []

for coder in self._coders:

try:

return coder.from_type_hint(typehint, registry)

except Exception as e:

msg = (

'%s could not provide a Coder for type %s: %s' %

(coder, typehint, e))

messages.append(msg)

raise ValueError(

'Cannot provide coder for %s: %s' % (typehint, ';'.join(messages)))

registry = CoderRegistry()