paimon-python/pypaimon/execution/python/write/blob_format_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 struct
from collections.abc import Mapping
from typing import BinaryIO, List, Optional

try:
    from isal import isal_zlib as crc_backend
except ImportError:
    import zlib as crc_backend

from pypaimon.execution.python.map_blob_key_serializer import create_map_blob_key_serializer
from pypaimon.schema.data_types import (
    MapType,
    is_array_blob_type,
    is_blob_file_type,
    is_map_blob_type,
)
from pypaimon.table.row.blob import (
    Blob,
    BlobConsumer,
    BlobData,
    BlobDescriptor,
    BlobDescriptorSerde,
)
from pypaimon.common.delta_varint_compressor import DeltaVarintCompressor


class BlobFormatWriter:
    VERSION = 1
    MAGIC_NUMBER = 1481511375
    ARRAY_VERSION = 1
    ARRAY_MAGIC_NUMBER = 1094861634
    MAP_VERSION = 1
    MAP_MAGIC_NUMBER = 0x4D424342
    NULL_LENGTH = -1
    PLACE_HOLDER_LENGTH = -2
    ARRAY_NULL_ELEMENT_LENGTH = -1
    MAP_NULL_KEY_LENGTH = -1
    MAP_NULL_VALUE_LENGTH = -1
    BUFFER_SIZE = 4096
    METADATA_SIZE = 12  # 8-byte length + 4-byte CRC

    def __init__(self, output_stream: BinaryIO,
                 blob_consumer: Optional[BlobConsumer] = None,
                 file_path: Optional[str] = None,
                 copy_buffer_size: int = BUFFER_SIZE):
        if copy_buffer_size <= 0:
            raise ValueError(
                f"BLOB copy buffer size must be positive, but was {copy_buffer_size}.")
        self.output_stream = output_stream
        self._blob_consumer = blob_consumer
        self._file_path = file_path
        self.copy_buffer_size = copy_buffer_size
        self.lengths: List[int] = []
        self.position = 0

    def add_element(self, row) -> None:
        if not hasattr(row, 'values') or len(row.values) != 1:
            raise ValueError("BlobFormatWriter only supports one field")

        blob_field_name = row.fields[0].name
        blob_field_type = row.fields[0].type
        blob_value = row.values[0]
        if isinstance(blob_field_type, MapType) and not is_map_blob_type(blob_field_type):
            raise ValueError(
                f"Map-Blob value type must be BLOB, but is {blob_field_type.value}."
            )
        if blob_value is None:
            self.lengths.append(self.NULL_LENGTH)
            if self._blob_consumer is not None:
                self._blob_consumer(blob_field_name, None)
            return

        if is_map_blob_type(blob_field_type):
            if blob_value is Blob.MAP_PLACE_HOLDER:
                self.lengths.append(self.PLACE_HOLDER_LENGTH)
                return
            self.add_blob_map(
                blob_field_name,
                blob_value,
                blob_field_type.key,
                value_nullable=blob_field_type.value.nullable,
            )
            return

        if is_array_blob_type(blob_field_type):
            if blob_value is Blob.ARRAY_PLACE_HOLDER:
                self.lengths.append(self.PLACE_HOLDER_LENGTH)
                return
            self.add_blob_array(
                blob_field_name,
                blob_value,
                element_nullable=blob_field_type.element.nullable,
            )
            return

        if blob_value is Blob.PLACE_HOLDER:
            self.lengths.append(self.PLACE_HOLDER_LENGTH)
            return

        self.add_blob(blob_field_name, blob_value)

    def add_blob(self, blob_field_name: str, blob_value: Blob) -> None:
        if not isinstance(blob_value, Blob):
            raise ValueError("Field must be Blob/BlobData instance")

        previous_pos = self.position
        crc32 = 0  # Initialize CRC32

        # Write magic number
        magic_bytes = struct.pack('<I', self.MAGIC_NUMBER)  # Little endian
        crc32 = self._write_with_crc(magic_bytes, crc32)

        blob_pos, blob_length, crc32 = self._write_blob_data(blob_value, crc32)

        # Calculate total length including magic + data + metadata (length + CRC)
        bin_length = self.position - previous_pos + self.METADATA_SIZE
        self.lengths.append(bin_length)

        # Write length (8 bytes, little endian)
        length_bytes = struct.pack('<Q', bin_length)
        crc32 = self._write_with_crc(length_bytes, crc32)

        # Write CRC32 (4 bytes, little endian)
        crc_bytes = struct.pack('<I', crc32 & 0xffffffff)
        self.output_stream.write(crc_bytes)
        self.position += 4

        if self._blob_consumer is not None:
            descriptor = BlobDescriptor(self._file_path, blob_pos, blob_length)
            flush = self._blob_consumer(blob_field_name, descriptor)
            if flush:
                self.output_stream.flush()

    def add_blob_array(
            self,
            blob_field_name: str,
            blob_values,
            element_nullable: bool = True) -> None:
        if isinstance(blob_values, (bytes, bytearray, str)) or not hasattr(blob_values, '__iter__'):
            raise ValueError("ARRAY<BLOB> field value must be a list or tuple of Blob values")

        blob_values = list(blob_values)
        for blob_value in blob_values:
            if blob_value is None:
                if not element_nullable:
                    raise ValueError(
                        f"ARRAY<BLOB> field '{blob_field_name}' does not allow null elements"
                    )
                continue
            if not isinstance(blob_value, Blob):
                raise ValueError("ARRAY<BLOB> elements must be Blob/BlobData instances or None")
            if (
                blob_value is Blob.PLACE_HOLDER
                or blob_value is Blob.ARRAY_PLACE_HOLDER
                or blob_value is Blob.MAP_PLACE_HOLDER
            ):
                raise ValueError("ARRAY<BLOB> elements do not support placeholders")

        previous_pos = self.position
        crc32 = 0

        crc32 = self._write_with_crc(struct.pack('<I', self.MAGIC_NUMBER), crc32)
        crc32 = self._write_with_crc(struct.pack('<I', self.ARRAY_MAGIC_NUMBER), crc32)
        crc32 = self._write_with_crc(struct.pack('<B', self.ARRAY_VERSION), crc32)
        crc32 = self._write_with_crc(struct.pack('<I', len(blob_values)), crc32)

        element_lengths = []
        flush = False
        for blob_value in blob_values:
            if blob_value is None:
                element_lengths.append(self.ARRAY_NULL_ELEMENT_LENGTH)
                continue

            blob_pos, blob_length, crc32 = self._write_blob_data(blob_value, crc32)
            element_lengths.append(blob_length)
            if self._blob_consumer is not None:
                descriptor = BlobDescriptor(self._file_path, blob_pos, blob_length)
                flush = self._blob_consumer(blob_field_name, descriptor) or flush

        element_index_bytes = DeltaVarintCompressor.compress(element_lengths)
        crc32 = self._write_with_crc(element_index_bytes, crc32)
        crc32 = self._write_with_crc(struct.pack('<I', len(element_index_bytes)), crc32)

        bin_length = self.position - previous_pos + self.METADATA_SIZE
        self.lengths.append(bin_length)
        crc32 = self._write_with_crc(struct.pack('<Q', bin_length), crc32)
        self.output_stream.write(struct.pack('<I', crc32 & 0xffffffff))
        self.position += 4

        if flush:
            self.output_stream.flush()

    def add_blob_map(
            self,
            blob_field_name: str,
            blob_values,
            key_type,
            value_nullable: bool = True) -> None:
        entries = self._map_entries(blob_values)
        if len(entries) > 0x7fffffff:
            raise ValueError(f"Invalid MAP<X, BLOB> entry count: {len(entries)}")

        key_serializer = create_map_blob_key_serializer(key_type)
        key_bytes = []
        seen_keys = set()
        for key, _ in entries:
            if key is None:
                serialized = None
            else:
                serialized = key_serializer.serialize(key)
                if len(serialized) > 0x7fffffff:
                    raise ValueError(f"MAP<X, BLOB> key is too large: {len(serialized)}")
            if serialized in seen_keys:
                raise ValueError("MAP<X, BLOB> keys must be unique.")
            seen_keys.add(serialized)
            key_bytes.append(serialized)

        for _, blob_value in entries:
            if blob_value is None:
                if not value_nullable:
                    raise ValueError(
                        f"MAP<X, BLOB> field '{blob_field_name}' does not allow null values"
                    )
                continue
            if (
                blob_value is Blob.PLACE_HOLDER
                or blob_value is Blob.ARRAY_PLACE_HOLDER
                or blob_value is Blob.MAP_PLACE_HOLDER
            ):
                raise ValueError("MAP<X, BLOB> values do not support placeholders")
            if not isinstance(blob_value, Blob):
                raise ValueError("MAP<X, BLOB> values must be Blob/BlobData instances or None")

        previous_pos = self.position
        crc32 = 0
        crc32 = self._write_with_crc(struct.pack('<I', self.MAGIC_NUMBER), crc32)
        crc32 = self._write_with_crc(struct.pack('<I', self.MAP_MAGIC_NUMBER), crc32)
        crc32 = self._write_with_crc(struct.pack('<B', self.MAP_VERSION), crc32)
        crc32 = self._write_with_crc(struct.pack('<I', len(entries)), crc32)

        key_lengths = []
        for serialized in key_bytes:
            if serialized is None:
                key_lengths.append(self.MAP_NULL_KEY_LENGTH)
            else:
                key_lengths.append(len(serialized))
                crc32 = self._write_with_crc(serialized, crc32)

        value_lengths = []
        flush = False
        for _, blob_value in entries:
            if blob_value is None:
                value_lengths.append(self.MAP_NULL_VALUE_LENGTH)
                continue
            blob_pos, blob_length, crc32 = self._write_blob_data(blob_value, crc32)
            value_lengths.append(blob_length)
            if self._blob_consumer is not None:
                descriptor = BlobDescriptor(self._file_path, blob_pos, blob_length)
                flush = self._blob_consumer(blob_field_name, descriptor) or flush

        key_index_bytes = DeltaVarintCompressor.compress(key_lengths)
        value_index_bytes = DeltaVarintCompressor.compress(value_lengths)
        crc32 = self._write_with_crc(key_index_bytes, crc32)
        crc32 = self._write_with_crc(value_index_bytes, crc32)
        crc32 = self._write_with_crc(struct.pack('<I', len(key_index_bytes)), crc32)
        crc32 = self._write_with_crc(struct.pack('<I', len(value_index_bytes)), crc32)

        bin_length = self.position - previous_pos + self.METADATA_SIZE
        self.lengths.append(bin_length)
        crc32 = self._write_with_crc(struct.pack('<Q', bin_length), crc32)
        self.output_stream.write(struct.pack('<I', crc32 & 0xffffffff))
        self.position += 4

        if flush:
            self.output_stream.flush()

    def _write_blob_data(self, blob_value: Blob, crc32: int):
        blob_pos = self.position

        if isinstance(blob_value, BlobData):
            data = blob_value.to_data()
            crc32 = self._write_with_crc(data, crc32)
        else:
            stream = blob_value.new_input_stream()
            try:
                chunk = stream.read(self.copy_buffer_size)
                while chunk:
                    crc32 = self._write_with_crc(chunk, crc32)
                    chunk = stream.read(self.copy_buffer_size)
            finally:
                stream.close()

        return blob_pos, self.position - blob_pos, crc32

    def _write_with_crc(self, data: bytes, crc32: int) -> int:
        crc32 = crc_backend.crc32(data, crc32)
        self.output_stream.write(data)
        self.position += len(data)
        return crc32

    def write_value(self, col_data, fields, uri_reader_factory=None) -> None:
        from pypaimon.table.row.generic_row import GenericRow
        from pypaimon.table.row.row_kind import RowKind

        is_map_blob = is_map_blob_type(fields[0].type)
        is_blob = is_blob_file_type(fields[0].type)

        if col_data is None:
            if not is_blob:
                raise RuntimeError("Null values are only supported for BLOB type fields")
            self.lengths.append(self.NULL_LENGTH)
            return

        if is_map_blob:
            row_values = [self._to_blob_map(col_data, uri_reader_factory)]
        elif is_array_blob_type(fields[0].type):
            row_values = [self._to_blob_array(col_data, uri_reader_factory)]
        elif is_blob:
            row_values = [self._to_blob(col_data, uri_reader_factory)]
        else:
            row_values = [col_data]

        row = GenericRow(row_values, fields, RowKind.INSERT)
        self.add_element(row)

    @staticmethod
    def _to_blob(col_data, uri_reader_factory=None) -> Blob:
        if col_data is Blob.PLACE_HOLDER:
            return Blob.PLACE_HOLDER
        if hasattr(col_data, 'as_py'):
            col_data = col_data.as_py()
        if isinstance(col_data, str):
            col_data = col_data.encode('utf-8')
        if isinstance(col_data, bytearray):
            col_data = bytes(col_data)
        if isinstance(col_data, Blob):
            return col_data
        if isinstance(col_data, bytes):
            if BlobDescriptorSerde.is_descriptor(col_data):
                if uri_reader_factory is None:
                    raise RuntimeError("uri_reader_factory is required for BlobDescriptor bytes.")
                descriptor = BlobDescriptorSerde.deserialize(col_data)
                uri_reader = uri_reader_factory.create(descriptor.uri)
                return Blob.from_descriptor(uri_reader, descriptor)
            return BlobData(col_data)
        raise RuntimeError(
            "Blob field value must be bytes/blob or serialized BlobDescriptor bytes."
        )

    @classmethod
    def _to_blob_array(cls, col_data, uri_reader_factory=None):
        if col_data is Blob.ARRAY_PLACE_HOLDER:
            return Blob.ARRAY_PLACE_HOLDER
        if hasattr(col_data, 'as_py'):
            col_data = col_data.as_py()
        if col_data is None:
            return None
        if isinstance(col_data, (bytes, bytearray, str)) or not hasattr(col_data, '__iter__'):
            raise RuntimeError("ARRAY<BLOB> field value must be a list or tuple.")
        result = []
        for element in col_data:
            if element is None:
                result.append(None)
                continue
            if (
                element is Blob.PLACE_HOLDER
                or element is Blob.ARRAY_PLACE_HOLDER
                or element is Blob.MAP_PLACE_HOLDER
            ):
                raise RuntimeError("ARRAY<BLOB> elements do not support placeholders.")
            result.append(cls._to_blob(element, uri_reader_factory))
        return result

    @classmethod
    def _to_blob_map(cls, col_data, uri_reader_factory=None):
        if col_data is Blob.MAP_PLACE_HOLDER:
            return Blob.MAP_PLACE_HOLDER
        if hasattr(col_data, 'as_py'):
            col_data = col_data.as_py()
        if col_data is None:
            return None

        result = []
        for key, value in cls._map_entries(col_data):
            if value is None:
                result.append((key, None))
                continue
            if (
                value is Blob.PLACE_HOLDER
                or value is Blob.ARRAY_PLACE_HOLDER
                or value is Blob.MAP_PLACE_HOLDER
            ):
                raise RuntimeError("MAP<X, BLOB> values do not support placeholders.")
            result.append((key, cls._to_blob(value, uri_reader_factory)))
        return result

    @staticmethod
    def _map_entries(blob_values):
        if isinstance(blob_values, Mapping):
            return list(blob_values.items())
        if isinstance(blob_values, (bytes, bytearray, str)) or not hasattr(
            blob_values, '__iter__'
        ):
            raise ValueError("MAP<X, BLOB> field value must be a mapping or key/value pairs")

        entries = []
        for entry in blob_values:
            if not isinstance(entry, (list, tuple)) or len(entry) != 2:
                raise ValueError("MAP<X, BLOB> entries must be key/value pairs")
            entries.append((entry[0], entry[1]))
        return entries

    def reach_target_size(self, target_size: int) -> bool:
        return self.position >= target_size

    def close(self) -> None:
        index_bytes = DeltaVarintCompressor.compress(self.lengths)
        self.output_stream.write(index_bytes)

        # Write header (index length + version)
        header = struct.pack('<I', len(index_bytes))  # Index length (4 bytes, little endian)
        header += struct.pack('<B', self.VERSION)  # Version (1 byte)
        self.output_stream.write(header)

        # Flush and close
        if hasattr(self.output_stream, 'flush'):
            self.output_stream.flush()
        if hasattr(self.output_stream, 'close'):
            self.output_stream.close()