split_read.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 os
from abc import ABC, abstractmethod
from functools import partial
from typing import Callable, Dict, List, Optional, Tuple

from pypaimon.common.merge_engine_dispatch import build_merge_function
from pypaimon.common.options.core_options import CoreOptions, MergeEngine
from pypaimon.common.predicate import Predicate
from pypaimon.deletionvectors import ApplyDeletionVectorReader
from pypaimon.deletionvectors.deletion_vector import DeletionVector
from pypaimon.globalindex import Range
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.read.interval_partition import IntervalPartition, SortedRun
from pypaimon.read.partition_info import PartitionInfo
from pypaimon.read.push_down_utils import rewrite_predicate_indices, trim_predicate_by_fields
from pypaimon.read.reader.concat_batch_reader import (
    BlobFallbackBatchReader, ConcatBatchReader,
    MergeAllBatchReader, DataEvolutionMergeReader)
from pypaimon.read.reader.concat_record_reader import ConcatRecordReader

from pypaimon.read.reader.data_file_batch_reader import DataFileBatchReader
from pypaimon.read.reader.drop_delete_reader import DropDeleteRecordReader
from pypaimon.read.reader.empty_record_reader import EmptyFileRecordReader
from pypaimon.read.reader.field_bunch import BlobBunch, DataBunch, FieldBunch, VectorBunch
from pypaimon.read.reader.filter_record_reader import FilterRecordReader
from pypaimon.read.reader.format_avro_reader import FormatAvroReader
from pypaimon.read.reader.blob_descriptor_convert_reader import BlobInlineConvertReader
from pypaimon.read.reader.filter_record_batch_reader import FilterRecordBatchReader
from pypaimon.read.reader.limited_record_reader import LimitedRecordBatchReader, LimitedRecordReader
from pypaimon.read.reader.row_range_filter_record_reader import RowIdFilterRecordBatchReader
from pypaimon.read.reader.format_blob_reader import FormatBlobReader
from pypaimon.read.reader.format_lance_reader import FormatLanceReader
from pypaimon.read.reader.format_pyarrow_reader import FormatPyArrowReader
from pypaimon.read.reader.format_row_reader import FormatRowReader
from pypaimon.read.reader.format_mosaic_reader import FormatMosaicReader
from pypaimon.read.reader.format_vortex_reader import FormatVortexReader
from pypaimon.read.reader.iface.record_batch_reader import (RecordBatchReader,
                                                            RowPositionReader, EmptyRecordBatchReader)
from pypaimon.read.reader.iface.record_reader import RecordReader
from pypaimon.read.reader.key_value_unwrap_reader import \
    KeyValueUnwrapRecordReader
from pypaimon.read.reader.key_value_wrap_reader import KeyValueWrapReader
from pypaimon.read.reader.shard_batch_reader import ShardBatchReader
from pypaimon.read.reader.aggregation_merge_function import (
    AggregateMergeFunction, build_field_aggregators)
from pypaimon.read.reader.sort_merge_reader import (SortMergeReaderWithMinHeap,
                                                    builtin_seq_comparator)
from pypaimon.read.push_down_utils import _get_all_fields
from pypaimon.read.split import Split
from pypaimon.read.sliced_split import SlicedSplit
from pypaimon.schema.data_types import DataField, PyarrowFieldParser
from pypaimon.table.special_fields import SpecialFields
from pypaimon.globalindex.indexed_split import IndexedSplit

KEY_PREFIX = "_KEY_"
KEY_FIELD_ID_START = 1000000
NULL_FIELD_INDEX = -1

_COMPRESS_EXTENSIONS = frozenset(['gz', 'bz2', 'deflate', 'snappy', 'lz4', 'zst'])


def _blob_field_indices(fields: List[DataField]) -> set:
    return {i for i, f in enumerate(fields)
            if hasattr(f.type, 'type') and f.type.type == 'BLOB'}


def _vector_field_indices(fields: List[DataField]) -> set:
    from pypaimon.schema.data_types import VectorType
    return {i for i, f in enumerate(fields) if isinstance(f.type, VectorType)}


def format_identifier(file_name):
    idx = file_name.rfind('.')
    assert idx != -1, "%s is not a legal file name." % file_name
    ext = file_name[idx + 1:]
    if ext.lower() in _COMPRESS_EXTENSIONS:
        second_idx = file_name.rfind('.', 0, idx)
        assert second_idx != -1, "%s is not a legal file name." % file_name
        return file_name[second_idx + 1:idx]
    return ext


