postpone_batch_table_write.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 Any, Dict, List, Optional

import pyarrow as pa

from pypaimon.common.options.core_options import CoreOptions
from pypaimon.write.commit_message import CommitMessage
from pypaimon.execution.python.write.file_store_write import PostponeFixedBucketFileStoreWrite
from pypaimon.execution.python.write.postpone_bucket import PythonPostponeBucketPlanner as PostponeBucketPlanner
from pypaimon.execution.python.write.table_write import PythonTableWrite
from pypaimon.execution.python.write.row_key_extractor import PostponeFixedBucketRowKeyExtractor
from pypaimon.write.row_utils import (
    normalize_inline_blob_row,
    require_columns,
    row_to_named_values,
    row_values_to_arrow_table,
)


class PythonPostponeFixedBucketWrite(PythonTableWrite):
    """Batch writer which plans postpone rows before routing new partitions."""

    def initialize(
        self,
        table,
        commit_user,
        static_partition: Optional[dict] = None,
        bucket_plan=None,
    ):
        self._planner = PostponeBucketPlanner(
            table,
            known_num_buckets=(
                bucket_plan.as_dict() if bucket_plan is not None
                else {} if static_partition is not None else None
            ),
        )
        self._bucket_plan = (
            bucket_plan
            if bucket_plan is not None
            else self._planner.current_plan()
        )
        self._plan_provided = bucket_plan is not None
        self._pending_inputs = []
        super().initialize(table, commit_user, static_partition)

    def _create_file_store_write(self, commit_user):
        return PostponeFixedBucketFileStoreWrite(self.table, commit_user)

    def _create_row_key_extractor(self, static_partition):
        return PostponeFixedBucketRowKeyExtractor(
            self.table, self._bucket_plan)

    def _write_partition_bucket_batch(self, partition, bucket, data):
        self.file_store_write.write(
            partition,
            bucket,
            data,
            self.row_key_extractor.num_buckets(partition),
        )

    def _write_partition_bucket_row(
        self, partition, bucket, row, values_by_name
    ):
        self.file_store_write.write_row(
            partition,
            bucket,
            row,
            values_by_name,
            self.row_key_extractor.num_buckets(partition),
        )

    def _buffer_input(self, data) -> bool:
        if self._plan_provided:
            return False
        return any(
            not self._bucket_plan.contains(partition)
            for partition in self._planner.input_partition_stats(data)
        )

    def write_arrow(self, table: pa.Table, bucket=None):
        if bucket is not None:
            raise ValueError('Precomputed bucket writes require HASH_FIXED or HASH_DYNAMIC tables')
        table = self._prepare_arrow_data(table)
        if not self._buffer_input(table):
            return super().write_arrow(table)
        self._pending_inputs.extend(
            ("batch", batch) for batch in table.to_batches())

    def write_arrow_batch(self, data: pa.RecordBatch, bucket=None):
        if bucket is not None:
            raise ValueError('Precomputed bucket writes require HASH_FIXED or HASH_DYNAMIC tables')
        data = self._prepare_arrow_data(data)
        if not self._buffer_input(data):
            return super().write_arrow_batch(data)
        self._pending_inputs.append(("batch", data))

    def write_row(self, row):
        if self._plan_provided:
            return super().write_row(row)

        values_by_name = row_to_named_values(
            row, self.table.table_schema.fields)
        column_names = (
            self.file_store_write.write_cols
            if self.file_store_write.write_cols is not None
            else list(self.table.field_names)
        )
        require_columns(values_by_name, column_names, "write_row")
        require_columns(values_by_name, self.table.partition_keys, "write_row")
        from pypaimon.write.row_kind import skip_write_row

        if skip_write_row(self.table, values_by_name, row.get_row_kind(), self.file_store_write.write_cols):
            return
        partition = tuple(
            values_by_name[key] for key in self.table.partition_keys)
        if self._bucket_plan.contains(partition):
            return super().write_row(row)

        row, values_by_name = normalize_inline_blob_row(
            row,
            values_by_name,
            self.table.field_dict,
            column_names,
            CoreOptions.blob_descriptor_fields(self.table.options),
            CoreOptions.blob_view_fields(self.table.options),
        )
        arrow_row = row_values_to_arrow_table(
            values_by_name, self.table.table_schema.fields, column_names)
        size = self._planner.input_partition_stats(arrow_row)[partition][1]
        self._pending_inputs.append(
            ("row", (row, partition, size)))

    def _flush_pending_inputs(self):
        if not self._pending_inputs:
            return

        partition_stats = {}
        for input_type, value in self._pending_inputs:
            if input_type == "batch":
                stats_by_partition = (
                    self._planner.input_partition_stats(value))
            else:
                _, partition, size = value
                stats_by_partition = {partition: (1, size)}
            for partition, stats in stats_by_partition.items():
                rows, size = partition_stats.get(partition, (0, 0))
                partition_stats[partition] = (
                    rows + stats[0], size + stats[1])

        self._bucket_plan = self._planner.plan(
            partition_stats,
        )
        self.row_key_extractor.with_bucket_plan(self._bucket_plan)
        inputs = self._pending_inputs
        self._pending_inputs = []
        for input_type, value in inputs:
            if input_type == "batch":
                super().write_arrow_batch(value)
            else:
                super().write_row(value[0])

    def prepare_commit(self, commit_identifier) -> List[CommitMessage]:
        self._flush_pending_inputs()
        return super().prepare_commit(commit_identifier)

    def _distributed_write_options(self, *, overwrite=False) -> Dict[str, Any]:
        if overwrite and not self._plan_provided:
            return {"postpone_bucket_planner": PostponeBucketPlanner(self.table, known_num_buckets={})}
        return {"postpone_bucket_planner": self._planner}

    def close(self):
        self._pending_inputs = []
        super().close()

    def abort(self):
        self._pending_inputs = []
        super().abort()