fix review

Author: HeartLinked

src/iceberg/json_internal.cc

@@ -1593,6 +1593,17 @@ Result<std::unique_ptr<TableUpdate>> TableUpdateFromJson(const nlohmann::json& j
     ICEBERG_ASSIGN_OR_RAISE(auto location, GetJsonValue<std::string>(json, kLocation));
     return std::make_unique<table::SetLocation>(std::move(location));
   }
+  if (action == kActionSetStatistics) {
+    ICEBERG_ASSIGN_OR_RAISE(auto statistics_json,
+                            GetJsonValue<nlohmann::json>(json, kStatistics));
+    ICEBERG_ASSIGN_OR_RAISE(auto statistics_file,
+                            StatisticsFileFromJson(statistics_json));
+    return std::make_unique<table::SetStatistics>(std::move(statistics_file));
+  }
+  if (action == kActionRemoveStatistics) {
+    ICEBERG_ASSIGN_OR_RAISE(auto snapshot_id, GetJsonValue<int64_t>(json, kSnapshotId));
+    return std::make_unique<table::RemoveStatistics>(snapshot_id);
+  }
 
   return JsonParseError("Unknown table update action: {}", action);
 }

src/iceberg/table_metadata.cc

@@ -613,7 +613,7 @@ class TableMetadataBuilder::Impl {
   Status SetCurrentSchema(int32_t schema_id);
   Status RemoveSchemas(const std::unordered_set<int32_t>& schema_ids);
   Result<int32_t> AddSchema(const Schema& schema, int32_t new_last_column_id);
-  Status SetLocation(std::string_view location);
+  void SetLocation(std::string_view location);
   Status AddSnapshot(std::shared_ptr<Snapshot> snapshot);
   Status SetBranchSnapshot(int64_t snapshot_id, const std::string& branch);
   Status SetBranchSnapshot(std::shared_ptr<Snapshot> snapshot, const std::string& branch);
@@ -1035,13 +1035,13 @@ Result<int32_t> TableMetadataBuilder::Impl::AddSchema(const Schema& schema,
   return new_schema_id;
 }
 
-Status TableMetadataBuilder::Impl::SetLocation(std::string_view location) {
+void TableMetadataBuilder::Impl::SetLocation(std::string_view location) {
   if (location == metadata_.location) {
-    return {};
+    return;
   }
   metadata_.location = std::string(location);
   changes_.push_back(std::make_unique<table::SetLocation>(std::string(location)));
-  return {};
+  return;
 }
 
 Status TableMetadataBuilder::Impl::AddSnapshot(std::shared_ptr<Snapshot> snapshot) {
@@ -1179,7 +1179,7 @@ Status TableMetadataBuilder::Impl::SetRef(const std::string& name,
 
 Status TableMetadataBuilder::Impl::SetStatistics(
     const std::shared_ptr<StatisticsFile>& statistics_file) {
-  ICEBERG_CHECK(statistics_file != nullptr, "Cannot set null statistics file");
+  ICEBERG_PRECHECK(statistics_file != nullptr, "Cannot set null statistics file");
 
   // Find and replace existing statistics for the same snapshot_id, or add new one
   auto it = std::ranges::find_if(
@@ -1194,24 +1194,19 @@ Status TableMetadataBuilder::Impl::SetStatistics(
     metadata_.statistics.push_back(statistics_file);
   }
 
-  changes_.push_back(std::make_unique<table::SetStatistics>(statistics_file));
+  changes_.push_back(std::make_unique<table::SetStatistics>(std::move(statistics_file)));
   return {};
 }
 
 Status TableMetadataBuilder::Impl::RemoveStatistics(int64_t snapshot_id) {
-  auto it = std::ranges::find_if(metadata_.statistics, [snapshot_id](const auto& stat) {
-    return stat && stat->snapshot_id == snapshot_id;
-  });
-
-  if (it == metadata_.statistics.end()) {
+  auto removed_count =
+      std::erase_if(metadata_.statistics, [snapshot_id](const auto& stat) {
+        return stat && stat->snapshot_id == snapshot_id;
+      });
+  if (removed_count == 0) {
     return {};
   }
 
-  // Remove statistics for the given snapshot_id
-  std::erase_if(metadata_.statistics, [snapshot_id](const auto& stat) {
-    return stat && stat->snapshot_id == snapshot_id;
-  });
-
   changes_.push_back(std::make_unique<table::RemoveStatistics>(snapshot_id));
   return {};
 }
@@ -1632,7 +1627,7 @@ TableMetadataBuilder& TableMetadataBuilder::SuppressHistoricalSnapshots() {
 }
 
 TableMetadataBuilder& TableMetadataBuilder::SetStatistics(
-    const std::shared_ptr<StatisticsFile>& statistics_file) {
+    std::shared_ptr<StatisticsFile> statistics_file) {
   ICEBERG_BUILDER_RETURN_IF_ERROR(impl_->SetStatistics(statistics_file));
   return *this;
 }
@@ -1665,7 +1660,7 @@ TableMetadataBuilder& TableMetadataBuilder::RemoveProperties(
 }
 
 TableMetadataBuilder& TableMetadataBuilder::SetLocation(std::string_view location) {
-  ICEBERG_BUILDER_RETURN_IF_ERROR(impl_->SetLocation(location));
+  impl_->SetLocation(location);
   return *this;
 }
 

src/iceberg/table_metadata.h

@@ -399,8 +399,7 @@ class ICEBERG_EXPORT TableMetadataBuilder : public ErrorCollector {
   ///
   /// \param statistics_file The statistics file to set
   /// \return Reference to this builder for method chaining
-  TableMetadataBuilder& SetStatistics(
-      const std::shared_ptr<StatisticsFile>& statistics_file);
+  TableMetadataBuilder& SetStatistics(std::shared_ptr<StatisticsFile> statistics_file);
 
   /// \brief Remove table statistics by snapshot ID
   ///

src/iceberg/table_update.cc

@@ -464,7 +464,13 @@ bool SetStatistics::Equals(const TableUpdate& other) const {
     return false;
   }
   const auto& other_set = static_cast<const SetStatistics&>(other);
-  return *statistics_file_ == *other_set.statistics_file_;
+  if (!statistics_file_ != !other_set.statistics_file_) {
+    return false;
+  }
+  if (statistics_file_ && !(*statistics_file_ == *other_set.statistics_file_)) {
+    return false;
+  }
+  return true;
 }
 
 std::unique_ptr<TableUpdate> SetStatistics::Clone() const {

src/iceberg/test/json_internal_test.cc

@@ -31,6 +31,7 @@
 #include "iceberg/snapshot.h"
 #include "iceberg/sort_field.h"
 #include "iceberg/sort_order.h"
+#include "iceberg/statistics_file.h"
 #include "iceberg/table_requirement.h"
 #include "iceberg/table_update.h"
 #include "iceberg/test/matchers.h"
@@ -508,6 +509,54 @@ TEST(JsonInternalTest, TableUpdateSetLocation) {
   EXPECT_EQ(*internal::checked_cast<table::SetLocation*>(parsed.value().get()), update);
 }
 
+TEST(JsonInternalTest, TableUpdateSetStatistics) {
+  auto stats_file = std::make_shared<StatisticsFile>();
+  stats_file->snapshot_id = 123456789;
+  stats_file->path = "s3://bucket/warehouse/table/metadata/stats-123456789.puffin";
+  stats_file->file_size_in_bytes = 1024;
+  stats_file->file_footer_size_in_bytes = 128;
+  stats_file->blob_metadata = {BlobMetadata{.type = "ndv",
+                                            .source_snapshot_id = 123456789,
+                                            .source_snapshot_sequence_number = 1,
+                                            .fields = {1, 2},
+                                            .properties = {{"prop1", "value1"}}}};
+
+  table::SetStatistics update(stats_file);
+  nlohmann::json expected = R"({
+    "action": "set-statistics",
+    "statistics": {
+      "snapshot-id": 123456789,
+      "statistics-path": "s3://bucket/warehouse/table/metadata/stats-123456789.puffin",
+      "file-size-in-bytes": 1024,
+      "file-footer-size-in-bytes": 128,
+      "blob-metadata": [{
+        "type": "ndv",
+        "snapshot-id": 123456789,
+        "sequence-number": 1,
+        "fields": [1, 2],
+        "properties": {"prop1": "value1"}
+      }]
+    }
+  })"_json;
+
+  EXPECT_EQ(ToJson(update), expected);
+  auto parsed = TableUpdateFromJson(expected);
+  ASSERT_THAT(parsed, IsOk());
+  EXPECT_EQ(*internal::checked_cast<table::SetStatistics*>(parsed.value().get()), update);
+}
+
+TEST(JsonInternalTest, TableUpdateRemoveStatistics) {
+  table::RemoveStatistics update(123456789);
+  nlohmann::json expected =
+      R"({"action":"remove-statistics","snapshot-id":123456789})"_json;
+
+  EXPECT_EQ(ToJson(update), expected);
+  auto parsed = TableUpdateFromJson(expected);
+  ASSERT_THAT(parsed, IsOk());
+  EXPECT_EQ(*internal::checked_cast<table::RemoveStatistics*>(parsed.value().get()),
+            update);
+}
+
 TEST(JsonInternalTest, TableUpdateUnknownAction) {
   nlohmann::json json = R"({"action":"unknown-action"})"_json;
   auto result = TableUpdateFromJson(json);

src/iceberg/test/update_statistics_test.cc

@@ -19,6 +19,7 @@
 
 #include "iceberg/update/update_statistics.h"
 
+#include <algorithm>
 #include <memory>
 
 #include <gmock/gmock.h>
@@ -54,6 +55,16 @@ class UpdateStatisticsTest : public UpdateTestBase {
 
     return stats_file;
   }
+
+  // Helper to find statistics file by snapshot_id in the result vector
+  std::shared_ptr<StatisticsFile> FindStatistics(
+      const std::vector<std::pair<int64_t, std::shared_ptr<StatisticsFile>>>& to_set,
+      int64_t snapshot_id) {
+    auto it = std::find_if(to_set.begin(), to_set.end(), [snapshot_id](const auto& p) {
+      return p.first == snapshot_id;
+    });
+    return it != to_set.end() ? it->second : nullptr;
+  }
 };
 
 TEST_F(UpdateStatisticsTest, EmptyUpdate) {
@@ -72,7 +83,7 @@ TEST_F(UpdateStatisticsTest, SetStatistics) {
   ICEBERG_UNWRAP_OR_FAIL(auto result, update->Apply());
   EXPECT_EQ(result.to_set.size(), 1);
   EXPECT_TRUE(result.to_remove.empty());
-  EXPECT_EQ(result.to_set.at(1), stats_file);
+  EXPECT_EQ(FindStatistics(result.to_set, 1), stats_file);
 }
 
 TEST_F(UpdateStatisticsTest, SetMultipleStatistics) {
@@ -87,8 +98,8 @@ TEST_F(UpdateStatisticsTest, SetMultipleStatistics) {
   ICEBERG_UNWRAP_OR_FAIL(auto result, update->Apply());
   EXPECT_EQ(result.to_set.size(), 2);
   EXPECT_TRUE(result.to_remove.empty());
-  EXPECT_EQ(result.to_set.at(1), stats_file_1);
-  EXPECT_EQ(result.to_set.at(2), stats_file_2);
+  EXPECT_EQ(FindStatistics(result.to_set, 1), stats_file_1);
+  EXPECT_EQ(FindStatistics(result.to_set, 2), stats_file_2);
 }
 
 TEST_F(UpdateStatisticsTest, RemoveStatistics) {
@@ -120,7 +131,7 @@ TEST_F(UpdateStatisticsTest, SetAndRemoveDifferentSnapshots) {
 
   ICEBERG_UNWRAP_OR_FAIL(auto result, update->Apply());
   EXPECT_EQ(result.to_set.size(), 1);
-  EXPECT_EQ(result.to_set.at(1), stats_file);
+  EXPECT_EQ(FindStatistics(result.to_set, 1), stats_file);
   EXPECT_EQ(result.to_remove.size(), 1);
   EXPECT_THAT(result.to_remove, ::testing::Contains(2));
 }
@@ -139,8 +150,8 @@ TEST_F(UpdateStatisticsTest, ReplaceStatistics) {
   EXPECT_EQ(result.to_set.size(), 1);
   EXPECT_TRUE(result.to_remove.empty());
   // Should have the second one (replacement)
-  EXPECT_EQ(result.to_set.at(1), stats_file_2);
-  EXPECT_NE(result.to_set.at(1), stats_file_1);
+  EXPECT_EQ(FindStatistics(result.to_set, 1), stats_file_2);
+  EXPECT_NE(FindStatistics(result.to_set, 1), stats_file_1);
 }
 
 TEST_F(UpdateStatisticsTest, SetThenRemoveSameSnapshot) {
@@ -168,7 +179,7 @@ TEST_F(UpdateStatisticsTest, RemoveThenSetSameSnapshot) {
   ICEBERG_UNWRAP_OR_FAIL(auto result, update->Apply());
   EXPECT_EQ(result.to_set.size(), 1);
   EXPECT_TRUE(result.to_remove.empty());
-  EXPECT_EQ(result.to_set.at(1), stats_file);
+  EXPECT_EQ(FindStatistics(result.to_set, 1), stats_file);
 }
 
 TEST_F(UpdateStatisticsTest, SetNullStatistics) {

src/iceberg/transaction.cc

@@ -195,8 +195,8 @@ Status Transaction::Apply(PendingUpdate& update) {
       auto& update_statistics = internal::checked_cast<UpdateStatistics&>(update);
       ICEBERG_ASSIGN_OR_RAISE(auto result, update_statistics.Apply());
       // Apply statistics changes to the metadata builder
-      for (const auto& [snapshot_id, stat_file] : result.to_set) {
-        metadata_builder_->SetStatistics(stat_file);
+      for (auto&& [_, stat_file] : result.to_set) {
+        metadata_builder_->SetStatistics(std::move(stat_file));
       }
       for (const auto& snapshot_id : result.to_remove) {
         metadata_builder_->RemoveStatistics(snapshot_id);

src/iceberg/update/update_statistics.cc

@@ -44,9 +44,9 @@ UpdateStatistics::UpdateStatistics(std::shared_ptr<Transaction> transaction)
 UpdateStatistics::~UpdateStatistics() = default;
 
 UpdateStatistics& UpdateStatistics::SetStatistics(
-    const std::shared_ptr<StatisticsFile>& statistics_file) {
+    std::shared_ptr<StatisticsFile> statistics_file) {
   ICEBERG_BUILDER_CHECK(statistics_file != nullptr, "Statistics file cannot be null");
-  statistics_to_set_[statistics_file->snapshot_id] = statistics_file;
+  statistics_to_set_[statistics_file->snapshot_id] = std::move(statistics_file);
   return *this;
 }
 
@@ -61,7 +61,7 @@ Result<UpdateStatistics::ApplyResult> UpdateStatistics::Apply() {
   ApplyResult result;
   for (const auto& [snapshot_id, stats] : statistics_to_set_) {
     if (stats) {
-      result.to_set[snapshot_id] = stats;
+      result.to_set.emplace_back(snapshot_id, stats);
     } else {
       result.to_remove.push_back(snapshot_id);
     }

src/iceberg/update/update_statistics.h

@@ -22,6 +22,7 @@
 #include <cstdint>
 #include <memory>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
 #include "iceberg/iceberg_export.h"
@@ -49,7 +50,7 @@ class ICEBERG_EXPORT UpdateStatistics : public PendingUpdate {
   ///
   /// \param statistics_file The statistics file to set
   /// \return Reference to this UpdateStatistics for chaining
-  UpdateStatistics& SetStatistics(const std::shared_ptr<StatisticsFile>& statistics_file);
+  UpdateStatistics& SetStatistics(std::shared_ptr<StatisticsFile> statistics_file);
 
   /// \brief Remove statistics for a snapshot.
   ///
@@ -62,7 +63,7 @@ class ICEBERG_EXPORT UpdateStatistics : public PendingUpdate {
   Kind kind() const final { return Kind::kUpdateStatistics; }
 
   struct ApplyResult {
-    std::unordered_map<int64_t, std::shared_ptr<StatisticsFile>> to_set;
+    std::vector<std::pair<int64_t, std::shared_ptr<StatisticsFile>>> to_set;
     std::vector<int64_t> to_remove;
   };