paimon-python/pypaimon/execution/python/write/writer/format_row_writer.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 datetime
import re
import struct
from decimal import Decimal, localcontext
from typing import Any, List

import pyarrow as pa

from pypaimon.common.delta_varint_compressor import DeltaVarintCompressor
from pypaimon.schema.data_types import (
    ArrayType, AtomicType, DataField, MapType, MultisetType, RowType, VectorType
)

FOOTER_SIZE = 32
MAGIC = 0x524F5753  # "ROWS"
VERSION = 1
DEFAULT_BLOCK_SIZE = 65536


class FormatRowWriter:

    def __init__(self, output_stream, fields: List[DataField],
                 block_size: int = DEFAULT_BLOCK_SIZE, zstd_level: int = 1):
        self._out = output_stream
        self._fields = fields
        self._block_size = block_size
        self._zstd_level = zstd_level

        self._block_buf = _BlockBuffer()
        self._row_offsets: List[int] = []

        self._block_compressed_sizes: List[int] = []
        self._block_uncompressed_sizes: List[int] = []
        self._block_row_starts: List[int] = []
        self._total_row_count = 0
        self._bytes_written = 0

    def write_table(self, data: pa.Table):
        columns = {field.name: data.column(field.name).to_pylist()
                   for field in self._fields if field.name in data.column_names}
        num_rows = data.num_rows

        for row_idx in range(num_rows):
            row_values = [columns[field.name][row_idx] for field in self._fields]
            self._write_row(row_values)
            self._total_row_count += 1

            if self._block_buf.position + len(self._row_offsets) * 4 + 4 >= self._block_size:
                self._flush_block()

    def close(self):
        self._flush_block()

        index_offset = self._bytes_written
        self._write_block_index()
        index_length = self._bytes_written - index_offset

        self._write_footer(index_offset, index_length)
        self._out.flush()

    def _write_row(self, values: List[Any]):
        self._row_offsets.append(self._block_buf.position)
        arity = len(self._fields)
        header_size = (arity + 7) // 8

        header_start = self._block_buf.position
        self._block_buf.write_zeros(header_size)

        for i, field in enumerate(self._fields):
            if values[i] is None:
                byte_idx = header_start + i // 8
                self._block_buf.buffer[byte_idx] |= (1 << (i % 8))
            else:
                _write_field(self._block_buf, values[i], field.type)

    def _flush_block(self):
        if not self._row_offsets:
            return

        import zstandard as zstd

        self._block_row_starts.append(self._total_row_count - len(self._row_offsets))

        for offset in self._row_offsets:
            self._block_buf.write_int_le(offset)
        self._block_buf.write_int_le(len(self._row_offsets))

        uncompressed = bytes(self._block_buf.buffer[:self._block_buf.position])
        self._block_uncompressed_sizes.append(len(uncompressed))

        compressor = zstd.ZstdCompressor(level=self._zstd_level)
        compressed = compressor.compress(uncompressed)
        self._block_compressed_sizes.append(len(compressed))

        self._out.write(compressed)
        self._bytes_written += len(compressed)

        self._block_buf.reset()
        self._row_offsets.clear()

    def _write_block_index(self):
        encoded_compressed = DeltaVarintCompressor.compress(self._block_compressed_sizes)
        encoded_uncompressed = DeltaVarintCompressor.compress(self._block_uncompressed_sizes)
        encoded_row_starts = DeltaVarintCompressor.compress(self._block_row_starts)

        self._write_varint_prefixed(encoded_compressed)
        self._write_varint_prefixed(encoded_uncompressed)
        self._write_varint_prefixed(encoded_row_starts)

    def _write_varint_prefixed(self, data: bytes):
        varint_bytes = _encode_var_int(len(data))
        self._out.write(varint_bytes)
        self._out.write(data)
        self._bytes_written += len(varint_bytes) + len(data)

    def _write_footer(self, index_offset: int, index_length: int):
        buf = bytearray(FOOTER_SIZE)
        struct.pack_into('<q', buf, 0, self._total_row_count)
        struct.pack_into('<i', buf, 8, len(self._block_compressed_sizes))
        struct.pack_into('<q', buf, 12, index_offset)
        struct.pack_into('<i', buf, 20, index_length)
        buf[24] = VERSION
        # bytes 25-27 reserved (zeros)
        struct.pack_into('<I', buf, 28, MAGIC)
        self._out.write(bytes(buf))
        self._bytes_written += FOOTER_SIZE


