feat: implement update stastics

Author: HeartLinked

src/iceberg/CMakeLists.txt

@@ -92,6 +92,7 @@ set(ICEBERG_SOURCES
     update/update_properties.cc
     update/update_schema.cc
     update/update_sort_order.cc
+    update/update_statistics.cc
     util/bucket_util.cc
     util/content_file_util.cc
     util/conversions.cc

src/iceberg/json_internal.cc

@@ -191,6 +191,8 @@ constexpr std::string_view kActionSetSnapshotRef = "set-snapshot-ref";
 constexpr std::string_view kActionSetProperties = "set-properties";
 constexpr std::string_view kActionRemoveProperties = "remove-properties";
 constexpr std::string_view kActionSetLocation = "set-location";
+constexpr std::string_view kActionSetStatistics = "set-statistics";
+constexpr std::string_view kActionRemoveStatistics = "remove-statistics";
 
 // TableUpdate field constants
 constexpr std::string_view kUUID = "uuid";
@@ -1399,6 +1401,18 @@ nlohmann::json ToJson(const TableUpdate& update) {
       json[kLocation] = u.location();
       break;
     }
+    case TableUpdate::Kind::kSetStatistics: {
+      const auto& u = internal::checked_cast<const table::SetStatistics&>(update);
+      json[kAction] = kActionSetStatistics;
+      json[kStatistics] = ToJson(*u.statistics_file());
+      break;
+    }
+    case TableUpdate::Kind::kRemoveStatistics: {
+      const auto& u = internal::checked_cast<const table::RemoveStatistics&>(update);
+      json[kAction] = kActionRemoveStatistics;
+      json[kSnapshotId] = u.snapshot_id();
+      break;
+    }
   }
   return json;
 }

src/iceberg/meson.build

