leaves12138 commented on code in PR #8826:
URL: https://github.com/apache/paimon/pull/8826#discussion_r3729372480
##########
paimon-python/pypaimon/ray/data_evolution_merge_join.py:
##########
@@ -503,112 +765,248 @@ def distributed_update_apply(
precomputed_info_ref = ray.put(files_info)
frid_col = "_FIRST_ROW_ID"
+ range_col = "_ROW_ID_RANGE_START"
captured_sorted = sorted_first_row_ids
captured_sorted_arr = np.asarray(captured_sorted, dtype=np.int64)
valid_ranges = planner.valid_row_id_ranges
range_starts = np.asarray([r.from_ for r in valid_ranges], dtype=np.int64)
range_ends = np.asarray([r.to for r in valid_ranges], dtype=np.int64)
+ row_id_range_starts = (
+ _row_id_range_starts(files_info, rows_per_range)
+ if rows_per_range is not None else None
+ )
def _assign_frid(batch: pa.Table) -> pa.Table:
if batch.num_rows == 0:
- return batch.append_column(
- frid_col, pa.array([], type=pa.int64())
- )
+ result = batch.append_column(
+ frid_col, pa.array([], type=pa.int64()))
+ if row_id_range_starts is None:
+ return result
+ return result.append_column(
+ range_col, pa.array([], type=pa.int64()))
rid_col = batch.column(row_id_name)
if rid_col.null_count:
raise ValueError(
"_ROW_ID is null; planner snapshot is stale "
"or matched rows come from a different table."
)
rids = rid_col.to_numpy(zero_copy_only=False)
- # Check each row_id belongs to a valid range (vectorized).
in_range = np.zeros(len(rids), dtype=bool)
- for s, e in zip(range_starts, range_ends):
- in_range |= (rids >= s) & (rids <= e)
+ for start, end in zip(range_starts, range_ends):
+ in_range |= (rids >= start) & (rids <= end)
if not in_range.all():
bad = rids[~in_range][0]
raise ValueError(
f"_ROW_ID {bad} does not belong to any valid range "
f"{[f'[{r.from_}, {r.to}]' for r in valid_ranges]}; "
- f"planner snapshot is stale or matched rows come "
- f"from a different table."
+ "planner snapshot is stale or row ids come from another table."
)
- idx = np.searchsorted(
- captured_sorted_arr, rids, side="right"
- ) - 1
- frids = captured_sorted_arr[idx]
- return batch.append_column(
- frid_col, pa.array(frids, type=pa.int64())
+ indexes = np.searchsorted(
+ captured_sorted_arr, rids, side="right") - 1
+ result = batch.append_column(
+ frid_col,
+ pa.array(captured_sorted_arr[indexes], type=pa.int64()))
+ if row_id_range_starts is None:
+ return result
+ return result.append_column(
+ range_col,
+ pa.array(row_id_range_starts[indexes], type=pa.int64()),
)
- map_kwargs = _map_kwargs(ray_remote_args)
- with_frid = update_ds.map_batches(_assign_frid, **map_kwargs)
+ worker_map_kwargs = _map_kwargs(ray_remote_args)
+ if transform is not None:
+ with_frid = ray.data.from_arrow(pa.table({
+ frid_col: pa.array(captured_sorted, type=pa.int64()),
+ range_col: pa.array(row_id_range_starts, type=pa.int64()),
+ }))
+ else:
+ with_frid = update_ds.map_batches(
+ _assign_frid, **worker_map_kwargs)
captured_table = table
captured_cols = cols
+ capture_group_errors = on_group_result is not None
def _apply_group(group: pa.Table) -> pa.Table:
- if group.num_rows == 0:
- return pa.Table.from_pydict({
- "msgs_blob": pa.array([], type=pa.binary()),
- "n_updated": pa.array([], type=pa.int64()),
- "row_ids_blob": pa.array([], type=pa.binary()),
- })
-
- if (
- pc.count_distinct(group.column(row_id_name)).as_py()
- != group.num_rows
- ):
- raise ValueError(
- "MERGE matched multiple source rows to the same "
- "target _ROW_ID. Deduplicate the source before "
- "merging."
+ worker = None
+ try:
+ if group.num_rows == 0:
+ return pa.Table.from_pydict({
+ "msgs_blob": pa.array([], type=pa.binary()),
+ "n_updated": pa.array([], type=pa.int64()),
+ "row_ids_blob": pa.array([], type=pa.binary()),
+ "error": pa.array([], type=pa.string()),
+ })
+
+ if (pc.count_distinct(group.column(row_id_name)).as_py()
+ != group.num_rows):
+ raise ValueError(
+ "MERGE matched multiple source rows to the same "
+ "target _ROW_ID. Deduplicate the source before merging.")
+
+ routing_columns = [frid_col]
+ if row_id_range_starts is not None:
+ routing_columns.append(range_col)
+ for_update = group.drop_columns(routing_columns)
+ row_ids = (
+ for_update.column(row_id_name).to_pylist()
+ if collect_row_ids else []
+ )
+ worker = TableUpdateByRowId(
+ captured_table,
+ "_merge_into_shard_" + uuid.uuid4().hex[:8],
+ BATCH_COMMIT_IDENTIFIER,
+ _precomputed_files_info=ray.get(precomputed_info_ref),
+ )
+ if capture_group_errors:
+ return _write_group_result(
+ captured_table, worker, for_update, captured_cols, row_ids)
+ messages = worker.update_columns(for_update, captured_cols)
+ return _group_result(messages, for_update.num_rows, row_ids)
+ except Exception as error:
+ if capture_group_errors:
+ return _failed_group_result(captured_table, worker, error)
+ raise
+
+ group_col = frid_col
+ num_groups = len(captured_sorted)
+ if row_id_range_starts is not None:
+ group_col = range_col
+ num_groups = len(set(row_id_range_starts.tolist()))
+ group_partitions = max(1, min(num_groups, num_partitions))
+ grouped = with_frid.groupby(
+ group_col, num_partitions=group_partitions)
+ if transform is None:
+ msgs_ds = grouped.map_groups(_apply_group, **worker_map_kwargs)
+ else:
+ retry_options = ray_remote_args or {}
+ worker_options = {
+ "table": scan_table,
Review Comment:
This prevents the documented new-column backfill use case. `scan_table` is
pinned to `base_snapshot_id`, so after adding `embedding` its schema still
lacks that column. `update_by_transform(..., update_cols=["embedding"])` then
fails in `TableUpdateByRowId.update_columns` with `ValueError: Column embedding
not found in table schema`. Please use separate tables: the pinned table for
reading, and the latest-schema `table` for writing, while keeping the
precomputed base-snapshot file metadata for row-id routing/conflict checks.
##########
paimon-python/pypaimon/ray/update_by_transform.py:
##########
@@ -0,0 +1,304 @@
+# 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 transform updates on Ray for data-evolution tables."""
+
+import logging
+import uuid
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from pypaimon.ray.data_evolution_merge_into import (
+ _reraise_inner,
+ _require_ray_join,
+ _resolve_num_partitions,
+)
+from pypaimon.ray.data_evolution_merge_join import (
+ GroupApplyError,
+ distributed_update_apply,
+)
+from pypaimon.ray.data_evolution_merge_transform import build_update_schema
+from pypaimon.ray.update_by_row_id import (
+ _abort_pending_update_messages,
+ _blob_col_names,
+)
+
+__all__ = ["update_by_transform"]
+
+logger = logging.getLogger(__name__)
+
+
+def _positive_int(name, value):
+ if (isinstance(value, bool)
+ or not isinstance(value, int)
+ or value <= 0):
+ raise ValueError("{} must be a positive integer.".format(name))
+
+
+def _prepare_transform_update(
+ target,
+ catalog_options,
+ read_projection,
+ transform,
+ update_cols,
+ rows_per_commit,
+ transform_filter,
+ transform_batch_size):
+ from pypaimon.catalog.catalog_factory import CatalogFactory
+ from pypaimon.common.where_parser import parse_where_clause
+ from pypaimon.schema.data_types import PyarrowFieldParser
+ from pypaimon.table.special_fields import SpecialFields
+
+ _positive_int("rows_per_commit", rows_per_commit)
+ _positive_int("transform_batch_size", transform_batch_size)
+ if not callable(transform):
+ raise ValueError("transform must be callable.")
+ if not read_projection:
+ raise ValueError("read_projection must be non-empty.")
+ if not update_cols:
+ raise ValueError("update_cols must be non-empty.")
+ if (transform_filter is not None
+ and not isinstance(transform_filter, str)
+ and not callable(transform_filter)):
+ raise ValueError("filter must be a SQL expression string or callable.")
+
+ update_cols = list(dict.fromkeys(update_cols))
+ read_cols = list(dict.fromkeys(read_projection))
+ table = CatalogFactory.create(catalog_options).get_table(target)
+ if table.is_primary_key_table:
+ raise ValueError(
+ "update_by_transform requires a non-primary-key table.")
+ if not table.options.data_evolution_enabled():
+ raise ValueError(
+ f"update_by_transform requires 'data-evolution.enabled'='true' "
+ f"on '{target}'.")
+ if not table.options.row_tracking_enabled():
+ raise ValueError(
+ f"update_by_transform requires 'row-tracking.enabled'='true' "
+ f"on '{target}'.")
+ if table.options.deletion_vectors_enabled():
+ raise ValueError(
+ "update_by_transform does not support deletion-vectors-enabled "
+ f"tables yet: '{target}'.")
+
+ rid = SpecialFields.ROW_ID.name
+ if rid in read_cols:
+ raise ValueError(
+ "update_by_transform keeps _ROW_ID internal; remove it from "
+ "read_projection.")
+ unknown = [col for col in read_cols if col not in table.field_names]
+ if unknown:
+ raise ValueError(
+ f"read column {unknown[0]!r} is not in target '{target}'.")
+ 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:
+ raise ValueError(
+ f"update_by_transform cannot update blob column {col!r}.")
+ if col in partition_keys:
+ raise ValueError(
+ f"update_by_transform cannot update partition column {col!r}.")
+
+ target_pa = PyarrowFieldParser.from_paimon_schema(
+ table.table_schema.fields)
+ update_schema = build_update_schema(target_pa, update_cols, rid)
+ predicate = (
+ parse_where_clause(transform_filter, table.table_schema.fields)
+ if isinstance(transform_filter, str) else None
+ )
+ filter_fn = transform_filter if callable(transform_filter) else None
+ return table, read_cols, update_cols, update_schema, predicate, filter_fn
+
+
+def update_by_transform(
+ target: str,
+ catalog_options: Dict[str, str],
+ *,
+ read_projection: List[str],
+ transform: Callable,
+ update_cols: List[str],
+ rows_per_commit: int,
+ filter: Optional[Union[str, Callable]] = None,
+ num_partitions: Optional[int] = None,
+ ray_remote_args: Optional[Dict[str, Any]] = None,
+ transform_batch_size: int = 1024,
+) -> Dict[str, int]:
+ """Transform matching rows and commit continuous row-id ranges.
+
+ The transform receives ``read_projection`` and returns ``update_cols``
+ with the same row count and order. ``filter`` defaults to the full target.
+ Row ids remain internal to this operation.
+
+ Returns ``{"num_updated": <rows>}``.
+ """
+ _require_ray_join()
+ (table, read_cols, update_cols, update_schema,
+ predicate, filter_fn) = _prepare_transform_update(
+ target,
+ catalog_options,
+ read_projection,
+ transform,
+ update_cols,
+ rows_per_commit,
+ filter,
+ transform_batch_size,
+ )
+ num_partitions = _resolve_num_partitions(num_partitions)
+ base = table.snapshot_manager().get_latest_snapshot()
+ if base is None or base.total_record_count == 0:
+ return {"num_updated": 0}
+
+ retention_tag = "pypaimon-transform-update-{}".format(uuid.uuid4().hex)
+ table.create_tag(retention_tag, snapshot_id=base.id, time_retained="30d")
+ committer = _IncrementalUpdateCommitter(
+ table, base, table.table_schema.id)
+ try:
+ _, num_updated, _ = distributed_update_apply(
+ None,
+ table,
+ update_cols,
+ num_partitions=num_partitions,
+ ray_remote_args=ray_remote_args,
+ base_snapshot_id=base.id,
+ on_group_result=committer.add_range,
+ rows_per_range=rows_per_commit,
+ read_projection=read_cols,
+ transform=transform,
+ transform_filter=filter_fn,
+ transform_predicate=predicate,
+ transform_update_schema=update_schema,
+ transform_batch_size=transform_batch_size,
+ )
+ committer.finish()
+ except GroupApplyError:
+ committer.finish()
+ raise
+ except Exception as error:
+ if committer.failed:
+ raise
+ _reraise_inner(error)
+ raise
+ finally:
+ committer.close()
+ try:
+ table.delete_tag(retention_tag)
+ except Exception as error:
+ logger.warning(
+ "Failed to delete transform retention tag %s: %s",
+ retention_tag,
+ error,
+ exc_info=error,
+ )
+ return {"num_updated": num_updated}
+
+
+class _IncrementalUpdateCommitter:
+
+ def __init__(self, table, base_snapshot=None, planned_schema_id=None):
+ self._table = table
+ self._table_commit = None
+ self._snapshot_callback = None
+ self._commit_user = None
+ self._checkpoint_snapshot = base_snapshot
+ self._planned_schema_id = planned_schema_id
+ self._next_commit_identifier = 1
+ self._deferred_commit_error = None
+
+ @property
+ def failed(self) -> bool:
+ return self._deferred_commit_error is not None
+
+ def add_range(self, commit_messages, _num_updated, _row_ids) -> None:
+ if self.failed:
+ _abort_pending_update_messages(self._table, commit_messages)
+ return
+ try:
+ self._commit(commit_messages)
+ except Exception as error:
+ self._deferred_commit_error = error
+
+ def finish(self) -> None:
+ if self.failed:
+ raise self._deferred_commit_error
+
+ def _commit(self, messages) -> None:
+ if not messages:
+ return
+ if self._table_commit is None:
+ builder = self._table.new_stream_write_builder()
+ if self._checkpoint_snapshot is not None:
+ self._commit_user = builder.commit_user
+ self._table_commit = builder.new_commit()
+ self._snapshot_callback = _SnapshotCallback()
+ self._table_commit.add_commit_callback(self._snapshot_callback)
+
+ identifier = self._next_commit_identifier
+ if self._checkpoint_snapshot is not None:
+ self._table_commit.protect_from_external_rewrites(
Review Comment:
The external-rewrite guard does not protect transform dependencies. Conflict
detection only sees the staged `update_cols`, so an overlapping concurrent
row-id update to a `read_projection` or filter column is accepted. Repro: read
`name="a"`, derive `age=len(name)`, concurrently update `name` to `"long"`
before this commit; the operation succeeds with `name="long", age=1` instead of
`age=4`. Please include read/filter dependencies in the overlapping-row
conflict set, or conservatively reject all concurrent row updates during an
incremental transform.
--
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]
