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 cd290a22ac [python] Prune source columns for Ray data-evolution merge
into (#8504)
cd290a22ac is described below
commit cd290a22acb42118fb37011e53f09fce6eb3f3fb
Author: umi <[email protected]>
AuthorDate: Thu Jul 9 13:08:37 2026 +0800
[python] Prune source columns for Ray data-evolution merge into (#8504)
---
.../pypaimon/ray/data_evolution_merge_into.py | 16 ++-
.../pypaimon/ray/data_evolution_merge_join.py | 32 ++++-
.../tests/ray_data_evolution_merge_into_test.py | 137 +++++++++++++++++++++
3 files changed, 179 insertions(+), 6 deletions(-)
diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_into.py
b/paimon-python/pypaimon/ray/data_evolution_merge_into.py
index a49b049af6..d43ac21f36 100644
--- a/paimon-python/pypaimon/ray/data_evolution_merge_into.py
+++ b/paimon-python/pypaimon/ray/data_evolution_merge_into.py
@@ -26,6 +26,7 @@ import pyarrow as pa
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.ray.data_evolution_merge_join import (
+ _resolve_source_projection,
build_matched_delete_ds,
build_matched_update_ds,
build_not_matched_insert_ds,
@@ -202,16 +203,20 @@ def _prepare(target, source, catalog_options,
when_matched, when_not_matched, on
source_col_names = set(full_target_field_names) | set(source_on_cols)
else:
source_snapshot_id = None
+ source_read_projection = None
if isinstance(source, str):
- source_snapshot = (
- catalog.get_table(source)
- .snapshot_manager()
- .get_latest_snapshot()
+ source_table = catalog.get_table(source)
+ source_read_projection = _resolve_source_projection(
+ matched_specs + not_matched_specs,
+ source_on_cols,
+ source_table.field_names,
)
+ source_snapshot =
source_table.snapshot_manager().get_latest_snapshot()
if source_snapshot is not None:
source_snapshot_id = source_snapshot.id
source_ds = _normalize_source(
source, catalog_options, source_snapshot_id=source_snapshot_id,
+ projection=source_read_projection,
)
_validate_source_on_cols(source_ds, source_on_cols)
source_col_names = set(_source_schema_or_raise(source_ds).names)
@@ -660,6 +665,7 @@ def _normalize_source(
source: Any,
catalog_options: Dict[str, str],
source_snapshot_id: Optional[int] = None,
+ projection: Optional[List[str]] = None,
):
import ray.data
@@ -670,6 +676,8 @@ def _normalize_source(
read_kwargs = {}
if source_snapshot_id is not None:
read_kwargs["snapshot_id"] = source_snapshot_id
+ if projection is not None:
+ read_kwargs["projection"] = projection
return read_paimon(source, catalog_options, **read_kwargs)
if isinstance(source, pa.Table):
return ray.data.from_arrow(source)
diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
index 5535b7ec1b..359af8e6ef 100644
--- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
+++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
@@ -44,6 +44,28 @@ def _map_kwargs(
return kwargs
+def _resolve_source_projection(
+ clauses: List[_NormalizedClause],
+ source_on: Sequence[str],
+ source_field_names: Sequence[str],
+) -> list:
+ needed = set(source_on)
+ source_set = set(source_field_names)
+
+ for clause in clauses:
+ for value in clause.spec.values():
+ if isinstance(value, SourceColumnRef):
+ needed.add(value.column)
+ if clause.condition is not None:
+ from pypaimon.ray.merge_condition import extract_columns
+ for ref in extract_columns(clause.condition):
+ prefix, col = ref.split(".", 1)
+ if prefix == "s" and col in source_set:
+ needed.add(col)
+
+ return [c for c in source_field_names if c in needed]
+
+
def _build_matched_transform(
clauses: List[_NormalizedClause],
on_map: Dict[str, str],
@@ -330,7 +352,10 @@ def build_matched_update_ds(
target_renamed = target_ds.rename_columns(
{c: f"t.{c}" for c in target_ds.schema().names}
)
- source_cols = list(source_ds.schema().names)
+ source_cols = _resolve_source_projection(
+ clauses, source_on, source_ds.schema().names,
+ )
+ source_ds = source_ds.select_columns(source_cols)
source_renamed = source_ds.rename_columns(
{c: f"s.{c}" for c in source_cols}
)
@@ -880,7 +905,10 @@ def build_not_matched_insert_ds(
captured_field_names = list(target_field_names)
out_schema = target_pa_schema
- source_cols = list(source_ds.schema().names)
+ source_cols = _resolve_source_projection(
+ clauses, source_on, source_ds.schema().names,
+ )
+ source_ds = source_ds.select_columns(source_cols)
source_renamed = source_ds.rename_columns(
{c: f"s.{c}" for c in source_cols}
)
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 36f10d1aff..cddb11ac37 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
@@ -140,6 +140,7 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
source,
self.catalog_options,
snapshot_id=expected_snapshot_id,
+ projection=['id', 'name', 'age'],
)
def test_no_clause_raises(self):
@@ -2372,6 +2373,142 @@ class TargetProjectionTest(unittest.TestCase):
self.assertIn('age', cols)
self.assertIn('id', cols)
+ def test_matched_source_projection_prunes_unneeded_cols(self):
+ from pypaimon.ray.data_evolution_merge_join import (
+ _resolve_source_projection,
+ )
+ from pypaimon.ray.data_evolution_merge_transform import (
+ LiteralValue,
+ SourceColumnRef,
+ TargetColumnRef,
+ )
+
+ cols = _resolve_source_projection(
+ [
+ self._clause(
+ {
+ 'age': SourceColumnRef('id'),
+ 'name': TargetColumnRef('name'),
+ 'note': LiteralValue('literal'),
+ },
+ condition="s.status = 't.fake' AND s.score > t.score",
+ )
+ ],
+ ['uid'],
+ ['uid', 'id', 'name', 'status', 'score', 'payload'],
+ )
+ self.assertEqual(['uid', 'id', 'status', 'score'], cols)
+
+ def test_literal_update_source_projection_keeps_only_join_key(self):
+ from pypaimon.ray.data_evolution_merge_join import (
+ _resolve_source_projection,
+ )
+ from pypaimon.ray.data_evolution_merge_transform import LiteralValue
+
+ cols = _resolve_source_projection(
+ [self._clause({'name': LiteralValue('updated')})],
+ ['id'],
+ ['id', 'name', 'age', 'payload'],
+ )
+ self.assertEqual(['id'], cols)
+
+ def test_matched_update_selects_needed_source_cols(self):
+ from pypaimon.ray.data_evolution_merge_join import
build_matched_update_ds
+ from pypaimon.ray.data_evolution_merge_transform import SourceColumnRef
+
+ source_ds = Mock()
+ source_ds.schema.return_value = pa.schema([
+ ('id', pa.int32()),
+ ('name', pa.string()),
+ ('payload', pa.string()),
+ ])
+ selected_ds = Mock()
+ source_renamed = Mock()
+ source_ds.select_columns.return_value = selected_ds
+ selected_ds.rename_columns.return_value = source_renamed
+
+ target_ds = Mock()
+ target_ds.schema.return_value = pa.schema([
+ ('_ROW_ID', pa.int64()),
+ ('id', pa.int32()),
+ ])
+ target_renamed = Mock()
+ joined = Mock()
+ result = object()
+ target_ds.rename_columns.return_value = target_renamed
+ target_renamed.join.return_value = joined
+ joined.map_batches.return_value = result
+
+ with patch(
+ 'pypaimon.ray.ray_paimon.read_paimon',
+ return_value=target_ds,
+ ):
+ out = build_matched_update_ds(
+ target_identifier='default.target',
+ source_ds=source_ds,
+ target_on=['id'],
+ source_on=['id'],
+ clauses=[self._clause({'name': SourceColumnRef('name')})],
+ target_field_names=['id', 'name'],
+ target_pa_schema=pa.schema([
+ ('id', pa.int32()),
+ ('name', pa.string()),
+ ]),
+ update_cols=['name'],
+ catalog_options={'warehouse': '/tmp/warehouse'},
+ num_partitions=1,
+ resolve_target_projection=lambda *args: ['id'],
+ )
+
+ self.assertIs(out, result)
+ source_ds.select_columns.assert_called_once_with(['id', 'name'])
+ selected_ds.rename_columns.assert_called_once_with({
+ 'id': 's.id',
+ 'name': 's.name',
+ })
+
+ def test_not_matched_insert_selects_needed_source_cols(self):
+ from pypaimon.ray.data_evolution_merge_join import (
+ build_not_matched_insert_ds,
+ )
+ from pypaimon.ray.data_evolution_merge_transform import SourceColumnRef
+
+ source_ds = Mock()
+ source_ds.schema.return_value = pa.schema([
+ ('id', pa.int32()),
+ ('name', pa.string()),
+ ('payload', pa.string()),
+ ])
+ selected_ds = Mock()
+ source_renamed = Mock()
+ result = object()
+ source_ds.select_columns.return_value = selected_ds
+ selected_ds.rename_columns.return_value = source_renamed
+ source_renamed.map_batches.return_value = result
+
+ out = build_not_matched_insert_ds(
+ target_identifier='default.target',
+ source_ds=source_ds,
+ target_on=['id'],
+ source_on=['id'],
+ clauses=[self._clause({'name': SourceColumnRef('name')})],
+ target_field_names=['id', 'name'],
+ target_pa_schema=pa.schema([
+ ('id', pa.int32()),
+ ('name', pa.string()),
+ ]),
+ catalog_options={'warehouse': '/tmp/warehouse'},
+ num_partitions=1,
+ target_empty=True,
+ )
+
+ self.assertIs(out, result)
+ source_ds.select_columns.assert_called_once_with(['id', 'name'])
+ selected_ds.rename_columns.assert_called_once_with({
+ 'id': 's.id',
+ 'name': 's.name',
+ })
+
class MergeConditionUnitTest(unittest.TestCase):
