datahub/batch/binary_record.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# 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.


from decimal import Decimal
from .utils import *
from .record_header import RecordHeader, RECORD_HEADER_SIZE
from ..models import FieldType
from ..utils.converters import to_binary, to_text
from ..exceptions import InvalidParameterException, DatahubException


BYTE_SIZE_ONE_FIELD = 8
FIELD_COUNT_BYTE_SIZE = 4
INT_BYTE_SIZE = 4
PADDING_BYTES = b'\x00'


class BinaryRecord:
    def __init__(self, schema=None, version_id=None, buffer=None, header=None):
        self._field_cnt = len(schema.field_list) if schema else 1
        self._attr_length = 0
        self._version_id = version_id
        self._schema = schema
        self._buffer = buffer if buffer else bytes()
        self._attr_map = dict()
        self._has_init_attr_map = False

        self._null_bit = [0] * (((self._field_cnt + 63) >> 6) << 3)  # N/8
        self._fields = [bytes()] * self._field_cnt  # N*8

        self._field_pos = RECORD_HEADER_SIZE + FIELD_COUNT_BYTE_SIZE + len(self._null_bit)
        self._next_pos = self._field_pos + self._field_cnt * BYTE_SIZE_ONE_FIELD

        # RecordHeader
        self._record_header = header

    # =======================
    # put record
    # =======================

    def serialize(self):
        # record header
        record_head_byte = RecordHeader.serialize(0, self._version_id, self.__get_record_size(), self._next_pos)
        if len(record_head_byte) != RECORD_HEADER_SIZE:
            raise DatahubException("Record header size is {}, should be {}".format(len(record_head_byte), RECORD_HEADER_SIZE))
        self.__append(record_head_byte)

        # field count
        field_count_byte = int2byte(self._field_cnt)
        self.__append(field_count_byte)

        # NullBit: N/8
        for null_bit in self._null_bit:
            null_bit_byte = int2byte(null_bit, size=1, unsigned=True)
            self.__append(null_bit_byte)

        # Field: N*8
        for field_data in self._fields:
            self.__append(field_data)

        # attribute map
        attr_map_len_byte = int2byte(len(self._attr_map))
        self.__append(attr_map_len_byte)

        for key, val in self._attr_map.items():
            key_len_byte = int2byte(len(key))
            self.__append(key_len_byte)
            key_byte = to_binary(key)
            self.__append(key_byte)
            val_len_byte = int2byte(len(val))
            self.__append(val_len_byte)
            val_byte = to_binary(val)
            self.__append(val_byte)
        return self._buffer

    def set_field(self, pos, value):
        self.__set_none(pos)
        if self._schema is None:    # Blob
            if not isinstance(value, bytes):
                raise InvalidParameterException("Only support write bytes for no schema")
            self._fields[0] = self.__set_byte_field(value)
        else:                       # Tuple
            field = self._schema.get_field(pos)
            field_type = field.type
            if field_type is FieldType.STRING or field_type is FieldType.DECIMAL:
                value_byte = self.__set_byte_field(to_binary(str(value)))
            elif field_type is FieldType.BOOLEAN:
                value_byte = bool2byte(value)
                value_byte += PADDING_BYTES * 7
            elif field_type is FieldType.FLOAT:
                value_byte = float2byte(value)
                value_byte += PADDING_BYTES * 4
            elif field_type is FieldType.DOUBLE:
                value_byte = double2byte(value)
            elif field_type in (FieldType.TINYINT, FieldType.SMALLINT, FieldType.INTEGER,
                                FieldType.BIGINT, FieldType.TIMESTAMP):
                value_byte = int2byte(value, size=8)
            else:
                raise DatahubException("Error field type. {}".format(field_type))
            self._fields[pos] = value_byte

    def add_attribute(self, key, value):
        self._attr_map[key] = value
        self._attr_length += (INT_BYTE_SIZE * 2 + len(key) + len(value))
        self._has_init_attr_map = False

    def __set_byte_field(self, value):
        value_byte_len = len(value)
        if value_byte_len <= 7:
            value += (PADDING_BYTES * (7 - value_byte_len))
            value += int2byte(value_byte_len | 0x80, size=1, unsigned=True)
            return value
        # offset from the end of RecordHeader
        value_offset = FIELD_COUNT_BYTE_SIZE + len(self._null_bit) + self._field_cnt * 8
        for i in range(self._field_cnt, len(self._fields)):
            value_offset += len(self._fields[i])
        # length(4 Byte) + offset(4 Byte)
        value_byte = int2byte(value_byte_len, size=4) + int2byte(value_offset, size=4)
        # add padding to 8N Byte
        value += (PADDING_BYTES * (BYTE_SIZE_ONE_FIELD - (len(value) % BYTE_SIZE_ONE_FIELD)))
        self._fields.append(value)
        self._next_pos += len(value)
        return value_byte

    def __append(self, src_buffer):
        self._buffer += src_buffer

    def __set_none(self, pos, none=False):
        if not none:
            self.__check_pos_valid(pos)
            index = pos >> 3
            self._null_bit[index] |= (1 << (pos & 0x07))

    # =======================
    # get record
    # =======================

    @classmethod
    def deserialize(cls, schema, buffer, record_header):
        binary_record = cls(schema, record_header.schema_version, buffer, record_header)

        # Deserialize null_bit
        null_bit_pos = RECORD_HEADER_SIZE + FIELD_COUNT_BYTE_SIZE
        for index in range(len(binary_record._null_bit)):
            binary_record._null_bit[index] = byte2int(binary_record._buffer[null_bit_pos:null_bit_pos + 1], size=1)
            null_bit_pos += 1

        # Deserialize filed
        for i in range(binary_record._field_cnt):
            binary_record._fields[i] = binary_record.__get_field(i)

        # Deserialize attribute map
        binary_record.__init_attr_map_if_need()

        return binary_record

    def get_field(self, pos):
        self.__check_pos_valid(pos)
        return self._fields[pos]

    def get_attribute(self):
        return self._attr_map

    def __get_field(self, pos):
        if self.__is_field_none(pos):
            return None
        field_byte = self.__read_field(pos)
        if self._schema is None:    # BLOB
            return self.__get_byte_field(field_byte)
        else:                       # TUPLE
            field_type = self._schema.get_field(pos).type
            value = None
            # length = 1
            if field_type is FieldType.BOOLEAN:
                value = byte2bool(field_byte[:1])
            elif field_type is FieldType.TINYINT:
                value = byte2int(field_byte[:1], size=1)
            # length = 2
            elif field_type is FieldType.SMALLINT:
                value = byte2int(field_byte[:2], size=2)
            # length = 4
            elif field_type is FieldType.INTEGER:
                value = byte2int(field_byte[:4], size=4)
            elif field_type is FieldType.FLOAT:
                value = byte2float(field_byte[:4])
            # length = 8
            elif field_type is FieldType.BIGINT or field_type is FieldType.TIMESTAMP:
                value = byte2int(field_byte, size=8)
            elif field_type is FieldType.DOUBLE:
                value = byte2double(field_byte)
            # length = ?
            elif field_type is FieldType.STRING or field_type is FieldType.DECIMAL:
                value = self.__get_str_field(field_byte)
                if field_type is FieldType.DECIMAL:
                    value = Decimal(value)
        return value

    def __get_byte_field(self, value_byte):
        data = byte2int(value_byte, size=8)
        is_little_str = (data & (0x80 << 56)) != 0
        if is_little_str:
            str_size = ((data >> 56) & 0x07)
            value = value_byte[:str_size]
        else:
            str_offset = RECORD_HEADER_SIZE + (data >> 32)
            str_size = byte2int(value_byte[:4], size=4)
            value = self._buffer[str_offset:str_offset + str_size]
            self._next_pos -= str_size
        return value

    def __get_str_field(self, value_byte):
        return to_text(self.__get_byte_field(value_byte))

    def __is_field_none(self, pos):
        self.__check_pos_valid(pos)
        index = pos >> 3
        value = self._null_bit[index] & (1 << (pos & 0x07))
        return value == 0

    def __init_attr_map_if_need(self):
        if self._has_init_attr_map:
            return
        offset = self._record_header.attr_offset
        attr_size = byte2int(self._buffer[offset:offset+4])
        if attr_size != 0 and self._attr_map is None:
            self._attr_map = dict()

        offset += INT_BYTE_SIZE
        for i in range(attr_size):
            key_size = byte2int(self._buffer[offset:offset+4])
            offset += INT_BYTE_SIZE
            key_str = to_text(self._buffer[offset:offset+key_size])
            offset += key_size
            val_size = byte2int(self._buffer[offset:offset+4])
            offset += INT_BYTE_SIZE
            value_str = to_text(self._buffer[offset:offset+val_size])
            offset += val_size
            self._attr_map[key_str] = value_str
            self._attr_length += (key_size + val_size + 2*INT_BYTE_SIZE)
        self._has_init_attr_map = True

    def __get_field_offset(self, pos):
        return self._field_pos + pos * BYTE_SIZE_ONE_FIELD

    def __read_field(self, pos):
        offset = self.__get_field_offset(pos)
        return self._buffer[offset:offset + BYTE_SIZE_ONE_FIELD]

    # =======================
    # common
    # =======================

    def __get_record_size(self):
        return INT_BYTE_SIZE + self._attr_length + self._next_pos

    def __check_pos_valid(self, pos):
        if pos < 0 or pos >= self._field_cnt:
            raise InvalidParameterException("Invalid position. position: {}, fieldCount: {}".format(pos, self._field_cnt))

    @property
    def field_cnt(self):
        return self._field_cnt

    @property
    def version_id(self):
        return self._version_id

    @property
    def schema(self):
        return self._schema

    @schema.setter
    def schema(self, schema):
        self._schema = schema