-
Notifications
You must be signed in to change notification settings - Fork 471
feat: Add metadata-only replace API to Table for REPLACE snapshot operations #3131
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 5 commits
19e8dee
6800bb4
e5e11b9
070c885
f12fa5d
94bd87e
a2f2b18
356d704
87f6848
41eb549
a91dfb4
ef9b84f
596df80
33aaef0
b0a770c
59f555e
c8162a8
d939b67
c3570d8
d7e89db
9681ec3
8f1f9b9
c60d5ad
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||
|---|---|---|---|---|
|
|
@@ -667,6 +667,82 @@ def _get_entries(manifest: ManifestFile) -> list[ManifestEntry]: | |||
| return [] | ||||
|
|
||||
|
|
||||
| class _RewriteFiles(_SnapshotProducer["_RewriteFiles"]): | ||||
| """A snapshot producer that rewrites data files.""" | ||||
|
|
||||
| def __init__(self, operation: Operation, transaction: Transaction, io: FileIO, snapshot_properties: dict[str, str]): | ||||
| super().__init__(operation, transaction, io, snapshot_properties=snapshot_properties) | ||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: we can remove the constructor here
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||||
|
|
||||
| def _commit(self) -> UpdatesAndRequirements: | ||||
| # Only produce a commit when there is something to rewrite | ||||
| if self._deleted_data_files or self._added_data_files: | ||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think we can replicate the
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||||
| return super()._commit() | ||||
| else: | ||||
| return (), () | ||||
|
|
||||
| def _deleted_entries(self) -> list[ManifestEntry]: | ||||
| """Check if we need to mark the files as deleted.""" | ||||
| if self._parent_snapshot_id is not None: | ||||
| previous_snapshot = self._transaction.table_metadata.snapshot_by_id(self._parent_snapshot_id) | ||||
| if previous_snapshot is None: | ||||
| raise ValueError(f"Could not find the previous snapshot: {self._parent_snapshot_id}") | ||||
|
|
||||
| executor = ExecutorFactory.get_or_create() | ||||
|
|
||||
| def _get_entries(manifest: ManifestFile) -> list[ManifestEntry]: | ||||
| return [ | ||||
| ManifestEntry.from_args( | ||||
| status=ManifestEntryStatus.DELETED, | ||||
| snapshot_id=entry.snapshot_id, | ||||
| sequence_number=entry.sequence_number, | ||||
| file_sequence_number=entry.file_sequence_number, | ||||
| data_file=entry.data_file, | ||||
| ) | ||||
| for entry in manifest.fetch_manifest_entry(self._io, discard_deleted=True) | ||||
| if entry.data_file.content == DataFileContent.DATA and entry.data_file in self._deleted_data_files | ||||
| ] | ||||
|
|
||||
| list_of_entries = executor.map(_get_entries, previous_snapshot.manifests(self._io)) | ||||
| return list(itertools.chain(*list_of_entries)) | ||||
| else: | ||||
| return [] | ||||
|
|
||||
| def _existing_manifests(self) -> list[ManifestFile]: | ||||
| """To determine if there are any existing manifests.""" | ||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This logic looks nearly identical to the
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||||
| existing_files = [] | ||||
| if snapshot := self._transaction.table_metadata.snapshot_by_name(name=self._target_branch): | ||||
| for manifest_file in snapshot.manifests(io=self._io): | ||||
| entries_to_write: set[ManifestEntry] = set() | ||||
| found_deleted_entries: set[ManifestEntry] = set() | ||||
|
|
||||
| for entry in manifest_file.fetch_manifest_entry(io=self._io, discard_deleted=True): | ||||
| if entry.data_file in self._deleted_data_files: | ||||
| found_deleted_entries.add(entry) | ||||
| else: | ||||
| entries_to_write.add(entry) | ||||
|
|
||||
| if len(found_deleted_entries) == 0: | ||||
| existing_files.append(manifest_file) | ||||
| continue | ||||
|
|
||||
| if len(entries_to_write) == 0: | ||||
| continue | ||||
|
|
||||
| with self.new_manifest_writer(self.spec(manifest_file.partition_spec_id)) as writer: | ||||
| for entry in entries_to_write: | ||||
| writer.add_entry( | ||||
| ManifestEntry.from_args( | ||||
| status=ManifestEntryStatus.EXISTING, | ||||
| snapshot_id=entry.snapshot_id, | ||||
| sequence_number=entry.sequence_number, | ||||
| file_sequence_number=entry.file_sequence_number, | ||||
| data_file=entry.data_file, | ||||
| ) | ||||
| ) | ||||
| existing_files.append(writer.to_manifest_file()) | ||||
| return existing_files | ||||
|
|
||||
|
|
||||
| class UpdateSnapshot: | ||||
| _transaction: Transaction | ||||
| _io: FileIO | ||||
|
|
@@ -724,6 +800,14 @@ def delete(self) -> _DeleteFiles: | |||
| snapshot_properties=self._snapshot_properties, | ||||
| ) | ||||
|
|
||||
| def replace(self) -> _RewriteFiles: | ||||
| return _RewriteFiles( | ||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'm sort of confused by the naming since we are introducing a user facing API
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Hi @geruh, you bring up a good point, and it's something I noticed seemed off along the way. The reason why we have this discrepancy is because we're mirroring what's found in the Java code itself.
I named the Python API That said, if you feel strongly about matching the Java API's user-facing method (
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yeah there is a bit of a distinction here, since rewrite is basically the rewrite of data files and replace is the logical change to your snapshot metadata. My thinking is that the users in java today are used to interacting with this api through: table.newRewrite()
.deleteFile(old)
.addFile(new)
.commit();So someone coming from Java Iceberg will look for rewrite, not replace. But ultimately maybe there is more of a history as to why the it follows this naming convention im missing on. WDYT @kevinjqliu? |
||||
| operation=Operation.REPLACE, | ||||
| transaction=self._transaction, | ||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I noticed that branch is missing here is there a reason for that?
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||||
| io=self._io, | ||||
| snapshot_properties=self._snapshot_properties, | ||||
| ) | ||||
|
|
||||
|
|
||||
| class _ManifestMergeManager(Generic[U]): | ||||
| _target_size_bytes: int | ||||
|
|
||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,109 @@ | ||
| # 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 pyiceberg.catalog import Catalog | ||
| from pyiceberg.manifest import DataFile, DataFileContent, FileFormat | ||
| from pyiceberg.schema import Schema | ||
| from pyiceberg.table.snapshots import Operation | ||
| from pyiceberg.typedef import Record | ||
|
|
||
|
|
||
| def test_replace_api(catalog: Catalog) -> None: | ||
| # Setup a basic table using the catalog fixture | ||
| catalog.create_namespace("default") | ||
| table = catalog.create_table( | ||
| identifier="default.test_replace", | ||
| schema=Schema(), | ||
| ) | ||
|
|
||
| # Create mock DataFiles for the test | ||
| file_to_delete = DataFile.from_args( | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: extract this into a helper or see if one exists already.
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||
| file_path="s3://bucket/test/data/deleted.parquet", | ||
| file_format=FileFormat.PARQUET, | ||
| partition=Record(), | ||
| record_count=100, | ||
| file_size_in_bytes=1024, | ||
| content=DataFileContent.DATA, | ||
| ) | ||
| file_to_delete.spec_id = 0 | ||
|
|
||
| file_to_add = DataFile.from_args( | ||
| file_path="s3://bucket/test/data/added.parquet", | ||
| file_format=FileFormat.PARQUET, | ||
| partition=Record(), | ||
| record_count=100, | ||
| file_size_in_bytes=1024, | ||
| content=DataFileContent.DATA, | ||
| ) | ||
| file_to_add.spec_id = 0 | ||
|
|
||
| # Initially append to have something to replace | ||
| with table.transaction() as tx: | ||
| with tx.update_snapshot().fast_append() as append_snapshot: | ||
| append_snapshot.append_data_file(file_to_delete) | ||
|
|
||
| # Verify initial append snapshot | ||
| assert len(table.history()) == 1 | ||
| snapshot = table.current_snapshot() | ||
| assert snapshot is not None | ||
| assert snapshot.summary is not None | ||
| assert snapshot.summary["operation"] == Operation.APPEND | ||
|
|
||
| # Call the replace API | ||
| table.replace(files_to_delete=[file_to_delete], files_to_add=[file_to_add]) | ||
|
|
||
| # Verify the replacement created a REPLACE snapshot | ||
| assert len(table.history()) == 2 | ||
| snapshot = table.current_snapshot() | ||
| assert snapshot is not None | ||
| assert snapshot.summary is not None | ||
| assert snapshot.summary["operation"] == Operation.REPLACE | ||
|
|
||
| # Verify the correct files are added and deleted | ||
| # The summary property tracks these counts | ||
| assert snapshot.summary["added-data-files"] == "1" | ||
| assert snapshot.summary["deleted-data-files"] == "1" | ||
| assert snapshot.summary["added-records"] == "100" | ||
| assert snapshot.summary["deleted-records"] == "100" | ||
|
|
||
| # Verify the new file exists in the new manifest | ||
| manifest_files = snapshot.manifests(table.io) | ||
| assert len(manifest_files) == 2 # One for ADDED, one for DELETED | ||
|
|
||
| # Check that sequence numbers were handled properly natively by verifying the manifest contents | ||
| entries = [] | ||
| for manifest in manifest_files: | ||
| for entry in manifest.fetch_manifest_entry(table.io, discard_deleted=False): | ||
| entries.append(entry) | ||
|
|
||
| # One entry for ADDED (new file), one for DELETED (old file) | ||
| assert len(entries) == 2 | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. we need to test a bit more on the status of each file. https://github.com/apache/iceberg/blob/main/core/src/test/java/org/apache/iceberg/TestRewriteFiles.java is a good reference
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Hi @kevinjqliu, I've addressed this in the latest 33aaef0, where the status of each file is tested. |
||
|
|
||
|
|
||
| def test_replace_empty_files(catalog: Catalog) -> None: | ||
| # Setup a basic table using the catalog fixture | ||
| catalog.create_namespace("default") | ||
| table = catalog.create_table( | ||
| identifier="default.test_replace_empty", | ||
| schema=Schema(), | ||
| ) | ||
|
|
||
| # Replacing empty lists should not throw errors, but should produce no changes. | ||
| table.replace([], []) | ||
|
|
||
| # History should be completely empty since no files were rewritten | ||
| assert len(table.history()) == 0 | ||
| assert table.current_snapshot() is None | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -398,8 +398,8 @@ def test_merge_snapshot_summaries_overwrite_summary() -> None: | |
|
|
||
| def test_invalid_operation() -> None: | ||
| with pytest.raises(ValueError) as e: | ||
| update_snapshot_summaries(summary=Summary(Operation.REPLACE)) | ||
| assert "Operation not implemented: Operation.REPLACE" in str(e.value) | ||
| update_snapshot_summaries(summary=Summary.model_construct(operation="unknown_operation")) | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Makes sense if we have support :) |
||
| assert "Operation not implemented: unknown_operation" in str(e.value) | ||
|
|
||
|
|
||
| def test_invalid_type() -> None: | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
we should not expose
replaceas a public function as we cannot guarantee that thefiles_to_deleteandfiles_to_addcontains the same records.I think we should start at
_RewriteFilesThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hi @kevinjqliu, I think that logic makes sense, as it would be dangerous for users to use these without being able to enforce the underlying expectations of the values input to these functions. Removed them in 94bd87e