class SplitRead(ABC):
    """Abstract base class for split reading operations."""

    def __init__(
            self,
            table,
            predicate: Optional[Predicate],
            read_type: List[DataField],
            split: Split,
            row_tracking_enabled: bool,
            nested_name_paths: Optional[List[List[str]]] = None,
            limit: Optional[int] = None):
        from pypaimon.table.file_store_table import FileStoreTable

        self.table: FileStoreTable = table
        self.predicate = predicate
        self.push_down_predicate = self._push_down_predicate()
        self.split = split
        self.row_tracking_enabled = row_tracking_enabled
        self.value_arity = len(read_type)
        self.nested_name_paths = nested_name_paths
        self.limit = limit
        # Snapshot the raw value-side schema before _create_key_value_fields
        # wraps it, so MergeFileSplitRead can hand per-value-field nullable
        # flags to merge functions that enforce NOT-NULL on every add().
        self.value_fields = list(read_type)

        self.trimmed_primary_key = self.table.trimmed_primary_keys
        self.read_fields = read_type
        if isinstance(self, MergeFileSplitRead):
            self.read_fields = self._create_key_value_fields(read_type)
        self._cached_nested_path_by_name = self._compute_nested_path_by_name()
        self.schema_id_2_fields = {}
        self.deletion_file_readers = {}
        # Apply filter only when all predicate columns are read by this scan,
        # AND remap predicate leaf indices into the row layout the reader sees.
        # Predicate leaves carry an `index` baked in by PredicateBuilder against
        # the *original* table schema; if `read_type` is narrower or reordered,
        # that index no longer matches the OffsetRow handed to
        # FilterRecordReader (which would otherwise raise IndexError).
        # We use `read_type` here, not `self.read_fields`: MergeFileSplitRead
        # augments `read_fields` with _KEY_*/_SEQ/_KIND prefixes, but
        # KeyValueUnwrapRecordReader returns kv.value whose arity equals
        # len(read_type) and whose coordinate space is read_type — that is
        # the space FilterRecordReader actually evaluates against.
        read_type_names = {f.name for f in read_type}
        if (
                self.predicate is not None
                and _get_all_fields(self.predicate).issubset(read_type_names)
        ):
            self.predicate_for_reader = rewrite_predicate_indices(
                self.predicate, read_type
            )
        else:
            self.predicate_for_reader = None

    def _compute_nested_path_by_name(self) -> Optional[Dict[str, List[str]]]:
        if not self.nested_name_paths:
            return None
        if not any(len(p) > 1 for p in self.nested_name_paths):
            return None
        out: Dict[str, List[str]] = {}
        for f, path in zip(self.read_fields[:self.value_arity],
                           self.nested_name_paths):
            out[f.name] = path
        return out

    def _nested_path_by_name(self) -> Optional[Dict[str, List[str]]]:
        return self._cached_nested_path_by_name

    def _resolve_schema(self, schema_id: int):
        """Resolve schema, short-circuiting current table schema id to avoid
        filesystem access (REST catalog would get 403).
        """
        if schema_id == self.table.table_schema.id:
            return self.table.table_schema
        return self.table.schema_manager.get_schema(schema_id)

    def _push_down_predicate(self) -> Optional[Predicate]:
        if self.predicate is None:
            return None
        elif self.table.is_primary_key_table:
            pk_predicate = trim_predicate_by_fields(self.predicate, self.table.primary_keys)
            if not pk_predicate:
                return None
            return pk_predicate
        else:
            return self.predicate

    @abstractmethod
    def create_reader(self) -> RecordReader:
        """Create a record reader for the given split."""

    # row_ranges: from IndexedSplit (ANN vector search), a list of discrete global row ID ranges.
    # shard_range: from SlicedSplit (parallel shard scan), a contiguous [start, end) row range within the file.
    def file_reader_supplier(self, file: DataFileMeta, for_merge_read: bool,
                             read_fields: List[str], row_tracking_enabled: bool,
                             row_ranges: Optional[List[Range]] = None,
                             shard_range: Optional[Tuple[int, int]] = None) -> RecordBatchReader:
        (read_file_fields, read_arrow_predicate) = self._get_fields_and_predicate(file.schema_id, read_fields)

        # Use external_path if available, otherwise use file_path
        file_path = file.external_path if file.external_path else file.file_path
        file_format = format_identifier(os.path.basename(file_path))

        batch_size = self.table.options.read_batch_size()

        # Convert global row_ranges (IndexedSplit) to local row_indices for native pushdown.
        row_indices = None
        if row_ranges is not None:
            effective_row_ranges = Range.and_(row_ranges, [file.row_id_range()])
            if len(effective_row_ranges) == 0:
                return EmptyRecordBatchReader()
            row_index_formats = (CoreOptions.FILE_FORMAT_BLOB,
                                 CoreOptions.FILE_FORMAT_VORTEX,
                                 CoreOptions.FILE_FORMAT_LANCE,
                                 CoreOptions.FILE_FORMAT_ROW)
            if file_format in row_index_formats:
                row_indices = []
                for r in effective_row_ranges:
                    start = r.from_ - file.first_row_id
                    end = r.to - file.first_row_id
                    row_indices.extend(range(start, end + 1))

        # Map nested paths into the order the format reader will see.
        nested_path_by_name = self._nested_path_by_name()
        has_nested = nested_path_by_name is not None

        # Field-id based per-file read (non-nested): select the file's OWN
        # physical fields by field id and read them under the file's original
        # names/types. A normalize step in DataFileBatchReader then aligns the
        # batch to the latest read schema by field id (not by name), so a
        # rename follows the id and a dropped-then-readded name cannot revive
        # stale data. Nested-projection reads stay on the legacy name path.
        file_read_fields = None if has_nested else self._file_read_fields(file)
        target_fields = None if has_nested else self._target_read_fields()
        if file_read_fields is not None:
            read_file_fields = [f.name for f in file_read_fields]
            name_to_field: Dict[str, DataField] = {
                f.name: f for f in file_read_fields}
        else:
            # Cover both the merge-internal aliases (``_KEY_id``) and the
            # bare user-facing PK name (``id``) the file actually stores.
            name_to_field = {f.name: f for f in self.read_fields}
            _, _trimmed_lookup_fields = self._get_trimmed_fields(
                self._get_read_data_fields(), self._get_all_data_fields()
            )
            for f in _trimmed_lookup_fields:
                name_to_field.setdefault(f.name, f)

        format_reader: RecordBatchReader
        if file_format == CoreOptions.FILE_FORMAT_AVRO:
            avro_nested_paths = (
                [nested_path_by_name[name] for name in read_file_fields]
                if has_nested else None
            )
            # Pass the alias-safe union so FormatAvroReader can resolve
            # the bare PK name (e.g. ``id``) requested by read_file_fields,
            # even when value projection drops it from self.read_fields.
            format_reader = FormatAvroReader(
                self.table.file_io, file_path, read_file_fields,
                list(name_to_field.values()),
                read_arrow_predicate, batch_size=batch_size,
                nested_name_paths=avro_nested_paths)
        elif file_format == CoreOptions.FILE_FORMAT_BLOB:
            if has_nested:
                raise NotImplementedError(
                    "Nested-field projection is not supported on BLOB files")
            blob_as_descriptor = CoreOptions.blob_as_descriptor(self.table.options)
            format_reader = FormatBlobReader(self.table.file_io, file_path, read_file_fields,
                                             self.read_fields, read_arrow_predicate, blob_as_descriptor,
                                             batch_size=batch_size,
                                             row_indices=row_indices)
        elif file_format == CoreOptions.FILE_FORMAT_LANCE:
            if has_nested:
                raise NotImplementedError(
                    "Nested-field projection is not supported on Lance files")
            ordered_read_fields = [name_to_field[n] for n in read_file_fields if n in name_to_field]
            format_reader = FormatLanceReader(self.table.file_io, file_path, ordered_read_fields,
                                              read_arrow_predicate, batch_size=batch_size,
                                              row_indices=row_indices,
                                              shard_range=shard_range)
        elif file_format == CoreOptions.FILE_FORMAT_VORTEX:
            if has_nested:
                raise NotImplementedError(
                    "Nested-field projection is not supported on Vortex files")
            ordered_read_fields = [name_to_field[n] for n in read_file_fields if n in name_to_field]
            predicate_fields = _get_all_fields(self.push_down_predicate) if self.push_down_predicate else set()
            format_reader = FormatVortexReader(self.table.file_io, file_path, ordered_read_fields,
                                               read_arrow_predicate, batch_size=batch_size,
                                               row_indices=row_indices,
                                               shard_range=shard_range,
                                               predicate_fields=predicate_fields)
        elif file_format == CoreOptions.FILE_FORMAT_MOSAIC:
            if has_nested:
                raise NotImplementedError(
                    "Nested-field projection is not supported on Mosaic files")
            ordered_read_fields = [name_to_field[n] for n in read_file_fields if n in name_to_field]
            format_reader = FormatMosaicReader(self.table.file_io, file_path, ordered_read_fields,
                                               read_arrow_predicate, batch_size=batch_size)
        elif file_format == CoreOptions.FILE_FORMAT_PARQUET or file_format == CoreOptions.FILE_FORMAT_ORC:
            ordered_read_fields = [name_to_field[n] for n in read_file_fields if n in name_to_field]
            ordered_nested_paths = (
                [nested_path_by_name[f.name] for f in ordered_read_fields]
                if has_nested else None
            )
            format_reader = FormatPyArrowReader(
                self.table.file_io, file_format, file_path,
                ordered_read_fields, read_arrow_predicate, batch_size=batch_size,
                options=self.table.options,
                nested_name_paths=ordered_nested_paths)
        elif file_format == CoreOptions.FILE_FORMAT_ROW:
            if has_nested:
                raise NotImplementedError(
                    "Nested-field projection is not supported on ROW files")
            file_schema = self._resolve_schema(file.schema_id)
            if file.write_cols:
                field_map = {f.name: f for f in file_schema.fields}
                row_full_fields = [field_map[n] for n in file.write_cols
                                   if n in field_map]
            elif self.table.is_primary_key_table:
                row_full_fields = self._create_key_value_fields(
                    file_schema.fields)
            else:
                row_full_fields = file_schema.fields
            format_reader = FormatRowReader(
                self.table.file_io, file_path, read_file_fields,
                row_full_fields,
                read_arrow_predicate, batch_size=batch_size,
                row_indices=row_indices)
        elif file_format in ('json', 'csv'):
            raise NotImplementedError(
                f"Reading '{file_format}' format is not yet supported in Python SDK. "
                f"Supported formats: parquet, orc, avro, lance, vortex, mosaic, blob, row.")
        else:
            raise ValueError(f"Unexpected file format: {file_format}")

        index_mapping = self.create_index_mapping()
        partition_info = self._create_partition_info()
        system_fields = SpecialFields.find_system_fields(self.read_fields)
        table_schema_fields = (
            SpecialFields.row_type_with_row_tracking(self.table.table_schema.fields)
            if row_tracking_enabled else self.table.table_schema.fields
        )

        # When native shard pushdown is used, the format reader only returns rows
        # starting from shard_range[0], so _ROW_ID must be offset accordingly.
        effective_first_row_id = file.first_row_id
        if (shard_range is not None and file.first_row_id is not None
                and file_format in (
                    CoreOptions.FILE_FORMAT_VORTEX, CoreOptions.FILE_FORMAT_LANCE)):
            effective_first_row_id = file.first_row_id + shard_range[0]

        if for_merge_read:
            reader = DataFileBatchReader(
                format_reader,
                index_mapping,
                partition_info,
                self.trimmed_primary_key,
                table_schema_fields,
                file.max_sequence_number,
                effective_first_row_id,
                row_tracking_enabled,
                system_fields,
                file_io=self.table.file_io,
                row_id_offsets=row_indices,
                file_data_fields=file_read_fields,
                target_data_fields=target_fields)
        else:
            reader = DataFileBatchReader(
                format_reader,
                index_mapping,
                partition_info,
                None,
                table_schema_fields,
                file.max_sequence_number,
                effective_first_row_id,
                row_tracking_enabled,
                system_fields,
                file_io=self.table.file_io,
                row_id_offsets=row_indices,
                file_data_fields=file_read_fields,
                target_data_fields=target_fields)

        # For non-Vortex formats, wrap with RowIdFilterRecordBatchReader
        if row_ranges is not None and row_indices is None:
            reader = RowIdFilterRecordBatchReader(reader, file.first_row_id, effective_row_ranges)

        # For formats without native shard support, wrap with ShardBatchReader
        if shard_range is not None and file_format not in (
                CoreOptions.FILE_FORMAT_VORTEX, CoreOptions.FILE_FORMAT_LANCE):
            reader = ShardBatchReader(reader, shard_range[0], shard_range[1])

        return reader

    def _get_fields_and_predicate(self, schema_id: int, read_fields):
        key = (schema_id, tuple(read_fields))
        if key not in self.schema_id_2_fields:
            nested_path_by_name = self._nested_path_by_name()
            schema = self._resolve_schema(schema_id)
            schema_fields = (
                SpecialFields.row_type_with_row_tracking(schema.fields)
                if self.row_tracking_enabled else schema.fields
            )
            schema_field_names = set(field.name for field in schema_fields)
            if self.table.is_primary_key_table:
                schema_field_names.add('_SEQUENCE_NUMBER')
                schema_field_names.add('_VALUE_KIND')

            def _is_reachable(name: str) -> bool:
                if name in schema_field_names:
                    return True
                if nested_path_by_name is not None:
                    path = nested_path_by_name.get(name)
                    if path:
                        return path[0] in schema_field_names
                return False

            read_file_fields = [
                read_field for read_field in read_fields
                if _is_reachable(read_field)
            ]
            read_predicate = trim_predicate_by_fields(self.push_down_predicate, read_file_fields)
            read_arrow_predicate = read_predicate.to_arrow() if read_predicate else None
            self.schema_id_2_fields[key] = (read_file_fields, read_arrow_predicate)
        return self.schema_id_2_fields[key]

    @abstractmethod
    def _all_data_fields_from(self, fields: List[DataField]) -> List[DataField]:
        """Apply this split-read's data-field shaping (row-tracking / kv
        wrapping) to the given base ``fields``. Called both for the latest
        table schema and for an older file schema."""

    def _get_all_data_fields(self):
        return self._all_data_fields_from(self.table.fields)

    def _get_read_data_fields(self):
        return self._read_data_fields_from(self._get_all_data_fields())

    def _read_data_fields_from(self, all_data_fields):
        read_field_ids = {field.id for field in self.read_fields}
        return [f for f in all_data_fields if f.id in read_field_ids]

    def _final_data_fields_from(self, all_data_fields: List[DataField]) -> List[DataField]:
        """The per-position target fields a batch must end up as: trimmed for
        kv (``_KEY_*``) duplicates and stripped of partition columns. The
        DataField analogue of ``_get_final_read_data_fields()``. Called with
        the latest-schema fields it yields the read target (latest names +
        types); called with a file's fields it yields what to physically read
        from that file (the file's own names + types)."""
        _, trimmed = self._get_trimmed_fields(
            self._read_data_fields_from(all_data_fields), all_data_fields)
        partition_keys = self.table.partition_keys
        if not partition_keys:
            return list(trimmed)
        return [f for f in trimmed if f.name not in partition_keys]

    def _target_read_fields(self) -> Optional[List[DataField]]:
        """Latest-schema target fields (names + types) that a normalized batch
        must align to, in order. None for nested-projection reads (kept on the
        legacy name-based path)."""
        if self._nested_path_by_name() is not None:
            return None
        return self._final_data_fields_from(self._get_all_data_fields())

    def _file_read_fields(self, file: DataFileMeta) -> Optional[List[DataField]]:
        """The fields to physically read from ``file``, in the file's own
        names/types, selected by field id against the read set. None for
        nested-projection reads."""
        if self._nested_path_by_name() is not None:
            return None
        file_schema = self._resolve_schema(file.schema_id)
        if file_schema is None:
            return None
        return self._final_data_fields_from(
            self._all_data_fields_from(file_schema.fields))

    def _create_key_value_fields(self, value_field: List[DataField]):
        all_fields: List[DataField] = self.table.fields
        all_data_fields = []

        for field in all_fields:
            if field.name in self.trimmed_primary_key:
                key_field_name = f"{KEY_PREFIX}{field.name}"
                key_field_id = field.id + KEY_FIELD_ID_START
                key_field = DataField(key_field_id, key_field_name, field.type)
                all_data_fields.append(key_field)

        all_data_fields.append(SpecialFields.SEQUENCE_NUMBER)
        all_data_fields.append(SpecialFields.VALUE_KIND)

        for field in value_field:
            all_data_fields.append(field)

        return all_data_fields

    def create_index_mapping(self):
        if self._nested_path_by_name() is not None:
            return None
        base_index_mapping = self._create_base_index_mapping(self.read_fields, self._get_read_data_fields())
        trimmed_key_mapping, _ = self._get_trimmed_fields(self._get_read_data_fields(), self._get_all_data_fields())
        if base_index_mapping is None:
            mapping = trimmed_key_mapping
        elif trimmed_key_mapping is None:
            mapping = base_index_mapping
        else:
            combined = [0] * len(base_index_mapping)
            for i in range(len(base_index_mapping)):
                if base_index_mapping[i] < 0:
                    combined[i] = base_index_mapping[i]
                else:
                    combined[i] = trimmed_key_mapping[base_index_mapping[i]]
            mapping = combined

        if mapping is not None:
            for i in range(len(mapping)):
                if mapping[i] != i:
                    return mapping

        return None

    def _create_base_index_mapping(self, table_fields: List[DataField], data_fields: List[DataField]):
        index_mapping = [0] * len(table_fields)
        field_id_to_index = {field.id: i for i, field in enumerate(data_fields)}

        for i, table_field in enumerate(table_fields):
            field_id = table_field.id
            data_field_index = field_id_to_index.get(field_id)
            if data_field_index is not None:
                index_mapping[i] = data_field_index
            else:
                index_mapping[i] = NULL_FIELD_INDEX

        for i in range(len(index_mapping)):
            if index_mapping[i] != i:
                return index_mapping

        return None

    def _get_final_read_data_fields(self) -> List[str]:
        if self._nested_path_by_name() is not None:
            return self._remove_partition_fields(list(self.read_fields))
        _, trimmed_fields = self._get_trimmed_fields(
            self._get_read_data_fields(), self._get_all_data_fields()
        )
        return self._remove_partition_fields(trimmed_fields)

    def _remove_partition_fields(self, fields: List[DataField]) -> List[str]:
        partition_keys = self.table.partition_keys
        if not partition_keys:
            return [field.name for field in fields]

        fields_without_partition = []
        for field in fields:
            if field.name not in partition_keys:
                fields_without_partition.append(field)

        return [field.name for field in fields_without_partition]

    def _get_trimmed_fields(self, read_data_fields: List[DataField],
                            all_data_fields: List[DataField]) -> Tuple[List[int], List[DataField]]:
        trimmed_mapping = [0] * len(read_data_fields)
        trimmed_fields = []

        field_id_to_field = {field.id: field for field in all_data_fields}
        position_map = {}
        for i, field in enumerate(read_data_fields):
            is_key_field = field.name.startswith(KEY_PREFIX)
            if is_key_field:
                original_id = field.id - KEY_FIELD_ID_START
            else:
                original_id = field.id
            original_field = field_id_to_field.get(original_id)

            if original_id in position_map:
                trimmed_mapping[i] = position_map[original_id]
            else:
                position = len(trimmed_fields)
                position_map[original_id] = position
                trimmed_mapping[i] = position
                if is_key_field:
                    trimmed_fields.append(original_field)
                else:
                    trimmed_fields.append(field)

        return trimmed_mapping, trimmed_fields

    def _create_partition_info(self):
        if not self.table.partition_keys:
            return None
        partition_mapping = self._construct_partition_mapping()
        if not partition_mapping:
            return None
        return PartitionInfo(partition_mapping, self.split.partition)

    def _construct_partition_mapping(self) -> List[int]:
        if self._nested_path_by_name() is not None:
            partition_names = self.table.partition_keys
            mapping = [0] * (len(self.read_fields) + 1)
            p_count = 0
            for i, field in enumerate(self.read_fields):
                if field.name in partition_names:
                    partition_index = partition_names.index(field.name)
                    mapping[i] = -(partition_index + 1)
                    p_count += 1
                else:
                    mapping[i] = (i - p_count) + 1
            return mapping
        _, trimmed_fields = self._get_trimmed_fields(
            self._get_read_data_fields(), self._get_all_data_fields()
        )
        partition_names = self.table.partition_keys

        mapping = [0] * (len(trimmed_fields) + 1)
        p_count = 0

        for i, field in enumerate(trimmed_fields):
            if field.name in partition_names:
                partition_index = partition_names.index(field.name)
                mapping[i] = -(partition_index + 1)
                p_count += 1
            else:
                mapping[i] = (i - p_count) + 1

        return mapping

    def _genarate_deletion_file_readers(self):
        self.deletion_file_readers = {}
        if self.split.data_deletion_files:
            for data_file, deletion_file in zip(self.split.files, self.split.data_deletion_files):
                if deletion_file is not None:
                    # Create a callable method to read the deletion vector
                    self.deletion_file_readers[data_file.file_name] = lambda df=deletion_file: DeletionVector.read(
                        self.table.file_io, df)