class _BlockBuffer:

    def __init__(self, initial_capacity: int = DEFAULT_BLOCK_SIZE):
        self.buffer = bytearray(initial_capacity)
        self.position = 0

    def reset(self):
        self.position = 0

    def _ensure_capacity(self, additional: int):
        required = self.position + additional
        if required > len(self.buffer):
            new_size = max(len(self.buffer) * 2, required)
            new_buf = bytearray(new_size)
            new_buf[:self.position] = self.buffer[:self.position]
            self.buffer = new_buf

    def write_zeros(self, count: int):
        self._ensure_capacity(count)
        for i in range(count):
            self.buffer[self.position + i] = 0
        self.position += count

    def write_boolean(self, value: bool):
        self._ensure_capacity(1)
        self.buffer[self.position] = 1 if value else 0
        self.position += 1

    def write_byte(self, value: int):
        self._ensure_capacity(1)
        self.buffer[self.position] = value & 0xFF
        self.position += 1

    def write_short_le(self, value: int):
        self._ensure_capacity(2)
        struct.pack_into('<h', self.buffer, self.position, value)
        self.position += 2

    def write_int_le(self, value: int):
        self._ensure_capacity(4)
        struct.pack_into('<i', self.buffer, self.position, value)
        self.position += 4

    def write_long_le(self, value: int):
        self._ensure_capacity(8)
        struct.pack_into('<q', self.buffer, self.position, value)
        self.position += 8

    def write_float_le(self, value: float):
        self._ensure_capacity(4)
        struct.pack_into('<f', self.buffer, self.position, value)
        self.position += 4

    def write_double_le(self, value: float):
        self._ensure_capacity(8)
        struct.pack_into('<d', self.buffer, self.position, value)
        self.position += 8

    def write_var_int(self, value: int):
        self._ensure_capacity(5)
        while (value & ~0x7F) != 0:
            self.buffer[self.position] = (value & 0x7F) | 0x80
            self.position += 1
            value >>= 7
        self.buffer[self.position] = value & 0x7F
        self.position += 1

    def write_bytes_with_length(self, data: bytes):
        self.write_var_int(len(data))
        self._ensure_capacity(len(data))
        self.buffer[self.position:self.position + len(data)] = data
        self.position += len(data)

    def write_string(self, value: str):
        encoded = value.encode('utf-8')
        self.write_bytes_with_length(encoded)


def _encode_var_int(value: int) -> bytes:
    result = bytearray()
    while (value & ~0x7F) != 0:
        result.append((value & 0x7F) | 0x80)
        value >>= 7
    result.append(value & 0x7F)
    return bytes(result)


