avro.py

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

from avro.errors import AvroTypeException, SchemaResolutionException
from avro.io import (
    BinaryDecoder,
    BinaryEncoder,
    DatumReader,
    DatumWriter,
    validate,
)

STRUCT_FLOAT = struct.Struct('>f')  # big-endian float
STRUCT_DOUBLE = struct.Struct('>d')  # big-endian double
STRUCT_INT = struct.Struct('>i')  # big-endian int
STRUCT_LONG_LONG = struct.Struct('>q')  # big-endian long long


class FlinkAvroBufferWrapper(object):

    def __init__(self):
        self._stream = None

    def switch_stream(self, stream):
        self._stream = stream

    def read(self, n=1):
        return self._stream.read(n)

    def write(self, data):
        return self._stream.write(data)


class FlinkAvroDecoder(BinaryDecoder):
    """
    Flink-customized basic type decoder, with some different encodings, e.g. Default avro encode int
    and long into variable-sized bytes, while Flink use fixed bytes in some places.
    """

    def __init__(self, reader):
        super().__init__(reader)

    def read_int(self):
        return STRUCT_INT.unpack(self.read(4))[0]

    def read_long(self):
        return STRUCT_LONG_LONG.unpack(self.read(8))[0]

    def read_var_long(self):
        """
        Flink implementation of variable-sized long serialization does not move sign to lower bit
        and flip the rest.
        """
        b = ord(self.read(1))
        n = b & 0x7F
        shift = 7
        while (b & 0x80) != 0:
            b = ord(self.read(1))
            n |= (b & 0x7F) << shift
            shift += 7
        return n

    def read_float(self):
        return STRUCT_FLOAT.unpack(self.read(4))[0]

    def read_double(self):
        return STRUCT_DOUBLE.unpack(self.read(8))[0]

    def read_bytes(self):
        nbytes = self.read_int()
        assert (nbytes >= 0), nbytes
        return self.read(nbytes)

    def skip_int(self):
        self.skip(4)

    def skip_long(self):
        self.skip(8)

    def skip_bytes(self):
        nbytes = self.read_int()
        assert (nbytes >= 0), nbytes
        self.skip(nbytes)


class FlinkAvroDatumReader(DatumReader):
    """
    Flink-customized datum reader to parse composite data structure, to fit with changed basic type
    serialization in FlinkAvroDecoder.
    """

    def __init__(self, writer_schema=None, reader_schema=None):
        super().__init__(writer_schema, reader_schema)

    def read_array(self, writer_schema, reader_schema, decoder: 'FlinkAvroDecoder'):
        read_items = []
        block_count = decoder.read_var_long()
        assert block_count >= 0
        if block_count == 0:
            return read_items
        for i in range(block_count):
            read_items.append(self.read_data(writer_schema.items,
                                             reader_schema.items, decoder))
        decoder.read_var_long()
        return read_items

    def skip_array(self, writer_schema, decoder: 'FlinkAvroDecoder'):
        block_count = decoder.read_var_long()
        assert block_count >= 0
        if block_count == 0:
            return
        for i in range(block_count):
            self.skip_data(writer_schema.items, decoder)
        decoder.read_var_long()

    def read_map(self, writer_schema, reader_schema, decoder: 'FlinkAvroDecoder'):
        read_items = {}
        block_count = decoder.read_var_long()
        assert block_count >= 0
        if block_count == 0:
            return read_items
        for i in range(block_count):
            key = decoder.read_utf8()
            read_items[key] = self.read_data(writer_schema.values,
                                             reader_schema.values, decoder)
        decoder.read_var_long()
        return read_items

    def skip_map(self, writer_schema, decoder: 'FlinkAvroDecoder'):
        block_count = decoder.read_var_long()
        assert block_count >= 0
        if block_count == 0:
            return
        for i in range(block_count):
            decoder.skip_utf8()
            self.skip_data(writer_schema.values, decoder)
        decoder.read_long()

    def read_union(self, writer_schema, reader_schema, decoder: 'FlinkAvroDecoder'):
        index_of_schema = int(decoder.read_int())
        if index_of_schema >= len(writer_schema.schemas):
            fail_msg = "Can't access branch index %d for union with %d branches" \
                       % (index_of_schema, len(writer_schema.schemas))
            raise SchemaResolutionException(fail_msg, writer_schema, reader_schema)
        selected_writer_schema = writer_schema.schemas[index_of_schema]

        return self.read_data(selected_writer_schema, reader_schema, decoder)

    def skip_union(self, writer_schema, decoder):
        index_of_schema = int(decoder.read_int())
        if index_of_schema >= len(writer_schema.schemas):
            fail_msg = "Can't access branch index %d for union with %d branches" \
                       % (index_of_schema, len(writer_schema.schemas))
            raise SchemaResolutionException(fail_msg, writer_schema)
        return self.skip_data(writer_schema.schemas[index_of_schema], decoder)


class FlinkAvroEncoder(BinaryEncoder):

    def __init__(self, writer):
        super().__init__(writer)

    def write_int(self, datum):
        self.write(STRUCT_INT.pack(datum))

    def write_long(self, datum):
        self.write(STRUCT_LONG_LONG.pack(datum))

    def write_var_long(self, datum):
        while datum & 0x80 != 0:
            self.write((datum & 0x80).to_bytes(1, 'big'))
            datum <<= 7
        self.write(datum.to_bytes(1, 'big'))

    def write_float(self, datum):
        self.write(STRUCT_FLOAT.pack(datum))

    def write_double(self, datum):
        self.write(STRUCT_DOUBLE.pack(datum))

    def write_bytes(self, datum):
        self.write_int(len(datum))
        self.write(datum)


class FlinkAvroDatumWriter(DatumWriter):

    def __init__(self, writer_schema=None):
        super().__init__(writers_schema=writer_schema)

    def write_array(self, writer_schema, datum, encoder):
        if len(datum) > 0:
            encoder.write_var_long(len(datum))
            for item in datum:
                self.write_data(writer_schema.items, item, encoder)
        encoder.write_long(0)

    def write_map(self, writer_schema, datum, encoder):
        if len(datum) > 0:
            encoder.write_var_long(len(datum))
            for key, val in datum.items():
                encoder.write_utf8(key)
                self.write_data(writer_schema.values, val, encoder)
        encoder.write_var_long(0)

    def write_union(self, writer_schema, datum, encoder):
        # resolve union
        index_of_schema = -1
        for i, candidate_schema in enumerate(writer_schema.schemas):
            if validate(candidate_schema, datum):
                index_of_schema = i
        if index_of_schema < 0:
            raise AvroTypeException(writer_schema, datum)

        # write data
        encoder.write_int(index_of_schema)
        self.write_data(writer_schema.schemas[index_of_schema], datum, encoder)