paimon-python/pypaimon/ray/update_by_row_id.py - paimon - Git at Google

 #  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.

 """Distributed row-id update on Ray for data-evolution tables.

 Update columns of a data-evolution table straight from a Ray Dataset that already
 carries ``_ROW_ID`` and the new values -- no full-target read and no big-table
 shuffle join (unlike ``merge_into(on=["_ROW_ID"])``, which reads and joins the whole
 target). Pairs with ``bucket_join``, which produces the row ids.
 """

 import logging
 from typing import Any, Dict, List, Optional

 import pyarrow as pa

 from pypaimon.ray.data_evolution_merge_into import (
     _normalize_source,
     _reraise_inner,
     _require_ray_join,
     _resolve_num_partitions,
 )
 from pypaimon.ray.data_evolution_merge_join import distributed_update_apply
 from pypaimon.ray.data_evolution_merge_transform import build_update_schema
 from pypaimon.schema.data_types import is_blob_file_field

 __all__ = ["update_by_row_id"]

 logger = logging.getLogger(__name__)


 def _blob_col_names(table: "FileStoreTable") -> set:
     return {f.name for f in table.table_schema.fields
             if is_blob_file_field(f)}


 def update_by_row_id(
     target: str,
     source: Any,
     catalog_options: Dict[str, str],
     *,
     update_cols: List[str],
     num_partitions: Optional[int] = None,
     ray_remote_args: Optional[Dict[str, Any]] = None,
 ) -> Dict[str, int]:
     """Update ``update_cols`` of a data-evolution table by ``_ROW_ID``.

     ``source`` (a ``ray.data.Dataset`` / ``pyarrow.Table`` / ``pandas.DataFrame``)
     must already carry the target ``_ROW_ID`` and the new values. Each row is
     routed to the data file owning its row id and only those files are rewritten --
     the target is never fully read and there is no join against it. Requires
     ``ray >= 2.50`` and a target with ``data-evolution.enabled`` + ``row-tracking.enabled``.

     Returns ``{"num_updated": <rows>}``.
     """
     from pypaimon.catalog.catalog_factory import CatalogFactory
     from pypaimon.schema.data_types import PyarrowFieldParser
     from pypaimon.table.special_fields import SpecialFields

     _require_ray_join()
     if not update_cols:
         raise ValueError("update_cols must be non-empty.")
     update_cols = list(dict.fromkeys(update_cols))  # de-dup, keep order
     num_partitions = _resolve_num_partitions(num_partitions)

     table = CatalogFactory.create(catalog_options).get_table(target)
     if not table.options.data_evolution_enabled():
         raise ValueError(
             f"update_by_row_id requires 'data-evolution.enabled'='true' on '{target}'.")
     if not table.options.row_tracking_enabled():
         raise ValueError(
             f"update_by_row_id requires 'row-tracking.enabled'='true' on '{target}'.")
     if table.options.deletion_vectors_enabled():
         # A DV-deleted row still lives in its data file, so row-id ranges can't tell it
         # apart without reading the target; refuse rather than update a deleted row.
         raise ValueError(
             f"update_by_row_id does not support deletion-vectors-enabled tables yet: "
             f"'{target}'.")

     rid = SpecialFields.ROW_ID.name
     blob_cols = _blob_col_names(table)
     partition_keys = set(table.partition_keys or [])
     for col in update_cols:
         if col not in table.field_names:
             raise ValueError(f"update column {col!r} is not in target '{target}'.")
         if col in blob_cols:
             # Update writes plain data files; blob deltas are a separate path.
             raise ValueError(f"update_by_row_id cannot update blob column {col!r}.")
         if col in partition_keys:
             # In-place rewrite can't move a row across partitions.
             raise ValueError(
                 f"update_by_row_id cannot update partition column {col!r}; "
                 "cross-partition row movement is not supported.")

     if isinstance(source, str):
         # A table's system _ROW_ID is its own, independent of the target's, so a
         # table-name source can't address target rows. Require in-memory data that
         # already carries the target row ids (e.g. produced by bucket_join).
         raise ValueError(
             "update_by_row_id does not accept a table-name source; pass a ray.data."
             f"Dataset / pyarrow.Table / pandas.DataFrame carrying the target {rid}.")
     source_ds = _normalize_source(source, catalog_options)
     src_cols = set(source_ds.schema().names)
     missing = [c for c in [rid] + update_cols if c not in src_cols]
     if missing:
         raise ValueError(
             f"source is missing columns {missing}; it must carry {rid} and {update_cols}.")

     # Cast to the on-disk schema (int64 _ROW_ID + target column types) so the writer
     # gets exactly the target types regardless of the source's arrow types.
     target_pa = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
     update_schema = build_update_schema(target_pa, update_cols, rid)

     def _project_cast(batch: pa.Table) -> pa.Table:
         return batch.select([rid] + update_cols).cast(update_schema)

     update_ds = source_ds.map_batches(_project_cast, batch_format="pyarrow")

     base = table.snapshot_manager().get_latest_snapshot()
     # Without deletion vectors (rejected above), total_record_count is the live row
     # count, so 0 means the target is empty (never written, or emptied by overwrite).
     if base is None or base.total_record_count == 0:
         # Every source row id is foreign; don't silently no-op non-empty input.
         if update_ds.limit(1).count() > 0:
             raise ValueError(
                 f"target '{target}' has no rows; every _ROW_ID in the source is foreign.")
         return {"num_updated": 0}
     try:
         msgs, num_updated, _ = distributed_update_apply(
             update_ds, table, update_cols,
             num_partitions=num_partitions,
             ray_remote_args=ray_remote_args,
             base_snapshot_id=base.id,
         )
     except Exception as e:
         _reraise_inner(e)
         raise  # _reraise_inner always raises; keeps msgs/num_updated defined for linters

     if msgs:
         _commit_update_messages(table, msgs)
     return {"num_updated": num_updated}


 def _commit_update_messages(table, commit_messages) -> None:
     pending_msgs: list = list(commit_messages)
     commit_started = False

     try:
         table_commit = None
         try:
             table_commit = table.new_batch_write_builder().new_commit()
             commit_started = True
             table_commit.commit(pending_msgs)
         finally:
             if table_commit is not None:
                 try:
                     table_commit.close()
                 except Exception as close_error:
                     logger.warning(
                         "Failed to close update_by_row_id commit: %s",
                         close_error,
                         exc_info=close_error,
                     )
     except Exception as e:
         if not commit_started:
             _abort_pending_update_messages(table, pending_msgs)
         _reraise_inner(e)


 def _abort_pending_update_messages(table, commit_messages) -> None:
     if not commit_messages:
         return

     table_commit = None
     try:
         table_commit = table.new_batch_write_builder().new_commit()
         table_commit.abort(commit_messages)
     except Exception as abort_error:
         logger.warning(
             "Failed to abort pending update_by_row_id commit messages: %s",
             abort_error,
             exc_info=abort_error,
         )
     finally:
         if table_commit is not None:
             try:
                 table_commit.close()
             except Exception as close_error:
                 logger.warning(
                     "Failed to close update_by_row_id abort commit: %s",
                     close_error,
                     exc_info=close_error,
                 )
	# 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.

	"""Distributed row-id update on Ray for data-evolution tables.

	Update columns of a data-evolution table straight from a Ray Dataset that already
	carries ``_ROW_ID`` and the new values -- no full-target read and no big-table
	shuffle join (unlike ``merge_into(on=["_ROW_ID"])``, which reads and joins the whole
	target). Pairs with ``bucket_join``, which produces the row ids.
	"""

	import logging
	from typing import Any, Dict, List, Optional

	import pyarrow as pa

	from pypaimon.ray.data_evolution_merge_into import (
	_normalize_source,
	_reraise_inner,
	_require_ray_join,
	_resolve_num_partitions,
	)
	from pypaimon.ray.data_evolution_merge_join import distributed_update_apply
	from pypaimon.ray.data_evolution_merge_transform import build_update_schema
	from pypaimon.schema.data_types import is_blob_file_field

	__all__ = ["update_by_row_id"]

	logger = logging.getLogger(__name__)


	def _blob_col_names(table: "FileStoreTable") -> set:
	return {f.name for f in table.table_schema.fields
	if is_blob_file_field(f)}


	def update_by_row_id(
	target: str,
	source: Any,
	catalog_options: Dict[str, str],
	*,
	update_cols: List[str],
	num_partitions: Optional[int] = None,
	ray_remote_args: Optional[Dict[str, Any]] = None,
	) -> Dict[str, int]:
	"""Update ``update_cols`` of a data-evolution table by ``_ROW_ID``.

	``source`` (a ``ray.data.Dataset`` / ``pyarrow.Table`` / ``pandas.DataFrame``)
	must already carry the target ``_ROW_ID`` and the new values. Each row is
	routed to the data file owning its row id and only those files are rewritten --
	the target is never fully read and there is no join against it. Requires
	``ray >= 2.50`` and a target with ``data-evolution.enabled`` + ``row-tracking.enabled``.

	Returns ``{"num_updated": <rows>}``.
	"""
	from pypaimon.catalog.catalog_factory import CatalogFactory
	from pypaimon.schema.data_types import PyarrowFieldParser
	from pypaimon.table.special_fields import SpecialFields

	_require_ray_join()
	if not update_cols:
	raise ValueError("update_cols must be non-empty.")
	update_cols = list(dict.fromkeys(update_cols)) # de-dup, keep order
	num_partitions = _resolve_num_partitions(num_partitions)

	table = CatalogFactory.create(catalog_options).get_table(target)
	if not table.options.data_evolution_enabled():
	raise ValueError(
	f"update_by_row_id requires 'data-evolution.enabled'='true' on '{target}'.")
	if not table.options.row_tracking_enabled():
	raise ValueError(
	f"update_by_row_id requires 'row-tracking.enabled'='true' on '{target}'.")
	if table.options.deletion_vectors_enabled():
	# A DV-deleted row still lives in its data file, so row-id ranges can't tell it
	# apart without reading the target; refuse rather than update a deleted row.
	raise ValueError(
	f"update_by_row_id does not support deletion-vectors-enabled tables yet: "
	f"'{target}'.")

	rid = SpecialFields.ROW_ID.name
	blob_cols = _blob_col_names(table)
	partition_keys = set(table.partition_keys or [])
	for col in update_cols:
	if col not in table.field_names:
	raise ValueError(f"update column {col!r} is not in target '{target}'.")
	if col in blob_cols:
	# Update writes plain data files; blob deltas are a separate path.
	raise ValueError(f"update_by_row_id cannot update blob column {col!r}.")
	if col in partition_keys:
	# In-place rewrite can't move a row across partitions.
	raise ValueError(
	f"update_by_row_id cannot update partition column {col!r}; "
	"cross-partition row movement is not supported.")

	if isinstance(source, str):
	# A table's system _ROW_ID is its own, independent of the target's, so a
	# table-name source can't address target rows. Require in-memory data that
	# already carries the target row ids (e.g. produced by bucket_join).
	raise ValueError(
	"update_by_row_id does not accept a table-name source; pass a ray.data."
	f"Dataset / pyarrow.Table / pandas.DataFrame carrying the target {rid}.")
	source_ds = _normalize_source(source, catalog_options)
	src_cols = set(source_ds.schema().names)
	missing = [c for c in [rid] + update_cols if c not in src_cols]
	if missing:
	raise ValueError(
	f"source is missing columns {missing}; it must carry {rid} and {update_cols}.")

	# Cast to the on-disk schema (int64 _ROW_ID + target column types) so the writer
	# gets exactly the target types regardless of the source's arrow types.
	target_pa = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
	update_schema = build_update_schema(target_pa, update_cols, rid)

	def _project_cast(batch: pa.Table) -> pa.Table:
	return batch.select([rid] + update_cols).cast(update_schema)

	update_ds = source_ds.map_batches(_project_cast, batch_format="pyarrow")

	base = table.snapshot_manager().get_latest_snapshot()
	# Without deletion vectors (rejected above), total_record_count is the live row
	# count, so 0 means the target is empty (never written, or emptied by overwrite).
	if base is None or base.total_record_count == 0:
	# Every source row id is foreign; don't silently no-op non-empty input.
	if update_ds.limit(1).count() > 0:
	raise ValueError(
	f"target '{target}' has no rows; every _ROW_ID in the source is foreign.")
	return {"num_updated": 0}
	try:
	msgs, num_updated, _ = distributed_update_apply(
	update_ds, table, update_cols,
	num_partitions=num_partitions,
	ray_remote_args=ray_remote_args,
	base_snapshot_id=base.id,
	)
	except Exception as e:
	_reraise_inner(e)
	raise # _reraise_inner always raises; keeps msgs/num_updated defined for linters

	if msgs:
	_commit_update_messages(table, msgs)
	return {"num_updated": num_updated}


	def _commit_update_messages(table, commit_messages) -> None:
	pending_msgs: list = list(commit_messages)
	commit_started = False

	try:
	table_commit = None
	try:
	table_commit = table.new_batch_write_builder().new_commit()
	commit_started = True
	table_commit.commit(pending_msgs)
	finally:
	if table_commit is not None:
	try:
	table_commit.close()
	except Exception as close_error:
	logger.warning(
	"Failed to close update_by_row_id commit: %s",
	close_error,
	exc_info=close_error,
	)
	except Exception as e:
	if not commit_started:
	_abort_pending_update_messages(table, pending_msgs)
	_reraise_inner(e)


	def _abort_pending_update_messages(table, commit_messages) -> None:
	if not commit_messages:
	return

	table_commit = None
	try:
	table_commit = table.new_batch_write_builder().new_commit()
	table_commit.abort(commit_messages)
	except Exception as abort_error:
	logger.warning(
	"Failed to abort pending update_by_row_id commit messages: %s",
	abort_error,
	exc_info=abort_error,
	)
	finally:
	if table_commit is not None:
	try:
	table_commit.close()
	except Exception as close_error:
	logger.warning(
	"Failed to close update_by_row_id abort commit: %s",
	close_error,
	exc_info=close_error,
	)