def _write_field(buf: _BlockBuffer, value: Any, data_type) -> None:
    if isinstance(data_type, AtomicType):
        type_name = data_type.type.upper()
        if type_name == 'BOOLEAN':
            buf.write_boolean(value)
        elif type_name == 'TINYINT':
            buf.write_byte(value)
        elif type_name == 'SMALLINT':
            buf.write_short_le(value)
        elif type_name in ('INT', 'INTEGER'):
            buf.write_int_le(value)
        elif type_name == 'DATE':
            if isinstance(value, datetime.date):
                epoch = datetime.date(1970, 1, 1)
                days = (value - epoch).days
                buf.write_int_le(days)
            else:
                buf.write_int_le(value)
        elif type_name == 'TIME' or (type_name.startswith('TIME') and not type_name.startswith('TIMESTAMP')):
            if isinstance(value, datetime.time):
                millis = (value.hour * 3600 + value.minute * 60 + value.second) * 1000 + value.microsecond // 1000
                buf.write_int_le(millis)
            else:
                buf.write_int_le(value)
        elif type_name == 'BIGINT':
            buf.write_long_le(value)
        elif type_name == 'FLOAT':
            buf.write_float_le(value)
        elif type_name == 'DOUBLE':
            buf.write_double_le(value)
        elif type_name == 'STRING' or type_name.startswith('CHAR') or type_name.startswith('VARCHAR'):
            buf.write_string(value)
        elif type_name == 'BYTES' or type_name.startswith('BINARY') or type_name.startswith('VARBINARY'):
            buf.write_bytes_with_length(value)
        elif type_name == 'BLOB':
            buf.write_bytes_with_length(value)
        elif type_name.startswith('DECIMAL'):
            precision, scale = _parse_decimal_params(type_name)
            dec = value if isinstance(value, Decimal) else Decimal(str(value))
            # Compute the unscaled value under a context wide enough for the column:
            # the default 28-digit precision would round a DECIMAL(p) value with more
            # than 28 significant digits, silently corrupting it. Mirrors the read path.
            with localcontext() as ctx:
                ctx.prec = max(precision + abs(scale), 38)
                unscaled = int(dec.scaleb(scale))
            if precision <= 18:
                buf.write_long_le(unscaled)
            else:
                raw = unscaled.to_bytes(
                    (unscaled.bit_length() + 8) // 8, byteorder='big', signed=True)
                buf.write_bytes_with_length(raw)
        elif type_name.startswith('TIMESTAMP'):
            precision = _parse_timestamp_precision(type_name)
            if isinstance(value, datetime.datetime):
                is_ltz = type_name.startswith('TIMESTAMP_LTZ') or 'WITH LOCAL TIME ZONE' in type_name
                if is_ltz and value.utcoffset() is not None:
                    value = value.astimezone(datetime.timezone.utc)
                # TIMESTAMP stores wall-clock fields; naive LTZ values use UTC.
                delta = value.replace(tzinfo=None) - datetime.datetime(1970, 1, 1)
                total_micros = ((delta.days * 86400 + delta.seconds) * 1_000_000
                                + delta.microseconds)
                if precision <= 3:
                    buf.write_long_le(total_micros // 1000)
                else:
                    millis = total_micros // 1000
                    nano_of_milli = (total_micros % 1000) * 1000
                    buf.write_long_le(millis)
                    buf.write_var_int(nano_of_milli)
            elif precision <= 3:
                buf.write_long_le(value)
            else:
                if isinstance(value, int):
                    millis = value // 1000
                    nano_of_milli = (value % 1000) * 1000
                else:
                    millis = int(value) // 1000
                    nano_of_milli = (int(value) % 1000) * 1000
                buf.write_long_le(millis)
                buf.write_var_int(nano_of_milli)
        elif type_name == 'VARIANT':
            if isinstance(value, dict):
                buf.write_bytes_with_length(value['value'])
                buf.write_bytes_with_length(value['metadata'])
            else:
                buf.write_bytes_with_length(value.value if hasattr(value, 'value') else b'')
                buf.write_bytes_with_length(value.metadata if hasattr(value, 'metadata') else b'')
        else:
            raise ValueError(f"Unsupported atomic type for writing: {type_name}")

    elif isinstance(data_type, ArrayType):
        _write_array_elements(buf, value, data_type.element)

    elif isinstance(data_type, VectorType):
        _write_vector(buf, value, data_type.element)

    elif isinstance(data_type, MapType):
        if isinstance(value, dict):
            keys = list(value.keys())
            values = list(value.values())
        else:
            keys = [pair[0] for pair in value]
            values = [pair[1] for pair in value]
        _write_array_elements(buf, keys, data_type.key)
        _write_array_elements(buf, values, data_type.value)

    elif isinstance(data_type, MultisetType):
        if isinstance(value, dict):
            keys = list(value.keys())
            counts = list(value.values())
        else:
            keys = [pair[0] for pair in value]
            counts = [pair[1] for pair in value]
        _write_array_elements(buf, keys, data_type.element)
        _write_array_elements(buf, counts, AtomicType("INT"))

    elif isinstance(data_type, RowType):
        _write_nested_row(buf, value, data_type)

    else:
        raise ValueError(f"Unsupported data type for writing: {data_type}")


def _write_array_elements(buf: _BlockBuffer, elements: list, element_type) -> None:
    size = len(elements)
    buf.write_var_int(size)
    null_bitmap_bytes = (size + 7) // 8

    null_start = buf.position
    buf.write_zeros(null_bitmap_bytes)

    for i, elem in enumerate(elements):
        if elem is None:
            buf.buffer[null_start + i // 8] |= (1 << (i % 8))
        else:
            _write_field(buf, elem, element_type)


def _write_vector(buf: _BlockBuffer, elements: list, element_type) -> None:
    size = len(elements)
    buf.write_var_int(size)
    for elem in elements:
        _write_field(buf, elem, element_type)


def _write_nested_row(buf: _BlockBuffer, value, row_type: RowType) -> None:
    fields = row_type.fields
    arity = len(fields)
    header_size = (arity + 7) // 8

    header_start = buf.position
    buf.write_zeros(header_size)

    for i, field in enumerate(fields):
        if isinstance(value, dict):
            field_value = value.get(field.name)
        else:
            field_value = value[i] if i < len(value) else None

        if field_value is None:
            buf.buffer[header_start + i // 8] |= (1 << (i % 8))
        else:
            _write_field(buf, field_value, field.type)


def _parse_decimal_params(type_name: str) -> tuple:
    match = re.fullmatch(r'DECIMAL\((\d+),\s*(\d+)\)', type_name)
    if match:
        return int(match.group(1)), int(match.group(2))
    match = re.fullmatch(r'DECIMAL\((\d+)\)', type_name)
    if match:
        return int(match.group(1)), 0
    return 10, 0


def _parse_timestamp_precision(type_name: str) -> int:
    match = re.search(r'\((\d+)\)', type_name)
    if match:
        return int(match.group(1))
    return 6