# 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 pyarrow as pa
import pyarrow.compute as pc
import uuid
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Tuple

from pypaimon.execution.python.write.format_file_writer import (
    write_blob_file, write_row_file, write_vortex_file
)
from pypaimon.common.options.core_options import CoreOptions, ChangelogProducer
from pypaimon.common.external_path_provider import ExternalPathProvider
from pypaimon.data.timestamp import Timestamp
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.manifest.schema.simple_stats import SimpleStats
from pypaimon.schema.data_types import PyarrowFieldParser
from pypaimon.table.bucket_mode import BucketMode
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.table.row.row_kind import RowKind
from pypaimon.table.special_fields import SpecialFields
from pypaimon.utils.file_store_path_factory import canonical_data_file_path
from pypaimon.execution.python.write.map_shared_shredding_writer import MapSharedShreddingWriter
from pypaimon.execution.python.write.writer.mosaic_writer_options import create_mosaic_writer_options
from pypaimon.execution.python.write.writer.parquet_writer_options import create_parquet_writer_options
from pypaimon.execution.python.write.writer.write_buffer import WriteBuffer


class DataWriter(ABC):
    """Base class for data writers that handle PyArrow tables directly."""

    ROW_SIDECAR_SUFFIX = ".row"

    def __init__(self, table, partition: Tuple, bucket: int, max_seq_number: int, options: CoreOptions = None,
                 write_cols: Optional[List[str]] = None,
                 changelog_producer: ChangelogProducer = ChangelogProducer.NONE):
        from pypaimon.table.file_store_table import FileStoreTable

        self.table: FileStoreTable = table
        self.partition = partition
        self.bucket = bucket

        self.file_io = self.table.file_io
        self.trimmed_primary_keys_fields = self.table.trimmed_primary_keys_fields
        self.trimmed_primary_keys = self.table.trimmed_primary_keys

        self.options = options
        self.target_file_size = self.options.target_file_size(self.table.is_primary_key_table)
        # Roll a file when it reaches target_file_row_num rows or target_file_size,
        # whichever comes first. Defaults to the max long (disabled), so plain
        # size-based rolling is unchanged unless the option is set.
        self.target_file_row_num = self.options.target_file_row_num()
        # POSTPONE_BUCKET uses AVRO format, otherwise default to PARQUET
        default_format = (
            CoreOptions.FILE_FORMAT_AVRO
            if self.bucket == BucketMode.POSTPONE_BUCKET.value
            else CoreOptions.FILE_FORMAT_PARQUET
        )
        self.file_format = self.options.file_format(default_format)
        self.compression = self.options.file_compression()
        self.zstd_level = self.options.file_compression_zstd_level()
        self.mosaic_writer_options = (
            create_mosaic_writer_options(self.options)
            if self.file_format == CoreOptions.FILE_FORMAT_MOSAIC
            else None
        )
        self.sequence_generator = SequenceGenerator(max_seq_number)

        self._buffer = WriteBuffer(self._merge_data)
        self.committed_files: List[DataFileMeta] = []
        self.committed_changelog_files: List[DataFileMeta] = []
        self.changelog_producer = changelog_producer
        self.changelog_file_format = (
            self.options.changelog_file_format()
            or self.file_format
        )
        self.parquet_writer_options = (
            create_parquet_writer_options(self.options)
            if self.file_format == CoreOptions.FILE_FORMAT_PARQUET
            or (changelog_producer == ChangelogProducer.INPUT
                and self.changelog_file_format == CoreOptions.FILE_FORMAT_PARQUET)
            else {}
        )
        self.write_cols = write_cols
        self.blob_as_descriptor = self.options.blob_as_descriptor()
        self.blob_descriptor_fields = CoreOptions.blob_descriptor_fields(self.options)
        self.blob_view_fields = CoreOptions.blob_view_fields(self.options)
        self.blob_inline_fields = self.blob_descriptor_fields.union(self.blob_view_fields)

        self.path_factory = self.table.path_factory()
        self.external_path_provider: Optional[ExternalPathProvider] = self.path_factory.create_external_path_provider(
            self.partition, self.bucket, canonical_partition=True
        )
        # Variant shredding (static mode) — col_name → (obj_fields, target_arrow_type)
        self._variant_shredding: Dict[str, Tuple] = {}
        if self.file_format == CoreOptions.FILE_FORMAT_PARQUET:
            shredding_json = self.options.variant_shredding_schema()
            if shredding_json:
                from pypaimon.data.variant_shredding import (
                    parse_shredding_schema_option, shredding_schema_to_arrow_type)
                col_schemas = parse_shredding_schema_option(shredding_json)
                for col_name, obj_fields in col_schemas.items():
                    target_type = shredding_schema_to_arrow_type(obj_fields)
                    self._variant_shredding[col_name] = (obj_fields, target_type)

        # Paimon field id map, used by _apply_variant_shredding; built once since
        # the table schema is fixed for the lifetime of this writer.
        self._paimon_field_id: Dict[str, int] = {pf.name: pf.id for pf in self.table.fields}
        self._map_shared_shredding = MapSharedShreddingWriter(
            self.table.fields,
            self.options,
            self.file_format,
            self.changelog_file_format
            if self.changelog_producer == ChangelogProducer.INPUT else None,
        )

    @property
    def pending_row_count(self) -> int:
        """Rows held for the file being written, not yet in a finished file.

        The composite writers override this: they keep their normal rows in
        ``_normal_buffer`` and hand the rest to their sub-writers, so the base
        ``_buffer`` stays empty for them.
        """
        return self._buffer.num_rows

    def write(self, data: pa.RecordBatch):
        try:
            self._validate_inline_stored_fields_input(data)
            processed_data = self._process_data(data)
            self._buffer.append(processed_data)
            self._check_and_roll_if_needed()
        except Exception as e:
            import logging
            logger = logging.getLogger(__name__)
            logger.warning("Exception occurs when writing data. Cleaning up.", exc_info=e)
            self.abort()
            raise e

    def _validate_inline_stored_fields_input(self, data: pa.RecordBatch):
        if not getattr(self, 'blob_inline_fields', None):
            return

        from pypaimon.table.row.blob import BlobDescriptor, BlobViewStruct

        for field_name in getattr(self, 'blob_descriptor_fields', ()):
            if field_name not in data.schema.names:
                continue
            values = data.column(data.schema.get_field_index(field_name)).to_pylist()
            for value in values:
                if value is None:
                    continue
                if hasattr(value, 'as_py'):
                    value = value.as_py()
                if isinstance(value, str):
                    value = value.encode('utf-8')
                if not isinstance(value, (bytes, bytearray)):
                    raise ValueError(
                        "blob-descriptor-field requires blob field value to be a serialized "
                        "BlobDescriptor."
                    )
                descriptor_bytes = bytes(value)
                # Like Java, schema-declared descriptors use the prefix parser.
                # Exact wire length is only needed when detecting descriptors
                # among arbitrary payload bytes.
                try:
                    BlobDescriptor.deserialize(descriptor_bytes)
                except Exception as e:
                    raise ValueError(
                        "blob-descriptor-field requires blob field value to be a serialized "
                        "BlobDescriptor."
                    ) from e

        for field_name in getattr(self, 'blob_view_fields', ()):
            if field_name not in data.schema.names:
                continue
            values = data.column(data.schema.get_field_index(field_name)).to_pylist()
            for value in values:
                if value is None:
                    continue
                if hasattr(value, 'as_py'):
                    value = value.as_py()
                if isinstance(value, str):
                    value = value.encode('utf-8')
                if not isinstance(value, (bytes, bytearray)):
                    raise ValueError(
                        "blob-view-field requires blob field value to be a serialized "
                        "BlobViewStruct."
                    )
                try:
                    view_bytes = bytes(value)
                    view_struct = BlobViewStruct.deserialize(view_bytes)
                    if view_struct.serialize() != view_bytes:
                        raise ValueError("BlobViewStruct payload contains trailing bytes.")
                except Exception as e:
                    raise ValueError(
                        "blob-view-field requires blob field value to be a serialized "
                        "BlobViewStruct."
                    ) from e

    def prepare_commit(self) -> List[DataFileMeta]:
        if self._buffer.num_rows > 0:
            # Clear only once the write lands: a caller that retries
            # prepare_commit after a failed write has to still find its rows.
            self._write_data_to_file(self._buffer.materialize())
            self._buffer.reset()

        return self.committed_files.copy()

    def _release_prepared_files(self) -> List[DataFileMeta]:
        """Clear metadata accepted by the parent and return files it must delete on abort.

        Call only after all sidecars are prepared and validated. Until then, this
        writer retains its metadata for retry and its responsibility for cleanup.
        """
        owned_files = (self.committed_files + self.committed_changelog_files
                       if self.delete_file_upon_abort() else [])
        self.committed_files.clear()
        self.committed_changelog_files.clear()
        return owned_files

    def delete_file_upon_abort(self) -> bool:
        return True

    def prepare_changelog_commit(self) -> List[DataFileMeta]:
        return self.committed_changelog_files.copy()

    def close(self):
        try:
            if self._buffer.num_rows > 0:
                self._write_data_to_file(self._buffer.materialize())
        except Exception as e:
            import logging
            logger = logging.getLogger(__name__)
            logger.warning("Exception occurs when closing writer. Cleaning up.", exc_info=e)
            self.abort()
            raise e
        finally:
            self._buffer.reset()
            # Note: Don't clear committed_files in close() - they should be returned by prepare_commit()

    def abort(self):
        """
        Abort all writers and clean up resources. This method should be called when an error occurs
        during writing. It deletes any files that were written and cleans up resources.
        """
        to_delete = self.committed_files + self.committed_changelog_files
        self._delete_committed_files(to_delete)

        # Clean up resources
        self._buffer.reset()
        self.committed_files.clear()
        self.committed_changelog_files.clear()

    def _delete_committed_files(self, file_metas: List[DataFileMeta]):
        for file_meta in file_metas:
            try:
                for path_to_delete in file_meta.collect_files():
                    self.file_io.delete_quietly(path_to_delete)
            except Exception as e:
                import logging
                logger = logging.getLogger(__name__)
                path_to_delete = file_meta.physical_path()
                logger.warning(f"Failed to delete file {path_to_delete} during abort: {e}")

    @abstractmethod
    def _process_data(self, data: pa.RecordBatch) -> pa.RecordBatch:
        """Process incoming data (e.g., add system fields, sort). Must be implemented by subclasses."""

    @abstractmethod
    def _merge_data(self, existing_data: pa.RecordBatch, new_data: pa.RecordBatch) -> pa.RecordBatch:
        """Merge existing data with new data. Must be implemented by subclasses."""

    def _numbers_append_rows(self) -> bool:
        # Unaware-bucket files keep the zero range that commit replaces with
        # the snapshot id on row-tracking tables.
        return (not self.table.is_primary_key_table
                and self.table.bucket_mode() != BucketMode.BUCKET_UNAWARE)

    def _append_file_sequence_range(self, row_count: int) -> Tuple[int, int]:
        if row_count <= 0:
            raise ValueError("row_count must be positive")

        if self.options.data_evolution_enabled(False):
            # Row-tracking commit stamps this sentinel range with the snapshot id.
            return 0, row_count - 1

        return -1, -1

    def _check_and_roll_if_needed(self):
        # Neither trigger can fire below these thresholds, so the running counts
        # rule out rolling -- the common case -- without concatenating anything.
        while (self._buffer.nbytes > self.target_file_size
                or self._buffer.num_rows > self.target_file_row_num):
            pending = self._buffer.materialize()
            num_rows = pending.num_rows
            # Row-count trigger: keep at most target_file_row_num rows per file.
            split_row = num_rows
            if num_rows > self.target_file_row_num:
                split_row = self.target_file_row_num
            # Size trigger: roll earlier if the size split point comes first.
            if pending.nbytes > self.target_file_size:
                size_split = self._find_optimal_split_point(
                    pending, self.target_file_size)
                # First row alone exceeds target_file_size: roll it by itself.
                if size_split <= 0:
                    size_split = 1
                if size_split < split_row:
                    split_row = size_split
            if split_row <= 0 or split_row >= num_rows:
                break
            self._write_data_to_file(pending.slice(0, split_row))
            self._buffer.reset(pending.slice(split_row))

    def _write_file(self, file_path: str, data: pa.Table, file_format: str) -> Dict:
        """Write contents without publishing metadata; return shredding statistics."""
        if file_format == CoreOptions.FILE_FORMAT_PARQUET:
            return self._write_parquet_data(file_path, data)
        elif file_format == CoreOptions.FILE_FORMAT_ORC:
            self.file_io.write_orc(file_path, data, compression=self.compression, zstd_level=self.zstd_level)
        elif file_format == CoreOptions.FILE_FORMAT_AVRO:
            self.file_io.write_avro(file_path, data, compression=self.compression, zstd_level=self.zstd_level)
        elif file_format == CoreOptions.FILE_FORMAT_BLOB:
            write_blob_file(self.file_io, file_path, data)
        elif file_format == CoreOptions.FILE_FORMAT_VORTEX:
            write_vortex_file(self.file_io, file_path, data)
        elif file_format == CoreOptions.FILE_FORMAT_MOSAIC:
            self.file_io.write_mosaic(file_path, data, options=self.mosaic_writer_options)
        elif file_format == CoreOptions.FILE_FORMAT_ROW:
            write_row_file(self.file_io, file_path, data, zstd_level=self.zstd_level)
        else:
            raise ValueError(f"Unsupported file format: {file_format}")
        return {}

    def _write_data_to_file(self, data: pa.Table):
        if data.num_rows == 0:
            return
        file_name = f"{CoreOptions.data_file_prefix(self.options)}{uuid.uuid4()}-0.{self.file_format}"
        file_path = self._generate_file_path(file_name)

        is_external_path = self.external_path_provider is not None
        external_path_str = file_path if is_external_path else None

        logical_data = data
        extra_files = []
        row_sidecar_path = None
        changelog_meta = None
        shared_shredding_stats = {}
        if self._variant_shredding:
            data = self._apply_variant_shredding(data)

        # One data file means up to three files on disk -- the data file, its row
        # sidecar and its changelog -- and none of them is committed until all of
        # them have landed. A caller that retries the flush still holds these
        # rows in its buffer, so publishing the data file before the changelog
        # exists would make the retry write a second copy of rows the first meta
        # already covers.
        try:
            shared_shredding_stats = self._write_file(file_path, data, self.file_format)

            if self._should_write_row_sidecar():
                row_sidecar_name = f"{file_name}{self.ROW_SIDECAR_SUFFIX}"
                row_sidecar_path = f"{file_path}{self.ROW_SIDECAR_SUFFIX}"
                write_row_file(
                    self.file_io,
                    row_sidecar_path,
                    logical_data,
                    fields=self._row_sidecar_fields(logical_data),
                    zstd_level=self.zstd_level)
                extra_files.append(row_sidecar_name)

            statistics = self._data_file_statistics(logical_data)
            data_meta = self._create_data_file_meta(
                file_name=file_name,
                file_path=file_path,
                row_count=data.num_rows,
                min_key=GenericRow(statistics['min_key'], self.trimmed_primary_keys_fields),
                max_key=GenericRow(statistics['max_key'], self.trimmed_primary_keys_fields),
                key_stats=statistics['key_stats'],
                value_stats=statistics['value_stats'],
                min_sequence_number=statistics['min_seq'],
                max_sequence_number=statistics['max_seq'],
                extra_files=extra_files,
                delete_row_count=self._count_delete_rows(data),
            )

            if (self.changelog_producer == ChangelogProducer.INPUT
                    and not self._writes_changelog_before_merge()):
                changelog_meta = self._write_changelog_file(
                    data, is_external=external_path_str is not None, **statistics)
        except Exception:
            self.file_io.delete_quietly(file_path)
            if row_sidecar_path is not None:
                self.file_io.delete_quietly(row_sidecar_path)
            raise

        self._finish_data_file(data_meta, changelog_meta, shared_shredding_stats)

    def _data_file_statistics(self, data):
        # min key & max key

        selected_table = data.select(self.trimmed_primary_keys)
        key_columns_batch = selected_table.to_batches()[0]
        min_key_row_batch = key_columns_batch.slice(0, 1)
        max_key_row_batch = key_columns_batch.slice(key_columns_batch.num_rows - 1, 1)
        min_key = [col.to_pylist()[0] for col in min_key_row_batch.columns]
        max_key = [col.to_pylist()[0] for col in max_key_row_batch.columns]

        # key stats & value stats
        value_stats_enabled = self.options.metadata_stats_enabled()
        if value_stats_enabled:
            stats_fields = self.table.fields if self.table.is_primary_key_table \
                else PyarrowFieldParser.to_paimon_schema(data.schema)
        else:
            stats_fields = self.table.trimmed_primary_keys_fields
        column_stats = {
            field.name: self._get_column_stats(data, field.name)
            for field in stats_fields
        }
        key_fields = self.trimmed_primary_keys_fields
        key_stats = self._collect_value_stats(
            data, key_fields, column_stats)
        if not self.options.primary_key_nullable() and not all(
                count == 0 for count in key_stats.null_counts):
            raise RuntimeError("Primary key should not be null")

        value_fields = stats_fields if value_stats_enabled else []
        value_stats = self._collect_value_stats(
            data, value_fields, column_stats)

        # Read the range without advancing it: the advance belongs with the
        # append below, so a retried flush derives the same range.
        if self.table.is_primary_key_table and '_SEQUENCE_NUMBER' in data.schema.names:
            # PK files can contain a deduplicated or split subset of the
            # buffer. Java records the range of the rows actually written.
            sequences = data.column('_SEQUENCE_NUMBER')
            minmax = pc.min_max(sequences)
            min_seq = minmax["min"].as_py()
            max_seq = minmax["max"].as_py()
        elif self._numbers_append_rows():
            # Like Java's AppendOnlyWriter, number the rows from the bucket's
            # restored max sequence number + 1.
            min_seq = self.sequence_generator.current + 1
            max_seq = self.sequence_generator.current + data.num_rows
        else:
            min_seq = self.sequence_generator.start
            max_seq = self.sequence_generator.current
        return dict(min_key=min_key, max_key=max_key, key_stats=key_stats,
                    value_stats=value_stats, min_seq=min_seq, max_seq=max_seq,
                    value_stats_enabled=value_stats_enabled)

    def _writes_changelog_before_merge(self):
        return False

    def _finish_data_file(self, data_meta, changelog_meta=None,
                          shared_shredding_stats=None):
        """Record a fully written data file and its optional changelog."""
        if self._numbers_append_rows():
            self.sequence_generator.current = data_meta.max_sequence_number
        self.sequence_generator.start = self.sequence_generator.current
        if shared_shredding_stats is not None:
            self._map_shared_shredding.file_completed(shared_shredding_stats)
        self.committed_files.append(data_meta)
        if changelog_meta is not None:
            self.committed_changelog_files.append(changelog_meta)

    def _write_parquet_data(self, path, data):
        from pypaimon.execution.python.write.format_file_writer import parquet_write_data
        data = parquet_write_data(data)
        if self._map_shared_shredding.is_active():
            return self._map_shared_shredding.write_parquet(
                self.file_io, path, data, self.compression, self.zstd_level)
        self.file_io.write_parquet(
            path, data, compression=self.compression, zstd_level=self.zstd_level,
            **self.parquet_writer_options)
        return {}

    def _create_data_file_meta(self, file_name, file_path, row_count,
                               min_key, max_key, key_stats, value_stats,
                               min_sequence_number, max_sequence_number,
                               extra_files=None, creation_time=None, delete_row_count=0):
        """Common metadata finalization for buffered and incremental files."""
        return DataFileMeta.create(
            file_name=file_name,
            file_size=self.file_io.get_file_size(file_path),
            row_count=row_count,
            min_key=min_key, max_key=max_key,
            key_stats=key_stats, value_stats=value_stats,
            min_sequence_number=min_sequence_number,
            max_sequence_number=max_sequence_number,
            schema_id=self.table.table_schema.id, level=0,
            extra_files=extra_files if extra_files is not None else [],
            creation_time=creation_time if creation_time is not None else Timestamp.now().to_millis_timestamp(),
            delete_row_count=delete_row_count, file_source=0,
            value_stats_cols=None if self.options.metadata_stats_enabled() else [],
            external_path=file_path if self.external_path_provider is not None else None,
            first_row_id=None, write_cols=self.write_cols, file_path=file_path,
        )

    def _count_delete_rows(self, data: pa.Table) -> int:
        # Count the final file contents, after merging and rolling the buffer.
        if (
            not self.table.is_primary_key_table
            or SpecialFields.VALUE_KIND.name not in data.schema.names
        ):
            return 0
        kind_counts = pc.value_counts(data.column(SpecialFields.VALUE_KIND.name))
        return sum(entry['counts'].as_py() for entry in kind_counts
                   if entry['values'].as_py() in (RowKind.UPDATE_BEFORE.value, RowKind.DELETE.value))

    def _apply_variant_shredding(self, data: pa.Table) -> pa.Table:
        """Transform VARIANT columns into shredded Parquet format.

        Each shredded parent column is tagged with a ``PARQUET:field_id`` so that
        the Java ``ParquetSchemaConverter.convertToPaimonField`` (called from
        ``VariantUtils.variantFileType``) can read ``parquetType.getId().intValue()``
        without a NullPointerException.
        """
        from pypaimon.execution.python.read.variant_shredding import shred_variant_column
        columns = list(data.columns)
        fields = list(data.schema)
        changed = False

        for i, f in enumerate(fields):
            if f.name in self._variant_shredding:
                obj_fields, target_type = self._variant_shredding[f.name]
                columns[i] = shred_variant_column(columns[i], obj_fields, target_type)
                pid = self._paimon_field_id.get(f.name)
                parent_meta = {b'PARQUET:field_id': str(pid).encode()} if pid is not None else None
                fields[i] = pa.field(f.name, target_type, nullable=f.nullable, metadata=parent_meta)
                changed = True
        if not changed:
            return data
        return pa.Table.from_arrays(columns, schema=pa.schema(fields))

    def _write_changelog_file(self, data, min_key, max_key, key_stats, value_stats,
                              min_seq, max_seq,
                              value_stats_enabled, is_external) -> DataFileMeta:
        """Write a complete changelog file and return its metadata.

        The caller publishes this metadata after success. A failure here leaves
        neither metadata to commit nor a partial changelog file on disk.
        """
        cl_fmt = self.changelog_file_format
        changelog_file_name = f"changelog-{uuid.uuid4()}-0.{cl_fmt}"
        changelog_file_path = self._generate_file_path(changelog_file_name)

        changelog_external_path = changelog_file_path if is_external else None

        try:
            if cl_fmt not in (CoreOptions.FILE_FORMAT_PARQUET,
                              CoreOptions.FILE_FORMAT_ORC,
                              CoreOptions.FILE_FORMAT_AVRO):
                raise ValueError(f"Unsupported changelog file format: {cl_fmt}. "
                                 f"Supported formats: parquet, orc, avro.")
            self._write_file(changelog_file_path, data, cl_fmt)

            return DataFileMeta.create(
                file_name=changelog_file_name,
                file_size=self.file_io.get_file_size(changelog_file_path),
                row_count=data.num_rows,
                min_key=GenericRow(min_key, self.trimmed_primary_keys_fields),
                max_key=GenericRow(max_key, self.trimmed_primary_keys_fields),
                key_stats=key_stats,
                value_stats=value_stats,
                min_sequence_number=min_seq,
                max_sequence_number=max_seq,
                schema_id=self.table.table_schema.id,
                level=0,
                extra_files=[],
                creation_time=Timestamp.now().to_millis_timestamp(),
                delete_row_count=self._count_delete_rows(data),
                file_source=0,
                value_stats_cols=None if value_stats_enabled else [],
                external_path=changelog_external_path,
                first_row_id=None,
                write_cols=self.write_cols,
                file_path=changelog_file_path,
            )
        except Exception:
            self.file_io.delete_quietly(changelog_file_path)
            raise

    def _generate_file_path(self, file_name: str) -> str:
        if self.external_path_provider:
            return self.external_path_provider.get_next_external_data_path(file_name)

        # New files use Java's escaped partition names, including on the pure
        # Python writer path. Never introduce another Python-specific layout.
        return canonical_data_file_path(self.table, self.partition, self.bucket, file_name)

    def _should_write_row_sidecar(self) -> bool:
        return (
            self.options.data_evolution_enabled(False)
            and self.options.data_evolution_row_sidecar_enabled(False)
        )

    def _row_sidecar_fields(self, data: pa.Table) -> List:
        if self.write_cols:
            field_map = {field.name: field for field in self.table.table_schema.fields}
            return [field_map[name] for name in self.write_cols if name in field_map]
        return PyarrowFieldParser.to_paimon_schema(data.schema)

    @staticmethod
    def _aligned_extra_file_path(file_meta: DataFileMeta, extra_file: str) -> str:
        return file_meta.aligned_file_path(extra_file)

    @staticmethod
    def _find_optimal_split_point(data: pa.RecordBatch, target_size: int) -> int:
        total_rows = data.num_rows
        if total_rows <= 1:
            return 0

        left, right = 1, total_rows
        best_split = 0

        while left <= right:
            mid = (left + right) // 2
            slice_data = data.slice(0, mid)
            slice_size = slice_data.nbytes

            if slice_size <= target_size:
                best_split = mid
                left = mid + 1
            else:
                right = mid - 1

        return best_split

    def _collect_value_stats(self, data: pa.Table, fields: List,
                             column_stats: Optional[Dict[str, Dict]] = None) -> SimpleStats:
        if not fields:
            return SimpleStats.empty_stats()
        
        if column_stats is None or not column_stats:
            column_stats = {
                field.name: self._get_column_stats(data, field.name)
                for field in fields
            }
        
        min_stats = [column_stats[field.name]['min_values'] for field in fields]
        max_stats = [column_stats[field.name]['max_values'] for field in fields]
        null_counts = [column_stats[field.name]['null_counts'] for field in fields]
        
        return SimpleStats(
            GenericRow(min_stats, fields),
            GenericRow(max_stats, fields),
            null_counts
        )

    @staticmethod
    def _get_column_stats(record_batch: pa.RecordBatch, column_name: str) -> Dict:
        column_array = record_batch.column(column_name)
        if column_array.null_count == len(column_array):
            return {
                "min_values": None,
                "max_values": None,
                "null_counts": column_array.null_count,
            }
        
        column_type = column_array.type
        supports_minmax = not (
            pa.types.is_nested(column_type) or pa.types.is_map(column_type) or pa.types.is_large_binary(column_type)
        )
        
        if not supports_minmax:
            return {
                "min_values": None,
                "max_values": None,
                "null_counts": column_array.null_count,
            }
        
        minmax = pc.min_max(column_array)
        min_values = minmax["min"].as_py()
        max_values = minmax["max"].as_py()
        null_counts = column_array.null_count
        return {
            "min_values": min_values,
            "max_values": max_values,
            "null_counts": null_counts,
        }


class SequenceGenerator:
    def __init__(self, start: int = 0):
        self.start = start
        self.current = start

    def next(self) -> int:
        self.current += 1
        return self.current