JingsongLi commented on code in PR #9357: URL: https://github.com/apache/paimon/pull/9357#discussion_r3837550831
########## paimon-python/pypaimon/ray/partitioning.py: ########## @@ -0,0 +1,144 @@ +################################################################################ +# 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, +) -> 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 + + target_size_bytes = int( + DataContext.get_current().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 = getattr( + operator, "can_modify_num_rows", None + ) + 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: Review Comment: [P2] Handle the callable cardinality flag on supported Ray versions On Ray 2.50 through 2.53, `can_modify_num_rows` is a method rather than a boolean field. In Ray 2.52 and 2.53, `MapBatches.can_modify_num_rows()` returns `udf_modifying_row_count`, so comparing the bound method directly with `False` stops traversal even when the caller explicitly used `udf_modifying_row_count=False`. The new cardinality test fails on Ray 2.53 with `None != 2`, and row-ID sizing can fall back from one partition to about 200. Please normalize the callable/boolean forms while retaining the conservative behavior for pre-2.52 MapBatches, and run this test in the Ray 2.53 compatibility lane. ########## paimon-python/pypaimon/ray/partitioning.py: ########## @@ -0,0 +1,144 @@ +################################################################################ +# 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, +) -> 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 + + target_size_bytes = int( + DataContext.get_current().target_max_block_size Review Comment: [P2] Size work from the Dataset sealed context Ray seals the `DataContext` when a Dataset is created, and later changes to `DataContext.get_current()` do not affect that Dataset or descendant operations. Reading the process-global context here (and again for `default_hash_shuffle_parallelism` below) can therefore choose a partition count from settings different from those used by the actual groupby or shuffle. For example, a 512 MiB Dataset captured with 128 MiB blocks should select four partitions, but changing the global setting to 1 GiB makes this helper select one. Please thread the relevant `source_ds.context` values into the row-ID and source-left paths; newly normalized Arrow/Pandas sources and explicit `num_partitions` can keep their current behavior. ########## paimon-python/pypaimon/ray/data_evolution_merge_into.py: ########## @@ -414,7 +437,7 @@ def _build_datasets( catalog_options=ctx.catalog_options, num_partitions=num_partitions, snapshot_id=base_snapshot_id, - target_empty=base_snapshot is None, + target_empty=target_empty, Review Comment: [P2] Preserve write parallelism for truncated targets For a never-written table this fast path already existed, but this change now sends truncated targets here as well. In this branch `build_not_matched_insert_ds` sets `unmatched = source_renamed`; no later operator consumes `num_partitions`, and `write_datasink` writes the existing blocks. Thus a one-block source remains one writer even with an explicit or resolved `num_partitions=37`, whereas the previous empty left-anti join produced 37 blocks. A large insert after truncate can therefore become a single-task bottleneck or OOM. Please repartition the fast-path dataset to `num_partitions` (or pass equivalent parallelism to the sink) and add a truncated-target block-count test. -- 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]