class RawFileSplitRead(SplitRead):
    def raw_reader_supplier(self, file: DataFileMeta, dv_factory: Optional[Callable] = None) -> Optional[RecordReader]:
        read_fields = self._get_final_read_data_fields()
        # Check if this is a SlicedSplit to get shard_file_idx_map
        shard_file_idx_map = (
            self.split.shard_file_idx_map() if isinstance(self.split, SlicedSplit) else {}
        )
        if file.file_name in shard_file_idx_map:
            (start_pos, end_pos) = shard_file_idx_map[file.file_name]
            if (start_pos, end_pos) == (-1, -1):
                return None
            file_batch_reader = self.file_reader_supplier(
                file=file,
                for_merge_read=False,
                read_fields=read_fields,
                row_tracking_enabled=True,
                shard_range=(start_pos, end_pos))
        else:
            file_batch_reader = self.file_reader_supplier(
                file=file,
                for_merge_read=False,
                read_fields=read_fields,
                row_tracking_enabled=True)
        dv = dv_factory() if dv_factory else None
        if dv:
            return ApplyDeletionVectorReader(RowPositionReader(file_batch_reader), dv)
        else:
            return file_batch_reader

    def create_reader(self) -> RecordReader:
        self._genarate_deletion_file_readers()
        data_readers = []
        for file in self.split.files:
            supplier = partial(
                self.raw_reader_supplier,
                file=file,
                dv_factory=self.deletion_file_readers.get(file.file_name, None)
            )
            data_readers.append(supplier)

        if not data_readers:
            return EmptyFileRecordReader()

        concat_reader = ConcatBatchReader(
            data_readers, file_io=self.table.file_io,
            blob_field_indices=_blob_field_indices(self.read_fields),
            vector_field_indices=_vector_field_indices(self.read_fields))
        # if the table is appendonly table, we don't need extra filter, all predicates has pushed down
        if self.table.is_primary_key_table and self.predicate_for_reader:
            reader = FilterRecordReader(concat_reader, self.predicate_for_reader)
            if self.limit is not None:
                reader = LimitedRecordReader(reader, self.limit)
        else:
            reader = concat_reader
            if self.limit is not None:
                reader = LimitedRecordBatchReader(reader, self.limit)
        return reader

    def _all_data_fields_from(self, fields):
        if self.row_tracking_enabled:
            return SpecialFields.row_type_with_row_tracking(fields)
        return fields


