| # 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. | |
| from abc import ABC, abstractmethod | |
| from typing import Callable, List, Optional, Dict, Tuple | |
| from pypaimon.common.options.core_options import CoreOptions | |
| from pypaimon.manifest.schema.data_file_meta import DataFileMeta | |
| from pypaimon.manifest.schema.manifest_entry import ManifestEntry | |
| from pypaimon.read.scan_distribution import shard_range, validate_shard, validate_slice | |
| from pypaimon.read.split import Split | |
| from pypaimon.read.split import DataSplit | |
| from pypaimon.table.row.generic_row import GenericRow | |
| from pypaimon.table.source.deletion_file import DeletionFile | |
| from pypaimon.utils.file_store_path_factory import canonical_data_file_path | |
| class AbstractSplitGenerator(ABC): | |
| """ | |
| Abstract base class for generating splits. | |
| """ | |
| # Special key for tracking file end position in split file index map | |
| NEXT_POS_KEY = '_next_pos' | |
| def __init__( | |
| self, | |
| table, | |
| target_split_size: int, | |
| open_file_cost: int, | |
| deletion_files_map: Optional[Dict] = None, | |
| snapshot_id: Optional[int] = None, | |
| ): | |
| self.table = table | |
| self.target_split_size = target_split_size | |
| self.open_file_cost = open_file_cost | |
| self.deletion_files_map = deletion_files_map or {} | |
| self.snapshot_id = snapshot_id | |
| self.default_part_value = table.options.options.get( | |
| CoreOptions.PARTITION_DEFAULT_NAME, "__DEFAULT_PARTITION__") | |
| # Shard configuration | |
| self.idx_of_this_subtask = None | |
| self.number_of_para_subtasks = None | |
| self.start_pos_of_this_subtask = None | |
| self.end_pos_of_this_subtask = None | |
| def with_shard(self, idx_of_this_subtask: int, number_of_para_subtasks: int): | |
| """Configure sharding for parallel processing.""" | |
| validate_shard(idx_of_this_subtask, number_of_para_subtasks) | |
| if self.start_pos_of_this_subtask is not None: | |
| raise ValueError("with_shard and with_slice cannot be used simultaneously") | |
| self.idx_of_this_subtask = idx_of_this_subtask | |
| self.number_of_para_subtasks = number_of_para_subtasks | |
| return self | |
| def with_slice(self, start_pos: int, end_pos: int): | |
| """Configure slice range for processing.""" | |
| validate_slice(start_pos, end_pos) | |
| if self.idx_of_this_subtask is not None: | |
| raise ValueError("with_slice and with_shard cannot be used simultaneously") | |
| self.start_pos_of_this_subtask = start_pos | |
| self.end_pos_of_this_subtask = end_pos | |
| return self | |
| @abstractmethod | |
| def create_splits(self, file_entries: List[ManifestEntry]) -> List[Split]: | |
| """ | |
| Create splits from manifest entries. | |
| """ | |
| pass | |
| def _build_split_from_pack( | |
| self, | |
| packed_files: List[List[DataFileMeta]], | |
| file_entries: List[ManifestEntry], | |
| for_primary_key_split: bool, | |
| use_optimized_path: bool = False | |
| ) -> List[Split]: | |
| """ | |
| Build splits from packed files. | |
| """ | |
| splits = [] | |
| if not packed_files or not file_entries: | |
| return splits | |
| for file_group in packed_files: | |
| if use_optimized_path: | |
| raw_convertible = True | |
| elif for_primary_key_split: | |
| raw_convertible = len(file_group) == 1 and self._without_delete_row(file_group[0]) | |
| else: | |
| raw_convertible = True | |
| self._set_data_file_paths( | |
| file_group, file_entries[0].partition, file_entries[0].bucket) | |
| if file_group: | |
| # Get deletion files for this split | |
| data_deletion_files = None | |
| if self.deletion_files_map: | |
| data_deletion_files = self._get_deletion_files_for_split( | |
| file_group, | |
| file_entries[0].partition, | |
| file_entries[0].bucket | |
| ) | |
| split = DataSplit( | |
| files=file_group, | |
| partition=file_entries[0].partition, | |
| bucket=file_entries[0].bucket, | |
| raw_convertible=raw_convertible, | |
| data_deletion_files=data_deletion_files, | |
| snapshot_id=self.snapshot_id, | |
| ) | |
| splits.append(split) | |
| return splits | |
| def _set_data_file_paths(self, files, partition: GenericRow, bucket: int): | |
| """Resolve Java/Rust escaped paths and legacy Python partition paths.""" | |
| values = tuple(partition.values) | |
| escaped_partition = False | |
| if values: | |
| path_factory = self.table.path_factory() | |
| escaped_partition = ( | |
| path_factory.bucket_path(values, bucket, canonical_partition=True) | |
| != path_factory.bucket_path(values, bucket)) | |
| for data_file in files: | |
| data_file.set_file_path( | |
| self.table.table_path, partition, bucket, self.default_part_value, | |
| self.table.options.data_file_path_directory()) | |
| if escaped_partition and not data_file.external_path: | |
| canonical_path = canonical_data_file_path( | |
| self.table, values, bucket, data_file.file_name) | |
| if self.table.file_io.exists(canonical_path): | |
| data_file.file_path = canonical_path | |
| def _get_deletion_files_for_split( | |
| self, | |
| data_files: List[DataFileMeta], | |
| partition: GenericRow, | |
| bucket: int | |
| ) -> Optional[List[DeletionFile]]: | |
| """Get deletion files for the given data files in a split.""" | |
| if not self.deletion_files_map: | |
| return None | |
| partition_key = (tuple(partition.values), bucket) | |
| file_deletion_map = self.deletion_files_map.get(partition_key, {}) | |
| if not file_deletion_map: | |
| return None | |
| deletion_files = [] | |
| for data_file in data_files: | |
| deletion_file = file_deletion_map.get(data_file.file_name) | |
| if deletion_file: | |
| deletion_files.append(deletion_file) | |
| else: | |
| deletion_files.append(None) | |
| return deletion_files if any(df is not None for df in deletion_files) else None | |
| @staticmethod | |
| def _without_delete_row(data_file_meta: DataFileMeta) -> bool: | |
| """Check if a data file has no deleted rows.""" | |
| if data_file_meta.delete_row_count is None: | |
| return True | |
| return data_file_meta.delete_row_count == 0 | |
| @staticmethod | |
| def _pack_for_ordered( | |
| items: List, | |
| weight_func: Callable, | |
| target_weight: int | |
| ) -> List[List]: | |
| """Pack items into groups based on target weight.""" | |
| packed = [] | |
| bin_items = [] | |
| bin_weight = 0 | |
| for item in items: | |
| weight = weight_func(item) | |
| if bin_weight + weight > target_weight and len(bin_items) > 0: | |
| packed.append(list(bin_items)) | |
| bin_items.clear() | |
| bin_weight = 0 | |
| bin_weight += weight | |
| bin_items.append(item) | |
| if len(bin_items) > 0: | |
| packed.append(bin_items) | |
| return packed | |
| def _compute_shard_range(self, total_row: int) -> Tuple[int, int]: | |
| """ | |
| Calculate start and end positions for this shard based on total rows. | |
| Uses balanced distribution to avoid last shard overload. | |
| """ | |
| return shard_range(total_row, self.idx_of_this_subtask, self.number_of_para_subtasks) | |
| @staticmethod | |
| def _compute_file_range( | |
| plan_start_pos: int, | |
| plan_end_pos: int, | |
| file_begin_pos: int, | |
| file_row_count: int | |
| ) -> Optional[Tuple[int, int]]: | |
| """ | |
| Compute the row range to read from a file given shard range and file position. | |
| Returns None if file is completely within shard range (no slicing needed). | |
| Returns (-1, -1) if file is completely outside shard range. | |
| """ | |
| file_end_pos = file_begin_pos + file_row_count | |
| if file_begin_pos <= plan_start_pos < plan_end_pos <= file_end_pos: | |
| return plan_start_pos - file_begin_pos, plan_end_pos - file_begin_pos | |
| elif file_begin_pos < plan_start_pos < file_end_pos: | |
| return plan_start_pos - file_begin_pos, file_row_count | |
| elif file_begin_pos < plan_end_pos < file_end_pos: | |
| return 0, plan_end_pos - file_begin_pos | |
| elif file_end_pos <= plan_start_pos or file_begin_pos >= plan_end_pos: | |
| return -1, -1 | |
| # File is completely within the shard range | |
| return None |