paimon-python/pypaimon/execution/python/read/reader/data_file_batch_reader.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.

from typing import List, Optional, Tuple

import pyarrow as pa
import pyarrow.compute as pc
from pyarrow import RecordBatch

from pypaimon.schema.arrow_schema import read_arrow_schema
from pypaimon.execution.python.casting.schema_evolution import cast_array_for_schema_evolution
from pypaimon.common.file_io import FileIO
from pypaimon.utils.arrow_utils import zero_column_batch
from pypaimon.data.map_shared_shredding import (
    is_map_selected_keys_field,
    map_selected_keys,
)
from pypaimon.execution.python.read.map_shared_shredding import assemble_normal_map_selected_keys
from pypaimon.read.partition_info import PartitionInfo
from pypaimon.execution.python.read.reader.format_blob_reader import FormatBlobReader
from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader
from pypaimon.schema.data_types import (ArrayType, AtomicType, DataField,
                                        MapType, PyarrowFieldParser, RowType)
from pypaimon.table.special_fields import SpecialFields


def _is_character_string_type(data_type) -> bool:
    if not isinstance(data_type, AtomicType):
        return False
    t = data_type.type.upper()
    return t == 'STRING' or t.startswith('VARCHAR') or t.startswith('CHAR')


def _unslice(array):
    """Re-materialize a sliced array so offsets/buffers start at zero.

    The list/map rebuilds below read ``offsets``/raw buffers directly; on a
    sliced array those still point into the parent storage, which either
    errors (list rebuild with a null mask) or silently misaligns rows (map
    rebuild from raw buffers)."""
    if array.offset == 0:
        return array
    return pa.concat_arrays([array])


def _to_string_values(array, data_type) -> list:
    """Render *array* as a list of per-row strings (None for NULL rows)."""
    if isinstance(data_type, (RowType, ArrayType, MapType)):
        return _constructed_to_string_array(array, data_type).to_pylist()
    return cast_array_for_schema_evolution(array, pa.string(), data_type, AtomicType('STRING')).to_pylist()


def _constructed_to_string_array(array, file_type):
    """Render a struct/list/map array in the engine's string form:
    ROW -> ``{v1, v2}``, ARRAY -> ``[e1, e2]``, MAP -> ``{k1 -> v1, k2 -> v2}``.
    Sub-values are rendered recursively; a NULL sub-value renders as the
    literal ``null`` while a NULL container row stays NULL."""
    array = _unslice(array)
    valid = pc.is_valid(array).to_pylist()
    out = []
    if isinstance(file_type, RowType):
        children = [
            _to_string_values(array.field(i), sub.type)
            for i, sub in enumerate(file_type.fields)
        ]
        for i in range(len(array)):
            if not valid[i]:
                out.append(None)
                continue
            vals = [c[i] if c[i] is not None else 'null' for c in children]
            out.append('{' + ', '.join(vals) + '}')
    elif isinstance(file_type, ArrayType):
        values = _to_string_values(array.values, file_type.element)
        offsets = array.offsets.to_pylist()
        for i in range(len(array)):
            if not valid[i]:
                out.append(None)
                continue
            elems = [v if v is not None else 'null'
                     for v in values[offsets[i]:offsets[i + 1]]]
            out.append('[' + ', '.join(elems) + ']')
    elif isinstance(file_type, MapType):
        keys = _to_string_values(array.keys, file_type.key)
        items = _to_string_values(array.items, file_type.value)
        offsets = array.offsets.to_pylist()
        for i in range(len(array)):
            if not valid[i]:
                out.append(None)
                continue
            entries = [
                '{} -> {}'.format(
                    keys[j] if keys[j] is not None else 'null',
                    items[j] if items[j] is not None else 'null')
                for j in range(offsets[i], offsets[i + 1])
            ]
            out.append('{' + ', '.join(entries) + '}')
    else:
        raise ValueError(
            'Unsupported constructed type for string rendering: {}'.format(file_type))
    return pa.array(out, type=pa.string())


class _RowIdRangeCursor:
    """Generate selected row offsets batch by batch from compact ranges."""

    def __init__(self, ranges: List[Tuple[int, int]]):
        self._ranges = ranges
        self._range_index = 0
        self._next_offset = ranges[0][0] if ranges else None

    def take(self, count: int) -> List[int]:
        offsets = []
        while len(offsets) < count:
            if self._range_index >= len(self._ranges):
                raise ValueError(
                    "Row range offsets were exhausted before the reader")
            lower, upper = self._ranges[self._range_index]
            current = max(lower, self._next_offset)
            take_count = min(count - len(offsets), upper - current + 1)
            offsets.extend(range(current, current + take_count))
            current += take_count
            if current > upper:
                self._range_index += 1
                self._next_offset = (
                    self._ranges[self._range_index][0]
                    if self._range_index < len(self._ranges) else None)
            else:
                self._next_offset = current
        return offsets


