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

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.read.map_shared_shredding import assemble_normal_map_selected_keys
from pypaimon.execution.python.read.reader.data_file_batch_reader import \
    cast_array_for_schema_evolution
from pypaimon.read.reader.field_indices import (
    blob_field_indices, descriptor_field_indices, vector_field_indices)
from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader
from pypaimon.schema.data_types import DataField
from pypaimon.utils.arrow_utils import as_blob_data, is_blob_data


def _struct_field(column, name):
    # Resolve literal names before Arrow can interpret them as field paths.
    field_index = [field.name for field in column.type].index(name)
    struct_field = getattr(pc, "struct_field", None)
    if struct_field is not None:
        return struct_field(column, field_index)

    return column.flatten()[field_index]


class NestedLeafBatchReader(RecordBatchReader):
    """Extract projected nested leaves from batches of full top-level columns.

    The inner reader yields batches carrying the widened top-level columns,
    already normalized to the latest schema by field id (renames followed,
    missing sub-fields padded NULL, types cast). Each requested name path is
    walked through the struct children (a NULL parent propagates to the
    leaf), producing the user's flat projected schema.
    """

    def __init__(self, inner: RecordBatchReader, name_paths: List[List[str]],
                 output_fields: List[DataField],
                 descriptor_field_names=None):
        if len(name_paths) != len(output_fields):
            raise ValueError(
                "name_paths length {} does not match output_fields length {}".format(
                    len(name_paths), len(output_fields)))
        self._inner = inner
        self._paths = name_paths
        self._schema = read_arrow_schema(output_fields)
        self.file_io = inner.file_io
        self.blob_field_indices = blob_field_indices(output_fields)
        self.descriptor_field_indices = descriptor_field_indices(
            output_fields, descriptor_field_names or ())
        self.vector_field_indices = vector_field_indices(output_fields)

    def read_arrow_batch(self) -> Optional[RecordBatch]:
        batch = self._inner.read_arrow_batch()
        if batch is None:
            return None
        arrays = []
        fields = list(self._schema)
        for i, path in enumerate(self._paths):
            column = batch.column(path[0])
            for name in path[1:]:
                if pa.types.is_map(column.type):
                    column = assemble_normal_map_selected_keys(
                        column, [name], column.type.item_type).field(0)
                else:
                    column = _struct_field(column, name)
            target_type = self._schema.field(i).type
            if column.type != target_type:
                column = cast_array_for_schema_evolution(column, target_type)
            arrays.append(column)
            if len(path) == 1 and is_blob_data(batch.schema.field(path[0])):
                fields[i] = as_blob_data(fields[i])
        return pa.RecordBatch.from_arrays(
            arrays, schema=pa.schema(fields, metadata=self._schema.metadata))

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