| # 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 time |
| from enum import Enum |
| from typing import ( |
| Any, |
| Dict, |
| List, |
| Mapping, |
| Optional, |
| ) |
| |
| from pydantic import Field, PrivateAttr, model_serializer |
| |
| from pyiceberg.io import FileIO |
| from pyiceberg.manifest import DataFile, DataFileContent, ManifestFile, read_manifest_list |
| from pyiceberg.typedef import IcebergBaseModel |
| |
| ADDED_DATA_FILES = 'added-data-files' |
| ADDED_DELETE_FILES = 'added-delete-files' |
| ADDED_EQUALITY_DELETES = 'added-equality-deletes' |
| ADDED_FILE_SIZE = 'added-files-size' |
| ADDED_POSITION_DELETES = 'added-position-deletes' |
| ADDED_POSITION_DELETE_FILES = 'added-position-delete-files' |
| ADDED_RECORDS = 'added-records' |
| DELETED_DATA_FILES = 'deleted-data-files' |
| DELETED_RECORDS = 'deleted-records' |
| ADDED_EQUALITY_DELETE_FILES = 'added-equality-delete-files' |
| REMOVED_DELETE_FILES = 'removed-delete-files' |
| REMOVED_EQUALITY_DELETES = 'removed-equality-deletes' |
| REMOVED_EQUALITY_DELETE_FILES = 'removed-equality-delete-files' |
| REMOVED_FILE_SIZE = 'removed-files-size' |
| REMOVED_POSITION_DELETES = 'removed-position-deletes' |
| REMOVED_POSITION_DELETE_FILES = 'removed-position-delete-files' |
| TOTAL_EQUALITY_DELETES = 'total-equality-deletes' |
| TOTAL_POSITION_DELETES = 'total-position-deletes' |
| TOTAL_DATA_FILES = 'total-data-files' |
| TOTAL_DELETE_FILES = 'total-delete-files' |
| TOTAL_RECORDS = 'total-records' |
| TOTAL_FILE_SIZE = 'total-files-size' |
| |
| |
| OPERATION = "operation" |
| |
| |
| class Operation(Enum): |
| """Describes the operation. |
| |
| Possible operation values are: |
| - append: Only data files were added and no files were removed. |
| - replace: Data and delete files were added and removed without changing table data; i.e., compaction, changing the data file format, or relocating data files. |
| - overwrite: Data and delete files were added and removed in a logical overwrite operation. |
| - delete: Data files were removed and their contents logically deleted and/or delete files were added to delete rows. |
| """ |
| |
| APPEND = "append" |
| REPLACE = "replace" |
| OVERWRITE = "overwrite" |
| DELETE = "delete" |
| |
| def __repr__(self) -> str: |
| """Return the string representation of the Operation class.""" |
| return f"Operation.{self.name}" |
| |
| |
| class Summary(IcebergBaseModel, Mapping[str, str]): |
| """A class that stores the summary information for a Snapshot. |
| |
| The snapshot summary’s operation field is used by some operations, |
| like snapshot expiration, to skip processing certain snapshots. |
| """ |
| |
| operation: Operation = Field() |
| _additional_properties: Dict[str, str] = PrivateAttr() |
| |
| def __init__(self, operation: Operation, **data: Any) -> None: |
| super().__init__(operation=operation, **data) |
| self._additional_properties = data |
| |
| def __getitem__(self, __key: str) -> Optional[Any]: # type: ignore |
| """Return a key as it is a map.""" |
| if __key.lower() == 'operation': |
| return self.operation |
| else: |
| return self._additional_properties.get(__key) |
| |
| def __setitem__(self, key: str, value: Any) -> None: |
| """Set a key as it is a map.""" |
| if key.lower() == 'operation': |
| self.operation = value |
| else: |
| self._additional_properties[key] = value |
| |
| def __len__(self) -> int: |
| """Return the number of keys in the summary.""" |
| # Operation is required |
| return 1 + len(self._additional_properties) |
| |
| @model_serializer |
| def ser_model(self) -> Dict[str, str]: |
| return { |
| "operation": str(self.operation.value), |
| **self._additional_properties, |
| } |
| |
| @property |
| def additional_properties(self) -> Dict[str, str]: |
| return self._additional_properties |
| |
| def __repr__(self) -> str: |
| """Return the string representation of the Summary class.""" |
| repr_properties = f", **{repr(self._additional_properties)}" if self._additional_properties else "" |
| return f"Summary({repr(self.operation)}{repr_properties})" |
| |
| def __eq__(self, other: Any) -> bool: |
| """Compare if the summary is equal to another summary.""" |
| return ( |
| self.operation == other.operation and self.additional_properties == other.additional_properties |
| if isinstance(other, Summary) |
| else False |
| ) |
| |
| |
| class Snapshot(IcebergBaseModel): |
| snapshot_id: int = Field(alias="snapshot-id") |
| parent_snapshot_id: Optional[int] = Field(alias="parent-snapshot-id", default=None) |
| sequence_number: Optional[int] = Field(alias="sequence-number", default=None) |
| timestamp_ms: int = Field(alias="timestamp-ms", default_factory=lambda: int(time.time() * 1000)) |
| manifest_list: Optional[str] = Field( |
| alias="manifest-list", description="Location of the snapshot's manifest list file", default=None |
| ) |
| summary: Optional[Summary] = Field(default=None) |
| schema_id: Optional[int] = Field(alias="schema-id", default=None) |
| |
| def __str__(self) -> str: |
| """Return the string representation of the Snapshot class.""" |
| operation = f"{self.summary.operation}: " if self.summary else "" |
| parent_id = f", parent_id={self.parent_snapshot_id}" if self.parent_snapshot_id else "" |
| schema_id = f", schema_id={self.schema_id}" if self.schema_id is not None else "" |
| result_str = f"{operation}id={self.snapshot_id}{parent_id}{schema_id}" |
| return result_str |
| |
| def manifests(self, io: FileIO) -> List[ManifestFile]: |
| if self.manifest_list is not None: |
| file = io.new_input(self.manifest_list) |
| return list(read_manifest_list(file)) |
| return [] |
| |
| |
| class MetadataLogEntry(IcebergBaseModel): |
| metadata_file: str = Field(alias="metadata-file") |
| timestamp_ms: int = Field(alias="timestamp-ms") |
| |
| |
| class SnapshotLogEntry(IcebergBaseModel): |
| snapshot_id: int = Field(alias="snapshot-id") |
| timestamp_ms: int = Field(alias="timestamp-ms") |
| |
| |
| class SnapshotSummaryCollector: |
| added_file_size: int |
| removed_file_size: int |
| added_data_files: int |
| removed_data_files: int |
| added_eq_delete_files: int |
| removed_eq_delete_files: int |
| added_pos_delete_files: int |
| removed_pos_delete_files: int |
| added_delete_files: int |
| removed_delete_files: int |
| added_records: int |
| deleted_records: int |
| added_pos_deletes: int |
| removed_pos_deletes: int |
| added_eq_deletes: int |
| removed_eq_deletes: int |
| |
| def __init__(self) -> None: |
| self.added_file_size = 0 |
| self.removed_file_size = 0 |
| self.added_data_files = 0 |
| self.removed_data_files = 0 |
| self.added_eq_delete_files = 0 |
| self.removed_eq_delete_files = 0 |
| self.added_pos_delete_files = 0 |
| self.removed_pos_delete_files = 0 |
| self.added_delete_files = 0 |
| self.removed_delete_files = 0 |
| self.added_records = 0 |
| self.deleted_records = 0 |
| self.added_pos_deletes = 0 |
| self.removed_pos_deletes = 0 |
| self.added_eq_deletes = 0 |
| self.removed_eq_deletes = 0 |
| |
| def add_file(self, data_file: DataFile) -> None: |
| self.added_file_size += data_file.file_size_in_bytes |
| |
| if data_file.content == DataFileContent.DATA: |
| self.added_data_files += 1 |
| self.added_records += data_file.record_count |
| elif data_file.content == DataFileContent.POSITION_DELETES: |
| self.added_delete_files += 1 |
| self.added_pos_delete_files += 1 |
| self.added_pos_deletes += data_file.record_count |
| elif data_file.content == DataFileContent.EQUALITY_DELETES: |
| self.added_delete_files += 1 |
| self.added_eq_delete_files += 1 |
| self.added_eq_deletes += data_file.record_count |
| else: |
| raise ValueError(f"Unknown data file content: {data_file.content}") |
| |
| def remove_file(self, data_file: DataFile) -> None: |
| self.removed_file_size += data_file.file_size_in_bytes |
| |
| if data_file.content == DataFileContent.DATA: |
| self.removed_data_files += 1 |
| self.deleted_records += data_file.record_count |
| elif data_file.content == DataFileContent.POSITION_DELETES: |
| self.removed_delete_files += 1 |
| self.removed_pos_delete_files += 1 |
| self.removed_pos_deletes += data_file.record_count |
| elif data_file.content == DataFileContent.EQUALITY_DELETES: |
| self.removed_delete_files += 1 |
| self.removed_eq_delete_files += 1 |
| self.removed_eq_deletes += data_file.record_count |
| else: |
| raise ValueError(f"Unknown data file content: {data_file.content}") |
| |
| def build(self) -> Dict[str, str]: |
| def set_when_positive(properties: Dict[str, str], num: int, property_name: str) -> None: |
| if num > 0: |
| properties[property_name] = str(num) |
| |
| properties: Dict[str, str] = {} |
| set_when_positive(properties, self.added_file_size, ADDED_FILE_SIZE) |
| set_when_positive(properties, self.removed_file_size, REMOVED_FILE_SIZE) |
| set_when_positive(properties, self.added_data_files, ADDED_DATA_FILES) |
| set_when_positive(properties, self.removed_data_files, DELETED_DATA_FILES) |
| set_when_positive(properties, self.added_eq_delete_files, ADDED_EQUALITY_DELETE_FILES) |
| set_when_positive(properties, self.removed_eq_delete_files, REMOVED_EQUALITY_DELETE_FILES) |
| set_when_positive(properties, self.added_pos_delete_files, ADDED_POSITION_DELETE_FILES) |
| set_when_positive(properties, self.removed_pos_delete_files, REMOVED_POSITION_DELETE_FILES) |
| set_when_positive(properties, self.added_delete_files, ADDED_DELETE_FILES) |
| set_when_positive(properties, self.removed_delete_files, REMOVED_DELETE_FILES) |
| set_when_positive(properties, self.added_records, ADDED_RECORDS) |
| set_when_positive(properties, self.deleted_records, DELETED_RECORDS) |
| set_when_positive(properties, self.added_pos_deletes, ADDED_POSITION_DELETES) |
| set_when_positive(properties, self.removed_pos_deletes, REMOVED_POSITION_DELETES) |
| set_when_positive(properties, self.added_eq_deletes, ADDED_EQUALITY_DELETES) |
| set_when_positive(properties, self.removed_eq_deletes, REMOVED_EQUALITY_DELETES) |
| |
| return properties |
| |
| |
| def _truncate_table_summary(summary: Summary, previous_summary: Mapping[str, str]) -> Summary: |
| for prop in { |
| TOTAL_DATA_FILES, |
| TOTAL_DELETE_FILES, |
| TOTAL_RECORDS, |
| TOTAL_FILE_SIZE, |
| TOTAL_POSITION_DELETES, |
| TOTAL_EQUALITY_DELETES, |
| }: |
| summary[prop] = '0' |
| |
| def get_prop(prop: str) -> int: |
| value = previous_summary.get(prop) or '0' |
| try: |
| return int(value) |
| except ValueError as e: |
| raise ValueError(f"Could not parse summary property {prop} to an int: {value}") from e |
| |
| if value := get_prop(TOTAL_DATA_FILES): |
| summary[DELETED_DATA_FILES] = str(value) |
| if value := get_prop(TOTAL_DELETE_FILES): |
| summary[REMOVED_DELETE_FILES] = str(value) |
| if value := get_prop(TOTAL_RECORDS): |
| summary[DELETED_RECORDS] = str(value) |
| if value := get_prop(TOTAL_FILE_SIZE): |
| summary[REMOVED_FILE_SIZE] = str(value) |
| if value := get_prop(TOTAL_POSITION_DELETES): |
| summary[REMOVED_POSITION_DELETES] = str(value) |
| if value := get_prop(TOTAL_EQUALITY_DELETES): |
| summary[REMOVED_EQUALITY_DELETES] = str(value) |
| |
| return summary |
| |
| |
| def update_snapshot_summaries( |
| summary: Summary, previous_summary: Optional[Mapping[str, str]] = None, truncate_full_table: bool = False |
| ) -> Summary: |
| if summary.operation not in {Operation.APPEND, Operation.OVERWRITE}: |
| raise ValueError(f"Operation not implemented: {summary.operation}") |
| |
| if truncate_full_table and summary.operation == Operation.OVERWRITE and previous_summary is not None: |
| summary = _truncate_table_summary(summary, previous_summary) |
| |
| if not previous_summary: |
| previous_summary = { |
| TOTAL_DATA_FILES: '0', |
| TOTAL_DELETE_FILES: '0', |
| TOTAL_RECORDS: '0', |
| TOTAL_FILE_SIZE: '0', |
| TOTAL_POSITION_DELETES: '0', |
| TOTAL_EQUALITY_DELETES: '0', |
| } |
| |
| def _update_totals(total_property: str, added_property: str, removed_property: str) -> None: |
| if previous_total_str := previous_summary.get(total_property): |
| try: |
| new_total = int(previous_total_str) |
| if new_total >= 0 and (added := summary.get(added_property)): |
| new_total += int(added) |
| if new_total >= 0 and (removed := summary.get(removed_property)): |
| new_total -= int(removed) |
| except ValueError as e: |
| raise ValueError(f"Could not parse summary property {total_property} to an int: {previous_total_str}") from e |
| |
| if new_total >= 0: |
| summary[total_property] = str(new_total) |
| |
| _update_totals( |
| total_property=TOTAL_DATA_FILES, |
| added_property=ADDED_DATA_FILES, |
| removed_property=DELETED_DATA_FILES, |
| ) |
| _update_totals( |
| total_property=TOTAL_DELETE_FILES, |
| added_property=ADDED_DELETE_FILES, |
| removed_property=REMOVED_DELETE_FILES, |
| ) |
| _update_totals( |
| total_property=TOTAL_RECORDS, |
| added_property=ADDED_RECORDS, |
| removed_property=DELETED_RECORDS, |
| ) |
| _update_totals( |
| total_property=TOTAL_FILE_SIZE, |
| added_property=ADDED_FILE_SIZE, |
| removed_property=REMOVED_FILE_SIZE, |
| ) |
| _update_totals( |
| total_property=TOTAL_POSITION_DELETES, |
| added_property=ADDED_POSITION_DELETES, |
| removed_property=REMOVED_POSITION_DELETES, |
| ) |
| _update_totals( |
| total_property=TOTAL_EQUALITY_DELETES, |
| added_property=ADDED_EQUALITY_DELETES, |
| removed_property=REMOVED_EQUALITY_DELETES, |
| ) |
| |
| return summary |
