| # 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. | |
| """Global-index coverage for data-evolution tables.""" | |
| from typing import Collection, Dict, List, Optional, Union | |
| from pypaimon.common.options.core_options import CoreOptions, GlobalIndexSearchMode | |
| from pypaimon.common.options.options import Options | |
| from pypaimon.common.predicate import Predicate | |
| from pypaimon.read.push_down_utils import _get_all_fields | |
| from pypaimon.schema.data_types import DataField | |
| from pypaimon.utils.range import Range | |
| class DataEvolutionGlobalIndexCoverage: | |
| """Computes global-index coverage by field id.""" | |
| def __init__( | |
| self, | |
| table, | |
| snapshot, | |
| partition_filter, | |
| index_files: Collection['IndexFileMeta'], | |
| ): | |
| self._table = table | |
| self._snapshot = snapshot | |
| self._partition_filter = partition_filter | |
| self._coverage_by_field: Dict[int, List[Range]] = {} | |
| for index_file in index_files: | |
| meta = index_file.global_index_meta | |
| if meta is None: | |
| continue | |
| if getattr(index_file, 'index_type', None) == 'btree' and meta.extra_field_ids: | |
| continue | |
| row_range = Range(meta.row_range_start, meta.row_range_end) | |
| self._add_coverage(meta.index_field_id, row_range) | |
| for field_id in meta.extra_field_ids or []: | |
| self._add_coverage(field_id, row_range) | |
| def unindexed_ranges( | |
| self, | |
| fields_or_field_id: Union[List[DataField], Collection[int], int], | |
| predicate: Optional[Predicate] = None, | |
| search_mode: Optional[GlobalIndexSearchMode] = None, | |
| ) -> List[Range]: | |
| if isinstance(fields_or_field_id, int): | |
| field_ids = {fields_or_field_id} | |
| elif _is_field_id_collection(fields_or_field_id): | |
| field_ids = set(fields_or_field_id) | |
| else: | |
| field_by_name = {f.name: f for f in fields_or_field_id} | |
| field_ids = set() | |
| for name in _get_all_fields(predicate): | |
| field = field_by_name.get(name) | |
| if field is not None: | |
| field_ids.add(field.id) | |
| return self._unindexed_ranges(field_ids, search_mode) | |
| def indexed_ranges(self, field_ids): | |
| return self._indexed_ranges(field_ids) | |
| def unindexed_ranges_from_indexed(self, indexed_ranges, search_mode=None): | |
| return self._unindexed_ranges((), search_mode, indexed_ranges) | |
| def _add_coverage(self, field_id: int, row_range: Range) -> None: | |
| self._coverage_by_field.setdefault(field_id, []).append(row_range) | |
| def _indexed_ranges(self, field_ids: Collection[int]) -> List[Range]: | |
| ranges = None | |
| for field_id in field_ids: | |
| field_ranges = self._coverage_by_field.get(field_id) | |
| if not field_ranges: | |
| return [] | |
| field_ranges = Range.sort_and_merge_overlap(field_ranges, True) | |
| ranges = field_ranges if ranges is None else Range.and_(ranges, field_ranges) | |
| if ranges is None: | |
| return [] | |
| return Range.sort_and_merge_overlap(ranges, True) | |
| def _unindexed_ranges( | |
| self, | |
| field_ids: Collection[int], | |
| search_mode: Optional[GlobalIndexSearchMode] = None, | |
| indexed_ranges=None, | |
| ) -> List[Range]: | |
| if search_mode is None: | |
| search_mode = _scalar_index_search_mode(self._table) | |
| if search_mode == GlobalIndexSearchMode.FAST: | |
| return [] | |
| next_row_id = getattr(self._snapshot, "next_row_id", None) | |
| if self._snapshot is None or next_row_id is None: | |
| return [] | |
| if next_row_id <= 0: | |
| return [] | |
| if search_mode == GlobalIndexSearchMode.DETAIL: | |
| data_ranges = self._data_ranges_by_data_files() | |
| else: | |
| data_ranges = [Range(0, next_row_id - 1)] | |
| indexed_ranges = Range.sort_and_merge_overlap( | |
| self._indexed_ranges(field_ids) if indexed_ranges is None | |
| else indexed_ranges, True) | |
| unindexed = [] | |
| for data_range in Range.sort_and_merge_overlap(data_ranges, True): | |
| unindexed.extend(data_range.exclude(indexed_ranges)) | |
| return Range.sort_and_merge_overlap(unindexed, True) | |
| def _data_ranges_by_data_files(self) -> List[Range]: | |
| if hasattr(self._table, "data_ranges_for_data_evolution_global_index_coverage"): | |
| return self._table.data_ranges_for_data_evolution_global_index_coverage( | |
| self._snapshot, | |
| self._partition_filter, | |
| ) | |
| manifest_list_manager = getattr(self._table, "manifest_list_manager", None) | |
| if manifest_list_manager is None: | |
| from pypaimon.manifest.manifest_list_manager import ManifestListManager | |
| manifest_list_manager = ManifestListManager(self._table) | |
| manifest_files = manifest_list_manager.read_all(self._snapshot) | |
| from pypaimon.manifest.manifest_file_manager import ManifestFileManager | |
| manager = ManifestFileManager(self._table) | |
| entries = manager.read_entries_parallel( | |
| manifest_files, | |
| self._entry_matches_partition, | |
| max_workers=self._table.options.scan_manifest_parallelism(), | |
| ) | |
| data_ranges = [] | |
| for entry in entries: | |
| row_range = entry.file.row_id_range() | |
| if row_range is not None: | |
| data_ranges.append(row_range) | |
| return data_ranges | |
| def _entry_matches_partition(self, entry) -> bool: | |
| if self._partition_filter is None: | |
| return True | |
| return self._partition_filter.test(entry.partition) | |
| def _scalar_index_search_mode(table): | |
| options = getattr(table, "options", None) | |
| if options is None: | |
| return GlobalIndexSearchMode.FULL | |
| if hasattr(options, "scalar_index_search_mode"): | |
| return options.scalar_index_search_mode() | |
| return CoreOptions(Options.from_none()).scalar_index_search_mode() | |
| def _is_field_id_collection(fields_or_field_id): | |
| if isinstance(fields_or_field_id, (str, bytes)): | |
| return False | |
| try: | |
| iterator = iter(fields_or_field_id) | |
| except TypeError: | |
| return False | |
| return all(isinstance(field_id, int) for field_id in iterator) |