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 8871cc4adf [python][ray] Adapt default partitions to input size (#9357)
8871cc4adf is described below
commit 8871cc4adfa80d5b0258362693eae3cebe4b7264
Author: XiaoHongbo <[email protected]>
AuthorDate: Thu Aug 27 14:23:56 2026 +0800
[python][ray] Adapt default partitions to input size (#9357)
## What changed
- Preserve an explicitly configured `num_partitions`.
- Size general MERGE from reliable in-memory source metadata. Unknown
inputs, including Paimon scans, use Ray's hash-shuffle default; a
nonempty target keeps it as a lower bound.
- Read block sizing and shuffle defaults from the source Dataset's
sealed `DataContext`.
- Handle Ray's property, callable, and boolean cardinality metadata
while keeping legacy `MapBatches` conservative.
- Treat never-written and truncated targets as empty, avoid target
joins, and repartition inserts to the resolved write parallelism.
- For row-ID update/read, combine reliable source metadata with
target-file fan-out.
- Keep self-merge on its previous CPU-based default.
- Avoid a separate target scan for partition estimation.
## Why
The previous default created 640 shuffle partitions on a 320-CPU cluster
even for small inputs. This follows Ray's own hash-shuffle default for
unknown sizes while retaining size-based scaling when in-memory metadata
is reliable.
Paimon split sizes are compressed file bytes, not Ray block memory
bytes, so they are not used for partition sizing.
## Validation
- Real Paimon source MERGE with highly compressed data: 200 partitions
instead of 640 or 1 on a mocked 320-CPU cluster.
- Truncated target with explicit parallelism: three materialized write
blocks and 30 inserted rows.
- `read_paimon(filter) -> map_batches -> update_by_row_id(None)`: 39
matching rows across 201 target files request 200 instead of 640.
- Dataset sealed at 128 MiB blocks still sizes a 512 MiB input to four
after the global context changes to 1 GiB.
- Cardinality compatibility verified on Ray 2.44, 2.50.1, 2.53, 2.54,
and 2.57; Ray 2.44 and 2.53 are covered by CI.
- Latest-master merge tree: 55 partition/update/read tests and 123
non-BLOB MERGE tests passed, with 18 subtests.
- Python 3.6 `py_compile`, flake8, and `git diff --check` passed.
---
.github/workflows/paimon-python-checks.yml | 6 +
docs/docs/pypaimon/ray-data.md | 17 +-
.../pypaimon/ray/data_evolution_merge_into.py | 113 +++++--
.../pypaimon/ray/data_evolution_merge_join.py | 65 ++++-
paimon-python/pypaimon/ray/partitioning.py | 177 +++++++++++
paimon-python/pypaimon/ray/read_by_row_id.py | 17 +-
paimon-python/pypaimon/ray/update_by_row_id.py | 17 +-
.../tests/ray_data_evolution_merge_into_test.py | 295 ++++++++++++++++++-
.../pypaimon/tests/ray_partitioning_test.py | 324 +++++++++++++++++++++
.../pypaimon/tests/ray_read_by_row_id_test.py | 12 +-
.../pypaimon/tests/ray_update_by_row_id_test.py | 213 +++++++++++++-
11 files changed, 1213 insertions(+), 43 deletions(-)
diff --git a/.github/workflows/paimon-python-checks.yml
b/.github/workflows/paimon-python-checks.yml
index d2a8ed781c..c66101d8e6 100755
--- a/.github/workflows/paimon-python-checks.yml
+++ b/.github/workflows/paimon-python-checks.yml
@@ -309,6 +309,12 @@ jobs:
if [ "$ray_version" = "2.50.1" ]; then
tests="$tests
pypaimon/tests/ray_update_by_row_id_test.py::RayUpdateByRowIdTest::test_empty_dataset_after_transform_is_noop"
fi
+ if [ "$ray_version" = "2.44.0" ]; then
+ tests="$tests
pypaimon/tests/ray_partitioning_test.py::RayPartitioningTest::test_ray_legacy_map_batches_cardinality_is_unknown"
+ fi
+ if [ "$ray_version" = "2.53.0" ]; then
+ tests="$tests
pypaimon/tests/ray_partitioning_test.py::RayPartitioningTest::test_ray_metadata_respects_transform_cardinality"
+ fi
python -m pytest $tests -v --tb=short || {
echo "Tests failed for Ray $ray_version"; python -m pip
uninstall -y ray; exit 1;
}
diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md
index 1f9194c92c..af6386e009 100644
--- a/docs/docs/pypaimon/ray-data.md
+++ b/docs/docs/pypaimon/ray-data.md
@@ -544,8 +544,11 @@ source-target merges.
- `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.
+- `num_partitions`: shuffle parallelism for the join and the write. When input
+ in-memory byte-size metadata is reliable, the default targets Ray's maximum
+ block size. Otherwise it uses Ray's hash-shuffle default. A nonempty target
+ keeps that default as a lower bound, and cluster CPUs cap the result.
Self-merge
+ keeps its CPU-based default. Set it explicitly to override the default.
- `ray_remote_args`: Ray remote options applied to the merge's map/group
tasks (update/delete transform, group write, insert transform).
- `concurrency`: scheduling for the insert sink.
@@ -595,8 +598,9 @@ print(metrics) # {"num_updated": 50}
values are cast to the target column types. A table-name source is not
accepted: a
table's system `_ROW_ID` is its own and cannot address the target's rows.
- `update_cols`: the non-blob columns to overwrite. Must be non-empty.
-- `num_partitions`: parallelism for grouping the update rows by target file;
- defaults to `max(1, cluster_cpus * 2)`.
+- `num_partitions`: parallelism for grouping the update rows by target file.
+ The default uses reliable source metadata when available; otherwise it uses
+ Ray's hash-shuffle default. Target file count and cluster CPUs bound the
result.
- `ray_remote_args`: Ray remote options applied to the update tasks.
**Returns:** `{"num_updated": <rows>}`.
@@ -646,8 +650,9 @@ ds = read_by_row_id(
(resolved later with `map_with_blobs`), or `scan.snapshot-id` /
`scan.tag-name` to read a
specific snapshot. Options that flip table invariants
(`data-evolution.enabled`,
`row-tracking.enabled`, `deletion-vectors.enabled`) are rejected.
-- `num_partitions`: parallelism for grouping the row ids by target file;
defaults to
- `max(1, cluster_cpus * 2)`.
+- `num_partitions`: parallelism for grouping the row ids by target file. The
+ default uses reliable source metadata when available; otherwise it uses Ray's
+ hash-shuffle default. Target file count and cluster CPUs bound the result.
- `ray_remote_args`: Ray remote options applied to the read tasks.
**Returns:** a `ray.data.Dataset` of `(*projection, _ROW_ID)`.
diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_into.py
b/paimon-python/pypaimon/ray/data_evolution_merge_into.py
index e0544f1890..3a280dd028 100644
--- a/paimon-python/pypaimon/ray/data_evolution_merge_into.py
+++ b/paimon-python/pypaimon/ray/data_evolution_merge_into.py
@@ -27,6 +27,7 @@ import pyarrow as pa
from pypaimon.common.predicate import Predicate
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.ray.data_evolution_merge_join import (
+ _resolve_matched_num_partitions,
_resolve_source_projection,
build_matched_delete_ds,
build_matched_update_ds,
@@ -49,6 +50,11 @@ from pypaimon.ray.data_evolution_merge_transform import (
WhenNotMatched,
_NormalizedClause,
)
+from pypaimon.ray.partitioning import (
+ _default_hash_shuffle_parallelism,
+ _estimate_dataset_size_bytes,
+ _resolve_num_partitions,
+)
__all__ = ["merge_into", "WhenMatched", "WhenNotMatched"]
@@ -84,7 +90,7 @@ def merge_into(
read_columns: Optional[Sequence[str]] = None,
) -> Dict[str, int]:
_require_ray_join()
- num_partitions = _resolve_num_partitions(num_partitions)
+ requested_num_partitions = num_partitions
table, source_ds, matched_specs, not_matched_specs, ctx = _prepare(
target, source, catalog_options,
@@ -92,16 +98,61 @@ def merge_into(
read_columns,
)
base_snapshot = table.snapshot_manager().get_latest_snapshot()
+ target_empty = _is_target_empty(base_snapshot)
+ estimated_size_bytes = None
+ if num_partitions is None:
+ estimated_size_bytes = _estimate_merge_input_size_bytes(
+ source_ds, ctx,
+ )
+ min_partitions = 1
+ unknown_num_partitions = None
+ data_context = (
+ None
+ if ctx.is_self_merge
+ else getattr(source_ds, "context", None)
+ )
+ if num_partitions is None and not ctx.is_self_merge:
+ unknown_num_partitions = _default_hash_shuffle_parallelism(
+ data_context
+ )
+ if not target_empty:
+ min_partitions = unknown_num_partitions
+ source_num_partitions = _resolve_num_partitions(
+ num_partitions,
+ estimated_size_bytes,
+ min_partitions=min_partitions,
+ unknown_num_partitions=unknown_num_partitions,
+ data_context=data_context,
+ )
update_ds, delete_ds, insert_ds, update_cols_union = _build_datasets(
table, target, source_ds, matched_specs, not_matched_specs,
- ctx, base_snapshot, num_partitions, ray_remote_args,
+ ctx, base_snapshot, source_num_partitions, ray_remote_args,
+ requested_num_partitions=requested_num_partitions,
+ estimated_size_bytes=estimated_size_bytes,
)
+ update_num_partitions = None
+ delete_num_partitions = None
+ if not ctx.is_self_merge:
+ if update_ds is not None:
+ update_num_partitions = _resolve_matched_num_partitions(
+ requested_num_partitions,
+ estimated_size_bytes,
+ update_ds,
+ )
+ if delete_ds is not None:
+ delete_num_partitions = _resolve_matched_num_partitions(
+ requested_num_partitions,
+ estimated_size_bytes,
+ delete_ds,
+ )
return _execute_and_commit(
table, update_ds, delete_ds, insert_ds, update_cols_union,
- base_snapshot, num_partitions,
+ base_snapshot, source_num_partitions,
ray_remote_args, concurrency,
+ update_num_partitions=update_num_partitions,
+ delete_num_partitions=delete_num_partitions,
)
@@ -326,7 +377,9 @@ def _is_self_merge(target, source, target_on_cols,
source_on_cols) -> bool:
def _build_datasets(
table, target, source_ds, matched_specs, not_matched_specs,
- ctx: "_PrepareCtx", base_snapshot, num_partitions, ray_remote_args,
+ ctx: "_PrepareCtx", base_snapshot, source_num_partitions, ray_remote_args,
+ requested_num_partitions: Optional[int] = None,
+ estimated_size_bytes: Optional[int] = None,
):
# Pin every target read to base_snapshot so all branches see the same
# snapshot the caller observed; otherwise concurrent commits in between
@@ -337,9 +390,10 @@ def _build_datasets(
delete_ds = None
insert_ds = None
update_cols_union: List[str] = []
+ target_empty = _is_target_empty(base_snapshot)
if ctx.is_self_merge:
- if matched_specs and base_snapshot is not None:
+ if matched_specs and not target_empty:
update_cols_union = _union_update_cols(matched_specs)
if update_cols_union:
update_ds = build_self_merge_update_plan(
@@ -369,7 +423,7 @@ def _build_datasets(
# Mirror Spark: matched/not-matched run as two independent joins
# (inner / left_anti). One unified left_outer join would force
# joined.materialize() to feed both branches, which can OOM on large
merges.
- if matched_specs and base_snapshot is not None:
+ if matched_specs and not target_empty:
update_cols_union = _union_update_cols(matched_specs)
if update_cols_union:
update_ds = build_matched_update_ds(
@@ -382,7 +436,8 @@ def _build_datasets(
target_pa_schema=ctx.update_pa_schema,
update_cols=update_cols_union,
catalog_options=ctx.catalog_options,
- num_partitions=num_partitions,
+ num_partitions=requested_num_partitions,
+ estimated_size_bytes=estimated_size_bytes,
resolve_target_projection=_resolve_target_projection,
snapshot_id=base_snapshot_id,
ray_remote_args=ray_remote_args,
@@ -396,7 +451,8 @@ def _build_datasets(
clauses=matched_specs,
target_field_names=ctx.settable_field_names,
catalog_options=ctx.catalog_options,
- num_partitions=num_partitions,
+ num_partitions=requested_num_partitions,
+ estimated_size_bytes=estimated_size_bytes,
resolve_target_projection=_resolve_target_projection,
snapshot_id=base_snapshot_id,
ray_remote_args=ray_remote_args,
@@ -412,9 +468,9 @@ def _build_datasets(
target_field_names=ctx.full_target_field_names,
target_pa_schema=ctx.full_pa_schema,
catalog_options=ctx.catalog_options,
- num_partitions=num_partitions,
+ num_partitions=source_num_partitions,
snapshot_id=base_snapshot_id,
- target_empty=base_snapshot is None,
+ target_empty=target_empty,
ray_remote_args=ray_remote_args,
)
@@ -425,6 +481,8 @@ def _execute_and_commit(
table, update_ds, delete_ds, insert_ds, update_cols_union,
base_snapshot, num_partitions,
ray_remote_args, concurrency,
+ update_num_partitions=None,
+ delete_num_partitions=None,
):
collect_action_row_ids = update_ds is not None and delete_ds is not None
commit_messages: list = []
@@ -438,6 +496,16 @@ def _execute_and_commit(
num_inserted = 0
insert_msgs: list = []
self_merge_update = isinstance(update_ds, _SelfMergeUpdatePlan)
+ update_num_partitions = (
+ num_partitions
+ if update_num_partitions is None
+ else update_num_partitions
+ )
+ delete_num_partitions = (
+ num_partitions
+ if delete_num_partitions is None
+ else delete_num_partitions
+ )
try:
if update_ds is not None:
@@ -445,7 +513,7 @@ def _execute_and_commit(
update_msgs, num_updated, update_row_ids = (
distributed_self_merge_update_apply(
update_ds,
- num_partitions=num_partitions,
+ num_partitions=update_num_partitions,
ray_remote_args=ray_remote_args,
collect_row_ids=collect_action_row_ids,
)
@@ -454,7 +522,7 @@ def _execute_and_commit(
update_msgs, num_updated, update_row_ids = (
distributed_update_apply(
update_ds, table, update_cols_union,
- num_partitions=num_partitions,
+ num_partitions=update_num_partitions,
ray_remote_args=ray_remote_args,
base_snapshot_id=(
base_snapshot.id
@@ -468,7 +536,7 @@ def _execute_and_commit(
if delete_ds is not None:
delete_msgs, num_deleted, delete_row_ids =
distributed_delete_apply(
delete_ds, table,
- num_partitions=num_partitions,
+ num_partitions=delete_num_partitions,
ray_remote_args=ray_remote_args,
base_snapshot_id=(
base_snapshot.id
@@ -545,16 +613,17 @@ def _normalize_on(on: OnSpec) -> Tuple[List[str],
List[str]]:
return target_cols, source_cols
-def _resolve_num_partitions(num_partitions: Optional[int]) -> int:
- if num_partitions is not None:
- return num_partitions
- try:
- import ray
+def _estimate_merge_input_size_bytes(
+ source_ds,
+ ctx: "_PrepareCtx",
+) -> Optional[int]:
+ if ctx.is_self_merge:
+ return None
+ return _estimate_dataset_size_bytes(source_ds)
+
- cpus = int(ray.cluster_resources().get("CPU", 4))
- return max(1, cpus * 2)
- except Exception:
- return 4
+def _is_target_empty(snapshot) -> bool:
+ return snapshot is None or snapshot.total_record_count == 0
def _require_ray_join() -> None:
diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
index af56d7ca30..5cb42fbf49 100644
--- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
+++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
@@ -31,6 +31,11 @@ from pypaimon.ray.data_evolution_merge_transform import (
vectorized_insert_transform,
vectorized_matched_transform,
)
+from pypaimon.ray.partitioning import (
+ _default_hash_shuffle_parallelism,
+ _resolve_num_partitions,
+ _resolve_row_id_num_partitions,
+)
def _map_kwargs(
@@ -45,6 +50,26 @@ def _map_kwargs(
return kwargs
+def _resolve_matched_num_partitions(
+ num_partitions: Optional[int],
+ estimated_size_bytes: Optional[int],
+ target_ds,
+) -> int:
+ """Resolve a target-left join from the context sealed on its left input."""
+ if num_partitions is not None:
+ return num_partitions
+
+ data_context = getattr(target_ds, "context", None)
+ default_shuffle = _default_hash_shuffle_parallelism(data_context)
+ return _resolve_num_partitions(
+ num_partitions,
+ estimated_size_bytes,
+ min_partitions=default_shuffle,
+ unknown_num_partitions=default_shuffle,
+ data_context=data_context,
+ )
+
+
@dataclass(frozen=True)
class _SelfMergeUpdatePlan:
"""Pinned target file groups for self-merge update execution."""
@@ -543,8 +568,9 @@ def build_matched_update_ds(
target_pa_schema: pa.Schema,
update_cols: Sequence[str],
catalog_options: Dict[str, str],
- num_partitions: int,
+ num_partitions: Optional[int],
resolve_target_projection,
+ estimated_size_bytes: Optional[int] = None,
snapshot_id: Optional[int] = None,
ray_remote_args: Optional[Dict[str, Any]] = None,
) -> Tuple:
@@ -566,6 +592,9 @@ def build_matched_update_ds(
target_renamed = target_ds.rename_columns(
{c: f"t.{c}" for c in target_ds.schema().names}
)
+ num_partitions = _resolve_matched_num_partitions(
+ num_partitions, estimated_size_bytes, target_renamed,
+ )
source_cols = _resolve_source_projection(
clauses, source_on, source_ds.schema().names,
)
@@ -602,8 +631,9 @@ def build_matched_delete_ds(
clauses: List[_NormalizedClause],
target_field_names: Sequence[str],
catalog_options: Dict[str, str],
- num_partitions: int,
+ num_partitions: Optional[int],
resolve_target_projection,
+ estimated_size_bytes: Optional[int] = None,
snapshot_id: Optional[int] = None,
ray_remote_args: Optional[Dict[str, Any]] = None,
) -> Tuple:
@@ -628,6 +658,9 @@ def build_matched_delete_ds(
target_renamed = target_ds.rename_columns(
{c: f"t.{c}" for c in target_ds.schema().names}
)
+ num_partitions = _resolve_matched_num_partitions(
+ num_partitions, estimated_size_bytes, target_renamed,
+ )
source_cols = list(source_ds.schema().names)
source_renamed = source_ds.rename_columns(
{c: f"s.{c}" for c in source_cols}
@@ -655,10 +688,13 @@ def distributed_update_apply(
table,
write_update_cols: Sequence[str],
*,
- num_partitions: int,
+ num_partitions: Optional[int],
ray_remote_args: Optional[Dict[str, Any]] = None,
base_snapshot_id: Optional[int] = None,
collect_row_ids: bool = False,
+ estimated_size_bytes: Optional[int] = None,
+ estimated_num_rows: Optional[int] = None,
+ data_context=None,
) -> Tuple[list, int, list]:
import numpy as np
import pickle
@@ -697,6 +733,14 @@ def distributed_update_apply(
if not sorted_first_row_ids:
return [], 0, []
+ num_partitions = _resolve_row_id_num_partitions(
+ num_partitions,
+ estimated_size_bytes,
+ estimated_num_rows,
+ len(sorted_first_row_ids),
+ data_context=data_context,
+ )
+
# Pin commit-time conflict check to the snapshot the join was built on,
# so concurrent commits between read and planner are detected.
check_from_snapshot = (
@@ -858,9 +902,12 @@ def distributed_read_by_row_id(
table,
projection: Sequence[str],
*,
- num_partitions: int,
+ num_partitions: Optional[int],
ray_remote_args: Optional[Dict[str, Any]] = None,
base_snapshot_id: Optional[int] = None,
+ estimated_size_bytes: Optional[int] = None,
+ estimated_num_rows: Optional[int] = None,
+ data_context=None,
):
"""Read ``projection`` for the ``_ROW_ID``s in ``row_ids_ds``, routing
each to its
owning file and reading only the matched rows via ``IndexedSplit`` slicing
(blob
@@ -902,6 +949,14 @@ def distributed_read_by_row_id(
if not sorted_first_row_ids:
return None
+ num_partitions = _resolve_row_id_num_partitions(
+ num_partitions,
+ estimated_size_bytes,
+ estimated_num_rows,
+ len(sorted_first_row_ids),
+ data_context=data_context,
+ )
+
precomputed_info_ref = ray.put(planner._snapshot_files_info())
frid_col = "_FIRST_ROW_ID"
sorted_arr = np.asarray(sorted_first_row_ids, dtype=np.int64)
@@ -1143,7 +1198,7 @@ def build_not_matched_insert_ds(
)
if target_empty:
- unmatched = source_renamed
+ unmatched = source_renamed.repartition(num_partitions)
else:
target_ds = read_paimon(
target_identifier, catalog_options,
diff --git a/paimon-python/pypaimon/ray/partitioning.py
b/paimon-python/pypaimon/ray/partitioning.py
new file mode 100644
index 0000000000..935992c30f
--- /dev/null
+++ b/paimon-python/pypaimon/ray/partitioning.py
@@ -0,0 +1,177 @@
+################################################################################
+# 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.
+################################################################################
+
+"""Best-effort partition sizing for PyPaimon Ray operations."""
+
+from typing import Optional
+
+
+def _resolve_num_partitions(
+ num_partitions: Optional[int],
+ estimated_size_bytes: Optional[int] = None,
+ min_partitions: int = 1,
+ unknown_num_partitions: Optional[int] = None,
+ data_context=None,
+) -> int:
+ """Resolve default shuffle partitions from input size and CPU count."""
+ if num_partitions is not None:
+ return num_partitions
+
+ try:
+ import ray
+
+ cpus = int(ray.cluster_resources().get("CPU", 4))
+ max_partitions = max(1, cpus * 2)
+ except Exception:
+ max_partitions = 4
+
+ if estimated_size_bytes is None:
+ if unknown_num_partitions is not None:
+ return min(
+ max_partitions,
+ max(min_partitions, int(unknown_num_partitions)),
+ )
+ return max_partitions
+
+ try:
+ from ray.data.context import DataContext
+
+ context = (
+ data_context
+ if data_context is not None
+ else DataContext.get_current()
+ )
+ target_size_bytes = int(context.target_max_block_size)
+ except Exception:
+ return max_partitions
+
+ if target_size_bytes <= 0:
+ return max_partitions
+ size_partitions = max(
+ 1,
+ (max(0, int(estimated_size_bytes)) + target_size_bytes - 1)
+ // target_size_bytes,
+ )
+ return min(max_partitions, max(min_partitions, size_partitions))
+
+
+def _estimate_dataset_size_bytes(dataset) -> Optional[int]:
+ """Read logical-plan size metadata without executing the Dataset."""
+ return _estimate_dataset_metadata(dataset, "size_bytes")
+
+
+def _estimate_dataset_num_rows(dataset) -> Optional[int]:
+ """Read logical-plan row metadata without executing the Dataset."""
+ return _estimate_dataset_metadata(dataset, "num_rows")
+
+
+def _estimate_dataset_metadata(dataset, field: str) -> Optional[int]:
+ try:
+ operator = getattr(getattr(dataset, "_logical_plan", None), "dag",
None)
+ while operator is not None:
+ infer_metadata = getattr(operator, "infer_metadata", None)
+ can_modify_num_rows = _can_modify_num_rows(operator)
+ if callable(infer_metadata):
+ value = getattr(infer_metadata(), field, None)
+ if (
+ value is not None
+ and int(value) >= 0
+ and not (
+ field == "size_bytes"
+ and can_modify_num_rows is not None
+ )
+ ):
+ return int(value)
+ if field == "size_bytes":
+ return None
+ # Only inherit row count through transforms Ray marks preserving.
+ if can_modify_num_rows is not False:
+ return None
+ dependencies = getattr(operator, "input_dependencies", ())
+ operator = dependencies[0] if len(dependencies) == 1 else None
+ except Exception:
+ pass
+ return None
+
+
+def _can_modify_num_rows(operator) -> Optional[bool]:
+ is_map_batches = type(operator).__name__ == "MapBatches"
+ class_value = getattr(type(operator), "can_modify_num_rows", None)
+ # Ray <= 2.45 exposes a False property even though the UDF may change rows.
+ if is_map_batches and isinstance(class_value, property):
+ return None
+
+ value = getattr(operator, "can_modify_num_rows", None)
+ if not callable(value):
+ return value if isinstance(value, bool) else None
+
+ # Ray 2.50/2.51 exposes a False method without a cardinality flag.
+ if is_map_batches and not hasattr(operator, "_udf_modifying_row_count"):
+ return None
+ try:
+ value = value()
+ except Exception:
+ return None
+ return value if isinstance(value, bool) else None
+
+
+def _default_hash_shuffle_parallelism(data_context=None) -> int:
+ try:
+ from ray.data.context import DataContext
+
+ context = (
+ data_context
+ if data_context is not None
+ else DataContext.get_current()
+ )
+ return max(
+ 1,
+ int(context.default_hash_shuffle_parallelism),
+ )
+ except Exception:
+ return 200
+
+
+def _resolve_row_id_num_partitions(
+ num_partitions: Optional[int],
+ estimated_size_bytes: Optional[int],
+ estimated_num_rows: Optional[int],
+ target_file_count: int,
+ data_context=None,
+) -> int:
+ """Resolve row-ID partitions from input size and target fan-out."""
+ if num_partitions is not None:
+ return num_partitions
+
+ default_shuffle = _default_hash_shuffle_parallelism(data_context)
+
+ possible_groups = max(1, target_file_count)
+ if estimated_num_rows is not None:
+ possible_groups = min(possible_groups, max(1, estimated_num_rows))
+ min_partitions = min(max(1, default_shuffle), possible_groups)
+ if estimated_size_bytes is None:
+ return min(
+ _resolve_num_partitions(None, data_context=data_context),
+ min_partitions,
+ )
+ return _resolve_num_partitions(
+ None,
+ estimated_size_bytes,
+ min_partitions=min_partitions,
+ data_context=data_context,
+ )
diff --git a/paimon-python/pypaimon/ray/read_by_row_id.py
b/paimon-python/pypaimon/ray/read_by_row_id.py
index 7c7a650803..d9358f3723 100644
--- a/paimon-python/pypaimon/ray/read_by_row_id.py
+++ b/paimon-python/pypaimon/ray/read_by_row_id.py
@@ -30,12 +30,15 @@ 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 (
_read_output_schema,
distributed_read_by_row_id,
)
+from pypaimon.ray.partitioning import (
+ _estimate_dataset_num_rows,
+ _estimate_dataset_size_bytes,
+)
__all__ = ["read_by_row_id"]
@@ -104,8 +107,6 @@ def read_by_row_id(
if not projection:
raise ValueError("projection must be non-empty.")
projection = list(dict.fromkeys(projection))
- num_partitions = _resolve_num_partitions(num_partitions)
-
table = CatalogFactory.create(catalog_options).get_table(target)
if not table.options.data_evolution_enabled():
raise ValueError(
@@ -141,6 +142,13 @@ def read_by_row_id(
"read_by_row_id does not accept a table-name source; pass a
ray.data."
"Dataset / pyarrow.Table / pandas.DataFrame carrying the target
row ids.")
source_ds = _normalize_source(row_ids, catalog_options)
+ estimated_size_bytes = None
+ estimated_num_rows = None
+ data_context = None
+ if num_partitions is None:
+ estimated_size_bytes = _estimate_dataset_size_bytes(source_ds)
+ estimated_num_rows = _estimate_dataset_num_rows(source_ds)
+ data_context = getattr(source_ds, "context", None)
# Only check now if the schema is free; fetching it would execute a lazy
source.
known_schema = source_ds.schema(fetch_if_missing=False)
if known_schema is not None and src_rid_col not in set(known_schema.names):
@@ -185,6 +193,9 @@ def read_by_row_id(
num_partitions=num_partitions,
ray_remote_args=ray_remote_args,
base_snapshot_id=base.id,
+ estimated_size_bytes=estimated_size_bytes,
+ estimated_num_rows=estimated_num_rows,
+ data_context=data_context,
)
except Exception as e:
_reraise_inner(e)
diff --git a/paimon-python/pypaimon/ray/update_by_row_id.py
b/paimon-python/pypaimon/ray/update_by_row_id.py
index bc1b437e8c..58679a8a98 100644
--- a/paimon-python/pypaimon/ray/update_by_row_id.py
+++ b/paimon-python/pypaimon/ray/update_by_row_id.py
@@ -32,10 +32,13 @@ 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
+from pypaimon.ray.partitioning import (
+ _estimate_dataset_num_rows,
+ _estimate_dataset_size_bytes,
+)
from pypaimon.schema.data_types import is_blob_file_field
__all__ = ["update_by_row_id"]
@@ -75,8 +78,6 @@ def update_by_row_id(
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(
@@ -114,6 +115,13 @@ def update_by_row_id(
"update_by_row_id does not accept a table-name source; pass a
ray.data."
f"Dataset / pyarrow.Table / pandas.DataFrame carrying the target
{rid}.")
source_ds = _normalize_source(source, catalog_options)
+ estimated_size_bytes = None
+ estimated_num_rows = None
+ data_context = None
+ if num_partitions is None:
+ estimated_size_bytes = _estimate_dataset_size_bytes(source_ds)
+ estimated_num_rows = _estimate_dataset_num_rows(source_ds)
+ data_context = getattr(source_ds, "context", None)
src_cols = set(source_ds.schema().names)
missing = [c for c in [rid] + update_cols if c not in src_cols]
if missing:
@@ -145,6 +153,9 @@ def update_by_row_id(
num_partitions=num_partitions,
ray_remote_args=ray_remote_args,
base_snapshot_id=base.id,
+ estimated_size_bytes=estimated_size_bytes,
+ estimated_num_rows=estimated_num_rows,
+ data_context=data_context,
)
except Exception as e:
_reraise_inner(e)
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 e0a35e1740..8a72558bb7 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
@@ -204,6 +204,140 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
projection=['id', 'name', 'age'],
)
+ def test_paimon_source_does_not_use_compressed_size(self):
+ from pypaimon.ray import data_evolution_merge_into as m
+
+ target = self._create_table()
+ source = self._create_table()
+ self._write(target, self._source(ids=(0,)))
+ value = 'x' * 1_000_000
+ source_data = pa.Table.from_pydict(
+ {
+ 'id': pa.array(list(range(16)), type=pa.int32()),
+ 'name': [value] * 16,
+ 'age': [10] * 16,
+ },
+ schema=self.pa_schema,
+ )
+ self._write(source, source_data)
+ source_table = self.catalog.get_table(source)
+ splits = source_table.new_read_builder().new_scan().plan().splits()
+ compressed_size = sum(split.file_size for split in splits)
+ self.assertLess(compressed_size * 10, source_data.nbytes)
+
+ real_resolve = m._resolve_num_partitions
+ resolved = []
+
+ def capture(*args, **kwargs):
+ result = real_resolve(*args, **kwargs)
+ resolved.append((args, kwargs, result))
+ return result
+
+ with patch(
+ 'ray.cluster_resources', return_value={'CPU': 320},
+ ), patch.object(
+ m, '_resolve_num_partitions', side_effect=capture,
+ ), patch.object(
+ m, '_build_datasets', return_value=(None, None, None, set()),
+ ), patch.object(
+ m, '_execute_and_commit', return_value={},
+ ):
+ merge_into(
+ target=target,
+ source=source,
+ catalog_options=self.catalog_options,
+ on=['id'],
+ when_matched=[WhenMatched.update('*')],
+ )
+
+ self.assertEqual(len(resolved), 1)
+ args, kwargs, result = resolved[0]
+ self.assertEqual(args, (None, None))
+ self.assertIsNotNone(kwargs.pop('data_context'))
+ self.assertEqual(kwargs, {
+ 'min_partitions': 200,
+ 'unknown_num_partitions': 200,
+ })
+ self.assertEqual(result, 200)
+
+ def test_matched_execution_uses_target_context(self):
+ from pypaimon.ray import data_evolution_merge_into as m
+
+ source_context = Mock(
+ target_max_block_size=512,
+ default_hash_shuffle_parallelism=7,
+ )
+ target_context = Mock(
+ target_max_block_size=128,
+ default_hash_shuffle_parallelism=3,
+ )
+ source_ds = Mock(context=source_context)
+ update_ds = Mock(context=target_context)
+ delete_ds = Mock(context=target_context)
+ ctx = Mock(is_self_merge=False)
+ snapshot = Mock(total_record_count=1)
+ table = Mock()
+ table.snapshot_manager().get_latest_snapshot.return_value = snapshot
+
+ with patch.object(
+ m, '_prepare',
+ return_value=(table, source_ds, [], [], ctx),
+ ), patch.object(
+ m, '_estimate_merge_input_size_bytes', return_value=512,
+ ), patch.object(
+ m, '_build_datasets',
+ return_value=(update_ds, delete_ds, None, ['age']),
+ ) as build_datasets, patch.object(
+ m, 'distributed_update_apply', return_value=([], 0, []),
+ ) as update_apply, patch.object(
+ m, 'distributed_delete_apply', return_value=([], 0, []),
+ ) as delete_apply, patch(
+ 'ray.cluster_resources', return_value={'CPU': 320},
+ ):
+ merge_into(
+ target='default.target',
+ source=source_ds,
+ catalog_options={'warehouse': '/tmp/warehouse'},
+ on=['id'],
+ when_matched=[WhenMatched.update('*')],
+ )
+ default_source_partitions = build_datasets.call_args.args[7]
+ default_update_partitions = (
+ update_apply.call_args.kwargs['num_partitions']
+ )
+ default_delete_partitions = (
+ delete_apply.call_args.kwargs['num_partitions']
+ )
+
+ build_datasets.reset_mock()
+ update_apply.reset_mock()
+ delete_apply.reset_mock()
+ merge_into(
+ target='default.target',
+ source=source_ds,
+ catalog_options={'warehouse': '/tmp/warehouse'},
+ on=['id'],
+ when_matched=[WhenMatched.update('*')],
+ num_partitions=11,
+ )
+
+ # The source-left branch would use 7 partitions.
+ self.assertEqual(default_source_partitions, 7)
+ # Both matched results inherit the target-left context: 512 / 128 = 4.
+ self.assertEqual(default_update_partitions, 4)
+ self.assertEqual(default_delete_partitions, 4)
+ # An explicit value still applies to every branch and execution stage.
+ self.assertEqual(build_datasets.call_args.args[7], 11)
+ self.assertEqual(
+ build_datasets.call_args.kwargs['requested_num_partitions'], 11
+ )
+ self.assertEqual(
+ update_apply.call_args.kwargs['num_partitions'], 11
+ )
+ self.assertEqual(
+ delete_apply.call_args.kwargs['num_partitions'], 11
+ )
+
def test_no_clause_raises(self):
target = self._create_table()
with self.assertRaises(ValueError):
@@ -566,6 +700,94 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
self.assertEqual(out['name'], ['a', 'b', 'c'])
self.assertEqual(out['age'], [10, 20, 30])
+ def test_insert_into_truncated_target_uses_empty_fast_path(self):
+ from pypaimon.ray import data_evolution_merge_into as m
+
+ target = self._create_table()
+ self._write(target, self._source(ids=(0,)))
+ table = self.catalog.get_table(target)
+ commit = table.new_batch_write_builder().new_commit()
+ commit.truncate_table()
+ commit.close()
+ snapshot = table.snapshot_manager().get_latest_snapshot()
+ self.assertIsNotNone(snapshot)
+ self.assertEqual(snapshot.total_record_count, 0)
+
+ real_resolve = m._resolve_num_partitions
+ resolved = []
+
+ def capture(*args, **kwargs):
+ result = real_resolve(*args, **kwargs)
+ resolved.append((args, kwargs, result))
+ return result
+
+ with patch(
+ 'ray.cluster_resources', return_value={'CPU': 320},
+ ), patch.object(
+ m, '_resolve_num_partitions', side_effect=capture,
+ ), patch.object(
+ ray.data.Dataset,
+ 'join',
+ side_effect=AssertionError('empty target must not be joined'),
+ ):
+ metrics = merge_into(
+ target=target,
+ source=self._source(ids=(1,)),
+ catalog_options=self.catalog_options,
+ on=['id'],
+ when_matched=[WhenMatched.update('*')],
+ when_not_matched=[WhenNotMatched(insert='*')],
+ )
+
+ self.assertEqual(metrics['num_inserted'], 1)
+ self.assertEqual(self._read_sorted(target)['id'], [1])
+ self.assertEqual(len(resolved), 1)
+ args, kwargs, result = resolved[0]
+ self.assertIsNone(args[0])
+ self.assertGreater(args[1], 0)
+ self.assertIsNotNone(kwargs.pop('data_context'))
+ self.assertEqual(kwargs, {
+ 'min_partitions': 1,
+ 'unknown_num_partitions': 200,
+ })
+ self.assertEqual(result, 1)
+
+ def test_insert_into_truncated_target_preserves_write_parallelism(self):
+ from pypaimon.ray import data_evolution_merge_into as m
+
+ target = self._create_table()
+ self._write(target, self._source(ids=(0,)))
+ table = self.catalog.get_table(target)
+ commit = table.new_batch_write_builder().new_commit()
+ commit.truncate_table()
+ commit.close()
+
+ real_write = m.distributed_write_collect_msgs
+ write_blocks = []
+
+ def capture(insert_ds, *args, **kwargs):
+ insert_ds = insert_ds.materialize()
+ write_blocks.append(insert_ds.num_blocks())
+ return real_write(insert_ds, *args, **kwargs)
+
+ with patch.object(
+ m, 'distributed_write_collect_msgs', side_effect=capture,
+ ):
+ metrics = merge_into(
+ target=target,
+ source=self._source(ids=range(1, 31)),
+ catalog_options=self.catalog_options,
+ on=['id'],
+ when_not_matched=[WhenNotMatched(insert='*')],
+ num_partitions=3,
+ )
+
+ self.assertEqual(metrics['num_inserted'], 30)
+ self.assertEqual(write_blocks, [3])
+ self.assertEqual(
+ self._read_sorted(target)['id'], list(range(1, 31))
+ )
+
def test_multi_source_match_raises_by_default(self):
# One target row matched by several source rows: the winning value is
# undefined (Spark DE's checkCardinality=false), so we refuse by
default.
@@ -4121,6 +4343,74 @@ class TargetProjectionTest(unittest.TestCase):
'id': 's.id',
'name': 's.name',
})
+ self.assertEqual(
+ target_renamed.join.call_args.kwargs['num_partitions'], 1
+ )
+
+ def test_matched_update_uses_target_left_context(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()),
+ ])
+ selected_ds = Mock()
+ source_renamed = Mock()
+ source_renamed.context = Mock(
+ target_max_block_size=512,
+ default_hash_shuffle_parallelism=7,
+ )
+ 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()
+ target_renamed.context = Mock(
+ target_max_block_size=128,
+ default_hash_shuffle_parallelism=3,
+ )
+ joined = Mock()
+ target_ds.rename_columns.return_value = target_renamed
+ target_renamed.join.return_value = joined
+
+ with patch(
+ 'pypaimon.ray.ray_paimon.read_paimon',
+ return_value=target_ds,
+ ), patch(
+ 'ray.cluster_resources', return_value={'CPU': 320},
+ ), patch(
+ 'ray.data.context.DataContext.get_current',
+ side_effect=AssertionError('must use target Dataset context'),
+ ):
+ 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=None,
+ estimated_size_bytes=512,
+ resolve_target_projection=lambda *args: ['id'],
+ )
+
+ join_kwargs = target_renamed.join.call_args.kwargs
+ # Target context: ceil(512 / 128) = 4. Source context would resolve 7.
+ self.assertEqual(join_kwargs['num_partitions'], 4)
def test_not_matched_insert_selects_needed_source_cols(self):
from pypaimon.ray.data_evolution_merge_join import (
@@ -4136,10 +4426,12 @@ class TargetProjectionTest(unittest.TestCase):
])
selected_ds = Mock()
source_renamed = Mock()
+ repartitioned = 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
+ source_renamed.repartition.return_value = repartitioned
+ repartitioned.map_batches.return_value = result
out = build_not_matched_insert_ds(
target_identifier='default.target',
@@ -4163,6 +4455,7 @@ class TargetProjectionTest(unittest.TestCase):
'id': 's.id',
'name': 's.name',
})
+ source_renamed.repartition.assert_called_once_with(1)
class MergeConditionUnitTest(unittest.TestCase):
diff --git a/paimon-python/pypaimon/tests/ray_partitioning_test.py
b/paimon-python/pypaimon/tests/ray_partitioning_test.py
new file mode 100644
index 0000000000..0d131fd1f7
--- /dev/null
+++ b/paimon-python/pypaimon/tests/ray_partitioning_test.py
@@ -0,0 +1,324 @@
+################################################################################
+# 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.
+################################################################################
+
+import types
+import unittest
+from unittest import mock
+
+import pyarrow as pa
+import pytest
+
+pytest.importorskip("ray")
+
+from pypaimon.ray.data_evolution_merge_into import (
+ _estimate_merge_input_size_bytes,
+)
+from pypaimon.ray.partitioning import (
+ _estimate_dataset_num_rows,
+ _estimate_dataset_size_bytes,
+ _resolve_num_partitions,
+ _resolve_row_id_num_partitions,
+)
+
+
+class RayPartitioningTest(unittest.TestCase):
+
+ def test_explicit_num_partitions_is_unchanged(self):
+ with mock.patch(
+ "ray.cluster_resources",
+ side_effect=AssertionError("must not inspect the cluster"),
+ ):
+ self.assertEqual(_resolve_num_partitions(37, 1), 37)
+
+ def test_known_size_uses_target_block_size_and_cpu_cap(self):
+ from ray.data.context import DataContext
+
+ context = DataContext.get_current()
+ previous_target = context.target_max_block_size
+ context.target_max_block_size = 128 * 1024 * 1024
+ try:
+ with mock.patch(
+ "ray.cluster_resources", return_value={"CPU": 320}
+ ):
+ self.assertEqual(_resolve_num_partitions(None, 0), 1)
+ self.assertEqual(
+ _resolve_num_partitions(None, 5 * 128 * 1024 * 1024),
+ 5,
+ )
+ self.assertEqual(
+ _resolve_num_partitions(None, 641 * 128 * 1024 * 1024),
+ 640,
+ )
+ # At the same average row width, a 10x larger input gets
+ # proportionally more partitions instead of both using 640.
+ self.assertEqual(
+ _resolve_num_partitions(None, 5_000_000 * 1024),
+ 39,
+ )
+ self.assertEqual(
+ _resolve_num_partitions(None, 50_000_000 * 1024),
+ 382,
+ )
+ self.assertEqual(
+ _resolve_num_partitions(None, 1, min_partitions=25),
+ 25,
+ )
+ self.assertEqual(
+ _resolve_num_partitions(None, 1, min_partitions=200),
+ 200,
+ )
+ finally:
+ context.target_max_block_size = previous_target
+
+ def test_known_size_uses_dataset_context(self):
+ import ray
+ from ray.data.context import DataContext
+
+ context = DataContext.get_current()
+ previous_target = context.target_max_block_size
+ try:
+ context.target_max_block_size = 128 * 1024 * 1024
+ data_context = ray.data.from_items([1]).context
+ context.target_max_block_size = 1024 * 1024 * 1024
+ with mock.patch(
+ "ray.cluster_resources", return_value={"CPU": 320}
+ ):
+ self.assertEqual(
+ _resolve_num_partitions(
+ None,
+ 512 * 1024 * 1024,
+ data_context=data_context,
+ ),
+ 4,
+ )
+ finally:
+ context.target_max_block_size = previous_target
+
+ def test_unknown_size_keeps_cpu_default(self):
+ with mock.patch(
+ "ray.cluster_resources", return_value={"CPU": 320}
+ ):
+ self.assertEqual(_resolve_num_partitions(None, None), 640)
+ self.assertEqual(
+ _resolve_num_partitions(
+ None,
+ None,
+ unknown_num_partitions=200,
+ ),
+ 200,
+ )
+
+ def test_dataset_estimate_does_not_call_size_bytes(self):
+ metadata = types.SimpleNamespace(size_bytes=1234)
+ dataset = types.SimpleNamespace(
+ _logical_plan=types.SimpleNamespace(
+ dag=types.SimpleNamespace(
+ infer_metadata=mock.Mock(return_value=metadata)
+ )
+ ),
+ size_bytes=mock.Mock(
+ side_effect=AssertionError("must not execute Dataset")
+ ),
+ )
+
+ self.assertEqual(_estimate_dataset_size_bytes(dataset), 1234)
+ dataset.size_bytes.assert_not_called()
+
+ def test_unknown_dataset_estimate_falls_back(self):
+ metadata = types.SimpleNamespace(size_bytes=None)
+ dataset = types.SimpleNamespace(
+ _logical_plan=types.SimpleNamespace(
+ dag=types.SimpleNamespace(
+ infer_metadata=mock.Mock(return_value=metadata)
+ )
+ )
+ )
+
+ self.assertIsNone(_estimate_dataset_size_bytes(dataset))
+
+ def test_ray_metadata_respects_transform_cardinality(self):
+ import inspect
+
+ import ray
+
+ dataset = ray.data.from_arrow(pa.table({"id": [1, 2]}))
+ self.assertEqual(_estimate_dataset_size_bytes(dataset), 16)
+ self.assertEqual(_estimate_dataset_num_rows(dataset), 2)
+
+ mapped = dataset.map_batches(lambda batch: batch)
+ self.assertIsNone(_estimate_dataset_size_bytes(mapped))
+ cardinality_parameter = inspect.signature(
+ ray.data.Dataset.map_batches
+ ).parameters.get("udf_modifying_row_count")
+ modifies_by_default = (
+ True
+ if cardinality_parameter is None
+ else cardinality_parameter.default
+ )
+ expected_rows = None if modifies_by_default else 2
+ self.assertEqual(_estimate_dataset_num_rows(mapped), expected_rows)
+
+ if "udf_modifying_row_count" in inspect.signature(
+ ray.data.Dataset.map_batches
+ ).parameters:
+ one_to_one = dataset.map_batches(
+ lambda batch: batch,
+ udf_modifying_row_count=False,
+ )
+ self.assertIsNone(_estimate_dataset_size_bytes(one_to_one))
+ self.assertEqual(_estimate_dataset_num_rows(one_to_one), 2)
+
+ widened = dataset.map_batches(
+ lambda batch: pa.table({
+ "payload": ["x" * 1_000_000] * batch.num_rows,
+ }),
+ batch_format="pyarrow",
+ udf_modifying_row_count=False,
+ )
+ self.assertIsNone(_estimate_dataset_size_bytes(widened))
+ self.assertEqual(_estimate_dataset_num_rows(widened), 2)
+ self.assertGreater(widened.materialize().size_bytes(), 1_000_000)
+
+ def test_legacy_map_batches_cardinality_is_unknown(self):
+ dependency = types.SimpleNamespace(
+ infer_metadata=mock.Mock(
+ return_value=types.SimpleNamespace(num_rows=2)
+ ),
+ can_modify_num_rows=None,
+ )
+
+ class MapBatches:
+ input_dependencies = (dependency,)
+
+ @staticmethod
+ def infer_metadata():
+ return types.SimpleNamespace(num_rows=None)
+
+ @property
+ def can_modify_num_rows(self):
+ return False
+
+ dataset = types.SimpleNamespace(
+ _logical_plan=types.SimpleNamespace(dag=MapBatches())
+ )
+ self.assertIsNone(_estimate_dataset_num_rows(dataset))
+
+ MapBatches.can_modify_num_rows = staticmethod(lambda: False)
+ self.assertIsNone(_estimate_dataset_num_rows(dataset))
+
+ def test_ray_legacy_map_batches_cardinality_is_unknown(self):
+ import inspect
+
+ import ray
+
+ dataset = ray.data.from_arrow(pa.table({"seed": [0]})).map_batches(
+ lambda _: pa.table({"id": list(range(500))}),
+ batch_format="pyarrow",
+ )
+ operator = dataset._logical_plan.dag
+ descriptor = inspect.getattr_static(
+ type(operator), "can_modify_num_rows", None
+ )
+ if not isinstance(descriptor, property):
+ self.skipTest("Ray does not use legacy MapBatches cardinality")
+
+ operator._input_dependencies = [
+ types.SimpleNamespace(
+ infer_metadata=mock.Mock(
+ return_value=types.SimpleNamespace(num_rows=1)
+ ),
+ can_modify_num_rows=None,
+ )
+ ]
+ self.assertIsNone(_estimate_dataset_num_rows(dataset))
+
+ def test_sparse_row_ids_keep_shuffle_parallelism(self):
+ with mock.patch(
+ "ray.cluster_resources", return_value={"CPU": 320}
+ ):
+ self.assertEqual(
+ _resolve_row_id_num_partitions(None, 4096, 500, 500),
+ 200,
+ )
+ self.assertEqual(
+ _resolve_row_id_num_partitions(None, None, None, 500),
+ 200,
+ )
+ self.assertEqual(
+ _resolve_row_id_num_partitions(None, 32, 1, 500),
+ 1,
+ )
+ self.assertEqual(
+ _resolve_row_id_num_partitions(37, 1, 500, 500),
+ 37,
+ )
+
+ def test_row_id_sizing_uses_dataset_context(self):
+ data_context = types.SimpleNamespace(
+ target_max_block_size=128 * 1024 * 1024,
+ default_hash_shuffle_parallelism=25,
+ )
+ with mock.patch(
+ "ray.cluster_resources", return_value={"CPU": 320}
+ ), mock.patch(
+ "ray.data.context.DataContext.get_current",
+ side_effect=AssertionError("must use the Dataset context"),
+ ):
+ self.assertEqual(
+ _resolve_row_id_num_partitions(
+ None,
+ 512 * 1024 * 1024,
+ None,
+ 500,
+ data_context=data_context,
+ ),
+ 25,
+ )
+
+ def test_merge_estimate_uses_only_reliable_source_size(self):
+ ctx = types.SimpleNamespace(
+ is_self_merge=False,
+ )
+
+ with mock.patch(
+ "pypaimon.ray.data_evolution_merge_into."
+ "_estimate_dataset_size_bytes",
+ return_value=100,
+ ):
+ result = _estimate_merge_input_size_bytes(
+ object(), ctx,
+ )
+
+ self.assertEqual(result, 100)
+
+ def test_merge_estimate_falls_back_when_source_size_is_unknown(self):
+ ctx = types.SimpleNamespace(is_self_merge=False)
+ with mock.patch(
+ "pypaimon.ray.data_evolution_merge_into."
+ "_estimate_dataset_size_bytes",
+ return_value=None,
+ ):
+ result = _estimate_merge_input_size_bytes(
+ object(), ctx,
+ )
+
+ self.assertIsNone(result)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/paimon-python/pypaimon/tests/ray_read_by_row_id_test.py
b/paimon-python/pypaimon/tests/ray_read_by_row_id_test.py
index c49d359042..d14392634b 100644
--- a/paimon-python/pypaimon/tests/ray_read_by_row_id_test.py
+++ b/paimon-python/pypaimon/tests/ray_read_by_row_id_test.py
@@ -283,13 +283,23 @@ class RayReadByRowIdTest(unittest.TestCase):
captured = {}
def fake_read(rid_ds, table, projection, *, num_partitions,
- ray_remote_args=None, base_snapshot_id=None):
+ ray_remote_args=None, base_snapshot_id=None,
+ estimated_size_bytes=None, estimated_num_rows=None,
+ data_context=None):
captured["base_snapshot_id"] = base_snapshot_id
+ captured["num_partitions"] = num_partitions
+ captured["estimated_size_bytes"] = estimated_size_bytes
+ captured["estimated_num_rows"] = estimated_num_rows
+ captured["data_context"] = data_context
return ray.data.from_arrow(pa.table({"_ROW_ID": pa.array([],
pa.int64())}))
with mock.patch.object(m, "distributed_read_by_row_id", fake_read):
read_by_row_id(target, src, self.catalog_options,
projection=["age"])
self.assertEqual(captured["base_snapshot_id"], expected_sid)
+ self.assertIsNone(captured["num_partitions"])
+ self.assertGreater(captured["estimated_size_bytes"], 0)
+ self.assertEqual(captured["estimated_num_rows"], 1)
+ self.assertIsNotNone(captured["data_context"])
def test_accepts_pyarrow_and_pandas_source(self):
target = self._create()
diff --git a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py
b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py
index e2bda92237..47dfe7f185 100644
--- a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py
+++ b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py
@@ -15,6 +15,7 @@
# specific language governing permissions and limitations
# under the License.
+import inspect
import os
import shutil
import tempfile
@@ -30,7 +31,7 @@ pypaimon = pytest.importorskip("pypaimon")
ray = pytest.importorskip("ray")
from pypaimon import CatalogFactory, Schema
-from pypaimon.ray import update_by_row_id
+from pypaimon.ray import read_paimon, update_by_row_id
class RayUpdateByRowIdTest(unittest.TestCase):
@@ -167,13 +168,221 @@ class RayUpdateByRowIdTest(unittest.TestCase):
captured = {}
def fake_apply(update_ds, table, cols, *, num_partitions,
- ray_remote_args=None, base_snapshot_id=None):
+ ray_remote_args=None, base_snapshot_id=None,
+ estimated_size_bytes=None, estimated_num_rows=None,
+ data_context=None):
captured["base_snapshot_id"] = base_snapshot_id
+ captured["num_partitions"] = num_partitions
+ captured["estimated_size_bytes"] = estimated_size_bytes
+ captured["estimated_num_rows"] = estimated_num_rows
+ captured["data_context"] = data_context
return [], 0, []
with mock.patch.object(m, "distributed_update_apply", fake_apply):
update_by_row_id(target, src, self.catalog_options,
update_cols=["age"])
self.assertEqual(captured["base_snapshot_id"], expected_sid)
+ self.assertIsNone(captured["num_partitions"])
+ self.assertGreater(captured["estimated_size_bytes"], 0)
+ self.assertEqual(captured["estimated_num_rows"], 1)
+ self.assertIsNotNone(captured["data_context"])
+
+ def test_transformed_paimon_source_adapts_partitions(self):
+ import pypaimon.ray.data_evolution_merge_join as merge_join
+
+ target = self._create()
+ for value in range(3):
+ self._write(target, pa.Table.from_pydict(
+ {"id": [value], "name": ["n{}".format(value)], "age": [0]},
+ schema=self.pa_schema,
+ ))
+
+ source = read_paimon(
+ target,
+ self.catalog_options,
+ projection=["_ROW_ID", "age"],
+ ).map_batches(
+ lambda batch: pa.table({
+ "_ROW_ID": batch.column("_ROW_ID"),
+ "age": pa.array([99] * batch.num_rows, type=pa.int32()),
+ }),
+ batch_format="pyarrow",
+ )
+ real_resolve = merge_join._resolve_row_id_num_partitions
+ resolved = []
+
+ def track_resolve(*args, **kwargs):
+ result = real_resolve(*args, **kwargs)
+ resolved.append((args, result))
+ return result
+
+ with mock.patch(
+ "ray.cluster_resources", return_value={"CPU": 320}
+ ), mock.patch.object(
+ merge_join,
+ "_resolve_row_id_num_partitions",
+ side_effect=track_resolve,
+ ):
+ stats = update_by_row_id(
+ target,
+ source,
+ self.catalog_options,
+ update_cols=["age"],
+ )
+
+ self.assertEqual(stats, {"num_updated": 3})
+ self.assertEqual([result for _, result in resolved], [3])
+ cardinality_parameter = inspect.signature(
+ ray.data.Dataset.map_batches
+ ).parameters.get("udf_modifying_row_count")
+ expected_num_rows = (
+ 3
+ if cardinality_parameter is not None
+ and cardinality_parameter.default is False
+ else None
+ )
+ self.assertEqual(
+ resolved[0][0][1:], (None, expected_num_rows, 3)
+ )
+ self.assertEqual(self._read(target).column("age").to_pylist(), [99] *
3)
+
+ def test_expanding_transform_keeps_target_parallelism(self):
+ import pypaimon.ray.data_evolution_merge_join as merge_join
+
+ target = self._create()
+ for value in range(3):
+ self._write(target, pa.Table.from_pydict(
+ {"id": [value], "name": ["n{}".format(value)], "age": [0]},
+ schema=self.pa_schema,
+ ))
+
+ rows = read_paimon(
+ target,
+ self.catalog_options,
+ projection=["_ROW_ID"],
+ ).take_all()
+ row_ids = [row["_ROW_ID"] for row in rows]
+ map_kwargs = {"batch_format": "pyarrow"}
+ if "udf_modifying_row_count" in inspect.signature(
+ ray.data.Dataset.map_batches
+ ).parameters:
+ map_kwargs["udf_modifying_row_count"] = True
+ source = ray.data.from_arrow(pa.table({"seed": [0]})).map_batches(
+ lambda batch: pa.table({
+ "_ROW_ID": pa.array(row_ids, type=pa.int64()),
+ "age": pa.array([77] * len(row_ids), type=pa.int32()),
+ }),
+ **map_kwargs,
+ )
+ real_resolve = merge_join._resolve_row_id_num_partitions
+ resolved = []
+
+ def track_resolve(*args, **kwargs):
+ result = real_resolve(*args, **kwargs)
+ resolved.append((args, result))
+ return result
+
+ with mock.patch(
+ "ray.cluster_resources", return_value={"CPU": 320}
+ ), mock.patch.object(
+ merge_join,
+ "_resolve_row_id_num_partitions",
+ side_effect=track_resolve,
+ ):
+ stats = update_by_row_id(
+ target,
+ source,
+ self.catalog_options,
+ update_cols=["age"],
+ )
+
+ self.assertEqual(stats, {"num_updated": 3})
+ self.assertEqual(resolved, [((None, None, None, 3), 3)])
+ self.assertEqual(self._read(target).column("age").to_pylist(), [77] *
3)
+
+ def test_filtered_paimon_source_uses_ray_shuffle_default(self):
+ import pypaimon.ray.data_evolution_merge_join as merge_join
+
+ target = "default.u_{}".format(uuid.uuid4().hex[:8])
+ schema = pa.schema([
+ ("content_key", pa.string()),
+ ("clip_id", pa.int32()),
+ ("age", pa.int32()),
+ ])
+ options = dict(self.de_options)
+ options["target-file-row-num"] = "1"
+ self.catalog.create_table(
+ target,
+ Schema.from_pyarrow_schema(schema, options=options),
+ False,
+ )
+ matched_rows = 39
+ total_rows = 201
+ self._write(target, pa.Table.from_pydict({
+ "content_key": ["wanted"] * matched_rows
+ + ["other"] * (total_rows - matched_rows),
+ "clip_id": [7] * matched_rows
+ + [8] * (total_rows - matched_rows),
+ "age": [0] * total_rows,
+ }, schema=schema))
+
+ table = self.catalog.get_table(target)
+ predicates = table.new_read_builder().new_predicate_builder()
+ predicate = predicates.and_predicates([
+ predicates.equal("content_key", "wanted"),
+ predicates.equal("clip_id", 7),
+ ])
+ source = read_paimon(
+ target,
+ self.catalog_options,
+ filter=predicate,
+ projection=["_ROW_ID", "content_key", "clip_id", "age"],
+ ).map_batches(
+ lambda batch: pa.table({
+ "_ROW_ID": batch.column("_ROW_ID"),
+ "age": pa.array([88] * batch.num_rows, type=pa.int32()),
+ }),
+ batch_format="pyarrow",
+ )
+ real_resolve = merge_join._resolve_row_id_num_partitions
+ real_groupby = ray.data.Dataset.groupby
+ resolved = []
+ group_partitions = []
+
+ def track_resolve(*args, **kwargs):
+ with mock.patch(
+ "ray.cluster_resources", return_value={"CPU": 320}
+ ):
+ result = real_resolve(*args, **kwargs)
+ resolved.append((args, result))
+ return result
+
+ def run_small_groupby(dataset, key, num_partitions=None):
+ # Assert the requested value but keep the test execution small.
+ group_partitions.append(num_partitions)
+ return real_groupby(dataset, key, num_partitions=4)
+
+ with mock.patch.object(
+ merge_join,
+ "_resolve_row_id_num_partitions",
+ side_effect=track_resolve,
+ ), mock.patch.object(
+ ray.data.Dataset,
+ "groupby",
+ new=run_small_groupby,
+ ):
+ stats = update_by_row_id(
+ target,
+ source,
+ self.catalog_options,
+ update_cols=["age"],
+ )
+
+ self.assertEqual(stats, {"num_updated": matched_rows})
+ self.assertEqual(resolved, [((None, None, None, total_rows), 200)])
+ self.assertEqual(group_partitions, [200])
+ ages = self._read(target).column("age").to_pylist()
+ self.assertEqual(ages.count(88), matched_rows)
+ self.assertEqual(ages.count(0), total_rows - matched_rows)
def test_new_commit_failure_preserves_pending_messages(self):
err = RuntimeError("new_commit failed")
