feat: add schema update to table metadata builder

Author: gty404

src/iceberg/partition_spec.cc

@@ -155,6 +155,42 @@ Status PartitionSpec::Validate(const Schema& schema, bool allow_missing_fields)
   return {};
 }
 
+Status PartitionSpec::ValidatePartitionName(const Schema& schema) const {
+  std::unordered_set<std::string> partition_names;
+  for (const auto& partition_field : fields_) {
+    auto name = std::string(partition_field.name());
+    if (name.empty()) {
+      return InvalidArgument("Cannot use empty partition name: {}", name);
+    }
+    if (partition_names.contains(name)) {
+      return InvalidArgument("Cannot use partition name more than once: {}", name);
+    }
+    partition_names.insert(name);
+
+    ICEBERG_ASSIGN_OR_RAISE(auto schema_field, schema.FindFieldByName(name));
+    auto transform_type = partition_field.transform()->transform_type();
+    if (transform_type == TransformType::kIdentity) {
+      // for identity transform case we allow conflicts between partition and schema field
+      // name as long as they are sourced from the same schema field
+      if (schema_field.has_value() &&
+          schema_field.value().get().field_id() != partition_field.source_id()) {
+        return InvalidArgument(
+            "Cannot create identity partition sourced from different field in schema: {}",
+            name);
+      }
+    } else {
+      // for all other transforms we don't allow conflicts between partition name and
+      // schema field name
+      if (schema_field.has_value()) {
+        return InvalidArgument(
+            "Cannot create partition from name that exists in schema: {}", name);
+      }
+    }
+  }
+
+  return {};
+}
+
 Result<std::vector<std::reference_wrapper<const PartitionField>>>
 PartitionSpec::GetFieldsBySourceId(int32_t source_id) const {
   ICEBERG_ASSIGN_OR_RAISE(auto source_id_to_fields, source_id_to_fields_.Get(*this));

src/iceberg/partition_spec.h

@@ -79,6 +79,10 @@ class ICEBERG_EXPORT PartitionSpec : public util::Formattable {
   /// \return Error status if the partition spec is invalid.
   Status Validate(const Schema& schema, bool allow_missing_fields) const;
 
+  // \brief Validates the partition field names are unique within the partition spec and
+  // schema.
+  Status ValidatePartitionName(const Schema& schema) const;
+
   /// \brief Get the partition fields by source ID.
   /// \param source_id The id of the source field.
   /// \return The partition fields by source ID, or NotFound if the source field is not

src/iceberg/schema.cc

@@ -30,6 +30,7 @@
 #include "iceberg/util/macros.h"
 #include "iceberg/util/type_util.h"
 #include "iceberg/util/visit_type.h"
+#include "table_metadata.h"
 
 namespace iceberg {
 
@@ -179,4 +180,43 @@ Result<std::unique_ptr<Schema>> Schema::Project(
                         std::nullopt);
 }
 
+Result<int32_t> Schema::HighestFieldId() const {
+  ICEBERG_ASSIGN_OR_RAISE(auto id_to_field, id_to_field_.Get(*this));
+
+  if (id_to_field.get().empty()) {
+    return kInitialColumnId;
+  }
+
+  auto max_it = std::ranges::max_element(
+      id_to_field.get(),
+      [](const auto& lhs, const auto& rhs) { return lhs.first < rhs.first; });
+
+  return max_it->first;
+}
+
+bool Schema::SameSchema(const Schema& other) const { return fields_ == other.fields_; }
+
+Status Schema::Validate(int32_t format_version) const {
+  // Get all fields including nested ones
+  ICEBERG_ASSIGN_OR_RAISE(auto id_to_field, id_to_field_.Get(*this));
+
+  // Check each field's type and defaults
+  for (const auto& [field_id, field_ref] : id_to_field.get()) {
+    const auto& field = field_ref.get();
+
+    // Check if the field's type requires a minimum format version
+    if (auto it = TableMetadata::kMinFormatVersions.find(field.type()->type_id());
+        it != TableMetadata::kMinFormatVersions.end()) {
+      if (int32_t min_format_version = it->second; format_version < min_format_version) {
+        return InvalidSchema("Invalid type for {}: {} is not supported until v{}",
+                             field.name(), *field.type(), min_format_version);
+      }
+    }
+
+    // TODO(GuoTao.yu): Check default values when they are supported
+  }
+
+  return {};
+}
+
 }  // namespace iceberg

src/iceberg/schema.h

@@ -46,6 +46,7 @@ namespace iceberg {
 class ICEBERG_EXPORT Schema : public StructType {
  public:
   static constexpr int32_t kInitialSchemaId = 0;
+  static constexpr int32_t kInitialColumnId = 0;
   static constexpr int32_t kInvalidColumnId = -1;
 
   explicit Schema(std::vector<SchemaField> fields,
@@ -103,6 +104,23 @@ class ICEBERG_EXPORT Schema : public StructType {
   Result<std::unique_ptr<Schema>> Project(
       const std::unordered_set<int32_t>& field_ids) const;
 
+  /// \brief Get the highest field ID in the schema.
+  /// \return The highest field ID.
+  Result<int32_t> HighestFieldId() const;
+
+  /// \brief Checks whether this schema is equivalent to another schema while ignoring the
+  /// schema id.
+  bool SameSchema(const Schema& other) const;
+
+  /// \brief Validate the schema for a given format version.
+  ///
+  /// This validates that the schema does not contain types that were released in later
+  /// format versions.
+  ///
+  /// \param format_version The format version to validate against.
+  /// \return Error status if the schema is invalid.
+  Status Validate(int32_t format_version) const;
+
   friend bool operator==(const Schema& lhs, const Schema& rhs) { return lhs.Equals(rhs); }
 
  private:

src/iceberg/table_metadata.cc

@@ -522,16 +522,166 @@ TableMetadataBuilder& TableMetadataBuilder::UpgradeFormatVersion(
 }
 
 TableMetadataBuilder& TableMetadataBuilder::SetCurrentSchema(
-    std::shared_ptr<Schema> schema, int32_t new_last_column_id) {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+    std::shared_ptr<Schema> const& schema, int32_t new_last_column_id) {
+  ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto schema_id,
+                                   AddSchemaInternal(*schema, new_last_column_id));
+  return SetCurrentSchema(schema_id);
 }
 
 TableMetadataBuilder& TableMetadataBuilder::SetCurrentSchema(int32_t schema_id) {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+  if (schema_id == kLastAdded) {
+    ICEBERG_BUILDER_CHECK(impl_->last_added_schema_id.has_value(),
+                          "Cannot set last added schema: no schema has been added");
+    return SetCurrentSchema(impl_->last_added_schema_id.value());
+  }
+
+  if (impl_->metadata.current_schema_id == schema_id) {
+    return *this;
+  }
+
+  auto it = impl_->schemas_by_id.find(schema_id);
+  ICEBERG_BUILDER_CHECK(it != impl_->schemas_by_id.end(),
+                        "Cannot set current schema to unknown schema: {}", schema_id);
+  const auto& schema = it->second;
+
+  // Rebuild all partition specs for the new current schema
+  std::vector<std::shared_ptr<PartitionSpec>> updated_specs;
+  for (const auto& spec : impl_->metadata.partition_specs) {
+    ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto updated_spec, UpdateSpecSchema(*schema, *spec));
+    updated_specs.push_back(std::move(updated_spec));
+  }
+  impl_->metadata.partition_specs = std::move(updated_specs);
+  impl_->specs_by_id.clear();
+  for (const auto& spec : impl_->metadata.partition_specs) {
+    impl_->specs_by_id.emplace(spec->spec_id(), spec);
+  }
+
+  // Rebuild all sort orders for the new current schema
+  std::vector<std::shared_ptr<SortOrder>> updated_orders;
+  for (const auto& order : impl_->metadata.sort_orders) {
+    ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto updated_order,
+                                     UpdateSortOrderSchema(*schema, *order));
+    updated_orders.push_back(std::move(updated_order));
+  }
+  impl_->metadata.sort_orders = std::move(updated_orders);
+  impl_->sort_orders_by_id.clear();
+  for (const auto& order : impl_->metadata.sort_orders) {
+    impl_->sort_orders_by_id.emplace(order->order_id(), order);
+  }
+
+  // Set the current schema ID
+  impl_->metadata.current_schema_id = schema_id;
+
+  // Record the change
+  if (impl_->last_added_schema_id.has_value() &&
+      impl_->last_added_schema_id.value() == schema_id) {
+    impl_->changes.push_back(std::make_unique<table::SetCurrentSchema>(kLastAdded));
+  } else {
+    impl_->changes.push_back(std::make_unique<table::SetCurrentSchema>(schema_id));
+  }
+
+  return *this;
 }
 
-TableMetadataBuilder& TableMetadataBuilder::AddSchema(std::shared_ptr<Schema> schema) {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+Result<std::shared_ptr<PartitionSpec>> TableMetadataBuilder::UpdateSpecSchema(
+    const Schema& schema, const PartitionSpec& partition_spec) {
+  // UpdateSpecSchema: Update partition spec to use the new schema
+  // This preserves the partition spec structure but rebinds it to the new schema
+
+  // Copy all fields from the partition spec. IDs should not change.
+  std::vector<PartitionField> fields;
+  fields.reserve(partition_spec.fields().size());
+  int32_t last_assigned_field_id = PartitionSpec::kLegacyPartitionDataIdStart;
+  for (const auto& field : partition_spec.fields()) {
+    fields.push_back(field);
+    last_assigned_field_id = std::max(last_assigned_field_id, field.field_id());
+  }
+
+  // Build without validation because the schema may have changed in a way that makes
+  // this spec invalid. The spec should still be preserved so that older metadata can
+  // be interpreted.
+  ICEBERG_ASSIGN_OR_RAISE(auto new_partition_spec,
+                          PartitionSpec::Make(partition_spec.spec_id(), std::move(fields),
+                                              last_assigned_field_id));
+
+  // Validate the new partition name against the new schema
+  ICEBERG_RETURN_UNEXPECTED(new_partition_spec->ValidatePartitionName(schema));
+  return new_partition_spec;
+}
+
+Result<std::unique_ptr<SortOrder>> TableMetadataBuilder::UpdateSortOrderSchema(
+    const Schema& schema, const SortOrder& sort_order) {
+  // Build without validation because the schema may have changed in a way that makes
+  // this order invalid. The order should still be preserved so that older metadata can
+  // be interpreted.
+  auto fields = sort_order.fields();
+  std::vector<SortField> new_fields{fields.begin(), fields.end()};
+  return SortOrder::Make(sort_order.order_id(), std::move(new_fields));
+}
+
+TableMetadataBuilder& TableMetadataBuilder::AddSchema(
+    std::shared_ptr<Schema> const& schema) {
+  ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto highest_field_id, schema->HighestFieldId());
+  AddSchemaInternal(*schema, std::max(impl_->metadata.last_column_id, highest_field_id));
+  return *this;
+}
+
+Result<int32_t> TableMetadataBuilder::AddSchemaInternal(const Schema& schema,
+                                                        int32_t new_last_column_id) {
+  if (new_last_column_id < impl_->metadata.last_column_id) {
+    return InvalidArgument("Invalid last column ID: {} < {} (previous last column ID)",
+                           new_last_column_id, impl_->metadata.last_column_id);
+  }
+
+  ICEBERG_RETURN_UNEXPECTED(schema.Validate(impl_->metadata.format_version));
+
+  auto new_schema_id = ReuseOrCreateNewSchemaId(schema);
+  if (impl_->schemas_by_id.find(new_schema_id) != impl_->schemas_by_id.end()) {
+    // update last_added_schema_id if the schema was added in this set of changes (since
+    // it is now the last)
+    bool is_new_schema =
+        impl_->last_added_schema_id.has_value() &&
+        std::ranges::find_if(impl_->changes, [new_schema_id](const auto& change) {
+          if (change->kind() != TableUpdate::Kind::kAddSchema) {
+            return false;
+          }
+          auto* add_schema = dynamic_cast<table::AddSchema*>(change.get());
+          return add_schema->schema()->schema_id().value_or(Schema::kInitialSchemaId) ==
+                 new_schema_id;
+        }) != impl_->changes.cend();
+    impl_->last_added_schema_id =
+        is_new_schema ? std::make_optional(new_schema_id) : std::nullopt;
+    return new_schema_id;
+  }
+
+  auto new_schema = std::make_shared<Schema>(
+      std::vector<SchemaField>(schema.fields().begin(), schema.fields().end()),
+      new_schema_id);
+
+  impl_->metadata.schemas.push_back(new_schema);
+  impl_->schemas_by_id.emplace(new_schema_id, new_schema);
+
+  impl_->changes.push_back(
+      std::make_unique<table::AddSchema>(new_schema, new_last_column_id));
+  impl_->metadata.last_column_id = new_last_column_id;
+  impl_->last_added_schema_id = new_schema_id;
+
+  return new_schema_id;
+}
+
+int32_t TableMetadataBuilder::ReuseOrCreateNewSchemaId(const Schema& new_schema) const {
+  // if the schema already exists, use its id; otherwise use the highest id + 1
+  auto new_schema_id =
+      impl_->metadata.current_schema_id.value_or(Schema::kInitialSchemaId);
+  for (auto& schema : impl_->metadata.schemas) {
+    auto schema_id = schema->schema_id().value_or(Schema::kInitialSchemaId);
+    if (schema->SameSchema(new_schema)) {
+      return schema_id;
+    } else if (new_schema_id <= schema_id) {
+      new_schema_id = schema_id + 1;
+    }
+  }
+  return new_schema_id;
 }
 
 TableMetadataBuilder& TableMetadataBuilder::SetDefaultPartitionSpec(
@@ -555,7 +705,23 @@ TableMetadataBuilder& TableMetadataBuilder::RemovePartitionSpecs(
 
 TableMetadataBuilder& TableMetadataBuilder::RemoveSchemas(
     const std::vector<int32_t>& schema_ids) {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+  std::unordered_set<int32_t> schema_ids_to_remove(schema_ids.begin(), schema_ids.end());
+  auto current_schema_id =
+      impl_->metadata.current_schema_id.value_or(Schema::kInitialSchemaId);
+  ICEBERG_BUILDER_CHECK(schema_ids_to_remove.contains(current_schema_id),
+                        "Cannot remove current schema: {}", current_schema_id);
+
+  if (!schema_ids_to_remove.empty()) {
+    impl_->metadata.schemas =
+        impl_->metadata.schemas | std::views::filter([&](const auto& schema) {
+          return !schema_ids_to_remove.contains(
+              schema->schema_id().value_or(Schema::kInitialSchemaId));
+        }) |
+        std::ranges::to<std::vector<std::shared_ptr<Schema>>>();
+    impl_->changes.push_back(
+        std::make_unique<table::RemoveSchemas>(std::move(schema_ids_to_remove)));
+  }
+  return *this;
 }
 
 TableMetadataBuilder& TableMetadataBuilder::SetDefaultSortOrder(
@@ -595,9 +761,11 @@ Result<int32_t> TableMetadataBuilder::AddSortOrderInternal(const SortOrder& orde
     bool is_new_order =
         impl_->last_added_order_id.has_value() &&
         std::ranges::find_if(impl_->changes, [new_order_id](const auto& change) {
+          if (change->kind() != TableUpdate::Kind::kAddSortOrder) {
+            return false;
+          }
           auto* add_sort_order = dynamic_cast<table::AddSortOrder*>(change.get());
-          return add_sort_order &&
-                 add_sort_order->sort_order()->order_id() == new_order_id;
+          return add_sort_order->sort_order()->order_id() == new_order_id;
         }) != impl_->changes.cend();
     impl_->last_added_order_id =
         is_new_order ? std::make_optional(new_order_id) : std::nullopt;
@@ -764,6 +932,27 @@ Result<std::unique_ptr<TableMetadata>> TableMetadataBuilder::Build() {
             std::chrono::system_clock::now().time_since_epoch())};
   }
 
+  auto current_schema_id =
+      impl_->metadata.current_schema_id.value_or(Schema::kInitialSchemaId);
+  auto schema_it = impl_->schemas_by_id.find(current_schema_id);
+  ICEBERG_PRECHECK(schema_it != impl_->schemas_by_id.end(),
+                   "Current schema ID {} not found in schemas", current_schema_id);
+  const auto& current_schema = schema_it->second;
+  {
+    auto spec_it = impl_->specs_by_id.find(impl_->metadata.default_spec_id);
+    // FIXME(GuoTao.yu): Default spec must exist after we support update partition spec
+    if (spec_it != impl_->specs_by_id.end()) {
+      ICEBERG_RETURN_UNEXPECTED(
+          spec_it->second->Validate(*current_schema, /*allow_missing_fields=*/false));
+    }
+    auto sort_order_it =
+        impl_->sort_orders_by_id.find(impl_->metadata.default_sort_order_id);
+    ICEBERG_PRECHECK(sort_order_it != impl_->sort_orders_by_id.end(),
+                     "Default sort order ID {} not found in sort orders",
+                     impl_->metadata.default_sort_order_id);
+    ICEBERG_RETURN_UNEXPECTED(sort_order_it->second->Validate(*current_schema));
+  }
+
   // 4. Buildup metadata_log from base metadata
   int32_t max_metadata_log_size =
       impl_->metadata.properties.Get(TableProperties::kMetadataPreviousVersionsMax);