This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new ba6a122499 [python] Support callable assignments in Ray self-merge
(#9327)
ba6a122499 is described below
commit ba6a122499e21efbe964c62a2923ac4c9146121b
Author: XiaoHongbo <[email protected]>
AuthorDate: Fri Aug 21 13:10:13 2026 +0800
[python] Support callable assignments in Ray self-merge (#9327)
---
docs/docs/pypaimon/ray-data.md | 26 +++
.../pypaimon/ray/data_evolution_merge_into.py | 43 +++-
.../pypaimon/ray/data_evolution_merge_join.py | 17 ++
.../pypaimon/ray/data_evolution_merge_transform.py | 30 ++-
.../tests/ray_data_evolution_merge_into_test.py | 248 +++++++++++++++++++++
5 files changed, 359 insertions(+), 5 deletions(-)
diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md
index 3f90068336..1f9194c92c 100644
--- a/docs/docs/pypaimon/ray-data.md
+++ b/docs/docs/pypaimon/ray-data.md
@@ -513,11 +513,37 @@ extra. Install the extra before using conditions: `pip
install pypaimon[sql]`.
]
```
+For self-merge (`source == target` and `on=["_ROW_ID"]`), update values may
+also be callables. A callable receives the matched `read_columns` plus
+`_ROW_ID` as a `pyarrow.Table` and must return one `pyarrow.Array` or
+`pyarrow.ChunkedArray` value per input row:
+
+```python
+import pyarrow.compute as pc
+
+merge_into(
+ target="db.table",
+ source="db.table",
+ catalog_options=catalog_options,
+ on=["_ROW_ID"],
+ read_columns=["age"],
+ when_matched=[WhenMatched.update({
+ "age": lambda rows: pc.add(rows["age"], 1),
+ }, condition="t.id IN (1, 3)")],
+)
+```
+
+Callables may run zero, one, or multiple times and must be deterministic,
+side-effect-free, and row-local. They are not supported for general
+source-target merges.
+
**Parameters:**
- `source`: a `ray.data.Dataset`, `pyarrow.Table`, `pandas.DataFrame`, or a
Paimon table identifier string. When a string is passed, it reads the table
from the same `catalog_options` at the latest snapshot.
- `on`: key columns, or `{target_col: source_col}` for renamed keys.
+- `read_columns`: columns passed to callable self-merge assignments. Required
+ when an update mapping contains a callable; otherwise it must be omitted.
- `num_partitions`: shuffle parallelism for the join and the write; defaults to
`max(1, cluster_cpus * 2)`. Raise it for large merges on big clusters.
- `ray_remote_args`: Ray remote options applied to the merge's map/group
diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_into.py
b/paimon-python/pypaimon/ray/data_evolution_merge_into.py
index 8695c07d77..3f8e3f3d16 100644
--- a/paimon-python/pypaimon/ray/data_evolution_merge_into.py
+++ b/paimon-python/pypaimon/ray/data_evolution_merge_into.py
@@ -65,6 +65,7 @@ class _PrepareCtx:
catalog_options: Dict[str, str]
is_self_merge: bool = False
self_merge_scan_predicate: Optional[Predicate] = None
+ read_columns: Tuple[str, ...] = ()
def merge_into(
@@ -78,6 +79,7 @@ def merge_into(
num_partitions: Optional[int] = None,
ray_remote_args: Optional[Dict[str, Any]] = None,
concurrency: Optional[int] = None,
+ read_columns: Optional[Sequence[str]] = None,
) -> Dict[str, int]:
_require_ray_join()
num_partitions = _resolve_num_partitions(num_partitions)
@@ -85,6 +87,7 @@ def merge_into(
table, source_ds, matched_specs, not_matched_specs, ctx = _prepare(
target, source, catalog_options,
list(when_matched), list(when_not_matched), on,
+ read_columns,
)
base_snapshot = table.snapshot_manager().get_latest_snapshot()
@@ -100,7 +103,10 @@ def merge_into(
)
-def _prepare(target, source, catalog_options, when_matched, when_not_matched,
on):
+def _prepare(
+ target, source, catalog_options, when_matched, when_not_matched, on,
+ read_columns=None,
+):
if not when_matched and not when_not_matched:
raise ValueError(
"At least one of when_matched or when_not_matched must be
non-empty."
@@ -139,12 +145,18 @@ def _prepare(target, source, catalog_options,
when_matched, when_not_matched, on
full_target_field_names = list(table.field_names)
settable_field_names = list(full_target_field_names)
on_map = dict(zip(target_on_cols, source_on_cols))
+ is_self_merge = _is_self_merge(
+ target, source, target_on_cols, source_on_cols
+ )
matched_specs = []
for c in when_matched:
spec = {}
if not c.delete:
spec = _normalize_set_spec(
- c.update, settable_field_names, on_map,
+ c.update,
+ settable_field_names,
+ on_map,
+ allow_callables=is_self_merge,
)
matched_specs.append(
_NormalizedClause(
@@ -193,13 +205,29 @@ def _prepare(target, source, catalog_options,
when_matched, when_not_matched, on
_NormalizedClause(spec=spec, condition=c.condition)
)
- is_self_merge = _is_self_merge(target, source, target_on_cols,
source_on_cols)
if is_self_merge and not_matched_specs:
raise ValueError(
"Self-merge (source == target with ON _ROW_ID) does not "
"support WHEN NOT MATCHED clauses."
)
+ read_columns = tuple(dict.fromkeys(read_columns or ()))
+ has_callable = any(
+ callable(value) and not isinstance(value, type)
+ for clause in matched_specs
+ for value in clause.spec.values()
+ )
+ if read_columns and not has_callable:
+ raise ValueError("read_columns requires a callable SET value.")
+ if has_callable:
+ if not read_columns:
+ raise ValueError("Callable SET values require read_columns.")
+ for col in read_columns:
+ if col not in full_target_field_names:
+ raise ValueError(
+ f"Read column {col!r} is not in target '{target}'."
+ )
+
if is_self_merge:
source_ds = None
source_col_names = set(full_target_field_names) | set(source_on_cols)
@@ -280,6 +308,7 @@ def _prepare(target, source, catalog_options, when_matched,
when_not_matched, on
catalog_options=catalog_options,
is_self_merge=is_self_merge,
self_merge_scan_predicate=self_merge_scan_predicate,
+ read_columns=read_columns,
)
return table, source_ds, matched_specs, not_matched_specs, ctx
@@ -321,6 +350,7 @@ def _build_datasets(
resolve_target_projection=_resolve_target_projection,
snapshot_id=base_snapshot_id,
scan_predicate=ctx.self_merge_scan_predicate,
+ read_columns=ctx.read_columns,
ray_remote_args=ray_remote_args,
)
if any(c.delete for c in matched_specs):
@@ -588,6 +618,7 @@ def _normalize_set_spec(
target_field_names: Sequence[str],
on_map: Optional[Mapping[str, str]] = None,
allow_target_refs: bool = True,
+ allow_callables: bool = False,
) -> Dict[str, Any]:
on_map = on_map or {}
if spec == "*":
@@ -610,9 +641,13 @@ def _normalize_set_spec(
result: Dict[str, Any] = {}
for key, val in spec.items():
if callable(val) and not isinstance(val, type):
+ if allow_callables:
+ result[key] = val
+ continue
raise TypeError(
"SET values must be source_col(), target_col(), "
- "lit(), or literals, not callables"
+ "lit(), or literals; callables are only supported "
+ "for self-merge"
)
if isinstance(val, SourceColumnRef):
result[key] = val
diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
index 863dfd5984..d7e381adc4 100644
--- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
+++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
@@ -73,6 +73,7 @@ def _build_matched_transform(
update_cols: List[str],
row_id_name: str,
update_schema: pa.Schema,
+ callable_input_columns: Optional[Sequence[str]] = None,
):
prepared_clauses = []
for clause in clauses:
@@ -105,10 +106,20 @@ def _build_matched_transform(
if matched.num_rows == 0:
continue
if not is_delete:
+ callable_input = None
+ if (callable_input_columns is not None
+ and any(callable(value) and not isinstance(value, type)
+ for value in spec.values())):
+ callable_input = pa.Table.from_arrays(
+ [matched.column(f"t.{col}")
+ for col in callable_input_columns],
+ names=list(callable_input_columns),
+ )
parts.append(vectorized_matched_transform(
matched, spec, on_pairs,
update_cols, row_id_name,
update_schema,
+ callable_input=callable_input,
))
if rewritten is not None and matched.num_rows < remaining.num_rows:
not_cond = f"COALESCE(NOT ({rewritten}), TRUE)"
@@ -189,6 +200,7 @@ def build_self_merge_update_ds(
resolve_target_projection,
snapshot_id: Optional[int] = None,
scan_predicate=None,
+ read_columns: Sequence[str] = (),
ray_remote_args: Optional[Dict[str, Any]] = None,
) -> Tuple:
from pypaimon.ray.ray_paimon import read_paimon
@@ -198,6 +210,7 @@ def build_self_merge_update_ds(
needed_cols = set(resolve_target_projection(
clauses, [row_id_name], update_cols, target_field_names,
))
+ needed_cols.update(read_columns)
for clause in clauses:
for value in clause.spec.values():
if isinstance(value, SourceColumnRef):
@@ -261,6 +274,10 @@ def build_self_merge_update_ds(
update_cols=list(update_cols),
row_id_name=row_id_name,
update_schema=update_schema,
+ callable_input_columns=(
+ list(read_columns) + [row_id_name]
+ if read_columns else None
+ ),
)
return aliased.map_batches(_transform, **_map_kwargs(ray_remote_args))
diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_transform.py
b/paimon-python/pypaimon/ray/data_evolution_merge_transform.py
index 80ae82732b..ad29368e51 100644
--- a/paimon-python/pypaimon/ray/data_evolution_merge_transform.py
+++ b/paimon-python/pypaimon/ray/data_evolution_merge_transform.py
@@ -101,6 +101,7 @@ def vectorized_matched_transform(
update_cols: Sequence[str],
row_id_name: str,
update_schema: pa.Schema,
+ callable_input: Optional[pa.Table] = None,
) -> pa.Table:
available = set(batch.schema.names)
arrays: list = [batch.column(f"t.{row_id_name}")]
@@ -109,7 +110,8 @@ def vectorized_matched_transform(
if col in spec:
arrays.append(
_resolve_spec_array(
- spec[col], batch, available, on_pairs, out_type
+ spec[col], batch, available, on_pairs, out_type,
+ callable_input=callable_input,
)
)
else:
@@ -173,7 +175,33 @@ def _resolve_spec_array(
available: set,
on_pairs: Sequence[Tuple[str, str]],
out_type: pa.DataType,
+ callable_input: Optional[pa.Table] = None,
):
+ if callable(val) and not isinstance(val, type):
+ if callable_input is None:
+ raise TypeError(
+ "Callable SET values are only supported for self-merge."
+ )
+ result = val(callable_input)
+ if not isinstance(result, (pa.Array, pa.ChunkedArray)):
+ raise ValueError(
+ "Callable SET values must return a pyarrow.Array or "
+ "pyarrow.ChunkedArray."
+ )
+ if len(result) != batch.num_rows:
+ raise ValueError(
+ "Callable SET result length must match matched row count: "
+ f"{len(result)} != {batch.num_rows}."
+ )
+ if result.type != out_type:
+ if isinstance(result, pa.ChunkedArray):
+ result = pa.chunked_array(
+ [chunk.cast(out_type) for chunk in result.chunks],
+ type=out_type,
+ )
+ else:
+ result = result.cast(out_type)
+ return result
if isinstance(val, LiteralValue):
return pa.array([val.value] * batch.num_rows, type=out_type)
if isinstance(val, SourceColumnRef):
diff --git a/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py
b/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py
index 70a0a0d189..e8a66a6061 100644
--- a/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py
+++ b/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py
@@ -25,6 +25,7 @@ import uuid
from unittest.mock import Mock, patch
import pyarrow as pa
+import pyarrow.compute as pc
import ray
from pypaimon import CatalogFactory, Schema
@@ -2037,6 +2038,253 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
self.assertEqual(out['age'], [99, 99, 99])
self.assertEqual(out['name'], ['a', 'b', 'c'])
+ @unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
+ def test_self_merge_callable_assignment(self):
+ from pypaimon.ray.ray_paimon import read_paimon as real_read_paimon
+
+ target = self._create_table()
+ self._write(
+ target,
+ pa.Table.from_pydict(
+ {
+ 'id': pa.array([1, 2, 3], type=pa.int32()),
+ 'name': ['a', 'b', 'c'],
+ 'age': pa.array([10, 20, 30], type=pa.int32()),
+ },
+ schema=self.pa_schema,
+ ),
+ )
+
+ def increment_age(rows):
+ if rows.column_names != ['age', '_ROW_ID']:
+ raise AssertionError(rows.column_names)
+ return pc.add(rows['age'], 1)
+
+ with patch(
+ 'pypaimon.ray.ray_paimon.read_paimon',
+ wraps=real_read_paimon,
+ ) as mock_read:
+ result = merge_into(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ read_columns=['age'],
+ when_matched=[WhenMatched.update(
+ {
+ 'age': increment_age,
+ 'name': lit('updated'),
+ },
+ condition='t.id IN (1, 3)',
+ )],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
+
+ self.assertEqual(result['num_matched'], 2)
+ self.assertEqual(
+ self._read_sorted(target),
+ {
+ 'id': [1, 2, 3],
+ 'name': ['updated', 'b', 'updated'],
+ 'age': [11, 20, 31],
+ },
+ )
+ self.assertEqual(
+ mock_read.call_args[1]['projection'], ['_ROW_ID', 'id', 'age']
+ )
+
+ @unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
+ def test_self_merge_callable_updates_variant_path(self):
+ from pypaimon.data.generic_variant import GenericVariant
+ from pypaimon.data.variant_path import variant_get, variant_replace
+
+ variant_type = pa.struct([
+ pa.field('value', pa.binary(), nullable=False),
+ pa.field('metadata', pa.binary(), nullable=False),
+ ])
+ pa_schema = pa.schema([
+ ('id', pa.int32()),
+ ('payload', variant_type),
+ ])
+ target = f'default.tbl_{uuid.uuid4().hex[:8]}'
+ self.catalog.create_table(
+ target,
+ Schema.from_pyarrow_schema(pa_schema, options=self.de_options),
+ False,
+ )
+ payload = GenericVariant.to_arrow_array([
+ GenericVariant.from_python({'value': 1.5}),
+ GenericVariant.from_python({'value': 2.5}),
+ ])
+ self._write(target, pa.table({
+ 'id': pa.array([1, 2], type=pa.int32()),
+ 'payload': payload,
+ }, schema=pa_schema))
+
+ def negate_value(rows):
+ values = variant_get(
+ rows['payload'], '$.value', pa.float64()
+ )
+ return variant_replace(
+ rows['payload'], '$.value', pc.negate(values), strict=True
+ )
+
+ result = merge_into(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ read_columns=['payload'],
+ when_matched=[WhenMatched.update(
+ {'payload': negate_value}, condition='t.id = 2',
+ )],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
+
+ self.assertEqual(result['num_matched'], 1)
+ output = self._read_sorted(target)['payload']
+ decoded = [
+ GenericVariant.from_arrow_struct(value).to_python()
+ for value in output
+ ]
+ self.assertEqual(decoded, [{'value': 1.5}, {'value': -2.5}])
+
+ @unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
+ def test_self_merge_no_match_does_not_invoke_callable(self):
+ target = self._create_table()
+ self._write(target, self._source())
+
+ def should_not_run(_rows):
+ raise AssertionError("Callable must not run without matched rows")
+
+ result = merge_into(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ read_columns=['age'],
+ when_matched=[WhenMatched.update(
+ {'age': should_not_run}, condition='t.id = 99',
+ )],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
+
+ self.assertEqual(result['num_matched'], 0)
+ self.assertEqual(self._read_sorted(target)['age'], [10])
+
+ def test_self_merge_callable_validation(self):
+ target = self._create_table()
+
+ with self.assertRaisesRegex(
+ ValueError, 'Callable SET values require read_columns'):
+ merge_into(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ when_matched=[WhenMatched.update({
+ 'age': lambda rows: rows['age'],
+ })],
+ )
+
+ with self.assertRaisesRegex(
+ ValueError, 'read_columns requires a callable SET value'):
+ merge_into(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ read_columns=['age'],
+ when_matched=[WhenMatched.update({'age': lit(99)})],
+ )
+
+ with self.assertRaisesRegex(
+ ValueError, "Read column 'missing' is not in target"):
+ merge_into(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ read_columns=['missing'],
+ when_matched=[WhenMatched.update({
+ 'age': lambda rows: rows['missing'],
+ })],
+ )
+
+ source = self._create_table()
+ with self.assertRaisesRegex(
+ TypeError, 'callables are only supported for self-merge'):
+ merge_into(
+ target=target,
+ source=source,
+ catalog_options=self.catalog_options,
+ on=['id'],
+ read_columns=['age'],
+ when_matched=[WhenMatched.update({
+ 'age': lambda rows: rows['age'],
+ })],
+ )
+
+ def test_self_merge_callable_rejects_invalid_result(self):
+ from pypaimon.ray.data_evolution_merge_transform import (
+ _resolve_spec_array,
+ )
+
+ batch = pa.table({'t.age': pa.array([10], type=pa.int32())})
+ callable_input = pa.table({
+ 'age': pa.array([10], type=pa.int32()),
+ })
+
+ with self.assertRaisesRegex(
+ ValueError, 'must return a pyarrow.Array'):
+ _resolve_spec_array(
+ lambda rows: rows['age'].to_pylist(),
+ batch,
+ set(batch.column_names),
+ [],
+ pa.int32(),
+ callable_input=callable_input,
+ )
+
+ with self.assertRaisesRegex(
+ ValueError, 'length must match matched row count'):
+ _resolve_spec_array(
+ lambda _rows: pa.array([], type=pa.int32()),
+ batch,
+ set(batch.column_names),
+ [],
+ pa.int32(),
+ callable_input=callable_input,
+ )
+
+ def test_self_merge_callable_preserves_chunked_result(self):
+ from pypaimon.ray.data_evolution_merge_transform import (
+ _resolve_spec_array,
+ )
+
+ batch = pa.table({'t.age': pa.array([10, 20], type=pa.int32())})
+ callable_input = pa.table({
+ 'age': pa.array([10, 20], type=pa.int32()),
+ })
+
+ def chunked_result(_rows):
+ return pa.chunked_array([
+ pa.array([10], type=pa.int64()),
+ pa.array([20], type=pa.int64()),
+ ])
+
+ result = _resolve_spec_array(
+ chunked_result,
+ batch,
+ set(batch.column_names),
+ [],
+ pa.int32(),
+ callable_input=callable_input,
+ )
+ self.assertIsInstance(result, pa.ChunkedArray)
+ self.assertEqual(result.num_chunks, 2)
+ self.assertEqual(result.type, pa.int32())
+
@unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
def test_self_merge_condition_pushes_down_predicate(self):
from pypaimon.common.options.core_options import (
