XiaoHongbo-Hope commented on code in PR #9357: URL: https://github.com/apache/paimon/pull/9357#discussion_r3837628445
########## 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: Fixed in e33d72589. Callable and boolean forms are normalized; Ray 2.50/2.51 MapBatches remains unknown because it has no reliable cardinality flag. The regression now runs in the Ray 2.53 compatibility lane, and I also verified it on Ray 2.50.1. ########## 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: Fixed in e33d72589. The empty-target fast path now repartitions to num_partitions before the insert transform. The new truncated-target regression materializes three write blocks and verifies all 30 rows are inserted. ########## 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: Fixed in e33d72589. General MERGE and row-ID update/read now pass the source Dataset context into partition sizing. A 512 MiB input sealed with 128 MiB blocks still resolves to four partitions after the global context changes to 1 GiB. -- 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]