class MergeFileSplitRead(SplitRead):
    def __init__(
            self,
            table,
            predicate: Optional[Predicate],
            read_type: List[DataField],
            split: Split,
            row_tracking_enabled: bool,
            outer_extract_name_paths: Optional[List[List[str]]] = None,
            limit: Optional[int] = None):
        # Merge functions need full ROW sub-structures, so nested paths
        # are not pushed down here; sub-path extraction happens above
        # the merge via OuterProjectionRecordReader.
        super().__init__(
            table=table,
            predicate=predicate,
            read_type=read_type,
            split=split,
            row_tracking_enabled=row_tracking_enabled,
            nested_name_paths=None,
            limit=limit,
        )
        self.outer_extract_name_paths = outer_extract_name_paths
        # Built once per split-read (value_fields and options are constant
        # for the object's life), not per section. ``None`` when
        # ``sequence.field`` is unset, in which case the heap falls back to
        # the file-level sequence number.
        self.seq_comparator = builtin_seq_comparator(
            self.value_fields,
            self.table.options.sequence_field(),
            self.table.options.sequence_field_sort_order_is_ascending(),
        )

    def kv_reader_supplier(self, file: DataFileMeta, dv_factory: Optional[Callable] = None) -> RecordReader:
        file_batch_reader = self.file_reader_supplier(file, True, self._get_final_read_data_fields(), False)
        dv = dv_factory() if dv_factory else None
        if dv:
            return ApplyDeletionVectorReader(
                KeyValueWrapReader(RowPositionReader(file_batch_reader),
                                   len(self.trimmed_primary_key), self.value_arity), dv)
        else:
            return KeyValueWrapReader(file_batch_reader, len(self.trimmed_primary_key), self.value_arity)

    def section_reader_supplier(self, section: List[SortedRun]) -> RecordReader:
        readers = []
        for sorter_run in section:
            data_readers = []
            for file in sorter_run.files:
                supplier = partial(self.kv_reader_supplier, file, self.deletion_file_readers.get(file.file_name, None))
                data_readers.append(supplier)
            readers.append(ConcatRecordReader(data_readers))
        merge_function = self._build_merge_function()
        return SortMergeReaderWithMinHeap(
            readers, self.table.table_schema, merge_function=merge_function,
            seq_comparator=self.seq_comparator)

    def _build_merge_function(self):
        """Pick the MergeFunction for the table's ``merge-engine`` option.

        Delegates to the shared dispatch in
        ``pypaimon.common.merge_engine_dispatch`` so the read path and
        the in-memory merge buffer on the write path cannot drift.
        ``AGGREGATE`` is special-cased here because building the per-
        field aggregators needs the full ``DataField`` objects, the
        full primary-key list and the parsed ``CoreOptions`` -- which
        sit outside the dispatch's raw-options contract. The writer-
        side merge buffer falls back to dedupe for aggregation anyway
        (see :meth:`FileStoreWrite._build_pk_merge_function`), so the
        two sides only need to share the simple engines.
        """
        engine = self.table.options.merge_engine()
        if engine == MergeEngine.AGGREGATE:
            # Use the full primary-key list, not ``trimmed_primary_key``:
            # ``value_fields`` still carries partition columns, so any PK
            # column that is also a partition column must be recognised
            # as PK here. Otherwise a table with
            # ``fields.default-aggregate-function`` would apply the
            # default aggregator to that partition-PK column.
            field_aggregators = build_field_aggregators(
                self.value_fields,
                self.table.primary_keys,
                self.table.options,
            )
            return AggregateMergeFunction(
                key_arity=len(self.trimmed_primary_key),
                value_arity=self.value_arity,
                field_aggregators=field_aggregators,
            )
        return build_merge_function(
            engine=engine,
            raw_options=self.table.options.options.to_map(),
            key_arity=len(self.trimmed_primary_key),
            value_arity=self.value_arity,
            value_field_nullables=[f.type.nullable for f in self.value_fields],
            value_field_names=[f.name for f in self.value_fields],
        )

    def create_reader(self) -> RecordReader:
        # Create a dict mapping data file name to deletion file reader method
        self._genarate_deletion_file_readers()
        section_readers = []
        sections = IntervalPartition(self.split.files).partition()
        for section in sections:
            supplier = partial(self.section_reader_supplier, section)
            section_readers.append(supplier)
        concat_reader = ConcatRecordReader(section_readers)
        kv_unwrap_reader = KeyValueUnwrapRecordReader(DropDeleteRecordReader(concat_reader))
        if self.predicate_for_reader:
            reader = FilterRecordReader(kv_unwrap_reader, self.predicate_for_reader)
        else:
            reader = kv_unwrap_reader
        if self.outer_extract_name_paths:
            from pypaimon.read.reader.outer_projection_record_reader import \
                OuterProjectionRecordReader
            inner_value_fields = self.read_fields[-self.value_arity:]
            reader = OuterProjectionRecordReader(
                reader, [f.name for f in inner_value_fields],
                self.outer_extract_name_paths,
                file_io=self.table.file_io,
                blob_field_indices=_blob_field_indices(inner_value_fields),
                vector_field_indices=_vector_field_indices(inner_value_fields))
        if self.limit is not None:
            reader = LimitedRecordReader(reader, self.limit)
        return reader

    def _all_data_fields_from(self, fields):
        return self._create_key_value_fields(fields)