@@ -110,6 +110,7 @@ iceberg_sources = files(
     'update/update_properties.cc',
     'update/update_schema.cc',
     'update/update_sort_order.cc',
+    'update/update_statistics.cc',
     'util/bucket_util.cc',
     'util/content_file_util.cc',
     'util/conversions.cc',

src/iceberg/table.cc

@@ -35,6 +35,7 @@
 #include "iceberg/update/update_partition_spec.h"
 #include "iceberg/update/update_properties.h"
 #include "iceberg/update/update_schema.h"
+#include "iceberg/update/update_statistics.h"
 #include "iceberg/util/macros.h"
 
 namespace iceberg {
@@ -199,6 +200,13 @@ Result<std::shared_ptr<UpdateLocation>> Table::NewUpdateLocation() {
   return transaction->NewUpdateLocation();
 }
 
+Result<std::shared_ptr<UpdateStatistics>> Table::NewUpdateStatistics() {
+  ICEBERG_ASSIGN_OR_RAISE(
+      auto transaction, Transaction::Make(shared_from_this(), Transaction::Kind::kUpdate,
+                                          /*auto_commit=*/true));
+  return transaction->NewUpdateStatistics();
+}
+
 Result<std::shared_ptr<StagedTable>> StagedTable::Make(
     TableIdentifier identifier, std::shared_ptr<TableMetadata> metadata,
     std::string metadata_location, std::shared_ptr<FileIO> io,

src/iceberg/table.h

@@ -151,6 +151,10 @@ class ICEBERG_EXPORT Table : public std::enable_shared_from_this<Table> {
   /// \brief Create a new ExpireSnapshots to remove expired snapshots and commit the
   /// changes.
   virtual Result<std::shared_ptr<ExpireSnapshots>> NewExpireSnapshots();
+  
+  /// \brief Create a new UpdateStatistics to update the table statistics and commit the
+  /// changes.
+  virtual Result<std::shared_ptr<UpdateStatistics>> NewUpdateStatistics();
 
   /// \brief Create a new UpdateLocation to update the table location and commit the
   /// changes.

src/iceberg/table_metadata.cc

@@ -612,14 +612,16 @@ 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);
-  void SetLocation(std::string_view location);
+  Status 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);
   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);
+  Status SetStatistics(const std::shared_ptr<StatisticsFile>& statistics_file);
+  Status RemoveStatistics(int64_t snapshot_id);
 
   Result<std::unique_ptr<TableMetadata>> Build();
 
@@ -1032,12 +1034,13 @@ Result<int32_t> TableMetadataBuilder::Impl::AddSchema(const Schema& schema,
   return new_schema_id;
 }
 
-void TableMetadataBuilder::Impl::SetLocation(std::string_view location) {
+Status 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 {};
 }
 
 Status TableMetadataBuilder::Impl::AddSnapshot(std::shared_ptr<Snapshot> snapshot) {
@@ -1173,6 +1176,45 @@ Status TableMetadataBuilder::Impl::SetRef(const std::string& name,
   return {};
 }
 
+Status TableMetadataBuilder::Impl::SetStatistics(
+    const std::shared_ptr<StatisticsFile>& statistics_file) {
+  ICEBERG_CHECK(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(
+      metadata_.statistics,
+      [snapshot_id = statistics_file->snapshot_id](const auto& stat) {
+        return stat && stat->snapshot_id == snapshot_id;
+      });
+
+  if (it != metadata_.statistics.end()) {
+    *it = statistics_file;
+  } else {
+    metadata_.statistics.push_back(statistics_file);
+  }
+
+  changes_.push_back(std::make_unique<table::SetStatistics>(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()) {
+    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 {};
+}
+
 std::unordered_set<int64_t> TableMetadataBuilder::Impl::IntermediateSnapshotIdSet(
     int64_t current_snapshot_id) const {
   std::unordered_set<int64_t> added_snapshot_ids;
@@ -1590,11 +1632,13 @@ TableMetadataBuilder& TableMetadataBuilder::SuppressHistoricalSnapshots() {
 
 TableMetadataBuilder& TableMetadataBuilder::SetStatistics(
     const std::shared_ptr<StatisticsFile>& statistics_file) {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+  impl_->SetStatistics(statistics_file);
+  return *this;
 }
 
 TableMetadataBuilder& TableMetadataBuilder::RemoveStatistics(int64_t snapshot_id) {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+  impl_->RemoveStatistics(snapshot_id);
+  return *this;
 }
 
 TableMetadataBuilder& TableMetadataBuilder::SetPartitionStatistics(

src/iceberg/table_update.cc

@@ -22,6 +22,7 @@
 #include "iceberg/exception.h"
 #include "iceberg/schema.h"
 #include "iceberg/sort_order.h"
+#include "iceberg/statistics_file.h"
 #include "iceberg/table_metadata.h"
 #include "iceberg/table_requirements.h"
 
@@ -446,4 +447,50 @@ std::unique_ptr<TableUpdate> SetLocation::Clone() const {
   return std::make_unique<SetLocation>(location_);
 }
 
+// SetStatistics
+
+int64_t SetStatistics::snapshot_id() const { return statistics_file_->snapshot_id; }
+
+void SetStatistics::ApplyTo(TableMetadataBuilder& builder) const {
+  builder.SetStatistics(statistics_file_);
+}
+
+void SetStatistics::GenerateRequirements(TableUpdateContext& context) const {
+  // SetStatistics doesn't generate any requirements
+}
+
+bool SetStatistics::Equals(const TableUpdate& other) const {
+  if (other.kind() != Kind::kSetStatistics) {
+    return false;
+  }
+  const auto& other_set = static_cast<const SetStatistics&>(other);
+  return *statistics_file_ == *other_set.statistics_file_;
+}
+
+std::unique_ptr<TableUpdate> SetStatistics::Clone() const {
+  return std::make_unique<SetStatistics>(statistics_file_);
+}
+
+// RemoveStatistics
+
+void RemoveStatistics::ApplyTo(TableMetadataBuilder& builder) const {
+  builder.RemoveStatistics(snapshot_id_);
+}
+
+void RemoveStatistics::GenerateRequirements(TableUpdateContext& context) const {
+  // RemoveStatistics doesn't generate any requirements
+}
+
+bool RemoveStatistics::Equals(const TableUpdate& other) const {
+  if (other.kind() != Kind::kRemoveStatistics) {
+    return false;
+  }
+  const auto& other_remove = static_cast<const RemoveStatistics&>(other);
+  return snapshot_id_ == other_remove.snapshot_id_;
+}
+
+std::unique_ptr<TableUpdate> RemoveStatistics::Clone() const {
+  return std::make_unique<RemoveStatistics>(snapshot_id_);
+}
+
 }  // namespace iceberg::table

src/iceberg/table_update.h

@@ -59,6 +59,8 @@ class ICEBERG_EXPORT TableUpdate {
     kSetProperties,
     kRemoveProperties,
     kSetLocation,
+    kSetStatistics,
+    kRemoveStatistics,
   };
 
   virtual ~TableUpdate();
@@ -509,6 +511,53 @@ class ICEBERG_EXPORT SetLocation : public TableUpdate {
   std::string location_;
 };
 
+/// \brief Represents setting statistics for a snapshot
+class ICEBERG_EXPORT SetStatistics : public TableUpdate {
+ public:
+  explicit SetStatistics(std::shared_ptr<StatisticsFile> statistics_file)
+      : statistics_file_(std::move(statistics_file)) {}
+
+  int64_t snapshot_id() const;
+
+  const std::shared_ptr<StatisticsFile>& statistics_file() const {
+    return statistics_file_;
+  }
+
+  void ApplyTo(TableMetadataBuilder& builder) const override;
+
+  void GenerateRequirements(TableUpdateContext& context) const override;
+
+  Kind kind() const override { return Kind::kSetStatistics; }
+
+  bool Equals(const TableUpdate& other) const override;
+
+  std::unique_ptr<TableUpdate> Clone() const override;
+
+ private:
+  std::shared_ptr<StatisticsFile> statistics_file_;
+};
+
+/// \brief Represents removing statistics for a snapshot
+class ICEBERG_EXPORT RemoveStatistics : public TableUpdate {
+ public:
+  explicit RemoveStatistics(int64_t snapshot_id) : snapshot_id_(snapshot_id) {}
+
+  int64_t snapshot_id() const { return snapshot_id_; }
+
+  void ApplyTo(TableMetadataBuilder& builder) const override;
+
+  void GenerateRequirements(TableUpdateContext& context) const override;
+
+  Kind kind() const override { return Kind::kRemoveStatistics; }
+
+  bool Equals(const TableUpdate& other) const override;
+
+  std::unique_ptr<TableUpdate> Clone() const override;
+
+ private:
+  int64_t snapshot_id_;
+};
+
 }  // namespace table
 
 }  // namespace iceberg

src/iceberg/test/CMakeLists.txt

@@ -177,7 +177,8 @@ if(ICEBERG_BUILD_BUNDLE)
                    update_partition_spec_test.cc
                    update_properties_test.cc
                    update_schema_test.cc
-                   update_sort_order_test.cc)
+                   update_sort_order_test.cc
+                   update_statistics_test.cc)
 
   add_iceberg_test(data_writer_test USE_BUNDLE SOURCES data_writer_test.cc)