zhjwpku commented on code in PR #774: URL: https://github.com/apache/iceberg-cpp/pull/774#discussion_r4443750734
########## src/iceberg/test/rewrite_manifests_test.cc: ########## @@ -0,0 +1,2802 @@ +/* + * 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. + */ + +#include "iceberg/update/rewrite_manifests.h" + +#include <algorithm> +#include <format> +#include <memory> +#include <mutex> +#include <optional> +#include <stdexcept> +#include <string> +#include <unordered_map> +#include <unordered_set> +#include <utility> +#include <vector> + +#include <gmock/gmock.h> +#include <gtest/gtest.h> + +#include "iceberg/avro/avro_register.h" +#include "iceberg/constants.h" +#include "iceberg/expression/expressions.h" +#include "iceberg/file_io.h" +#include "iceberg/file_reader.h" +#include "iceberg/partition_spec.h" +#include "iceberg/schema.h" +#include "iceberg/table_metadata.h" +#include "iceberg/table_properties.h" +#include "iceberg/test/executor.h" +#include "iceberg/test/mock_catalog.h" +#include "iceberg/test/update_test_base.h" +#include "iceberg/transaction.h" +#include "iceberg/transform.h" +#include "iceberg/update/delete_files.h" +#include "iceberg/update/fast_append.h" +#include "iceberg/update/row_delta.h" +#include "iceberg/update/update_partition_spec.h" +#include "iceberg/update/update_properties.h" +#include "iceberg/util/macros.h" + +namespace iceberg { + +namespace { + +class TestRewriteManifests : public RewriteManifests { + public: + static Result<std::shared_ptr<TestRewriteManifests>> Make( + const std::shared_ptr<Table>& table) { + ICEBERG_ASSIGN_OR_RAISE(auto ctx, + TransactionContext::Make(table, TransactionKind::kUpdate)); + return std::shared_ptr<TestRewriteManifests>( + new TestRewriteManifests(table->full_name(), std::move(ctx))); + } + + Result<std::vector<ManifestFile>> ApplyForTest( + const TableMetadata& metadata, const std::shared_ptr<Snapshot>& snapshot) { + return RewriteManifests::Apply(metadata, snapshot); + } + + private: + TestRewriteManifests(std::string table_name, std::shared_ptr<TransactionContext> ctx) + : RewriteManifests(std::move(table_name), std::move(ctx)) {} +}; + +struct ManifestReadState { + std::unordered_set<std::string> source_paths; + std::mutex mutex; + int open_readers = 0; + int max_open_readers = 0; + int opened_readers = 0; + int64_t unread_entries = 0; + int64_t max_unread_entries = 0; + int64_t processed_entries = 0; + bool fail_after_first_batch = false; +}; + +class TrackingManifestReader : public Reader { + public: + TrackingManifestReader(std::unique_ptr<Reader> delegate, + std::shared_ptr<ManifestReadState> state) + : delegate_(std::move(delegate)), state_(std::move(state)) {} + + ~TrackingManifestReader() override { Release(); } + + Status Open(const ReaderOptions& options) override { + auto read_options = options; + const bool track = state_->source_paths.contains(options.path); + if (track) { + read_options.properties.Set(ReaderProperties::kBatchSize, int64_t{2}); + } + ICEBERG_RETURN_UNEXPECTED(delegate_->Open(read_options)); + if (track) { + std::lock_guard lock(state_->mutex); + tracked_ = true; + ++state_->opened_readers; + state_->max_open_readers = + std::max(state_->max_open_readers, ++state_->open_readers); + } + return {}; + } + + Status Close() override { + Release(); + return delegate_->Close(); + } + + Result<std::optional<ArrowArray>> Next() override { + if (tracked_ && state_->fail_after_first_batch && read_batch_) { + return IOError("Injected streaming manifest read failure"); + } + ICEBERG_ASSIGN_OR_RAISE(auto batch, delegate_->Next()); + if (tracked_ && batch.has_value()) { + read_batch_ = true; + std::lock_guard lock(state_->mutex); + state_->unread_entries += batch->length; + state_->max_unread_entries = + std::max(state_->max_unread_entries, state_->unread_entries); + } + return batch; + } + + Result<ArrowSchema> Schema() override { return delegate_->Schema(); } + + Result<std::unordered_map<std::string, std::string>> Metadata() override { + return delegate_->Metadata(); + } + + private: + void Release() { + if (tracked_) { + std::lock_guard lock(state_->mutex); + --state_->open_readers; + tracked_ = false; + } + } + + std::unique_ptr<Reader> delegate_; + std::shared_ptr<ManifestReadState> state_; + bool tracked_ = false; + bool read_batch_ = false; +}; + +class ScopedManifestReadTracking { + public: + explicit ScopedManifestReadTracking(std::shared_ptr<ManifestReadState> state) + : previous_(ReaderFactoryRegistry::GetFactory(FileFormatType::kAvro)) { + ReaderFactoryRegistry::GetFactory(FileFormatType::kAvro) = + [factory = previous_, + state = std::move(state)]() -> Result<std::unique_ptr<Reader>> { + ICEBERG_ASSIGN_OR_RAISE(auto reader, factory()); + return std::make_unique<TrackingManifestReader>(std::move(reader), state); + }; + } + + ~ScopedManifestReadTracking() { + ReaderFactoryRegistry::GetFactory(FileFormatType::kAvro) = std::move(previous_); + } + + private: + ReaderFactory previous_; +}; + +struct ManifestWriteFailureState { + int manifest_output_count = 0; + int fail_on_manifest_output = 0; + int manifest_close_count = 0; + int fail_on_manifest_close = 0; + int throw_on_manifest_close = 0; + std::vector<std::string> manifest_outputs; + std::vector<std::string> manifest_close_attempts; + std::vector<std::string> deleted_files; + std::unordered_set<std::string> open_manifest_paths; + std::vector<std::string> deleted_open_files; +}; + +class ManifestWriteFailureOutputStream : public PositionOutputStream { + public: + ManifestWriteFailureOutputStream(std::unique_ptr<PositionOutputStream> delegate, + std::string location, + std::shared_ptr<ManifestWriteFailureState> state) + : delegate_(std::move(delegate)), + location_(std::move(location)), + state_(std::move(state)) { + state_->open_manifest_paths.insert(location_); + } + + Result<int64_t> Position() const override { return delegate_->Position(); } + + Result<int64_t> StoredLength() const override { return delegate_->StoredLength(); } + + Status Write(std::span<const std::byte> data) override { + return delegate_->Write(data); + } + + Status Flush() override { return delegate_->Flush(); } + + Status Close() override { + if (closed_) { + return {}; + } + closed_ = true; + state_->manifest_close_attempts.push_back(location_); + ++state_->manifest_close_count; + ICEBERG_RETURN_UNEXPECTED(delegate_->Close()); + state_->open_manifest_paths.erase(location_); + if (state_->manifest_close_count == state_->throw_on_manifest_close) { + throw std::runtime_error("Injected rewrite manifest close exception"); + } + if (state_->manifest_close_count == state_->fail_on_manifest_close) { + return IOError("Injected rewrite manifest close failure"); + } + return {}; + } + + private: + std::unique_ptr<PositionOutputStream> delegate_; + std::string location_; + std::shared_ptr<ManifestWriteFailureState> state_; + bool closed_{false}; +}; + +class ManifestWriteFailureOutputFile : public OutputFile { + public: + ManifestWriteFailureOutputFile(std::unique_ptr<OutputFile> delegate, + std::string location, + std::shared_ptr<ManifestWriteFailureState> state) + : delegate_(std::move(delegate)), + location_(std::move(location)), + state_(std::move(state)) {} + + std::string_view location() const override { return location_; } + + Result<std::unique_ptr<PositionOutputStream>> Create() override { + ICEBERG_ASSIGN_OR_RAISE(auto stream, delegate_->Create()); + return std::make_unique<ManifestWriteFailureOutputStream>(std::move(stream), + location_, state_); + } + + Result<std::unique_ptr<PositionOutputStream>> CreateOrOverwrite() override { + ICEBERG_ASSIGN_OR_RAISE(auto stream, delegate_->CreateOrOverwrite()); + return std::make_unique<ManifestWriteFailureOutputStream>(std::move(stream), + location_, state_); + } + + private: + std::unique_ptr<OutputFile> delegate_; + std::string location_; + std::shared_ptr<ManifestWriteFailureState> state_; +}; + +class ManifestWriteFailureFileIO : public FileIO { + public: + ManifestWriteFailureFileIO(std::shared_ptr<FileIO> delegate, + std::string metadata_location, + std::shared_ptr<ManifestWriteFailureState> state) + : delegate_(std::move(delegate)), + metadata_location_(std::move(metadata_location)), + state_(std::move(state)) {} + + Result<std::unique_ptr<InputFile>> NewInputFile(std::string file_location) override { + return delegate_->NewInputFile(std::move(file_location)); + } + + Result<std::unique_ptr<InputFile>> NewInputFile(std::string file_location, + size_t length) override { + return delegate_->NewInputFile(std::move(file_location), length); + } + + Result<std::unique_ptr<OutputFile>> NewOutputFile(std::string file_location) override { + const bool is_manifest_output = IsManifestOutput(file_location); + if (is_manifest_output) { + state_->manifest_outputs.push_back(file_location); + ++state_->manifest_output_count; + if (state_->manifest_output_count == state_->fail_on_manifest_output) { + return IOError("Injected rewrite manifest writer failure"); + } + } + ICEBERG_ASSIGN_OR_RAISE(auto output_file, delegate_->NewOutputFile(file_location)); + if (is_manifest_output) { + return std::make_unique<ManifestWriteFailureOutputFile>( + std::move(output_file), std::move(file_location), state_); + } + return output_file; + } + + Status DeleteFile(const std::string& file_location) override { + if (state_->open_manifest_paths.contains(file_location)) { + state_->deleted_open_files.push_back(file_location); + } + state_->deleted_files.push_back(file_location); + return delegate_->DeleteFile(file_location); + } + + Status DeleteFiles(const std::vector<std::string>& file_locations) override { + state_->deleted_files.insert(state_->deleted_files.end(), file_locations.begin(), + file_locations.end()); + return delegate_->DeleteFiles(file_locations); + } + + private: + bool IsManifestOutput(const std::string& file_location) const { + return file_location.starts_with(metadata_location_) && + file_location.ends_with(".avro"); + } + + std::shared_ptr<FileIO> delegate_; + std::string metadata_location_; + std::shared_ptr<ManifestWriteFailureState> state_; +}; + +} // namespace + +class RewriteManifestsTest : public MinimalUpdateTestBase, + public ::testing::WithParamInterface<int8_t> { + protected: + static void SetUpTestSuite() { avro::RegisterAll(); } + + void SetUp() override { + MinimalUpdateTestBase::SetUp(); + + ICEBERG_UNWRAP_OR_FAIL(spec_, table_->spec()); + ICEBERG_UNWRAP_OR_FAIL(schema_, table_->schema()); + file_a_ = MakeDataFile("a", 10, 100, 1); + file_b_ = MakeDataFile("b", 20, 200, 2); + file_c_ = MakeDataFile("c", 30, 300, 3); + file_d_ = MakeDataFile("d", 40, 400, 4); + } + + int8_t format_version() const override { return GetParam(); } + + std::shared_ptr<DataFile> MakeDataFile(const std::string& name, int64_t record_count, + int64_t size, int64_t partition_value) { + return MakeDataFile(name, spec_, record_count, size, + {Literal::Long(partition_value)}); + } + + std::shared_ptr<DataFile> MakeDataFile(const std::string& name, + std::shared_ptr<PartitionSpec> spec, + int64_t record_count, int64_t size, + std::vector<Literal> partition_values) { + auto file = std::make_shared<DataFile>(); + file->content = DataFile::Content::kData; + file->file_path = std::format("{}/data/{}.parquet", table_location_, name); + file->file_format = FileFormatType::kParquet; + file->partition = PartitionValues(std::move(partition_values)); + file->record_count = record_count; + file->file_size_in_bytes = size; + file->partition_spec_id = spec->spec_id(); + return file; + } + + std::shared_ptr<DataFile> MakeDeleteFile( + const std::string& name, int64_t partition_value, + std::optional<std::string> referenced_data_file = std::nullopt) { + auto file = MakeDataFile(name, 7, 50, partition_value); + file->content = DataFile::Content::kPositionDeletes; + file->file_path = std::format("{}/delete/{}.parquet", table_location_, name); + if (table_->metadata()->format_version >= 3) { + file->file_format = FileFormatType::kPuffin; + file->referenced_data_file = referenced_data_file.value_or(file_a_->file_path); + file->content_offset = 0; + file->content_size_in_bytes = 10; + } + return file; + } + + std::shared_ptr<DataFile> MakeEqualityDeleteFile(const std::string& name, + int64_t partition_value) { + auto file = MakeDataFile(name, /*record_count=*/1, /*size=*/10, partition_value); + file->content = DataFile::Content::kEqualityDeletes; + file->file_path = std::format("{}/delete/{}.parquet", table_location_, name); + file->equality_ids = {1}; + return file; + } + + Status AppendFiles(std::vector<std::shared_ptr<DataFile>> files) { + ICEBERG_ASSIGN_OR_RAISE(auto append, table_->NewFastAppend()); + for (const auto& file : files) { + append->AppendFile(file); + } + ICEBERG_RETURN_UNEXPECTED(append->Commit()); + return table_->Refresh(); + } + + Status AppendDeleteFiles(std::vector<std::shared_ptr<DataFile>> files) { + ICEBERG_ASSIGN_OR_RAISE(auto row_delta, table_->NewRowDelta()); + for (const auto& file : files) { + row_delta->AddDeletes(file); + } + ICEBERG_RETURN_UNEXPECTED(row_delta->Commit()); + return table_->Refresh(); + } + + Status RemoveDeleteFiles(std::vector<std::shared_ptr<DataFile>> files) { + ICEBERG_ASSIGN_OR_RAISE(auto row_delta, table_->NewRowDelta()); + for (const auto& file : files) { + row_delta->RemoveDeletes(file); + } + ICEBERG_RETURN_UNEXPECTED(row_delta->Commit()); + return table_->Refresh(); + } + + void SetSnapshotIdInheritanceEnabled() { + ICEBERG_UNWRAP_OR_FAIL(auto props, table_->NewUpdateProperties()); + props->Set(std::string(TableProperties::kSnapshotIdInheritanceEnabled.key()), "true"); + EXPECT_THAT(props->Commit(), IsOk()); + EXPECT_THAT(table_->Refresh(), IsOk()); + } + + void ReplayV1ReplacementAfterUpgrade(bool inherit_snapshot_id) { + if (format_version() != 1) { + GTEST_SKIP() << "This regression starts from a v1 table"; + } + + ASSERT_THAT(AppendFiles({file_a_, file_b_}), IsOk()); + if (inherit_snapshot_id) { + SetSnapshotIdInheritanceEnabled(); + } + ICEBERG_UNWRAP_OR_FAIL(auto snapshot, table_->current_snapshot()); + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 1U); + ICEBERG_UNWRAP_OR_FAIL(auto replacement, WriteExistingManifest("legacy-replacement", + snapshot->snapshot_id, + {file_a_, file_b_})); + if (inherit_snapshot_id) { + replacement.added_snapshot_id.reset(); + } + + ICEBERG_UNWRAP_OR_FAIL(auto txn, table_->NewTransaction()); + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, txn->NewRewriteManifests()); + rewrite->DeleteManifest(before[0]).AddManifest(replacement); + ASSERT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto props, table_->NewUpdateProperties()); + props->Set(std::string(TableProperties::kFormatVersion.key()), "2"); + ASSERT_THAT(props->Commit(), IsOk()); + ASSERT_THAT(table_->Refresh(), IsOk()); + ASSERT_THAT(txn->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto after, CurrentManifests()); + ASSERT_EQ(after.size(), 1U); + EXPECT_NE(after[0].manifest_path, replacement.manifest_path); + EXPECT_TRUE(FileExists(replacement.manifest_path)); + ICEBERG_UNWRAP_OR_FAIL(auto entries, ReadManifestEntries(after[0])); + ASSERT_EQ(entries.size(), 2U); + for (const auto& entry : entries) { + EXPECT_EQ(entry.status, ManifestStatus::kExisting); + EXPECT_EQ(entry.snapshot_id, snapshot->snapshot_id); + EXPECT_EQ(entry.sequence_number, std::make_optional<int64_t>(0)); + EXPECT_EQ(entry.file_sequence_number, std::make_optional<int64_t>(0)); + } + ICEBERG_UNWRAP_OR_FAIL(auto rewritten_snapshot, table_->current_snapshot()); + ValidateSummary(rewritten_snapshot, /*replaced=*/1, /*kept=*/0, /*created=*/1, + /*entry_count=*/0); + + ICEBERG_UNWRAP_OR_FAIL(auto remove, table_->NewDeleteFiles()); + remove->DeleteFile(file_a_); + ASSERT_THAT(remove->Commit(), IsOk()); + } + + void SetManifestMergeEnabled(bool enabled) { + ICEBERG_UNWRAP_OR_FAIL(auto props, table_->NewUpdateProperties()); + props->Set(std::string(TableProperties::kManifestMergeEnabled.key()), + enabled ? "true" : "false"); + EXPECT_THAT(props->Commit(), IsOk()); + EXPECT_THAT(table_->Refresh(), IsOk()); + } + + void SetCommitRetryProperties(int32_t retries) { + ICEBERG_UNWRAP_OR_FAIL(auto props, table_->NewUpdateProperties()); + props->Set(std::string(TableProperties::kCommitNumRetries.key()), + std::to_string(retries)); + props->Set(std::string(TableProperties::kCommitMinRetryWaitMs.key()), "1"); + props->Set(std::string(TableProperties::kCommitMaxRetryWaitMs.key()), "1"); + props->Set(std::string(TableProperties::kCommitTotalRetryTimeMs.key()), "1000"); + EXPECT_THAT(props->Commit(), IsOk()); + EXPECT_THAT(table_->Refresh(), IsOk()); + } + + void BindTableWithFailingCommits(int failures) { + auto mock_catalog = std::make_shared<::testing::NiceMock<MockCatalog>>(); + std::weak_ptr<::testing::NiceMock<MockCatalog>> weak_catalog = mock_catalog; + + ON_CALL(*mock_catalog, LoadTable(::testing::_)) + .WillByDefault([this, weak_catalog](const TableIdentifier& identifier) + -> Result<std::shared_ptr<Table>> { + ICEBERG_ASSIGN_OR_RAISE(auto loaded, catalog_->LoadTable(identifier)); + auto catalog = weak_catalog.lock(); + ICEBERG_PRECHECK(catalog != nullptr, "Mock catalog expired"); + return Table::Make(loaded->name(), loaded->metadata(), + std::string(loaded->metadata_file_location()), loaded->io(), + catalog); + }); + + ON_CALL(*mock_catalog, UpdateTable(::testing::_, ::testing::_, ::testing::_)) + .WillByDefault( + [this, weak_catalog, failures]( + const TableIdentifier& identifier, + const std::vector<std::unique_ptr<TableRequirement>>& requirements, + const std::vector<std::unique_ptr<TableUpdate>>& updates) mutable + -> Result<std::shared_ptr<Table>> { + if (failures-- > 0) { + return CommitFailed("Injected failure"); + } + ICEBERG_ASSIGN_OR_RAISE( + auto updated, catalog_->UpdateTable(identifier, requirements, updates)); + auto catalog = weak_catalog.lock(); + ICEBERG_PRECHECK(catalog != nullptr, "Mock catalog expired"); + return Table::Make(updated->name(), updated->metadata(), + std::string(updated->metadata_file_location()), + updated->io(), catalog); + }); + + ICEBERG_UNWRAP_OR_FAIL(auto bound_table, + Table::Make(table_->name(), table_->metadata(), + std::string(table_->metadata_file_location()), + table_->io(), mock_catalog)); + table_ = std::move(bound_table); + mock_catalog_ = std::move(mock_catalog); + } + + void BindTableWithFileIO(std::shared_ptr<FileIO> file_io) { + ICEBERG_UNWRAP_OR_FAIL(auto bound_table, + Table::Make(table_->name(), table_->metadata(), + std::string(table_->metadata_file_location()), + std::move(file_io), catalog_)); + table_ = std::move(bound_table); + } + + void ValidateSummary(const std::shared_ptr<Snapshot>& snapshot, int replaced, int kept, + int created, int entry_count) { + ASSERT_NE(snapshot, nullptr); + const auto& summary = snapshot->summary; + EXPECT_THAT(summary, ::testing::Contains(::testing::Pair( + std::string(SnapshotSummaryFields::kManifestsReplaced), + std::to_string(replaced)))); + EXPECT_THAT(summary, ::testing::Contains(::testing::Pair( + std::string(SnapshotSummaryFields::kManifestsKept), + std::to_string(kept)))); + EXPECT_THAT(summary, ::testing::Contains(::testing::Pair( + std::string(SnapshotSummaryFields::kManifestsCreated), + std::to_string(created)))); + EXPECT_THAT(summary, ::testing::Contains(::testing::Pair( + std::string(SnapshotSummaryFields::kEntriesProcessed), + std::to_string(entry_count)))); + } + + void MatchNumberOfManifestFileWithSpecId(std::span<const ManifestFile> manifests, + int32_t partition_spec_id, + int expected_manifest_count) { + const auto matched_manifest_count = std::ranges::count_if( + manifests, [partition_spec_id](const ManifestFile& manifest) { + return manifest.partition_spec_id == partition_spec_id; + }); + EXPECT_EQ(matched_manifest_count, expected_manifest_count); + } + + bool FileExists(const std::string& path) { + auto input_file = file_io_->NewInputFile(path); + if (!input_file.has_value()) { + return false; + } + return input_file.value()->Size().has_value(); + } + + Result<const ManifestFile*> ManifestContainingPath( + std::span<const ManifestFile> manifests, const std::string& path) { + for (const auto& manifest : manifests) { + ICEBERG_ASSIGN_OR_RAISE(auto entries, ReadManifestEntries(manifest)); + for (const auto& entry : entries) { + if (entry.data_file != nullptr && entry.data_file->file_path == path) { + return &manifest; + } + } + } + return NotFound("Manifest containing path {} is not found", path); + } + + void ExpectManifestEntryContents(const ManifestFile& manifest, + std::vector<std::string> expected_paths, + std::vector<DataFile::Content> expected_contents) { + ASSERT_EQ(expected_contents.size(), expected_paths.size()); + + ICEBERG_UNWRAP_OR_FAIL(auto entries, ReadManifestEntries(manifest)); + ASSERT_EQ(entries.size(), expected_paths.size()); + + std::vector<std::pair<std::string, DataFile::Content>> actual; + actual.reserve(entries.size()); + for (const auto& entry : entries) { + ASSERT_NE(entry.data_file, nullptr); + actual.emplace_back(entry.data_file->file_path, entry.data_file->content); + } + + std::vector<std::pair<std::string, DataFile::Content>> expected; + expected.reserve(expected_paths.size()); + for (size_t i = 0; i < expected_paths.size(); ++i) { + expected.emplace_back(std::move(expected_paths[i]), expected_contents[i]); + } + EXPECT_THAT(actual, ::testing::UnorderedElementsAreArray(expected)); + } + + Result<std::vector<ManifestFile>> CurrentManifests() { + ICEBERG_RETURN_UNEXPECTED(table_->Refresh()); + ICEBERG_ASSIGN_OR_RAISE(auto snapshot, table_->current_snapshot()); + SnapshotReader reader(snapshot.get()); + ICEBERG_ASSIGN_OR_RAISE(auto manifests, reader.DataManifests(file_io_)); + return std::vector<ManifestFile>(manifests.begin(), manifests.end()); + } + + Result<std::vector<ManifestFile>> CurrentDeleteManifests() { + ICEBERG_RETURN_UNEXPECTED(table_->Refresh()); + ICEBERG_ASSIGN_OR_RAISE(auto snapshot, table_->current_snapshot()); + SnapshotReader reader(snapshot.get()); + ICEBERG_ASSIGN_OR_RAISE(auto manifests, reader.DeleteManifests(file_io_)); + return std::vector<ManifestFile>(manifests.begin(), manifests.end()); + } + + Result<ManifestFile> WriteExistingManifest( + const std::string& name, int64_t snapshot_id, + const std::vector<std::shared_ptr<DataFile>>& files, + ManifestContent content = ManifestContent::kData, + std::shared_ptr<PartitionSpec> manifest_spec = nullptr, + int64_t data_sequence_number = TableMetadata::kInitialSequenceNumber, + std::optional<int64_t> file_sequence_number = std::nullopt) { + auto path = std::format("{}/metadata/{}.avro", table_location_, name); + auto spec = manifest_spec != nullptr ? std::move(manifest_spec) : spec_; + std::vector<std::shared_ptr<DataFile>> files_to_write; + files_to_write.reserve(files.size()); + std::optional<int64_t> first_row_id = std::nullopt; + if (table_->metadata()->format_version >= 3 && content == ManifestContent::kData) { + first_row_id = TableMetadata::kInitialRowId; + ICEBERG_ASSIGN_OR_RAISE(auto current_manifests, CurrentManifests()); + std::unordered_map<std::string, std::shared_ptr<DataFile>> current_files_by_path; + for (const auto& manifest : current_manifests) { + ICEBERG_ASSIGN_OR_RAISE(auto entries, ReadManifestEntries(manifest)); + for (const auto& entry : entries) { + if (entry.IsAlive() && entry.data_file != nullptr) { + current_files_by_path.emplace(entry.data_file->file_path, + std::make_shared<DataFile>(*entry.data_file)); + } + } + } + for (const auto& file : files) { + auto file_to_write = file; + if (!file_to_write->first_row_id.has_value()) { + if (auto it = current_files_by_path.find(file->file_path); + it != current_files_by_path.end()) { + file_to_write = it->second; + } + } + files_to_write.push_back(std::move(file_to_write)); + } + } else { + files_to_write = files; + } + ICEBERG_ASSIGN_OR_RAISE( + auto writer, + ManifestWriter::MakeWriter(table_->metadata()->format_version, kInvalidSnapshotId, + path, file_io_, spec, schema_, content, first_row_id)); + for (const auto& file : files_to_write) { + ICEBERG_RETURN_UNEXPECTED(writer->WriteExistingEntry( + file, snapshot_id, data_sequence_number, file_sequence_number)); + } + ICEBERG_RETURN_UNEXPECTED(writer->Close()); + return writer->ToManifestFile(); + } + + struct ExistingManifestEntry { + std::shared_ptr<DataFile> file; + int64_t snapshot_id; + int64_t data_sequence_number; + int64_t file_sequence_number; + }; + + Result<ManifestFile> WriteExistingDeleteManifest( + const std::string& name, const std::vector<ExistingManifestEntry>& entries) { + auto path = std::format("{}/metadata/{}.avro", table_location_, name); + ICEBERG_ASSIGN_OR_RAISE( + auto writer, ManifestWriter::MakeWriter(table_->metadata()->format_version, + kInvalidSnapshotId, path, file_io_, spec_, + schema_, ManifestContent::kDeletes)); + for (const auto& entry : entries) { + ICEBERG_RETURN_UNEXPECTED(writer->WriteExistingEntry(entry.file, entry.snapshot_id, + entry.data_sequence_number, + entry.file_sequence_number)); + } + ICEBERG_RETURN_UNEXPECTED(writer->Close()); + return writer->ToManifestFile(); + } + + Result<ManifestFile> WriteExistingManifest( + const std::string& name, + const std::vector<std::pair<std::shared_ptr<DataFile>, int64_t>>& files) { + auto path = std::format("{}/metadata/{}.avro", table_location_, name); + std::vector<std::pair<std::shared_ptr<DataFile>, int64_t>> files_to_write; + files_to_write.reserve(files.size()); + std::optional<int64_t> first_row_id = std::nullopt; + if (table_->metadata()->format_version >= 3) { + first_row_id = TableMetadata::kInitialRowId; + ICEBERG_ASSIGN_OR_RAISE(auto current_manifests, CurrentManifests()); + std::unordered_map<std::string, std::shared_ptr<DataFile>> current_files_by_path; + for (const auto& manifest : current_manifests) { + ICEBERG_ASSIGN_OR_RAISE(auto entries, ReadManifestEntries(manifest)); + for (const auto& entry : entries) { + if (entry.IsAlive() && entry.data_file != nullptr) { + current_files_by_path.emplace(entry.data_file->file_path, + std::make_shared<DataFile>(*entry.data_file)); + } + } + } + for (const auto& [file, snapshot_id] : files) { + auto file_to_write = file; + if (!file_to_write->first_row_id.has_value()) { + if (auto it = current_files_by_path.find(file->file_path); + it != current_files_by_path.end()) { + file_to_write = it->second; + } + } + files_to_write.emplace_back(std::move(file_to_write), snapshot_id); + } + } else { + files_to_write = files; + } + ICEBERG_ASSIGN_OR_RAISE( + auto writer, ManifestWriter::MakeWriter( + table_->metadata()->format_version, kInvalidSnapshotId, path, + file_io_, spec_, schema_, ManifestContent::kData, first_row_id)); + for (const auto& [file, snapshot_id] : files_to_write) { + ICEBERG_RETURN_UNEXPECTED(writer->WriteExistingEntry( + file, snapshot_id, TableMetadata::kInitialSequenceNumber)); + } + ICEBERG_RETURN_UNEXPECTED(writer->Close()); + return writer->ToManifestFile(); + } + + Result<ManifestFile> WriteAddedManifest( + const std::string& name, const std::vector<std::shared_ptr<DataFile>>& files, + std::optional<int64_t> entry_snapshot_id = std::nullopt, + std::optional<int64_t> data_sequence_number = std::nullopt) { + auto path = std::format("{}/metadata/{}.avro", table_location_, name); + ICEBERG_ASSIGN_OR_RAISE( + auto writer, ManifestWriter::MakeWriter(table_->metadata()->format_version, + entry_snapshot_id, path, file_io_, spec_, + schema_, ManifestContent::kData)); + for (const auto& file : files) { + ICEBERG_RETURN_UNEXPECTED(writer->WriteAddedEntry(file, data_sequence_number)); + } + ICEBERG_RETURN_UNEXPECTED(writer->Close()); + ICEBERG_ASSIGN_OR_RAISE(auto manifest, writer->ToManifestFile()); + manifest.added_snapshot_id = kInvalidSnapshotId; + return manifest; + } + + Result<ManifestFile> WriteDeletedManifest( + const std::string& name, const std::vector<std::shared_ptr<DataFile>>& files, + std::optional<int64_t> entry_snapshot_id = std::nullopt, + int64_t data_sequence_number = TableMetadata::kInitialSequenceNumber) { + auto path = std::format("{}/metadata/{}.avro", table_location_, name); + ICEBERG_ASSIGN_OR_RAISE( + auto writer, ManifestWriter::MakeWriter(table_->metadata()->format_version, + entry_snapshot_id, path, file_io_, spec_, + schema_, ManifestContent::kData)); + for (const auto& file : files) { + ICEBERG_RETURN_UNEXPECTED(writer->WriteDeletedEntry(file, data_sequence_number)); + } + ICEBERG_RETURN_UNEXPECTED(writer->Close()); + return writer->ToManifestFile(); + } + + std::shared_ptr<PartitionSpec> spec_; + std::shared_ptr<Schema> schema_; + std::shared_ptr<DataFile> file_a_; + std::shared_ptr<DataFile> file_b_; + std::shared_ptr<DataFile> file_c_; + std::shared_ptr<DataFile> file_d_; + std::shared_ptr<::testing::NiceMock<MockCatalog>> mock_catalog_; +}; + +TEST_P(RewriteManifestsTest, RewriteManifestsAppendedDirectly) { + SetSnapshotIdInheritanceEnabled(); + ICEBERG_UNWRAP_OR_FAIL(auto new_manifest, + WriteAddedManifest("manifest-file-1", {file_a_})); + + ICEBERG_UNWRAP_OR_FAIL(auto append, table_->NewFastAppend()); + append->AppendManifest(new_manifest); + EXPECT_THAT(append->Commit(), IsOk()); + EXPECT_THAT(table_->Refresh(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot, table_->current_snapshot()); + const int64_t append_snapshot_id = append_snapshot->snapshot_id; + + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 1U); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile&) { return ""; }); + EXPECT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 1U); + ExpectManifestEntriesWithSnapshotIds(manifests[0], {file_a_->file_path}, + {ManifestStatus::kExisting}, {append_snapshot_id}); +} + +TEST_P(RewriteManifestsTest, FactoryResultCanCommitDirectly) { + ASSERT_THAT(AppendFiles({file_a_, file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto original, table_->current_snapshot()); + ICEBERG_UNWRAP_OR_FAIL(auto ctx, + TransactionContext::Make(table_, TransactionKind::kUpdate)); + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, + RewriteManifests::Make(table_->name().name, std::move(ctx))); + rewrite->ClusterBy([](const DataFile&) { return ""; }); + ASSERT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 1U); + ExpectManifestEntriesWithSnapshotIds( + manifests[0], {file_a_->file_path, file_b_->file_path}, + {ManifestStatus::kExisting, ManifestStatus::kExisting}, + {original->snapshot_id, original->snapshot_id}); +} + +TEST_P(RewriteManifestsTest, StreamsEntriesWithBoundedReaders) { + ICEBERG_UNWRAP_OR_FAIL(auto append, table_->NewFastAppend()); + std::vector<std::string> expected_paths; + for (int i = 0; i < 40; ++i) { + std::vector<std::shared_ptr<DataFile>> files; + for (int j = 0; j < 3; ++j) { + auto file = MakeDataFile(std::format("stream-{}-{}", i, j), 1, 10, i); + files.push_back(file); + expected_paths.push_back(file->file_path); + } + ICEBERG_UNWRAP_OR_FAIL(auto manifest, + WriteAddedManifest(std::format("stream-{}", i), files)); + append->AppendManifest(manifest); + } + ASSERT_THAT(append->Commit(), IsOk()); + + for (bool use_executor : {false, true}) { + SCOPED_TRACE(use_executor); + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + auto state = std::make_shared<ManifestReadState>(); + for (const auto& manifest : before) { + state->source_paths.insert(manifest.manifest_path); + } + { + test::ThreadExecutor executor; + ScopedManifestReadTracking tracking(state); + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([state](const DataFile& file) { + std::lock_guard lock(state->mutex); + --state->unread_entries; + ++state->processed_entries; + return file.file_path.substr(0, file.file_path.find_last_of('-')); + }); + if (use_executor) { + rewrite->ScanManifestsWith(executor); + } + ASSERT_THAT(rewrite->Commit(), IsOk()); + if (use_executor) { + EXPECT_GT(executor.submit_count(), 0); + } + } + EXPECT_EQ(state->processed_entries, expected_paths.size()); + EXPECT_EQ(state->unread_entries, 0); + EXPECT_LE(state->max_unread_entries, 2); + EXPECT_EQ(state->opened_readers, before.size()); + EXPECT_EQ(state->open_readers, 0); + EXPECT_LT(state->max_open_readers, before.size()); + if (!use_executor) { + EXPECT_EQ(state->max_open_readers, 1); + } + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + std::vector<std::string> actual_paths; + for (const auto& manifest : manifests) { + ICEBERG_UNWRAP_OR_FAIL(auto entries, ReadManifestEntries(manifest)); + for (const auto& entry : entries) { + EXPECT_EQ(entry.status, ManifestStatus::kExisting); + actual_paths.push_back(entry.data_file->file_path); + } + } + EXPECT_THAT(actual_paths, ::testing::UnorderedElementsAreArray(expected_paths)); + } +} + +TEST_P(RewriteManifestsTest, CleansWritersAfterStreamingReadFailure) { + ASSERT_THAT(AppendFiles({file_a_, file_b_, file_c_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + for (bool use_executor : {false, true}) { + SCOPED_TRACE(use_executor); + auto reads = std::make_shared<ManifestReadState>(); + reads->source_paths.insert(before[0].manifest_path); + reads->fail_after_first_batch = true; + auto writes = std::make_shared<ManifestWriteFailureState>(); + BindTableWithFileIO(std::make_shared<ManifestWriteFailureFileIO>( + file_io_, table_location_ + "/metadata/", writes)); + { + test::ThreadExecutor executor; + ScopedManifestReadTracking tracking(reads); + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile&) { return ""; }); + if (use_executor) { + rewrite->ScanManifestsWith(executor); + } + auto result = rewrite->Commit(); + EXPECT_THAT(result, IsError(ErrorKind::kIOError)); + EXPECT_THAT(result, HasErrorMessage("Injected streaming manifest read failure")); + } + EXPECT_EQ(reads->open_readers, 0); + ASSERT_EQ(writes->manifest_outputs.size(), 1U); + EXPECT_TRUE(writes->open_manifest_paths.empty()); + EXPECT_TRUE(writes->deleted_open_files.empty()); + EXPECT_THAT(writes->deleted_files, + ::testing::Contains(writes->manifest_outputs.front())); + EXPECT_FALSE(FileExists(writes->manifest_outputs.front())); + ICEBERG_UNWRAP_OR_FAIL(auto after, CurrentManifests()); + EXPECT_EQ(after, before); + } +} + +TEST_P(RewriteManifestsTest, RewriteManifestsWithScanExecutor) { + SetSnapshotIdInheritanceEnabled(); + ICEBERG_UNWRAP_OR_FAIL(auto new_manifest, + WriteAddedManifest("manifest-file-1", {file_a_})); + + ICEBERG_UNWRAP_OR_FAIL(auto append, table_->NewFastAppend()); + append->AppendManifest(new_manifest); + EXPECT_THAT(append->Commit(), IsOk()); + EXPECT_THAT(table_->Refresh(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 1U); + + test::ThreadExecutor executor; + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile&) { return ""; }).ScanManifestsWith(executor); + EXPECT_THAT(rewrite->Commit(), IsOk()); + + EXPECT_GT(executor.submit_count(), 0); + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 1U); + ExpectManifestEntries(manifests[0], {file_a_->file_path}, {ManifestStatus::kExisting}); +} + +TEST_P(RewriteManifestsTest, RewriteManifestsGeneratedAndAppendedDirectly) { + SetSnapshotIdInheritanceEnabled(); + ICEBERG_UNWRAP_OR_FAIL(auto new_manifest, + WriteAddedManifest("manifest-file-1", {file_a_})); + + ICEBERG_UNWRAP_OR_FAIL(auto append_manifest, table_->NewFastAppend()); + append_manifest->AppendManifest(new_manifest); + EXPECT_THAT(append_manifest->Commit(), IsOk()); + EXPECT_THAT(table_->Refresh(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto manifest_append_snapshot, table_->current_snapshot()); + const int64_t manifest_append_id = manifest_append_snapshot->snapshot_id; + + ASSERT_THAT(AppendFiles({file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto file_append_snapshot, table_->current_snapshot()); + const int64_t file_append_id = file_append_snapshot->snapshot_id; + + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 2U); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile&) { return ""; }); + EXPECT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 1U); + ExpectManifestEntriesWithSnapshotIds( + manifests[0], {file_a_->file_path, file_b_->file_path}, + {ManifestStatus::kExisting, ManifestStatus::kExisting}, + {manifest_append_id, file_append_id}); +} + +TEST_P(RewriteManifestsTest, RewriteManifestsWithoutClusterBy) { + ASSERT_THAT(AppendFiles({file_a_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot, table_->current_snapshot()); + const int64_t append_id = append_snapshot->snapshot_id; + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 1U); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + EXPECT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 1U); + EXPECT_EQ(manifests[0].manifest_path, before[0].manifest_path); + ExpectManifestEntriesWithSnapshotIds(manifests[0], {file_a_->file_path}, + {ManifestStatus::kAdded}, {append_id}); +} + +TEST_P(RewriteManifestsTest, RewriteAfterV2ToV3UpgradeAssignsRowIds) { + if (format_version() != 2) { + GTEST_SKIP() << "This regression starts from a v2 table"; + } + + ASSERT_THAT(AppendFiles({file_a_}), IsOk()); + ASSERT_THAT(AppendFiles({file_b_}), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto original_manifests, CurrentManifests()); + ASSERT_EQ(original_manifests.size(), 2U); + for (const auto& manifest : original_manifests) { + EXPECT_FALSE(manifest.first_row_id.has_value()); + } + + ICEBERG_UNWRAP_OR_FAIL(auto props, table_->NewUpdateProperties()); + props->Set(std::string(TableProperties::kFormatVersion.key()), "3"); + EXPECT_THAT(props->Commit(), IsOk()); + EXPECT_THAT(table_->Refresh(), IsOk()); + EXPECT_EQ(table_->metadata()->format_version, 3); + EXPECT_EQ(table_->metadata()->next_row_id, TableMetadata::kInitialRowId); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile&) { return ""; }); + EXPECT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 1U); + EXPECT_EQ(manifests[0].first_row_id, + std::make_optional<int64_t>(TableMetadata::kInitialRowId)); + + ICEBERG_UNWRAP_OR_FAIL(auto entries, ReadManifestEntries(manifests[0])); + ASSERT_EQ(entries.size(), 2U); + + int64_t next_row_id = TableMetadata::kInitialRowId; + int64_t assigned_rows = 0; + for (const auto& entry : entries) { + ASSERT_NE(entry.data_file, nullptr); + ASSERT_TRUE(entry.data_file->first_row_id.has_value()); + EXPECT_EQ(entry.data_file->first_row_id.value(), next_row_id); + next_row_id += entry.data_file->record_count; + assigned_rows += entry.data_file->record_count; + } + EXPECT_EQ(table_->metadata()->next_row_id, + TableMetadata::kInitialRowId + assigned_rows); +} + +TEST_P(RewriteManifestsTest, RewritesWhenClusterByChangesAfterApply) { + ASSERT_THAT(AppendFiles({file_a_, file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto snapshot, table_->current_snapshot()); + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, TestRewriteManifests::Make(table_)); + rewrite->ClusterBy([](const DataFile&) { return "all"; }); + ICEBERG_UNWRAP_OR_FAIL(auto prepared, + rewrite->ApplyForTest(*table_->metadata(), snapshot)); + ASSERT_EQ(prepared.size(), 1U); + EXPECT_TRUE(FileExists(prepared[0].manifest_path)); + + rewrite->ClusterBy([](const DataFile& file) { return file.file_path; }); + ASSERT_THAT(rewrite->Commit(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto after, CurrentManifests()); + ASSERT_EQ(after.size(), 2U); + for (const auto& manifest : after) { + ICEBERG_UNWRAP_OR_FAIL(auto entries, ReadManifestEntries(manifest)); + EXPECT_EQ(entries.size(), 1U); + } + EXPECT_FALSE(FileExists(prepared[0].manifest_path)); + ICEBERG_UNWRAP_OR_FAIL(auto rewritten_snapshot, table_->current_snapshot()); + ValidateSummary(rewritten_snapshot, /*replaced=*/1, /*kept=*/0, /*created=*/2, + /*entry_count=*/2); +} + +TEST_P(RewriteManifestsTest, RewritesWhenPredicateChangesAfterApply) { + ASSERT_THAT(AppendFiles({file_a_}), IsOk()); + ASSERT_THAT(AppendFiles({file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 2U); + ICEBERG_UNWRAP_OR_FAIL(auto snapshot, table_->current_snapshot()); + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, TestRewriteManifests::Make(table_)); + rewrite->ClusterBy([](const DataFile&) { return "all"; }); + rewrite->RewriteIf([path = before[0].manifest_path](const ManifestFile& manifest) { + return manifest.manifest_path == path; + }); + ICEBERG_UNWRAP_OR_FAIL(auto prepared, + rewrite->ApplyForTest(*table_->metadata(), snapshot)); + + rewrite->RewriteIf([](const ManifestFile&) { return false; }); + ASSERT_THAT(rewrite->Commit(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto after, CurrentManifests()); + EXPECT_THAT(after, ::testing::UnorderedElementsAreArray(before)); + for (const auto& manifest : prepared) { + if (!std::ranges::contains(before, manifest)) { + EXPECT_FALSE(FileExists(manifest.manifest_path)); + } + } + ICEBERG_UNWRAP_OR_FAIL(auto rewritten_snapshot, table_->current_snapshot()); + ValidateSummary(rewritten_snapshot, /*replaced=*/0, /*kept=*/2, /*created=*/0, + /*entry_count=*/0); +} + +TEST_P(RewriteManifestsTest, ReplaysV1ReplacementAfterUpgradeWithSnapshotIdInheritance) { + ReplayV1ReplacementAfterUpgrade(/*inherit_snapshot_id=*/true); +} + +TEST_P(RewriteManifestsTest, + ReplaysV1ReplacementAfterUpgradeWithoutSnapshotIdInheritance) { + ReplayV1ReplacementAfterUpgrade(/*inherit_snapshot_id=*/false); +} + +TEST_P(RewriteManifestsTest, CopiesV1ReplacementAddedAfterFormatUpgrade) { + if (format_version() != 1) { + GTEST_SKIP() << "This regression starts from a v1 manifest"; + } + ASSERT_THAT(AppendFiles({file_a_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto snapshot, table_->current_snapshot()); + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ICEBERG_UNWRAP_OR_FAIL( + auto replacement, + WriteExistingManifest("legacy-replacement", snapshot->snapshot_id, {file_a_})); + replacement.added_snapshot_id.reset(); + + ICEBERG_UNWRAP_OR_FAIL(auto props, table_->NewUpdateProperties()); + props->Set(std::string(TableProperties::kFormatVersion.key()), "2"); + ASSERT_THAT(props->Commit(), IsOk()); + ASSERT_THAT(table_->Refresh(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->DeleteManifest(before[0]).AddManifest(replacement); + ASSERT_THAT(rewrite->Commit(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto after, CurrentManifests()); + ASSERT_EQ(after.size(), 1U); + EXPECT_NE(after[0].manifest_path, replacement.manifest_path); + EXPECT_TRUE(FileExists(replacement.manifest_path)); + ICEBERG_UNWRAP_OR_FAIL(auto entries, ReadManifestEntries(after[0])); + ASSERT_EQ(entries.size(), 1U); + EXPECT_EQ(entries[0].snapshot_id, snapshot->snapshot_id); + EXPECT_EQ(entries[0].sequence_number, std::make_optional<int64_t>(0)); + EXPECT_EQ(entries[0].file_sequence_number, std::make_optional<int64_t>(0)); +} + +TEST_P(RewriteManifestsTest, ReplaceManifestsSeparate) { + ASSERT_THAT(AppendFiles({file_a_, file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot, table_->current_snapshot()); + const int64_t append_id = append_snapshot->snapshot_id; + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 1U); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile& file) { return file.file_path; }); + EXPECT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 2U); + ICEBERG_UNWRAP_OR_FAIL(auto manifest_a, + ManifestContainingPath(manifests, file_a_->file_path)); + ICEBERG_UNWRAP_OR_FAIL(auto manifest_b, + ManifestContainingPath(manifests, file_b_->file_path)); + ExpectManifestEntriesWithSnapshotIds(*manifest_a, {file_a_->file_path}, + {ManifestStatus::kExisting}, {append_id}); + ExpectManifestEntriesWithSnapshotIds(*manifest_b, {file_b_->file_path}, + {ManifestStatus::kExisting}, {append_id}); +} + +TEST_P(RewriteManifestsTest, ReplaceManifestsConsolidate) { + ASSERT_THAT(AppendFiles({file_a_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot_a, table_->current_snapshot()); + const int64_t append_id_a = append_snapshot_a->snapshot_id; + ASSERT_THAT(AppendFiles({file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot_b, table_->current_snapshot()); + const int64_t append_id_b = append_snapshot_b->snapshot_id; + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 2U); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile&) { return "file"; }); + EXPECT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 1U); + ExpectManifestEntriesWithSnapshotIds( + manifests[0], {file_a_->file_path, file_b_->file_path}, + {ManifestStatus::kExisting, ManifestStatus::kExisting}, {append_id_a, append_id_b}); +} + +TEST_P(RewriteManifestsTest, ReplaceManifestsWithFilter) { + ASSERT_THAT(AppendFiles({file_a_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot_a, table_->current_snapshot()); + const int64_t append_id_a = append_snapshot_a->snapshot_id; + ASSERT_THAT(AppendFiles({file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot_b, table_->current_snapshot()); + const int64_t append_id_b = append_snapshot_b->snapshot_id; + ASSERT_THAT(AppendFiles({file_c_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot_c, table_->current_snapshot()); + const int64_t append_id_c = append_snapshot_c->snapshot_id; + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 3U); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile&) { return "file"; }); + rewrite->RewriteIf([&](const ManifestFile& manifest) { + auto entries_result = ReadManifestEntries(manifest); + if (!entries_result.has_value() || entries_result.value().empty()) { + return false; + } + return entries_result.value()[0].data_file->file_path != file_a_->file_path; + }); + EXPECT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 2U); + + const ManifestFile* kept_manifest = nullptr; + const ManifestFile* rewritten_manifest = nullptr; + for (const auto& manifest : manifests) { + ICEBERG_UNWRAP_OR_FAIL(auto entries, ReadManifestEntries(manifest)); + ASSERT_FALSE(entries.empty()); + ASSERT_NE(entries[0].data_file, nullptr); + if (entries[0].data_file->file_path == file_a_->file_path) { + kept_manifest = &manifest; + } else { + rewritten_manifest = &manifest; + } + } + + ASSERT_NE(rewritten_manifest, nullptr); + ExpectManifestEntriesWithSnapshotIds( + *rewritten_manifest, {file_b_->file_path, file_c_->file_path}, + {ManifestStatus::kExisting, ManifestStatus::kExisting}, {append_id_b, append_id_c}); + + ASSERT_NE(kept_manifest, nullptr); + ExpectManifestEntriesWithSnapshotIds(*kept_manifest, {file_a_->file_path}, + {ManifestStatus::kAdded}, {append_id_a}); +} + +TEST_P(RewriteManifestsTest, ReplaceManifestsMaxSize) { + ASSERT_THAT(AppendFiles({file_a_, file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot, table_->current_snapshot()); + const int64_t append_snapshot_id = append_snapshot->snapshot_id; + + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 1U); + + ICEBERG_UNWRAP_OR_FAIL(auto props, table_->NewUpdateProperties()); + props->Set(std::string(TableProperties::kManifestTargetSizeBytes.key()), "1"); + EXPECT_THAT(props->Commit(), IsOk()); + EXPECT_THAT(table_->Refresh(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile&) { return "file"; }); + EXPECT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 2U); + std::ranges::sort(manifests, {}, &ManifestFile::manifest_path); + ExpectManifestEntriesWithSnapshotIds(manifests[0], {file_a_->file_path}, + {ManifestStatus::kExisting}, {append_snapshot_id}); + ExpectManifestEntriesWithSnapshotIds(manifests[1], {file_b_->file_path}, + {ManifestStatus::kExisting}, {append_snapshot_id}); +} + +TEST_P(RewriteManifestsTest, CleansRolledRewriteManifestsAfterApplyFailure) { + ASSERT_THAT(AppendFiles({file_a_, file_b_, file_c_}), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 1U); + + ICEBERG_UNWRAP_OR_FAIL(auto props, table_->NewUpdateProperties()); + props->Set(std::string(TableProperties::kManifestTargetSizeBytes.key()), "1"); + EXPECT_THAT(props->Commit(), IsOk()); + EXPECT_THAT(table_->Refresh(), IsOk()); + + auto failure_state = std::make_shared<ManifestWriteFailureState>(); + failure_state->fail_on_manifest_output = 2; + BindTableWithFileIO(std::make_shared<ManifestWriteFailureFileIO>( + file_io_, table_location_ + "/metadata/", failure_state)); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile&) { return "file"; }); + + auto result = rewrite->Commit(); + EXPECT_THAT(result, IsError(ErrorKind::kIOError)); + EXPECT_THAT(result, HasErrorMessage("Injected rewrite manifest writer failure")); + + ASSERT_GE(failure_state->manifest_outputs.size(), 2U); + const auto& closed_manifest_path = failure_state->manifest_outputs[0]; + EXPECT_THAT(failure_state->deleted_files, ::testing::Contains(closed_manifest_path)); + EXPECT_FALSE(FileExists(closed_manifest_path)); +} + +TEST_P(RewriteManifestsTest, ClosesOpenRewriteWritersAfterApplyFailure) { + ASSERT_THAT(AppendFiles({file_a_, file_b_, file_c_}), IsOk()); + + auto failure_state = std::make_shared<ManifestWriteFailureState>(); + failure_state->fail_on_manifest_output = 3; + BindTableWithFileIO(std::make_shared<ManifestWriteFailureFileIO>( + file_io_, table_location_ + "/metadata/", failure_state)); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile& file) { return file.file_path; }); + + auto result = rewrite->Commit(); + EXPECT_THAT(result, IsError(ErrorKind::kIOError)); + EXPECT_THAT(result, HasErrorMessage("Injected rewrite manifest writer failure")); + + ASSERT_EQ(failure_state->manifest_outputs.size(), 3U); + EXPECT_THAT(failure_state->manifest_close_attempts, + ::testing::UnorderedElementsAre(failure_state->manifest_outputs[0], + failure_state->manifest_outputs[1])); + for (size_t i = 0; i < 2; ++i) { + EXPECT_THAT(failure_state->deleted_files, + ::testing::Contains(failure_state->manifest_outputs[i])); + EXPECT_FALSE(FileExists(failure_state->manifest_outputs[i])); + } +} + +TEST_P(RewriteManifestsTest, ClosesAllRewriteWritersAfterCloseFailure) { + ASSERT_THAT(AppendFiles({file_a_, file_b_, file_c_}), IsOk()); + + auto failure_state = std::make_shared<ManifestWriteFailureState>(); + failure_state->fail_on_manifest_close = 1; + BindTableWithFileIO(std::make_shared<ManifestWriteFailureFileIO>( + file_io_, table_location_ + "/metadata/", failure_state)); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile& file) { return file.file_path; }); + + auto result = rewrite->Commit(); + EXPECT_THAT(result, IsError(ErrorKind::kIOError)); + EXPECT_THAT(result, HasErrorMessage("Injected rewrite manifest close failure")); + + ASSERT_EQ(failure_state->manifest_outputs.size(), 3U); + ASSERT_EQ(failure_state->manifest_close_attempts.size(), 3U); + EXPECT_THAT(failure_state->manifest_close_attempts, + ::testing::UnorderedElementsAre(failure_state->manifest_outputs[0], + failure_state->manifest_outputs[1], + failure_state->manifest_outputs[2])); + for (size_t i = 0; i < failure_state->manifest_close_attempts.size(); ++i) { + const auto& closed_manifest_path = failure_state->manifest_close_attempts[i]; + EXPECT_THAT(failure_state->deleted_files, ::testing::Contains(closed_manifest_path)); + EXPECT_FALSE(FileExists(closed_manifest_path)); + } +} + +TEST_P(RewriteManifestsTest, PreservesRewriteFailureWhenClosingWritersFails) { + ASSERT_THAT(AppendFiles({file_a_, file_b_, file_c_}), IsOk()); + + auto failure_state = std::make_shared<ManifestWriteFailureState>(); + failure_state->fail_on_manifest_output = 3; + failure_state->fail_on_manifest_close = 1; + BindTableWithFileIO(std::make_shared<ManifestWriteFailureFileIO>( + file_io_, table_location_ + "/metadata/", failure_state)); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile& file) { return file.file_path; }); + + auto result = rewrite->Commit(); + EXPECT_THAT(result, IsError(ErrorKind::kIOError)); + EXPECT_THAT(result, HasErrorMessage("Injected rewrite manifest writer failure")); + EXPECT_THAT(result, HasErrorMessage("additionally failed to close manifest writer")); + EXPECT_THAT(result, HasErrorMessage("Injected rewrite manifest close failure")); + + ASSERT_EQ(failure_state->manifest_outputs.size(), 3U); + ASSERT_EQ(failure_state->manifest_close_attempts.size(), 2U); + for (const auto& manifest_path : failure_state->manifest_outputs) { + EXPECT_THAT(failure_state->deleted_files, ::testing::Contains(manifest_path)); + EXPECT_FALSE(FileExists(manifest_path)); + } +} + +TEST_P(RewriteManifestsTest, ClosesWritersWhenClusterByThrows) { + ASSERT_THAT(AppendFiles({file_a_, file_b_, file_c_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + + for (bool explicit_transaction : {false, true}) { + for (int close_failure : {0, 1, 2}) { + SCOPED_TRACE(std::format("explicit_transaction={}, close_failure={}", + explicit_transaction, close_failure)); + auto state = std::make_shared<ManifestWriteFailureState>(); + state->fail_on_manifest_close = close_failure == 1 ? 1 : 0; + state->throw_on_manifest_close = close_failure == 2 ? 1 : 0; + BindTableWithFileIO(std::make_shared<ManifestWriteFailureFileIO>( + file_io_, table_location_ + "/metadata/", state)); + + std::shared_ptr<Transaction> txn; + std::shared_ptr<RewriteManifests> rewrite; + if (explicit_transaction) { + ICEBERG_UNWRAP_OR_FAIL(txn, table_->NewTransaction()); + ICEBERG_UNWRAP_OR_FAIL(rewrite, txn->NewRewriteManifests()); + } else { + ICEBERG_UNWRAP_OR_FAIL(rewrite, table_->NewRewriteManifests()); + } + rewrite->ClusterBy([calls = 0](const DataFile& file) mutable -> std::string { + if (++calls == 3) { + throw std::runtime_error("Injected cluster failure"); + } + return file.file_path; + }); + auto result = rewrite->Commit(); + EXPECT_THAT(result, IsError(ErrorKind::kValidationFailed)); + EXPECT_THAT(result, HasErrorMessage("Injected cluster failure")); + if (txn) { + EXPECT_EQ(txn->state(), TransactionState::kFailed); + } + ASSERT_EQ(state->manifest_outputs.size(), 2U); + EXPECT_THAT(state->manifest_close_attempts, + ::testing::UnorderedElementsAreArray(state->manifest_outputs)); + EXPECT_TRUE(state->open_manifest_paths.empty()); + EXPECT_TRUE(state->deleted_open_files.empty()); + for (const auto& path : state->manifest_outputs) { + EXPECT_THAT(state->deleted_files, ::testing::Contains(path)); + EXPECT_FALSE(FileExists(path)); + } + ICEBERG_UNWRAP_OR_FAIL(auto after, CurrentManifests()); + EXPECT_EQ(after, before); + } + } +} + +TEST_P(RewriteManifestsTest, CleansWritersAfterTransactionApplyFailure) { + ASSERT_THAT(AppendFiles({file_a_, file_b_, file_c_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + auto state = std::make_shared<ManifestWriteFailureState>(); + state->fail_on_manifest_output = 3; + BindTableWithFileIO(std::make_shared<ManifestWriteFailureFileIO>( + file_io_, table_location_ + "/metadata/", state)); + + ICEBERG_UNWRAP_OR_FAIL(auto txn, table_->NewTransaction()); + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, txn->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile& file) { return file.file_path; }); + EXPECT_THAT(rewrite->Commit(), IsError(ErrorKind::kIOError)); + EXPECT_EQ(txn->state(), TransactionState::kFailed); + ASSERT_EQ(state->manifest_outputs.size(), 3U); + EXPECT_THAT(state->manifest_close_attempts, + ::testing::UnorderedElementsAre(state->manifest_outputs[0], + state->manifest_outputs[1])); + EXPECT_TRUE(state->open_manifest_paths.empty()); + EXPECT_TRUE(state->deleted_open_files.empty()); + for (const auto& path : state->manifest_outputs) { + EXPECT_THAT(state->deleted_files, ::testing::Contains(path)); + EXPECT_FALSE(FileExists(path)); + } + EXPECT_THAT(txn->Commit(), IsError(ErrorKind::kValidationFailed)); + ASSERT_THAT(txn->Abort(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto after, CurrentManifests()); + EXPECT_EQ(after, before); +} + +TEST_P(RewriteManifestsTest, ClosesAndCleansCopiedManifestOnFailure) { + ASSERT_THAT(AppendFiles({file_a_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto snapshot, table_->current_snapshot()); + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 1U); + ICEBERG_UNWRAP_OR_FAIL( + auto added_manifest, + WriteAddedManifest("copy-added-entry", {file_a_}, snapshot->snapshot_id)); + ICEBERG_UNWRAP_OR_FAIL( + auto deleted_manifest, + WriteDeletedManifest("copy-deleted-entry", {file_a_}, snapshot->snapshot_id)); + // Counts pass AddManifest validation; reading the entries must still reject it. + added_manifest.added_files_count = 0; + added_manifest.existing_files_count = 1; + deleted_manifest.deleted_files_count = 0; + deleted_manifest.existing_files_count = 1; + deleted_manifest.added_snapshot_id = kInvalidSnapshotId; + ICEBERG_UNWRAP_OR_FAIL( + auto valid_manifest, + WriteExistingManifest("copy-close-failure", snapshot->snapshot_id, {file_a_})); + + for (bool explicit_transaction : {false, true}) { + for (const auto& [source, error_message] : + {std::pair{&valid_manifest, "Injected rewrite manifest close failure"}, + std::pair{&added_manifest, + "Invalid manifest entry status: ADDED (allowed status: EXISTING)"}, + std::pair{ + &deleted_manifest, + "Invalid manifest entry status: DELETED (allowed status: EXISTING)"}}) { + const bool invalid_entry = source != &valid_manifest; + SCOPED_TRACE(std::format("explicit_transaction={}, error={}", explicit_transaction, + error_message)); + auto state = std::make_shared<ManifestWriteFailureState>(); + state->fail_on_manifest_close = invalid_entry ? 0 : 1; + BindTableWithFileIO(std::make_shared<ManifestWriteFailureFileIO>( + file_io_, table_location_ + "/metadata/", state)); + std::shared_ptr<Transaction> txn; + std::shared_ptr<RewriteManifests> rewrite; + if (explicit_transaction) { + ICEBERG_UNWRAP_OR_FAIL(txn, table_->NewTransaction()); + ICEBERG_UNWRAP_OR_FAIL(rewrite, txn->NewRewriteManifests()); + } else { + ICEBERG_UNWRAP_OR_FAIL(rewrite, table_->NewRewriteManifests()); + } + rewrite->DeleteManifest(before[0]).AddManifest(*source); + auto result = rewrite->Commit(); + EXPECT_THAT(result, IsError(invalid_entry ? ErrorKind::kValidationFailed + : ErrorKind::kIOError)); + EXPECT_THAT(result, HasErrorMessage(error_message)); + if (txn) { + EXPECT_EQ(txn->state(), TransactionState::kFailed); + } + ASSERT_EQ(state->manifest_outputs.size(), 1U); + EXPECT_THAT(state->manifest_close_attempts, + ::testing::ElementsAre(state->manifest_outputs[0])); + EXPECT_TRUE(state->open_manifest_paths.empty()); + EXPECT_TRUE(state->deleted_open_files.empty()); + EXPECT_THAT(state->deleted_files, ::testing::Contains(state->manifest_outputs[0])); + EXPECT_FALSE(FileExists(state->manifest_outputs[0])); + EXPECT_TRUE(FileExists(source->manifest_path)); + ICEBERG_UNWRAP_OR_FAIL(auto after, CurrentManifests()); + EXPECT_EQ(after, before); + } + } +} + +TEST_P(RewriteManifestsTest, ConcurrentRewriteManifest) { + ASSERT_THAT(AppendFiles({file_a_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot_a, table_->current_snapshot()); + const int64_t append_id_a = append_snapshot_a->snapshot_id; + ASSERT_THAT(AppendFiles({file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot_b, table_->current_snapshot()); + const int64_t append_id_b = append_snapshot_b->snapshot_id; + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 2U); + + ICEBERG_UNWRAP_OR_FAIL(auto txn, table_->NewTransaction()); + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, txn->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile&) { return "file"; }); + ASSERT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto concurrent, table_->NewRewriteManifests()); + concurrent->ClusterBy([](const DataFile&) { return "file"; }); + concurrent->RewriteIf([&](const ManifestFile& manifest) { + auto entries_result = ReadManifestEntries(manifest); + if (!entries_result.has_value() || entries_result.value().empty()) { + return false; + } + return entries_result.value()[0].data_file->file_path != file_a_->file_path; + }); + EXPECT_THAT(concurrent->Commit(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto after_concurrent, CurrentManifests()); + ASSERT_EQ(after_concurrent.size(), 2U); + + ASSERT_THAT(txn->Commit(), IsOk()); + ASSERT_THAT(table_->Refresh(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 1U); + ExpectManifestEntriesWithSnapshotIds( + manifests[0], {file_a_->file_path, file_b_->file_path}, + {ManifestStatus::kExisting, ManifestStatus::kExisting}, {append_id_a, append_id_b}); +} + +TEST_P(RewriteManifestsTest, AppendDuringRewriteManifest) { + ASSERT_THAT(AppendFiles({file_a_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot_a, table_->current_snapshot()); + const int64_t append_id_a = append_snapshot_a->snapshot_id; + + ICEBERG_UNWRAP_OR_FAIL(auto txn, table_->NewTransaction()); + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, txn->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile&) { return "file"; }); + ASSERT_THAT(rewrite->Commit(), IsOk()); + + ASSERT_THAT(AppendFiles({file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot_b, table_->current_snapshot()); + const int64_t append_id_b = append_snapshot_b->snapshot_id; + ICEBERG_UNWRAP_OR_FAIL(auto after_append, CurrentManifests()); + ASSERT_EQ(after_append.size(), 2U); + + // commit the rewrite manifests in progress + ASSERT_THAT(txn->Commit(), IsOk()); + ASSERT_THAT(table_->Refresh(), IsOk()); + + // Retry cleanup discards the staged rewrite and reclusters the latest manifests. + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 1U); + ExpectManifestEntriesWithSnapshotIds( + manifests[0], {file_a_->file_path, file_b_->file_path}, + {ManifestStatus::kExisting, ManifestStatus::kExisting}, {append_id_a, append_id_b}); +} + +TEST_P(RewriteManifestsTest, RewriteManifestDuringAppend) { + ASSERT_THAT(AppendFiles({file_a_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto append_snapshot, table_->current_snapshot()); + const int64_t append_snapshot_id = append_snapshot->snapshot_id; + + ICEBERG_UNWRAP_OR_FAIL(auto txn, table_->NewTransaction()); + ICEBERG_UNWRAP_OR_FAIL(auto append, txn->NewFastAppend()); + append->AppendFile(file_b_); + ASSERT_THAT(append->Commit(), IsOk()); + + // rewrite the manifests - only affects the first + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->ClusterBy([](const DataFile&) { return "file"; }); + EXPECT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto rewritten_manifests, CurrentManifests()); + ASSERT_EQ(rewritten_manifests.size(), 1U); + + // commit the append in progress + ASSERT_THAT(txn->Commit(), IsOk()); + ASSERT_THAT(table_->Refresh(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto file_append_snapshot, table_->current_snapshot()); + const int64_t file_append_id = file_append_snapshot->snapshot_id; + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 2U); + ExpectManifestEntriesWithSnapshotIds(manifests[0], {file_b_->file_path}, + {ManifestStatus::kAdded}, {file_append_id}); + ExpectManifestEntriesWithSnapshotIds(manifests[1], {file_a_->file_path}, + {ManifestStatus::kExisting}, {append_snapshot_id}); +} + +TEST_P(RewriteManifestsTest, BasicManifestReplacement) { + ASSERT_THAT(AppendFiles({file_a_, file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto first_snapshot, table_->current_snapshot()); + const int64_t first_snapshot_id = first_snapshot->snapshot_id; + ICEBERG_UNWRAP_OR_FAIL(auto first_snapshot_manifests, CurrentManifests()); + ASSERT_EQ(first_snapshot_manifests.size(), 1U); + + ASSERT_THAT(AppendFiles({file_c_, file_d_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto second_snapshot, table_->current_snapshot()); + const int64_t second_snapshot_id = second_snapshot->snapshot_id; + + ICEBERG_UNWRAP_OR_FAIL( + auto manifest_a, WriteExistingManifest("rewrite-a", first_snapshot_id, {file_a_})); + ICEBERG_UNWRAP_OR_FAIL( + auto manifest_b, WriteExistingManifest("rewrite-b", first_snapshot_id, {file_b_})); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->DeleteManifest(first_snapshot_manifests[0]) + .AddManifest(manifest_a) + .AddManifest(manifest_b); + EXPECT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 3U); + EXPECT_NE(manifests[0].manifest_path, manifest_a.manifest_path); + EXPECT_NE(manifests[1].manifest_path, manifest_b.manifest_path); + ExpectManifestEntriesWithSnapshotIds(manifests[0], {file_a_->file_path}, + {ManifestStatus::kExisting}, {first_snapshot_id}); + ExpectManifestEntriesWithSnapshotIds(manifests[1], {file_b_->file_path}, + {ManifestStatus::kExisting}, {first_snapshot_id}); + ExpectManifestEntriesWithSnapshotIds(manifests[2], + {file_c_->file_path, file_d_->file_path}, + {ManifestStatus::kAdded, ManifestStatus::kAdded}, + {second_snapshot_id, second_snapshot_id}); + + ICEBERG_UNWRAP_OR_FAIL(auto after, table_->current_snapshot()); + ValidateSummary(after, /*replaced=*/1, /*kept=*/1, /*created=*/2, + /*entry_count=*/0); +} + +TEST_P(RewriteManifestsTest, ManifestReplacementWithUnassignedSnapshotId) { + ASSERT_THAT(AppendFiles({file_a_, file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto original_snapshot, table_->current_snapshot()); + + for (bool inherit_snapshot_id : {false, true}) { + SCOPED_TRACE(inherit_snapshot_id); + if (inherit_snapshot_id) { + SetSnapshotIdInheritanceEnabled(); + } + + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 1U); + ICEBERG_UNWRAP_OR_FAIL(auto entries, ReadManifestEntries(before[0])); + auto path = std::format("{}/metadata/unassigned-{}.avro", table_location_, + inherit_snapshot_id); + ICEBERG_UNWRAP_OR_FAIL( + auto writer, ManifestWriter::MakeWriter( + format_version(), std::nullopt, path, file_io_, spec_, schema_, + ManifestContent::kData, before[0].first_row_id)); + for (const auto& entry : entries) { + ASSERT_THAT(writer->WriteExistingEntry(entry), IsOk()); + } + ASSERT_THAT(writer->Close(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto replacement, writer->ToManifestFile()); + ASSERT_FALSE(replacement.added_snapshot_id.has_value()); + + auto assigned_manifest = replacement; + assigned_manifest.added_snapshot_id = original_snapshot->snapshot_id; + ICEBERG_UNWRAP_OR_FAIL(auto invalid_rewrite, table_->NewRewriteManifests()); + invalid_rewrite->DeleteManifest(before[0]).AddManifest(assigned_manifest); + EXPECT_THAT(invalid_rewrite->Commit(), + HasErrorMessage("Snapshot id must be assigned during commit")); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->DeleteManifest(before[0]).AddManifest(replacement); + ASSERT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 1U); + ICEBERG_UNWRAP_OR_FAIL(auto snapshot, table_->current_snapshot()); + EXPECT_NE(snapshot->snapshot_id, original_snapshot->snapshot_id); + EXPECT_EQ(manifests[0].added_snapshot_id, snapshot->snapshot_id); + if (inherit_snapshot_id || format_version() > 1) { + EXPECT_EQ(manifests[0].manifest_path, replacement.manifest_path); + } else { + EXPECT_NE(manifests[0].manifest_path, replacement.manifest_path); + } + ExpectManifestEntriesWithSnapshotIds( + manifests[0], {file_a_->file_path, file_b_->file_path}, + {ManifestStatus::kExisting, ManifestStatus::kExisting}, + {original_snapshot->snapshot_id, original_snapshot->snapshot_id}); + } +} + +TEST_P(RewriteManifestsTest, BasicManifestReplacementWithSnapshotIdInheritance) { + SetSnapshotIdInheritanceEnabled(); + + ASSERT_THAT(AppendFiles({file_a_, file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto first_snapshot, table_->current_snapshot()); + const int64_t first_snapshot_id = first_snapshot->snapshot_id; + ICEBERG_UNWRAP_OR_FAIL(auto first_snapshot_manifests, CurrentManifests()); + ASSERT_EQ(first_snapshot_manifests.size(), 1U); + + ASSERT_THAT(AppendFiles({file_c_, file_d_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto second_snapshot, table_->current_snapshot()); + const int64_t second_snapshot_id = second_snapshot->snapshot_id; + + ICEBERG_UNWRAP_OR_FAIL( + auto manifest_a, + WriteExistingManifest("rewrite-inherit-a", first_snapshot_id, {file_a_})); + ICEBERG_UNWRAP_OR_FAIL( + auto manifest_b, + WriteExistingManifest("rewrite-inherit-b", first_snapshot_id, {file_b_})); + manifest_a.added_snapshot_id.reset(); + manifest_b.added_snapshot_id.reset(); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->DeleteManifest(first_snapshot_manifests[0]) + .AddManifest(manifest_a) + .AddManifest(manifest_b); + EXPECT_THAT(rewrite->Commit(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 3U); + EXPECT_EQ(manifests[0].manifest_path, manifest_a.manifest_path); + EXPECT_EQ(manifests[1].manifest_path, manifest_b.manifest_path); + ExpectManifestEntriesWithSnapshotIds(manifests[0], {file_a_->file_path}, + {ManifestStatus::kExisting}, {first_snapshot_id}); + ExpectManifestEntriesWithSnapshotIds(manifests[1], {file_b_->file_path}, + {ManifestStatus::kExisting}, {first_snapshot_id}); + ExpectManifestEntriesWithSnapshotIds(manifests[2], + {file_c_->file_path, file_d_->file_path}, + {ManifestStatus::kAdded, ManifestStatus::kAdded}, + {second_snapshot_id, second_snapshot_id}); + + ICEBERG_UNWRAP_OR_FAIL(auto after, table_->current_snapshot()); + ValidateSummary(after, /*replaced=*/1, /*kept=*/1, /*created=*/2, + /*entry_count=*/0); + + // validate that any subsequent operation does not fail + ICEBERG_UNWRAP_OR_FAIL(auto delete_files, table_->NewDeleteFiles()); + delete_files->DeleteFromRowFilter(Expressions::AlwaysTrue()); + EXPECT_THAT(delete_files->Commit(), IsOk()); +} + +TEST_P(RewriteManifestsTest, ManifestReplacementWithExplicitInvalidSnapshotIdIsCopied) { + SetSnapshotIdInheritanceEnabled(); + ASSERT_THAT(AppendFiles({file_a_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto original_snapshot, table_->current_snapshot()); + ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests()); + ASSERT_EQ(before.size(), 1U); + ICEBERG_UNWRAP_OR_FAIL( + auto replacement, + WriteExistingManifest("explicit-invalid-id", original_snapshot->snapshot_id, + {file_a_}, ManifestContent::kData, nullptr, + original_snapshot->sequence_number, + original_snapshot->sequence_number)); + ASSERT_EQ(replacement.added_snapshot_id, kInvalidSnapshotId); + + ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests()); + rewrite->DeleteManifest(before[0]).AddManifest(replacement); + ASSERT_THAT(rewrite->Commit(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests()); + ASSERT_EQ(manifests.size(), 1U); + EXPECT_NE(manifests[0].manifest_path, replacement.manifest_path); + EXPECT_TRUE(FileExists(replacement.manifest_path)); + ICEBERG_UNWRAP_OR_FAIL(auto snapshot, table_->current_snapshot()); + EXPECT_EQ(manifests[0].added_snapshot_id, snapshot->snapshot_id); + ICEBERG_UNWRAP_OR_FAIL(auto entries, ReadManifestEntries(manifests[0])); + ASSERT_EQ(entries.size(), 1U); + EXPECT_EQ(entries[0].status, ManifestStatus::kExisting); + EXPECT_EQ(entries[0].snapshot_id, original_snapshot->snapshot_id); + EXPECT_EQ(entries[0].sequence_number, original_snapshot->sequence_number); + EXPECT_EQ(entries[0].file_sequence_number, original_snapshot->sequence_number); +} + +TEST_P(RewriteManifestsTest, WithMultiplePartitionSpec) { + ASSERT_THAT(AppendFiles({file_a_, file_b_}), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto initial_manifests, CurrentManifests()); + ASSERT_EQ(initial_manifests.size(), 1U); + const int32_t initial_spec_id = initial_manifests[0].partition_spec_id; + ICEBERG_UNWRAP_OR_FAIL(auto initial_spec, + table_->metadata()->PartitionSpecById(initial_spec_id)); + ASSERT_EQ(initial_spec->fields().size(), 1U); + const int32_t initial_partition_field_id = initial_spec->fields()[0].field_id(); + + // Build the new spec using the table's schema, which uses fresh partition IDs. + ICEBERG_UNWRAP_OR_FAIL(auto update_spec, table_->NewUpdatePartitionSpec()); + update_spec->AddField(Expressions::Bucket("y", 16)) + .AddField(Expressions::Bucket("z", 4)); + EXPECT_THAT(update_spec->Commit(), IsOk()); + EXPECT_THAT(table_->Refresh(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto new_spec, table_->spec()); + ASSERT_EQ(new_spec->fields().size(), 3U); + EXPECT_EQ(new_spec->fields()[1].name(), "y_bucket_16"); + EXPECT_EQ(new_spec->fields()[1].transform()->ToString(), "bucket[16]"); + EXPECT_GT(new_spec->fields()[1].field_id(), initial_partition_field_id); + EXPECT_EQ(new_spec->fields()[2].name(), "z_bucket_4"); + EXPECT_EQ(new_spec->fields()[2].transform()->ToString(), "bucket[4]"); + EXPECT_GT(new_spec->fields()[2].field_id(), new_spec->fields()[1].field_id()); + + std::vector<Literal> partition_y{Literal::Long(/*x=*/5L), + Literal::Int(/*y_bucket_16=*/2), + Literal::Int(/*z_bucket_4=*/3)}; + auto file_y = MakeDataFile("y", new_spec, /*record_count=*/1, /*size=*/10, + std::move(partition_y)); + std::vector<Literal> partition_z{Literal::Long(/*x=*/7L), + Literal::Int(/*y_bucket_16=*/2), + Literal::Int(/*z_bucket_4=*/4)}; Review Comment: OK, will change 4 to 1 in both tests. -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
