JingsongLi commented on code in PR #8445: URL: https://github.com/apache/paimon/pull/8445#discussion_r3518588123
########## paimon-python/pypaimon/ray/update_by_row_id.py: ########## @@ -0,0 +1,149 @@ +# 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. +""" + +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 + +__all__ = ["update_by_row_id"] + + +def _blob_col_names(table: "FileStoreTable") -> set: + return {f.name for f in table.table_schema.fields + if getattr(f.type, "type", None) == "BLOB"} + + +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`` / + table-name str) must already carry ``_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.") + + source_ds = _normalize_source(source, catalog_options) Review Comment: This currently breaks the documented table-name source path. When `source` is a Paimon table identifier, `_normalize_source` calls `read_paimon(source, catalog_options)` without a projection, and the default read schema does not include row-tracking system fields. `_ROW_ID` is only returned when it is explicitly projected, so `src_cols` will miss `_ROW_ID` and this API will reject a string source even though the docs say it is supported. Please special-case string sources to read `projection=[rid] + update_cols` (and ideally pin the source snapshot, like `merge_into` already does) and add a regression test that calls `update_by_row_id(..., source="db.table", ...)`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
