attempt to fix inconsistency

Author: wgtmac

src/iceberg/snapshot.cc

@@ -52,6 +52,19 @@ SnapshotRefType SnapshotRef::type() const noexcept {
       retention);
 }
 
+std::optional<int64_t> SnapshotRef::max_ref_age_ms() const noexcept {
+  return std::visit(
+      [&](const auto& retention) -> std::optional<int64_t> {
+        using T = std::remove_cvref_t<decltype(retention)>;
+        if constexpr (std::is_same_v<T, Branch>) {
+          return retention.max_ref_age_ms;
+        } else {
+          return retention.max_ref_age_ms;
+        }
+      },
+      retention);
+}
+
 Status SnapshotRef::Validate() const {
   if (type() == SnapshotRefType::kBranch) {
     const auto& branch = std::get<Branch>(this->retention);

src/iceberg/snapshot.h

@@ -113,6 +113,8 @@ struct ICEBERG_EXPORT SnapshotRef {
 
   SnapshotRefType type() const noexcept;
 
+  std::optional<int64_t> max_ref_age_ms() const noexcept;
+
   /// \brief Create a branch reference
   ///
   /// \param snapshot_id The snapshot ID for the branch

src/iceberg/table_metadata.cc

@@ -617,7 +617,6 @@ class TableMetadataBuilder::Impl {
   Status SetBranchSnapshot(int64_t snapshot_id, const std::string& branch);
   Status SetBranchSnapshot(std::shared_ptr<Snapshot> snapshot, const std::string& branch);
   Status SetRef(const std::string& name, std::shared_ptr<SnapshotRef> ref);
-
   Status RemoveRef(const std::string& name);
   Status RemoveSnapshots(const std::vector<int64_t>& snapshot_ids);
   Status RemovePartitionSpecs(const std::vector<int32_t>& spec_ids);
@@ -1339,15 +1338,11 @@ int32_t TableMetadataBuilder::Impl::ReuseOrCreateNewSchemaId(
 }
 
 Status TableMetadataBuilder::Impl::RemoveRef(const std::string& name) {
-  // Handle main branch specially
   if (name == SnapshotRef::kMainBranch) {
     metadata_.current_snapshot_id = kInvalidSnapshotId;
   }
 
-  // Remove the ref from the map
-  auto it = metadata_.refs.find(name);
-  if (it != metadata_.refs.end()) {
-    metadata_.refs.erase(it);
+  if (metadata_.refs.erase(name) != 0) {
     changes_.push_back(std::make_unique<table::RemoveSnapshotRef>(name));
   }
 
@@ -1360,47 +1355,39 @@ Status TableMetadataBuilder::Impl::RemoveSnapshots(
     return {};
   }
 
-  std::unordered_set<int64_t> snapshot_ids_set(snapshot_ids.begin(), snapshot_ids.end());
-
-  // Build a map of snapshot IDs for quick lookup
-  std::unordered_map<int64_t, std::shared_ptr<Snapshot>> snapshots_by_id;
-  for (const auto& snapshot : metadata_.snapshots) {
-    if (snapshot) {
-      snapshots_by_id[snapshot->snapshot_id] = snapshot;
-    }
-  }
-
-  // Filter snapshots to retain
+  std::unordered_set<int64_t> ids_to_remove(snapshot_ids.begin(), snapshot_ids.end());
   std::vector<std::shared_ptr<Snapshot>> retained_snapshots;
-  retained_snapshots.reserve(metadata_.snapshots.size());
-
-  for (const auto& snapshot : metadata_.snapshots) {
-    if (!snapshot) continue;
-
-    int64_t snapshot_id = snapshot->snapshot_id;
-    if (snapshot_ids_set.contains(snapshot_id)) {
-      // Remove from the map
-      snapshots_by_id.erase(snapshot_id);
-      // Record the removal
-      changes_.push_back(
-          std::make_unique<table::RemoveSnapshots>(std::vector<int64_t>{snapshot_id}));
-      // Note: Statistics and partition statistics removal would be handled here
-      // if those features were implemented
+  retained_snapshots.reserve(metadata_.snapshots.size() - snapshot_ids.size());
+  std::vector<int64_t> snapshot_ids_to_remove;
+  snapshot_ids_to_remove.reserve(snapshot_ids.size());
+
+  for (auto& snapshot : metadata_.snapshots) {
+    ICEBERG_CHECK(snapshot != nullptr, "Encountered null snapshot in metadata");
+    const int64_t snapshot_id = snapshot->snapshot_id;
+    if (ids_to_remove.contains(snapshot_id)) {
+      snapshots_by_id_.erase(snapshot_id);
+      snapshot_ids_to_remove.push_back(snapshot_id);
+      // FIXME: implement statistics removal and uncomment below
+      // ICEBERG_RETURN_UNEXPECTED(RemoveStatistics(snapshot_id));
+      // ICEBERG_RETURN_UNEXPECTED(RemovePartitionStatistics(snapshot_id));
     } else {
-      retained_snapshots.push_back(snapshot);
+      retained_snapshots.push_back(std::move(snapshot));
     }
   }
 
+  if (!snapshot_ids_to_remove.empty()) {
+    changes_.push_back(std::make_unique<table::RemoveSnapshots>(snapshot_ids_to_remove));
+  }
+
   metadata_.snapshots = std::move(retained_snapshots);
 
   // Remove any refs that are no longer valid (dangling refs)
   std::vector<std::string> dangling_refs;
   for (const auto& [ref_name, ref] : metadata_.refs) {
-    if (!snapshots_by_id.contains(ref->snapshot_id)) {
+    if (!snapshots_by_id_.contains(ref->snapshot_id)) {
       dangling_refs.push_back(ref_name);
     }
   }
-
   for (const auto& ref_name : dangling_refs) {
     ICEBERG_RETURN_UNEXPECTED(RemoveRef(ref_name));
   }
@@ -1414,25 +1401,15 @@ Status TableMetadataBuilder::Impl::RemovePartitionSpecs(
     return {};
   }
 
-  std::unordered_set<int32_t> spec_ids_set(spec_ids.begin(), spec_ids.end());
-
-  // Validate that we're not removing the default spec
-  ICEBERG_PRECHECK(!spec_ids_set.contains(metadata_.default_spec_id),
+  std::unordered_set<int32_t> spec_ids_to_remove(spec_ids.begin(), spec_ids.end());
+  ICEBERG_PRECHECK(!spec_ids_to_remove.contains(metadata_.default_spec_id),
                    "Cannot remove the default partition spec");
 
-  // Filter partition specs to retain
   metadata_.partition_specs =
       metadata_.partition_specs | std::views::filter([&](const auto& spec) {
-        return !spec_ids_set.contains(spec->spec_id());
+        return !spec_ids_to_remove.contains(spec->spec_id());
       }) |
-      std::ranges::to<std::vector<std::shared_ptr<iceberg::PartitionSpec>>>();
-
-  // Update the specs_by_id_ index
-  for (int32_t spec_id : spec_ids) {
-    specs_by_id_.erase(spec_id);
-  }
-
-  // Record the change
+      std::ranges::to<std::vector<std::shared_ptr<PartitionSpec>>>();
   changes_.push_back(std::make_unique<table::RemovePartitionSpecs>(spec_ids));
 
   return {};

src/iceberg/test/expire_snapshots_test.cc

@@ -24,71 +24,43 @@
 
 namespace iceberg {
 
-class ExpireSnapshotsTest : public UpdateTestBase {
- protected:
-};
+class ExpireSnapshotsTest : public UpdateTestBase {};
 
-TEST_F(ExpireSnapshotsTest, Empty) {
+TEST_F(ExpireSnapshotsTest, DefaultExpireByAge) {
   ICEBERG_UNWRAP_OR_FAIL(auto update, table_->NewExpireSnapshots());
   ICEBERG_UNWRAP_OR_FAIL(auto result, update->Apply());
-  EXPECT_THAT(result.snapshot_ids_to_remove.size(), 1);
-  EXPECT_THAT(result.snapshot_ids_to_remove.at(0), 3051729675574597004);
-  EXPECT_THAT(result.ref_to_remove.empty(), true);
-  EXPECT_THAT(result.schema_ids_to_remove.empty(), true);
-  EXPECT_THAT(result.partition_spec_ids_to_remove.empty(), true);
+  EXPECT_EQ(result.snapshot_ids_to_remove.size(), 1);
+  EXPECT_EQ(result.snapshot_ids_to_remove.at(0), 3051729675574597004);
 }
 
-TEST_F(ExpireSnapshotsTest, Keep2) {
+TEST_F(ExpireSnapshotsTest, KeepAll) {
   ICEBERG_UNWRAP_OR_FAIL(auto update, table_->NewExpireSnapshots());
   update->RetainLast(2);
   ICEBERG_UNWRAP_OR_FAIL(auto result, update->Apply());
-  EXPECT_THAT(result.snapshot_ids_to_remove.empty(), true);
-  EXPECT_THAT(result.ref_to_remove.empty(), true);
-  EXPECT_THAT(result.schema_ids_to_remove.empty(), true);
-  EXPECT_THAT(result.partition_spec_ids_to_remove.empty(), true);
+  EXPECT_TRUE(result.snapshot_ids_to_remove.empty());
+  EXPECT_TRUE(result.refs_to_remove.empty());
 }
 
 TEST_F(ExpireSnapshotsTest, ExpireById) {
   ICEBERG_UNWRAP_OR_FAIL(auto update, table_->NewExpireSnapshots());
   update->ExpireSnapshotId(3051729675574597004);
-  update->RetainLast(2);
   ICEBERG_UNWRAP_OR_FAIL(auto result, update->Apply());
-  EXPECT_THAT(result.snapshot_ids_to_remove.size(), 1);
-  EXPECT_THAT(result.snapshot_ids_to_remove.at(0), 3051729675574597004);
-  EXPECT_THAT(result.ref_to_remove.empty(), true);
-  EXPECT_THAT(result.schema_ids_to_remove.empty(), true);
-  EXPECT_THAT(result.partition_spec_ids_to_remove.empty(), true);
-}
-
-TEST_F(ExpireSnapshotsTest, ExpireByIdNotExist) {
-  ICEBERG_UNWRAP_OR_FAIL(auto update, table_->NewExpireSnapshots());
-  update->ExpireSnapshotId(3055729675574597003);
-  update->RetainLast(2);
-  auto result = update->Apply();
-  EXPECT_THAT(result.has_value(), false);
-  EXPECT_THAT(result.error().message.contains("Snapshot:3055729675574597003 not exist"),
-              true);
+  EXPECT_EQ(result.snapshot_ids_to_remove.size(), 1);
+  EXPECT_EQ(result.snapshot_ids_to_remove.at(0), 3051729675574597004);
 }
 
-TEST_F(ExpireSnapshotsTest, ExpireOlderThan1) {
-  ICEBERG_UNWRAP_OR_FAIL(auto update, table_->NewExpireSnapshots());
-  update->ExpireOlderThan(1515100955770 - 1);
-  ICEBERG_UNWRAP_OR_FAIL(auto result, update->Apply());
-  EXPECT_THAT(result.snapshot_ids_to_remove.empty(), true);
-  EXPECT_THAT(result.ref_to_remove.empty(), true);
-  EXPECT_THAT(result.schema_ids_to_remove.empty(), true);
-  EXPECT_THAT(result.partition_spec_ids_to_remove.empty(), true);
-}
-
-TEST_F(ExpireSnapshotsTest, ExpireOlderThan2) {
-  ICEBERG_UNWRAP_OR_FAIL(auto update, table_->NewExpireSnapshots());
-  update->ExpireOlderThan(1515100955770 + 1);
-  ICEBERG_UNWRAP_OR_FAIL(auto result, update->Apply());
-  EXPECT_THAT(result.snapshot_ids_to_remove.size(), 1);
-  EXPECT_THAT(result.snapshot_ids_to_remove.at(0), 3051729675574597004);
-  EXPECT_THAT(result.ref_to_remove.empty(), true);
-  EXPECT_THAT(result.schema_ids_to_remove.empty(), true);
-  EXPECT_THAT(result.partition_spec_ids_to_remove.empty(), true);
+TEST_F(ExpireSnapshotsTest, ExpireOlderThan) {
+  struct TestCase {
+    int64_t expire_older_than;
+    size_t expected_num_expired;
+  };
+  std::vector<TestCase> test_cases = {{1515100955770 - 1, 0}, {1515100955770 + 1, 1}};
+  for (const auto& test_case : test_cases) {
+    ICEBERG_UNWRAP_OR_FAIL(auto update, table_->NewExpireSnapshots());
+    update->ExpireOlderThan(test_case.expire_older_than);
+    ICEBERG_UNWRAP_OR_FAIL(auto result, update->Apply());
+    EXPECT_EQ(result.snapshot_ids_to_remove.size(), test_case.expected_num_expired);
+  }
 }
 
 }  // namespace iceberg

src/iceberg/test/table_metadata_builder_test.cc

@@ -1133,28 +1133,19 @@ TEST(TableMetadataBuilderTest, RemoveSchemasAfterSchemaChange) {
   ASSERT_THAT(builder->Build(), HasErrorMessage("Cannot remove current schema: 1"));
 }
 
-// Test RemoveSnapshotRef
-TEST(TableMetadataBuilderTest, RemoveSnapshotRefBasic) {
+TEST(TableMetadataBuilderTest, RemoveSnapshotRef) {
   auto base = CreateBaseMetadata();
   auto builder = TableMetadataBuilder::BuildFrom(base.get());
 
   // Add multiple snapshots
-  auto snapshot1 = std::make_shared<Snapshot>();
-  snapshot1->snapshot_id = 1;
-  builder->AddSnapshot(snapshot1);
-  auto snapshot2 = std::make_shared<Snapshot>();
-  snapshot2->snapshot_id = 2;
-  builder->AddSnapshot(snapshot2);
+  builder->AddSnapshot(std::make_shared<Snapshot>(Snapshot{.snapshot_id = 1}));
+  builder->AddSnapshot(std::make_shared<Snapshot>(Snapshot{.snapshot_id = 2}));
 
   // Add multiple refs
-  auto ref1 = std::make_shared<SnapshotRef>();
-  ref1->snapshot_id = 1;
-  ref1->retention = SnapshotRef::Branch{};
-  builder->SetRef("ref1", ref1);
-  auto ref2 = std::make_shared<SnapshotRef>();
-  ref2->snapshot_id = 2;
-  ref2->retention = SnapshotRef::Tag{};
-  builder->SetRef("ref2", ref2);
+  ICEBERG_UNWRAP_OR_FAIL(auto ref1, SnapshotRef::MakeBranch(1));
+  ICEBERG_UNWRAP_OR_FAIL(auto ref2, SnapshotRef::MakeBranch(2));
+  builder->SetRef("ref1", std::move(ref1));
+  builder->SetRef("ref2", std::move(ref2));
 
   ICEBERG_UNWRAP_OR_FAIL(auto metadata, builder->Build());
   ASSERT_EQ(metadata->refs.size(), 2);
@@ -1167,68 +1158,50 @@ TEST(TableMetadataBuilderTest, RemoveSnapshotRefBasic) {
   EXPECT_TRUE(metadata->refs.contains("ref1"));
 }
 
-// Test RemoveSnapshot
-TEST(TableMetadataBuilderTest, RemoveSnapshotBasic) {
+TEST(TableMetadataBuilderTest, RemoveSnapshot) {
   auto base = CreateBaseMetadata();
   auto builder = TableMetadataBuilder::BuildFrom(base.get());
 
   // Add multiple snapshots
-  auto snapshot1 = std::make_shared<Snapshot>();
-  snapshot1->snapshot_id = 1;
-  builder->AddSnapshot(snapshot1);
-  auto snapshot2 = std::make_shared<Snapshot>();
-  snapshot2->snapshot_id = 2;
-  builder->AddSnapshot(snapshot2);
+  builder->AddSnapshot(std::make_shared<Snapshot>(Snapshot{.snapshot_id = 1}));
+  builder->AddSnapshot(std::make_shared<Snapshot>(Snapshot{.snapshot_id = 2}));
 
   ICEBERG_UNWRAP_OR_FAIL(auto metadata, builder->Build());
   ASSERT_EQ(metadata->snapshots.size(), 2);
 
   // Remove one snapshot
   builder = TableMetadataBuilder::BuildFrom(metadata.get());
-  std::vector<int64_t> to_remove{snapshot2->snapshot_id};
+  std::vector<int64_t> to_remove{2};
   builder->RemoveSnapshots(to_remove);
   ICEBERG_UNWRAP_OR_FAIL(metadata, builder->Build());
   ASSERT_EQ(metadata->snapshots.size(), 1);
-  for (const auto& s : metadata->snapshots) {
-    std::cout << s->snapshot_id << std::endl;
-  }
-  EXPECT_TRUE(
-      std::ranges::find_if(metadata->snapshots, [&](const std::shared_ptr<Snapshot>& s) {
-        return s->snapshot_id == snapshot1->snapshot_id;
-      }) != metadata->snapshots.end());
+  ASSERT_THAT(metadata->SnapshotById(2), IsError(ErrorKind::kNotFound));
 }
 
 TEST(TableMetadataBuilderTest, RemoveSnapshotNotExist) {
   auto base = CreateBaseMetadata();
   auto builder = TableMetadataBuilder::BuildFrom(base.get());
 
   // Add multiple snapshots
-  auto snapshot1 = std::make_shared<Snapshot>();
-  snapshot1->snapshot_id = 1;
-  builder->AddSnapshot(snapshot1);
-  auto snapshot2 = std::make_shared<Snapshot>();
-  snapshot2->snapshot_id = 2;
-  builder->AddSnapshot(snapshot2);
+  builder->AddSnapshot(std::make_shared<Snapshot>(Snapshot{.snapshot_id = 1}));
+  builder->AddSnapshot(std::make_shared<Snapshot>(Snapshot{.snapshot_id = 2}));
 
   ICEBERG_UNWRAP_OR_FAIL(auto metadata, builder->Build());
   ASSERT_EQ(metadata->snapshots.size(), 2);
 
   // Remove one snapshot
   builder = TableMetadataBuilder::BuildFrom(metadata.get());
-  std::vector<int64_t> to_remove{3};
-  builder->RemoveSnapshots(to_remove);
+  builder->RemoveSnapshots(std::vector<int64_t>{3});
   ICEBERG_UNWRAP_OR_FAIL(metadata, builder->Build());
   ASSERT_EQ(metadata->snapshots.size(), 2);
 
   builder = TableMetadataBuilder::BuildFrom(metadata.get());
-  to_remove = {1, 2};
-  builder->RemoveSnapshots(to_remove);
+  builder->RemoveSnapshots(std::vector<int64_t>{1, 2});
   ICEBERG_UNWRAP_OR_FAIL(metadata, builder->Build());
   ASSERT_EQ(metadata->snapshots.size(), 0);
 }
 
-// Test RemovePartitionSpec
-TEST(TableMetadataBuilderTest, RemovePartitionSpecBasic) {
+TEST(TableMetadataBuilderTest, RemovePartitionSpec) {
   // Add multiple specs
   PartitionField field1(2, 4, "field1", Transform::Identity());
   PartitionField field2(3, 5, "field2", Transform::Identity());
@@ -1248,10 +1221,7 @@ TEST(TableMetadataBuilderTest, RemovePartitionSpecBasic) {
   builder->RemovePartitionSpecs({2});
   ICEBERG_UNWRAP_OR_FAIL(metadata, builder->Build());
   ASSERT_EQ(metadata->partition_specs.size(), 1);
-  EXPECT_TRUE(std::ranges::find_if(metadata->partition_specs,
-                                   [&](const std::shared_ptr<PartitionSpec>& s) {
-                                     return s->spec_id() == 1;
-                                   }) != metadata->partition_specs.end());
+  ASSERT_THAT(metadata->PartitionSpecById(2), IsError(ErrorKind::kNotFound));
 }
 
 TEST(TableMetadataBuilderTest, RemovePartitionSpecNotExist) {
@@ -1269,7 +1239,7 @@ TEST(TableMetadataBuilderTest, RemovePartitionSpecNotExist) {
   ICEBERG_UNWRAP_OR_FAIL(auto metadata, builder->Build());
   ASSERT_EQ(metadata->partition_specs.size(), 2);
 
-  // Remove one not exist spec
+  // Remove one non-existing spec
   builder = TableMetadataBuilder::BuildFrom(metadata.get());
   builder->RemovePartitionSpecs({3});
   ICEBERG_UNWRAP_OR_FAIL(metadata, builder->Build());

src/iceberg/transaction.cc

@@ -170,6 +170,11 @@ Status Transaction::Apply(PendingUpdate& update) {
       if (!result.snapshot_ids_to_remove.empty()) {
         metadata_builder_->RemoveSnapshots(std::move(result.snapshot_ids_to_remove));
       }
+      if (!result.refs_to_remove.empty()) {
+        for (const auto& ref_name : result.refs_to_remove) {
+          metadata_builder_->RemoveRef(ref_name);
+        }
+      }
       if (!result.partition_spec_ids_to_remove.empty()) {
         metadata_builder_->RemovePartitionSpecs(
             std::move(result.partition_spec_ids_to_remove));