class DataFileBatchReader(RecordBatchReader):
    """
    Reads record batch from files of different formats
    """

    def __init__(self, format_reader: RecordBatchReader, index_mapping: List[int], partition_info: PartitionInfo,
                 system_primary_key: Optional[List[str]], fields: List[DataField],
                 max_sequence_number: int,
                 first_row_id: int,
                 row_tracking_enabled: bool,
                 system_fields: dict,
                 file_io: Optional[FileIO] = None,
                 row_id_offsets: Optional[List[int]] = None,
                 file_data_fields: Optional[List[DataField]] = None,
                 target_data_fields: Optional[List[DataField]] = None,
                 row_id_offset_ranges: Optional[List[Tuple[int, int]]] = None):
        if (row_id_offsets is not None
                and row_id_offset_ranges is not None):
            raise ValueError(
                "row_id_offsets and row_id_offset_ranges cannot both "
                "be provided")
        self.format_reader = format_reader
        self.index_mapping = index_mapping
        self.partition_info = partition_info
        self.system_primary_key = system_primary_key
        source_fields = (
            file_data_fields if file_data_fields is not None else fields)
        self._file_data_field_map = {
            field.name: field for field in source_fields
        }
        self.schema_map = {field.name: field for field in read_arrow_schema(fields)}
        self.row_tracking_enabled = row_tracking_enabled
        self.first_row_id = first_row_id
        self.row_id_offsets = row_id_offsets
        self._row_id_cursor = 0
        self._row_id_range_cursor = (
            _RowIdRangeCursor(row_id_offset_ranges)
            if row_id_offset_ranges is not None else None)
        self.max_sequence_number = max_sequence_number
        self.system_fields = system_fields
        self.file_io = file_io
        # Per-file field-id normalization: map the physically-read columns
        # (the file's own field order/names) onto the latest read target by
        # field id, padding missing ids with NULL and recursing into nested
        # ROW / ARRAY<ROW> / MAP<.,ROW> sub-fields the same way. ``None`` when
        # there is no evolution to reconcile -- the common path stays zero-copy.
        self._normalize_plan = self._build_normalize_plan(file_data_fields, target_data_fields)

    @staticmethod
    def _build_normalize_plan(file_data_fields, target_data_fields):
        """Build a per-file field-id alignment plan.

        Returns a list of ``(pos, file_field, target_field)`` -- one per target
        field, in target order -- where ``pos`` is the column index in the
        physically-read batch carrying ``target_field`` (matched by field id),
        or -1 if the file does not contain that id (pad NULL). Returns ``None``
        when the file already matches the target exactly (no evolution), so the
        caller stays zero-copy.
        """
        if file_data_fields is None or target_data_fields is None:
            return None
        # Recursive equality covers nested sub-field changes too: any rename /
        # add / drop / type change at any depth makes the file != target.
        if file_data_fields == target_data_fields:
            return None
        file_id_to_pos = {f.id: i for i, f in enumerate(file_data_fields)}
        plan = []
        for target in target_data_fields:
            pos = file_id_to_pos.get(target.id, -1)
            file_field = file_data_fields[pos] if pos >= 0 else None
            plan.append((pos, file_field, target))
        return plan

    def _normalize_batch(self, record_batch: RecordBatch) -> RecordBatch:
        """Reorder/pad the physically-read batch onto the latest read target by
        field id, relabel columns to the latest names, and align nested ROW
        sub-fields by id. Missing ids become typed all-NULL columns."""
        if self._normalize_plan is None:
            return record_batch
        if not self._normalize_plan:
            return record_batch.select([])
        num_rows = record_batch.num_rows
        arrays = []
        names = []
        for pos, file_field, target_field in self._normalize_plan:
            target_pa_type = PyarrowFieldParser.from_paimon_type(target_field.type)
            if pos < 0:
                arrays.append(pa.nulls(num_rows, type=target_pa_type))
            else:
                arrays.append(self._align_array_by_id(
                    record_batch.column(pos), file_field.type, target_field.type))
            names.append(target_field.name)
        return pa.RecordBatch.from_arrays(arrays, names=names)

    def _assemble_selected_map_keys(
            self, record_batch: RecordBatch) -> RecordBatch:
        columns = list(record_batch.columns)
        fields = list(record_batch.schema)
        changed = False
        for index, field in enumerate(fields):
            data_field = self._file_data_field_map.get(field.name)
            if (data_field is None
                    or not is_map_selected_keys_field(data_field)
                    or not pa.types.is_map(columns[index].type)):
                continue
            value_type = PyarrowFieldParser.from_paimon_type(
                data_field.type.fields[0].type)
            columns[index] = assemble_normal_map_selected_keys(
                columns[index],
                map_selected_keys(data_field.description),
                value_type)
            fields[index] = pa.field(
                field.name, columns[index].type, nullable=field.nullable)
            changed = True
        if not changed:
            return record_batch
        return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields))

    def _align_array_by_id(self, array, file_type, target_type):
        """Return *array* converted to *target_type*, matching ROW sub-fields by
        field id (reorder, pad missing with NULL, follow renames, cast changed
        types) recursively, transparently through ARRAY/MAP wrappers."""
        if isinstance(target_type, RowType) and isinstance(file_type, RowType):
            n = len(array)
            file_id_to_pos = {f.id: i for i, f in enumerate(file_type.fields)}
            children = []
            pa_fields = []
            for tsub in target_type.fields:
                p = file_id_to_pos.get(tsub.id, -1)
                if p < 0:
                    child = pa.nulls(n, type=PyarrowFieldParser.from_paimon_type(tsub.type))
                else:
                    child = self._align_array_by_id(
                        array.field(p), file_type.fields[p].type, tsub.type)
                children.append(child)
                pa_fields.append(pa.field(tsub.name, child.type, nullable=tsub.type.nullable))
            # Preserve the struct's own null mask; child values under a null
            # struct are irrelevant.
            return pa.StructArray.from_arrays(
                children, fields=pa_fields, mask=pc.is_null(array))
        if isinstance(target_type, ArrayType) and isinstance(file_type, ArrayType):
            array = _unslice(array)
            aligned_values = self._align_array_by_id(
                array.values, file_type.element, target_type.element)
            return pa.ListArray.from_arrays(
                array.offsets, aligned_values, mask=pc.is_null(array))
        if isinstance(target_type, MapType) and isinstance(file_type, MapType):
            array = _unslice(array)
            aligned_items = self._align_array_by_id(
                array.items, file_type.value, target_type.value)
            # MapArray.from_arrays cannot carry a null mask (a null map would
            # collapse to an empty one), so rebuild from buffers, reusing the
            # original validity/offset buffers and only swapping the value child.
            target_pa = PyarrowFieldParser.from_paimon_type(target_type)
            entries = pa.StructArray.from_arrays(
                [array.keys, aligned_items],
                fields=[target_pa.key_field, target_pa.item_field])
            return pa.Array.from_buffers(
                target_pa, len(array), array.buffers()[:2], children=[entries])
        # A constructed type changed to a character string: pyarrow cannot
        # cast struct/list/map to utf8 directly, so render the engine's
        # string form instead.
        if (isinstance(file_type, (RowType, ArrayType, MapType))
                and _is_character_string_type(target_type)):
            return _constructed_to_string_array(array, file_type)
        # Leaf / non-nested: cast to the target type when it differs.
        target_pa_type = PyarrowFieldParser.from_paimon_type(target_type)
        return cast_array_for_schema_evolution(array, target_pa_type, file_type, target_type)

    def read_arrow_batch(self, start_idx=None, end_idx=None) -> Optional[RecordBatch]:
        if isinstance(self.format_reader, FormatBlobReader):
            record_batch = self.format_reader.read_arrow_batch(start_idx, end_idx)
        else:
            record_batch = self.format_reader.read_arrow_batch()
        if record_batch is None:
            return None
        record_batch = self._assemble_selected_map_keys(record_batch)
        record_batch = self._normalize_batch(record_batch)

        if self.partition_info is None and self.index_mapping is None:
            # A file written under an older schema (e.g. before an INT -> BIGINT
            # promotion or a DECIMAL precision change) yields columns in the
            # data file's original types. Without reordering or partition padding
            # to rebuild the batch, those old types would otherwise leak through
            # here -- returning a type that depends on whether this read happens
            # to span newer-schema files, and failing to concatenate when it
            # does. Align them to the current read schema, mirroring the rebuild
            # path below.
            record_batch = self._align_batch_to_read_schema(
                record_batch.schema.names, record_batch.columns, record_batch.num_rows)
            if self.row_tracking_enabled and self.system_fields:
                record_batch = self._assign_row_tracking(record_batch)
            return record_batch

        inter_arrays = []
        inter_names = []
        num_rows = record_batch.num_rows

        if self.partition_info is not None:
            for i in range(self.partition_info.size()):
                if self.partition_info.is_partition_row(i):
                    partition_value, partition_field = self.partition_info.get_partition_value(i)
                    const_array = pa.repeat(partition_value, num_rows)
                    inter_arrays.append(const_array)
                    inter_names.append(partition_field.name)
                else:
                    real_index = self.partition_info.get_real_index(i)
                    if real_index < record_batch.num_columns:
                        inter_arrays.append(record_batch.column(real_index))
                        inter_names.append(record_batch.schema.field(real_index).name)
        else:
            inter_arrays = record_batch.columns
            inter_names = record_batch.schema.names

        if self.index_mapping is not None:
            mapped_arrays = []
            mapped_names = []
            for i, real_index in enumerate(self.index_mapping):
                if 0 <= real_index < len(inter_arrays):
                    mapped_arrays.append(inter_arrays[real_index])
                    mapped_names.append(inter_names[real_index])
                else:
                    null_array = pa.nulls(num_rows)
                    mapped_arrays.append(null_array)
                    mapped_names.append(f"null_col_{i}")

            if self.system_primary_key:
                for i in range(len(self.system_primary_key)):
                    if not mapped_names[i].startswith("_KEY_"):
                        mapped_names[i] = f"_KEY_{mapped_names[i]}"

            inter_arrays = mapped_arrays
            inter_names = mapped_names

        # Rebuild the batch typed by the read schema (carries 'not null' and
        # aligns old-schema column types).
        record_batch = self._align_batch_to_read_schema(inter_names, inter_arrays, num_rows)

        # Handle row tracking fields
        if self.row_tracking_enabled and self.system_fields:
            record_batch = self._assign_row_tracking(record_batch)

        return record_batch

    def _align_batch_to_read_schema(self, names: List[str], arrays: list, num_rows: int) -> RecordBatch:
        """Build a record batch for ``names``/``arrays`` typed by the read schema.

        Each known field is cast to the current read schema's type, which also
        carries the 'not null' property; unknown columns keep the array's own
        type. Columns whose type already matches are reused as-is, keeping the
        common (non-evolution) path zero-copy.

        DECIMAL scale reductions use Java's HALF_UP and overflow-to-NULL
        semantics; numeric narrowing and timestamp precision also follow Java casts.
        Evolution legality is the writer's concern (``DataTypeCasts``).
        """
        out_arrays = []
        out_fields = []
        for name, array in zip(names, arrays):
            target_field = self.schema_map.get(name)
            if target_field is None:
                target_field = pa.field(name, array.type)
            elif array.type != target_field.type:
                array = cast_array_for_schema_evolution(array, target_field.type)
            out_arrays.append(array)
            out_fields.append(target_field)
        if not out_arrays:
            return zero_column_batch(num_rows)
        return pa.RecordBatch.from_arrays(out_arrays, schema=pa.schema(out_fields))

    def _assign_row_tracking(self, record_batch: RecordBatch) -> RecordBatch:
        """Assign row tracking meta fields (_ROW_ID and _SEQUENCE_NUMBER)."""
        arrays = list(record_batch.columns)

        # Handle _ROW_ID field
        if SpecialFields.ROW_ID.name in self.system_fields.keys():
            idx = self.system_fields[SpecialFields.ROW_ID.name]
            if self.first_row_id is None:
                arrays[idx] = pa.nulls(record_batch.num_rows, type=pa.int64())
            elif self.row_id_offsets is not None:
                end = self._row_id_cursor + record_batch.num_rows
                row_ids = [self.first_row_id + o for o in self.row_id_offsets[self._row_id_cursor:end]]
                arrays[idx] = pa.array(row_ids, type=pa.int64())
                self._row_id_cursor = end
            elif self._row_id_range_cursor is not None:
                row_ids = [
                    self.first_row_id + offset
                    for offset in self._row_id_range_cursor.take(
                        record_batch.num_rows)
                ]
                arrays[idx] = pa.array(row_ids, type=pa.int64())
            else:
                row_id_range = range(self.first_row_id, self.first_row_id + record_batch.num_rows)
                arrays[idx] = pa.array(row_id_range, type=pa.int64())
                self.first_row_id += record_batch.num_rows

        # Handle _SEQUENCE_NUMBER field
        if SpecialFields.SEQUENCE_NUMBER.name in self.system_fields.keys():
            idx = self.system_fields[SpecialFields.SEQUENCE_NUMBER.name]
            # Java's tracking vector retains physical versions and only fills
            # NULL entries from the manifest's max sequence number.
            arrays[idx] = pc.fill_null(arrays[idx], self.max_sequence_number)

        return pa.RecordBatch.from_arrays(arrays, schema=record_batch.schema)

    def close(self) -> None:
        self.format_reader.close()