[python][ray] Fix Ray 2.60 shuffle v2 compatibility

paimon-python/pypaimon/ray/data_evolution_merge_join.py

@@ -26,12 +26,14 @@
     _NormalizedClause,
     build_delete_schema,
     build_update_schema,
+    cast_to_schema,
 )
 from pypaimon.ray.partitioning import (
     _default_hash_shuffle_parallelism,
     _resolve_num_partitions,
     _resolve_row_id_num_partitions,
 )
+from pypaimon.ray.shuffle import _group_by_shuffle_key
 
 
 def _map_kwargs(
@@ -667,15 +669,18 @@ def distributed_update_apply(
     sentinel_row_id = -1
     captured_sorted = sorted_first_row_ids
 
+    from pypaimon.schema.data_types import PyarrowFieldParser
+    target_pa = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
+    update_schema = build_update_schema(target_pa, cols, row_id_name)
+
     def _assign_frid(batch: pa.Table) -> pa.Table:
+        # Ray's shuffle reducer concatenates these blocks with the sentinel.
+        batch = cast_to_schema(batch.select(update_schema.names), update_schema)
         return batch.append_column(frid_col, planner.route(batch.column(row_id_name)))
 
     map_kwargs = _map_kwargs(ray_remote_args)
     with_frid = update_ds.map_batches(_assign_frid, **map_kwargs)
 
-    from pypaimon.schema.data_types import PyarrowFieldParser
-    target_pa = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
-    update_schema = build_update_schema(target_pa, cols, row_id_name)
     # Ray drops the schema of an empty shuffle input. A negative row id keeps
     # the group-by schema in a separate group that is never written.
     sentinel = pa.Table.from_arrays(
@@ -734,8 +739,8 @@ def _apply_group(group: pa.Table) -> pa.Table:
     group_partitions = max(
         1, min(len(captured_sorted), num_partitions)
     )
-    msgs_ds = with_frid.groupby(
-        frid_col, num_partitions=group_partitions
+    msgs_ds = _group_by_shuffle_key(
+        with_frid, frid_col, num_partitions=group_partitions
     ).map_groups(_apply_group, **map_kwargs)
 
     all_msgs: list = []
@@ -829,7 +834,7 @@ def _read_group(group: pa.Table) -> pa.Table:
     map_kwargs = _map_kwargs(ray_remote_args)
     with_frid = row_ids_ds.map_batches(_assign_frid, **map_kwargs)
     group_partitions = max(1, min(len(sorted_first_row_ids), num_partitions))
-    return with_frid.groupby(frid_col, num_partitions=group_partitions).map_groups(
+    return _group_by_shuffle_key(with_frid, frid_col, num_partitions=group_partitions).map_groups(
         _read_group, **map_kwargs
     )
 

paimon-python/pypaimon/ray/shuffle.py

@@ -31,6 +31,7 @@
 """
 
 import uuid
+from functools import partial
 from typing import TYPE_CHECKING, List
 
 import pyarrow as pa
@@ -144,7 +145,45 @@ def _group_by_partition_bucket(
         bucket_udf, batch_format="pyarrow", zero_copy_batch=True,
     )
     group_keys: List[str] = partition_keys + [bucket_col]
-    return ds_with_bucket.groupby(group_keys), bucket_col
+    return _group_by_shuffle_key(ds_with_bucket, group_keys), bucket_col
+
+
+def _group_by_shuffle_key(dataset, keys, **kwargs):
+    """Group routed rows while tolerating schema-less shards in Ray 2.60.0."""
+    context = getattr(dataset, "context", None)
+    if getattr(context, "shuffle_strategy", None) == "shuffle_v2":
+        import ray
+
+        if ray.__version__ == "2.60.0":
+            from ray.data._internal.execution.execution_callback import ExecutionCallback
+            from ray.data._internal.execution.operators.shuffle_operators.shuffle_reduce_operator import (
+                ShuffleReduceOp,
+            )
+
+            group_keys = (keys,) if isinstance(keys, str) else tuple(keys)
+            reduce_prefix = "HashShuffleReduce(keys={}, ".format(group_keys)
+
+            class IgnoreSchemalessShuffleShards(ExecutionCallback):
+                def before_execution_starts(self, executor):
+                    for operator in executor._topology:
+                        if isinstance(operator, ShuffleReduceOp) and operator.name.startswith(reduce_prefix):
+                            reduce = operator._reduce_fn
+                            if getattr(reduce, "__qualname__", None) == "_sort_reduce.<locals>._reduce":
+                                operator._reduce_fn = partial(_reduce_without_schemaless_shards, reduce)
+
+            # Derived datasets otherwise share their parent's DataContext.
+            dataset = ray.data.Dataset.copy(dataset, _deep_copy=True)
+            dataset.context.custom_execution_callback_classes.append(IgnoreSchemalessShuffleShards)
+    return dataset.groupby(keys, **kwargs)
+
+
+def _reduce_without_schemaless_shards(reduce, partition_id, tables_by_input):
+    # A map with no batches may emit a zero-column placeholder. Preserve
+    # typed empties so the sorted reducer can still see the grouping keys.
+    return reduce(partition_id, [
+        [table for table in tables if table.num_columns or table.num_rows]
+        for tables in tables_by_input
+    ])
 
 
 def _identity_batch(batch: pa.Table) -> pa.Table: