| # 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. |
| |
| """ |
| Module to write a Paimon table from a Ray Dataset, by using the Ray Datasink API. |
| """ |
| |
| import logging |
| import traceback |
| from dataclasses import dataclass |
| from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional |
| |
| from ray.data.datasource.datasink import Datasink |
| |
| from ray.util.annotations import DeveloperAPI |
| from ray.data.block import BlockAccessor, Block |
| from ray.data._internal.execution.interfaces import TaskContext |
| import pyarrow as pa |
| |
| from pypaimon.write.commit_callback import CommitCallback |
| |
| if TYPE_CHECKING: |
| from pypaimon.table.table import Table |
| from pypaimon.write.write_builder import WriteBuilder |
| from pypaimon.write.commit_message import CommitMessage |
| |
| logger = logging.getLogger(__name__) |
| |
| |
| @dataclass(frozen=True) |
| class PaimonWriteResult: |
| """Metadata reported by the exact coordinator commit.""" |
| |
| row_count: int |
| snapshot_id: int |
| |
| |
| class _SnapshotIdRecorder(CommitCallback): |
| |
| def __init__(self): |
| self.snapshot_id = None |
| |
| def call(self, context): |
| self.snapshot_id = context.snapshot.id |
| |
| |
| class _TaskCommitMessages(list): |
| """Commit messages plus the logical rows written by one Ray task.""" |
| |
| def __init__(self, messages=(), row_count=0): |
| super().__init__(messages) |
| self.row_count = row_count |
| |
| |
| def _cast_binary_to_table_schema(table: pa.Table, target_schema: pa.Schema) -> pa.Table: |
| """Cast binary to large_binary for BLOB fields. |
| |
| When map_batches returns Python dicts, PyArrow infers bytes as binary, |
| losing the original large_binary (BLOB) type. Cast back before writing. |
| """ |
| cast_indices = [] |
| for i, field in enumerate(table.schema): |
| target_field = target_schema.field(field.name) if field.name in target_schema.names else None |
| if target_field and pa.types.is_binary(field.type) and pa.types.is_large_binary(target_field.type): |
| cast_indices.append(i) |
| |
| if not cast_indices: |
| return table |
| |
| columns = table.columns |
| for i in cast_indices: |
| columns[i] = columns[i].cast(pa.large_binary()) |
| fields = [target_schema.field(f.name) if i in cast_indices else f |
| for i, f in enumerate(table.schema)] |
| return pa.table(columns, schema=pa.schema(fields)) |
| |
| # Python 3.8 / Ray 2.10: Datasink is not subscriptable at runtime |
| try: |
| _DatasinkBase = Datasink[List["CommitMessage"]] |
| except TypeError: |
| _DatasinkBase = Datasink |
| |
| |
| @DeveloperAPI |
| class PaimonDatasink(_DatasinkBase): |
| def __init__( |
| self, |
| table: "Table", |
| overwrite: bool = False, |
| static_partition: Optional[Dict[str, Any]] = None, |
| postpone_bucket_plan=None, |
| ): |
| self.table = table |
| self.overwrite = overwrite |
| self.static_partition = static_partition |
| self._postpone_bucket_plan = postpone_bucket_plan |
| self._table_name = table.identifier.get_full_name() |
| self._writer_builder: Optional["WriteBuilder"] = None |
| self.commit_result: Optional[PaimonWriteResult] = None |
| |
| def _is_overwrite(self) -> bool: |
| return self.overwrite or self.static_partition is not None |
| |
| def __getstate__(self) -> dict: |
| state = self.__dict__.copy() |
| return state |
| |
| def __setstate__(self, state: dict) -> None: |
| self.__dict__.update(state) |
| writer_builder = getattr(self, '_writer_builder', None) |
| if writer_builder is not None and not hasattr(writer_builder, 'table'): |
| self._writer_builder = None |
| if not hasattr(self, '_table_name'): |
| self._table_name = self.table.identifier.get_full_name() |
| if not hasattr(self, 'static_partition'): |
| self.static_partition = None |
| if not hasattr(self, '_postpone_bucket_plan'): |
| self._postpone_bucket_plan = None |
| |
| def on_write_start(self, schema=None) -> None: |
| logger.info(f"Starting write job for table {self._table_name}") |
| |
| self._writer_builder = self.table.new_batch_write_builder() |
| if self._is_overwrite(): |
| self._writer_builder = self._writer_builder.overwrite(self.static_partition) |
| |
| def write( |
| self, |
| blocks: Iterable[Block], |
| ctx: TaskContext, |
| ) -> List["CommitMessage"]: |
| commit_messages_list: List["CommitMessage"] = [] |
| row_count = 0 |
| table_write = None |
| |
| try: |
| writer_builder = ( |
| self.table.new_postpone_fixed_bucket_write_builder() |
| if self._postpone_bucket_plan is not None |
| else self.table.new_batch_write_builder() |
| ) |
| if self._is_overwrite(): |
| writer_builder = writer_builder.overwrite(self.static_partition) |
| |
| if self._postpone_bucket_plan is not None: |
| writer_builder.with_bucket_plan(self._postpone_bucket_plan) |
| table_write = writer_builder.new_write() |
| |
| table_schema = self.table.table_schema |
| from pypaimon.schema.data_types import PyarrowFieldParser |
| target_pa_schema = PyarrowFieldParser.from_paimon_schema(table_schema.fields) |
| |
| for block in blocks: |
| block_arrow: pa.Table = BlockAccessor.for_block(block).to_arrow() |
| |
| if block_arrow.num_rows == 0: |
| continue |
| row_count += block_arrow.num_rows |
| |
| block_arrow = _cast_binary_to_table_schema(block_arrow, target_pa_schema) |
| |
| table_write.write_arrow(block_arrow) |
| |
| commit_messages = table_write.prepare_commit() |
| commit_messages_list.extend(commit_messages) |
| |
| table_write.close() |
| table_write = None |
| return _TaskCommitMessages(commit_messages_list, row_count) |
| except Exception: |
| if table_write is not None: |
| try: |
| table_write.abort() |
| except Exception as abort_error: |
| logger.warning( |
| f"Error aborting worker-side table_write: {abort_error}", |
| exc_info=abort_error |
| ) |
| raise |
| |
| @staticmethod |
| def _extract_write_returns(write_result: Any): |
| """Normalize WriteResult.write_returns (Ray 2.44+) vs list of returns |
| (older Ray) into a list of per-task commit-message lists.""" |
| if hasattr(write_result, "write_returns"): |
| return write_result.write_returns |
| if isinstance(write_result, list): |
| return write_result |
| raise TypeError( |
| f"Unexpected write_result type {type(write_result).__name__}: " |
| "expected object with .write_returns or list of commit message " |
| "lists. Refusing to proceed to avoid silent data loss." |
| ) |
| |
| @staticmethod |
| def _extract_row_count(write_result: Any, write_returns) -> int: |
| if hasattr(write_result, "num_rows"): |
| return int(write_result.num_rows) |
| return sum( |
| int(getattr(messages, "row_count", 0)) |
| for messages in write_returns |
| ) |
| |
| def on_write_complete( |
| self, write_result: Any |
| ): |
| table_commit = None |
| try: |
| write_returns = self._extract_write_returns(write_result) |
| all_commit_messages = [ |
| commit_message |
| for commit_messages in write_returns |
| for commit_message in commit_messages |
| ] |
| |
| non_empty_messages = [ |
| msg for msg in all_commit_messages if not msg.is_empty() |
| ] |
| |
| if not non_empty_messages and not self._is_overwrite(): |
| logger.info("No data to commit (all commit messages are empty)") |
| return |
| |
| # Ray does not call on_write_start when the input has no blocks. |
| if self._writer_builder is None: |
| self.on_write_start() |
| |
| logger.info( |
| f"Committing {len(non_empty_messages)} commit messages " |
| f"for table {self._table_name}" |
| ) |
| |
| table_commit = self._writer_builder.new_commit() |
| recorder = _SnapshotIdRecorder() |
| table_commit.add_commit_callback(recorder) |
| table_commit.commit(non_empty_messages) |
| if recorder.snapshot_id is not None: |
| self.commit_result = PaimonWriteResult( |
| row_count=self._extract_row_count( |
| write_result, write_returns), |
| snapshot_id=recorder.snapshot_id, |
| ) |
| |
| logger.info(f"Successfully committed write job for table {self._table_name}") |
| except Exception as e: |
| logger.error( |
| f"Error committing write job for table {self._table_name}: {e}", |
| exc_info=e |
| ) |
| raise |
| finally: |
| if table_commit is not None: |
| try: |
| table_commit.close() |
| except Exception as e: |
| logger.warning( |
| f"Error closing table_commit: {e}", |
| exc_info=e |
| ) |
| |
| def on_write_failed(self, error: Exception) -> None: |
| logger.error( |
| f"Write job failed for table {self._table_name}. Error: {error}", |
| exc_info=error |
| ) |
| # Do not abort files returned by completed write tasks. Ray or an outer |
| # scheduler may replay those commit messages in another attempt. |
| |
| |
| def write_paimon_dataset( |
| dataset, |
| table, |
| *, |
| overwrite: bool = False, |
| static_partition: Optional[Dict[str, Any]] = None, |
| concurrency: Optional[int] = None, |
| ray_remote_args: Optional[Dict[str, Any]] = None, |
| hash_fixed_precluster: str = "auto", |
| postpone_bucket_planner=None, |
| ) -> Optional[PaimonWriteResult]: |
| """Write a Ray Dataset through the safe path for the table's bucket mode.""" |
| from pypaimon.ray.shuffle import ( |
| HASH_FIXED_PRECLUSTER_MAP_GROUPS, |
| HASH_FIXED_PRECLUSTER_MODES, |
| HASH_FIXED_PRECLUSTER_OFF, |
| maybe_apply_repartition, |
| ) |
| from pypaimon.table.bucket_mode import BucketMode |
| |
| if hash_fixed_precluster not in HASH_FIXED_PRECLUSTER_MODES: |
| raise ValueError( |
| "hash_fixed_precluster must be one of {}, got {!r}".format( |
| sorted(HASH_FIXED_PRECLUSTER_MODES), |
| hash_fixed_precluster, |
| ) |
| ) |
| |
| if ( |
| table.bucket_mode() == BucketMode.POSTPONE_MODE |
| and ( |
| postpone_bucket_planner is not None |
| or ( |
| table.options.postpone_batch_write_fixed_bucket() |
| and hash_fixed_precluster != HASH_FIXED_PRECLUSTER_OFF |
| ) |
| ) |
| ): |
| from pypaimon.write.postpone_bucket import ( |
| PostponeBucketPlanner, |
| ) |
| from pypaimon.write.row_key_extractor import ( |
| PostponeFixedBucketRowKeyExtractor, |
| ) |
| |
| planner = ( |
| postpone_bucket_planner |
| if postpone_bucket_planner is not None |
| else PostponeBucketPlanner(table) |
| ) |
| plan = planner.current_plan() |
| if table.partition_keys or not plan.contains(()): |
| dataset, partition_stats = _collect_partition_stats( |
| dataset, planner |
| ) |
| plan = planner.plan( |
| partition_stats, |
| include_postpone_rows=not ( |
| overwrite or static_partition is not None |
| ), |
| ) |
| return _write_postpone_primary_key_blocks( |
| dataset, |
| table, |
| overwrite=overwrite, |
| static_partition=static_partition, |
| concurrency=concurrency, |
| ray_remote_args=ray_remote_args, |
| bucket_extractor=PostponeFixedBucketRowKeyExtractor(table, plan), |
| postpone_bucket_plan=plan, |
| ) |
| |
| if ( |
| hash_fixed_precluster == HASH_FIXED_PRECLUSTER_MAP_GROUPS |
| and table.bucket_mode() == BucketMode.HASH_FIXED |
| and getattr(table, "is_primary_key_table", False) |
| ): |
| return _write_primary_key_groups( |
| dataset, |
| table, |
| overwrite=overwrite, |
| static_partition=static_partition, |
| concurrency=concurrency, |
| ray_remote_args=ray_remote_args, |
| ) |
| |
| dataset = maybe_apply_repartition(dataset, table, hash_fixed_precluster) |
| datasink = PaimonDatasink( |
| table, |
| overwrite=overwrite, |
| static_partition=static_partition, |
| ) |
| dataset.write_datasink( |
| datasink, |
| concurrency=concurrency, |
| ray_remote_args=ray_remote_args, |
| ) |
| return datasink.commit_result |
| |
| |
| def _write_postpone_primary_key_blocks( |
| dataset, |
| table, |
| *, |
| overwrite: bool, |
| static_partition: Optional[Dict[str, Any]], |
| concurrency: Optional[int], |
| ray_remote_args: Optional[Dict[str, Any]], |
| bucket_extractor, |
| postpone_bucket_plan, |
| ) -> Optional[PaimonWriteResult]: |
| import pickle |
| |
| from pypaimon.ray.shuffle import ( |
| _coerce_large_string_types, |
| _sort_by_partition_bucket_primary_key, |
| ) |
| |
| sorted_dataset, routing_columns = ( |
| _sort_by_partition_bucket_primary_key( |
| dataset, table, bucket_extractor |
| ) |
| ) |
| message_col = "__paimon_commit_messages__" |
| error_col = "__paimon_write_error__" |
| captured_table = table |
| |
| def _write_block(batch: pa.Table) -> pa.Table: |
| if batch.num_rows == 0: |
| return pa.table({ |
| message_col: pa.array([], type=pa.binary()), |
| error_col: pa.array([], type=pa.string()), |
| }) |
| |
| rows = _coerce_large_string_types( |
| batch.drop_columns(routing_columns) |
| ) |
| worker_sink = PaimonDatasink( |
| captured_table, |
| overwrite=overwrite, |
| static_partition=static_partition, |
| postpone_bucket_plan=postpone_bucket_plan, |
| ) |
| try: |
| commit_messages = worker_sink.write([rows], None) |
| error = None |
| except Exception: |
| commit_messages = [] |
| error = traceback.format_exc() |
| return pa.table({ |
| message_col: pa.array( |
| [pickle.dumps(commit_messages)], type=pa.binary() |
| ), |
| error_col: pa.array([error], type=pa.string()), |
| }) |
| |
| map_kwargs = _ray_map_kwargs( |
| sorted_dataset.map_batches, |
| concurrency, |
| ray_remote_args, |
| batch_size=None, |
| batch_format="pyarrow", |
| zero_copy_batch=True, |
| ) |
| |
| results = sorted_dataset.map_batches(_write_block, **map_kwargs) |
| coordinator = PaimonDatasink( |
| table, |
| overwrite=overwrite, |
| static_partition=static_partition, |
| ) |
| coordinator.on_write_start() |
| return _consume_write_results( |
| results, coordinator, message_col, error_col |
| ) |
| |
| |
| def _ray_map_kwargs(method, concurrency, ray_remote_args, **kwargs): |
| import inspect |
| |
| if concurrency is not None: |
| concurrency_param = inspect.signature( |
| method |
| ).parameters.get("concurrency") |
| if ( |
| concurrency_param is not None |
| and concurrency_param.kind != inspect.Parameter.VAR_KEYWORD |
| ): |
| kwargs["concurrency"] = concurrency |
| else: |
| from ray.data._internal.compute import TaskPoolStrategy |
| kwargs["compute"] = TaskPoolStrategy(size=concurrency) |
| if ray_remote_args: |
| kwargs.update(ray_remote_args) |
| return kwargs |
| |
| |
| def _consume_write_results( |
| results, |
| coordinator, |
| message_col, |
| error_col=None, |
| ) -> Optional[PaimonWriteResult]: |
| import pickle |
| |
| write_returns = [] |
| errors = [] |
| try: |
| for batch in results.iter_batches(batch_format="pyarrow"): |
| messages = batch.column(message_col).to_pylist() |
| batch_errors = ( |
| batch.column(error_col).to_pylist() |
| if error_col is not None else [None] * len(messages) |
| ) |
| for blob, error in zip(messages, batch_errors): |
| commit_messages = pickle.loads(blob) |
| write_returns.append(commit_messages) |
| if error is not None: |
| errors.append(error) |
| if errors: |
| raise RuntimeError( |
| "One or more Ray write tasks failed:\n{}".format( |
| "\n".join(errors) |
| ) |
| ) |
| coordinator.on_write_complete(write_returns) |
| return coordinator.commit_result |
| except Exception as error: |
| coordinator.on_write_failed(error) |
| raise |
| |
| |
| def _collect_partition_stats(dataset, planner): |
| import pickle |
| |
| partition_col = "__paimon_partition__" |
| rows_col = "__paimon_rows__" |
| size_col = "__paimon_size__" |
| |
| def _stats(batch: pa.Table) -> pa.Table: |
| stats = planner.input_partition_stats(batch) |
| items = list(stats.items()) |
| return pa.table({ |
| partition_col: pa.array( |
| [pickle.dumps(partition) for partition, _ in items], |
| type=pa.binary(), |
| ), |
| rows_col: pa.array( |
| [value[0] for _, value in items], type=pa.int64() |
| ), |
| size_col: pa.array( |
| [value[1] for _, value in items], type=pa.int64() |
| ), |
| }) |
| |
| materialized = dataset.materialize() |
| stats_dataset = materialized.map_batches( |
| _stats, batch_format="pyarrow", zero_copy_batch=True |
| ) |
| combined = {} |
| for batch in stats_dataset.iter_batches(batch_format="pyarrow"): |
| for partition, rows, size in zip( |
| batch.column(partition_col).to_pylist(), |
| batch.column(rows_col).to_pylist(), |
| batch.column(size_col).to_pylist(), |
| ): |
| key = pickle.loads(partition) |
| previous_rows, previous_size = combined.get(key, (0, 0)) |
| combined[key] = ( |
| previous_rows + rows, |
| previous_size + size, |
| ) |
| return materialized, combined |
| |
| |
| def _write_primary_key_groups( |
| dataset, |
| table, |
| *, |
| overwrite: bool, |
| static_partition: Optional[Dict[str, Any]], |
| concurrency: Optional[int], |
| ray_remote_args: Optional[Dict[str, Any]], |
| bucket_extractor=None, |
| postpone_bucket_plan=None, |
| ) -> Optional[PaimonWriteResult]: |
| import pickle |
| |
| from pypaimon.ray.shuffle import ( |
| _coerce_large_string_types, |
| _group_by_partition_bucket, |
| ) |
| |
| grouped, bucket_col = _group_by_partition_bucket( |
| dataset, table, extractor=bucket_extractor |
| ) |
| message_col = "__paimon_commit_messages__" |
| error_col = "__paimon_write_error__" |
| captured_table = table |
| |
| # Keep the writer inside the group UDF. Ray may split the UDF output |
| # into multiple blocks, so only serialized commit messages leave it. |
| def _write_group(group: pa.Table) -> pa.Table: |
| if group.num_rows == 0: |
| return pa.table({ |
| message_col: pa.array([], type=pa.binary()), |
| error_col: pa.array([], type=pa.string()), |
| }) |
| |
| rows = _coerce_large_string_types( |
| group.drop_columns([bucket_col]) |
| ) |
| worker_sink = PaimonDatasink( |
| captured_table, |
| overwrite=overwrite, |
| static_partition=static_partition, |
| postpone_bucket_plan=postpone_bucket_plan, |
| ) |
| try: |
| commit_messages = worker_sink.write([rows], None) |
| error = None |
| except Exception: |
| commit_messages = [] |
| error = traceback.format_exc() |
| return pa.table({ |
| message_col: pa.array( |
| [pickle.dumps(commit_messages)], type=pa.binary() |
| ), |
| error_col: pa.array([error], type=pa.string()), |
| }) |
| |
| map_kwargs = _ray_map_kwargs( |
| grouped.map_groups, |
| concurrency, |
| ray_remote_args, |
| batch_format="pyarrow", |
| ) |
| |
| messages = grouped.map_groups(_write_group, **map_kwargs) |
| coordinator = PaimonDatasink( |
| table, |
| overwrite=overwrite, |
| static_partition=static_partition, |
| ) |
| coordinator.on_write_start() |
| return _consume_write_results( |
| messages, coordinator, message_col, error_col |
| ) |