class DataEvolutionSplitRead(SplitRead):

    def __init__(
            self,
            table,
            predicate: Optional[Predicate],
            read_type: List[DataField],
            split: Split,
            row_tracking_enabled: bool,
            nested_name_paths: Optional[List[List[str]]] = None,
            limit: Optional[int] = None):
        self.row_ranges = None
        actual_split = split
        if isinstance(split, IndexedSplit):
            self.row_ranges = split.row_ranges()
            actual_split = split.data_split()
        super().__init__(
            table, predicate, read_type, actual_split, row_tracking_enabled,
            nested_name_paths=nested_name_paths,
            limit=limit,
        )

    def _push_down_predicate(self) -> Optional[Predicate]:
        # Data evolution: files may have different schemas, so we don't push predicate
        # to file readers; filtering is done in FilterRecordBatchReader after merge.
        return None

    def create_reader(self) -> RecordReader:
        reader = self._create_raw_reader()

        if ((CoreOptions.blob_view_fields(self.table.options) and CoreOptions.blob_view_resolve_enabled(
                self.table.options))
                or (not CoreOptions.blob_as_descriptor(self.table.options)
                    and CoreOptions.blob_descriptor_fields(self.table.options))):
            reader = BlobInlineConvertReader(
                reader, self.table,
                prescan_reader_factory=lambda names: self._create_prescan_reader(names))

        return reader

    def _create_raw_reader(self) -> RecordReader:
        """Core read logic: split_by_row_id -> suppliers -> ConcatBatchReader -> filter."""
        files = self.split.files
        suppliers = []

        # Split files by row ID
        split_by_row_id = self._split_by_row_id(files)

        for need_merge_files in split_by_row_id:
            if len(need_merge_files) == 1 or not self.read_fields:
                # No need to merge fields, just create a single file reader
                suppliers.append(
                    lambda f=need_merge_files[0]: self._create_file_reader(f, self._get_final_read_data_fields())
                )
            else:
                suppliers.append(
                    lambda files=need_merge_files: self._create_union_reader(files)
                )

        merge_reader = ConcatBatchReader(
            suppliers, file_io=self.table.file_io,
            blob_field_indices=_blob_field_indices(self.read_fields),
            vector_field_indices=_vector_field_indices(self.read_fields))
        if self.predicate_for_reader is not None:
            reader = FilterRecordBatchReader(
                merge_reader,
                self.predicate_for_reader,
                field_names=[f.name for f in self.read_fields],
                schema_fields=self.read_fields,
            )
        else:
            reader = merge_reader

        if self.limit is not None:
            reader = LimitedRecordBatchReader(reader, self.limit)

        return reader

    def _create_prescan_reader(self, field_names):
        """Create a prescan reader by constructing a new DataEvolutionSplitRead
        instance that only projects the specified field names.
        
        Align with Java's configureBlobViewPrescanRead: pass limit to prescan reader
        to avoid scanning entire split when there's a LIMIT clause.
        """
        from pypaimon.read.reader.iface.record_batch_reader import EmptyRecordBatchReader

        prescan_fields = [f for f in self.read_fields if f.name in field_names]
        if not prescan_fields:
            return EmptyRecordBatchReader()

        # When there's a normal field predicate, don't push down limit to prescan reader
        # because the outer reader will apply predicate+limit filtering,
        # while prescan reader would only apply limit without normal field predicate
        # TODO support limit+predicate push down
        prescan_read = DataEvolutionSplitRead(
            table=self.table,
            predicate=self.predicate,
            read_type=prescan_fields,
            split=self.split,
            row_tracking_enabled=False,
            limit=None if self.predicate else self.limit,
        )
        prescan_read.row_ranges = self.row_ranges
        return prescan_read._create_raw_reader()

    def _split_by_row_id(self, files: List[DataFileMeta]) -> List[List[DataFileMeta]]:
        """Split files by firstRowId for data evolution."""

        # Sort files by firstRowId and then by maxSequenceNumber
        def sort_key(file: DataFileMeta) -> tuple:
            first_row_id = file.first_row_id if file.first_row_id is not None else float('-inf')
            is_special = 1 if (DataFileMeta.is_blob_file(file.file_name)
                               or DataFileMeta.is_vector_file(file.file_name)) else 0
            max_seq = file.max_sequence_number
            return (first_row_id, is_special, -max_seq)

        sorted_files = sorted(files, key=sort_key)

        # Split files by firstRowId
        split_by_row_id = []
        last_row_id = -1
        check_row_id_start = 0
        current_split = []

        for file in sorted_files:
            first_row_id = file.first_row_id
            if first_row_id is None:
                split_by_row_id.append([file])
                continue

            if (not DataFileMeta.is_blob_file(file.file_name)
                    and not DataFileMeta.is_vector_file(file.file_name)
                    and first_row_id != last_row_id):
                if current_split:
                    split_by_row_id.append(current_split)
                if first_row_id < check_row_id_start:
                    raise ValueError(
                        f"There are overlapping files in the split: {files}, "
                        f"the wrong file is: {file}"
                    )
                current_split = []
                last_row_id = first_row_id
                check_row_id_start = first_row_id + file.row_count
            current_split.append(file)

        if current_split:
            split_by_row_id.append(current_split)

        return split_by_row_id

    def _create_union_reader(self, need_merge_files: List[DataFileMeta]) -> RecordReader:
        """Create a DataEvolutionFileReader for merging multiple files."""
        # Split field bunches
        fields_files = self._split_field_bunches(need_merge_files)

        # Validate row counts and first row IDs (skip when row ranges are pushed down)
        row_count = fields_files[0].row_count()
        first_row_id = fields_files[0].files()[0].first_row_id

        if self.row_ranges is None:
            for bunch in fields_files:
                if bunch.row_count() != row_count:
                    raise ValueError(
                        "All files in a field merge split should have the same row count.")
                if bunch.files()[0].first_row_id != first_row_id:
                    raise ValueError(
                        "All files in a field merge split should have the same "
                        "first row id and could not be null."
                    )

        # Create the union reader
        all_read_fields = self.read_fields
        file_record_readers = [None] * len(fields_files)
        read_field_index = [field.id for field in all_read_fields]

        # Initialize offsets
        row_offsets = [-1] * len(all_read_fields)
        field_offsets = [-1] * len(all_read_fields)
        schema_pos = {f.id: p for p, f in enumerate(self.table.fields)}

        for i, bunch in enumerate(fields_files):
            first_file = bunch.files()[0]

            # Get field IDs for this bunch
            if DataFileMeta.is_blob_file(first_file.file_name):
                field_ids = [self._get_field_id_from_write_cols(first_file)]
            elif DataFileMeta.is_vector_file(first_file.file_name):
                field_ids = self._get_field_ids_from_write_cols(first_file.write_cols)
            elif first_file.write_cols:
                field_ids = self._get_field_ids_from_write_cols(first_file.write_cols)
            else:
                # For regular files without write_cols, derive field IDs from
                # the file's schema version, not the current table schema.
                # The file only contains columns from when it was written.
                file_schema = self.table.schema_manager.get_schema(first_file.schema_id)
                field_ids = [field.id for field in file_schema.fields]
                field_ids.append(SpecialFields.ROW_ID.id)
                field_ids.append(SpecialFields.SEQUENCE_NUMBER.id)

            read_fields = []
            for j, read_field_id in enumerate(read_field_index):
                for field_id in field_ids:
                    if read_field_id == field_id:
                        if row_offsets[j] == -1:
                            row_offsets[j] = i
                            read_fields.append(all_read_fields[j])
                        break

            if not read_fields:
                file_record_readers[i] = None
            else:
                read_fields.sort(key=lambda f: schema_pos.get(f.id, float('inf')))
                id_to_pos = {f.id: p for p, f in enumerate(read_fields)}
                for j in range(len(read_field_index)):
                    if row_offsets[j] == i:
                        field_offsets[j] = id_to_pos[read_field_index[j]]

                read_field_names = self._remove_partition_fields(read_fields)
                table_fields = self.read_fields
                self.read_fields = read_fields  # create reader based on read_fields
                batch_size = self.table.options.read_batch_size()
                # Create reader for this bunch
                if len(bunch.files()) == 1:
                    suppliers = [lambda r=self._create_file_reader(
                        bunch.files()[0], read_field_names
                    ): r]
                    file_record_readers[i] = MergeAllBatchReader(suppliers, batch_size=batch_size)
                elif DataFileMeta.is_blob_file(first_file.file_name):
                    file_reader_suppliers = [
                        (
                            file,
                            partial(
                                self._create_raw_blob_file_reader,
                                file=file,
                                read_fields=read_field_names,
                            ),
                        )
                        for file in bunch.files()
                    ]
                    file_record_readers[i] = BlobFallbackBatchReader(
                        file_reader_suppliers,
                        read_fields[0].name,
                        PyarrowFieldParser.from_paimon_schema(
                            [read_fields[0]]
                        ).field(0).type,
                        self.row_ranges,
                        CoreOptions.blob_as_descriptor(self.table.options),
                    )
                else:
                    # Create concatenated reader for multiple files
                    suppliers = [
                        partial(self._create_file_reader, file=file,
                                read_fields=read_field_names) for file in bunch.files()
                    ]
                    file_record_readers[i] = MergeAllBatchReader(suppliers, batch_size=batch_size)
                self.read_fields = table_fields

        # Validate that all required fields are found
        for i, field in enumerate(all_read_fields):
            if row_offsets[i] == -1:
                if not field.type.nullable:
                    raise ValueError(f"Field {field} is not null but can't find any file contains it.")

        output_schema = PyarrowFieldParser.from_paimon_schema(all_read_fields)
        return DataEvolutionMergeReader(row_offsets, field_offsets, file_record_readers, schema=output_schema)

    def _create_file_reader(self, file: DataFileMeta, read_fields: [str]) -> Optional[RecordReader]:
        """Create a file reader for a single file."""
        return self.file_reader_supplier(
            file=file,
            for_merge_read=False,
            read_fields=read_fields,
            row_tracking_enabled=True,
            row_ranges=self.row_ranges)

    def _create_raw_blob_file_reader(
            self, file: DataFileMeta, read_fields: [str]) -> Optional[FormatBlobReader]:
        row_indices = None
        if self.row_ranges is not None:
            row_indices = [
                row_id - file.first_row_id
                for row_range in Range.and_([file.row_id_range()], self.row_ranges)
                for row_id in range(row_range.from_, row_range.to + 1)
            ]
            if not row_indices:
                return None

        file_path = file.external_path if file.external_path else file.file_path
        return FormatBlobReader(
            self.table.file_io,
            file_path,
            read_fields,
            self.read_fields,
            None,
            CoreOptions.blob_as_descriptor(self.table.options),
            batch_size=self.table.options.read_batch_size(),
            row_indices=row_indices,
        )

    def _split_field_bunches(self, need_merge_files: List[DataFileMeta]) -> List[FieldBunch]:
        """Split files into field bunches."""

        fields_files = []
        blob_bunch_map = {}
        vector_bunch_map = {}
        row_count = -1
        row_id_push_down = self.row_ranges is not None

        for file in need_merge_files:
            if DataFileMeta.is_blob_file(file.file_name):
                field_id = self._get_field_id_from_write_cols(file)
                if field_id not in blob_bunch_map:
                    blob_bunch_map[field_id] = BlobBunch(row_count, row_id_push_down)
                blob_bunch_map[field_id].add(file)
            elif DataFileMeta.is_vector_file(file.file_name):
                field_id = self._get_field_id_from_write_cols(file)
                if field_id not in vector_bunch_map:
                    vector_bunch_map[field_id] = VectorBunch(row_count, row_id_push_down)
                vector_bunch_map[field_id].add(file)
            else:
                fields_files.append(DataBunch(file))
                row_count = file.row_count

        fields_files.extend(blob_bunch_map.values())
        fields_files.extend(vector_bunch_map.values())
        return fields_files

    def _get_field_id_from_write_cols(self, file: DataFileMeta) -> int:
        """Get field ID from write columns for blob/vector files."""
        if not file.write_cols or len(file.write_cols) == 0:
            raise ValueError("Blob/vector file must have write columns")

        # Find the field by name in the table schema
        field_name = file.write_cols[0]
        for field in self.table.fields:
            if field.name == field_name:
                return field.id
        raise ValueError(f"Field {field_name} not found in table schema")

    def _get_field_ids_from_write_cols(self, write_cols: List[str]) -> List[int]:
        field_ids = []
        for field_name in write_cols:
            for field in self.table.fields:
                if field.name == field_name:
                    field_ids.append(field.id)
        field_ids.append(SpecialFields.ROW_ID.id)
        field_ids.append(SpecialFields.SEQUENCE_NUMBER.id)
        return field_ids

    def _all_data_fields_from(self, fields):
        return SpecialFields.row_type_with_row_tracking(fields)