resolve some comments

Author: gty404

src/iceberg/partition_spec.cc

@@ -159,9 +159,8 @@ 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);
-    }
+    ICEBERG_PRECHECK(!name.empty(), "Cannot use empty partition name: {}", name);
+
     if (partition_names.contains(name)) {
       return InvalidArgument("Cannot use partition name more than once: {}", name);
     }

src/iceberg/schema.cc

@@ -25,12 +25,12 @@
 #include "iceberg/result.h"
 #include "iceberg/row/struct_like.h"
 #include "iceberg/schema_internal.h"
+#include "iceberg/table_metadata.h"
 #include "iceberg/type.h"
 #include "iceberg/util/formatter.h"  // IWYU pragma: keep
 #include "iceberg/util/macros.h"
 #include "iceberg/util/type_util.h"
 #include "iceberg/util/visit_type.h"
-#include "table_metadata.h"
 
 namespace iceberg {
 
@@ -148,6 +148,20 @@ Result<std::unordered_map<int32_t, std::vector<size_t>>> Schema::InitIdToPositio
   return visitor.Finish();
 }
 
+Result<int32_t> Schema::InitHighestFieldId(const Schema& self) {
+  ICEBERG_ASSIGN_OR_RAISE(auto id_to_field, self.id_to_field_.Get(self));
+
+  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;
+}
+
 Result<std::unique_ptr<StructLikeAccessor>> Schema::GetAccessorById(
     int32_t field_id) const {
   ICEBERG_ASSIGN_OR_RAISE(auto id_to_position_path, id_to_position_path_.Get(*this));
@@ -229,22 +243,12 @@ Result<std::vector<std::string>> Schema::IdentifierFieldNames() const {
   return names;
 }
 
-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;
-  }
+Result<int32_t> Schema::HighestFieldId() const { return highest_field_id_.Get(*this); }
 
-  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_ && identifier_field_ids_ == other.identifier_field_ids_;
 }
 
-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));

src/iceberg/schema.h

@@ -173,6 +173,7 @@ class ICEBERG_EXPORT Schema : public StructType {
   InitLowerCaseNameToIdMap(const Schema&);
   static Result<std::unordered_map<int32_t, std::vector<size_t>>> InitIdToPositionPath(
       const Schema&);
+  static Result<int32_t> InitHighestFieldId(const Schema&);
 
   const std::optional<int32_t> schema_id_;
   /// Field IDs that uniquely identify rows in the table.
@@ -185,6 +186,8 @@ class ICEBERG_EXPORT Schema : public StructType {
   Lazy<InitLowerCaseNameToIdMap> lowercase_name_to_id_;
   /// Mapping from field id to (nested) position path to access the field.
   Lazy<InitIdToPositionPath> id_to_position_path_;
+  /// Highest field ID in the schema.
+  Lazy<InitHighestFieldId> highest_field_id_;
 };
 
 }  // namespace iceberg

src/iceberg/table_metadata.cc

@@ -429,7 +429,7 @@ class TableMetadataBuilder::Impl {
   Status SetProperties(const std::unordered_map<std::string, std::string>& updated);
   Status RemoveProperties(const std::unordered_set<std::string>& removed);
   Status SetCurrentSchema(int32_t schema_id);
-  Status RemoveSchemas(const std::vector<int32_t>& schema_ids);
+  Status RemoveSchemas(const std::unordered_set<int32_t>& schema_ids);
   Result<int32_t> AddSchema(const Schema& schema, int32_t new_last_column_id);
 
   Result<std::unique_ptr<TableMetadata>> Build();
@@ -634,7 +634,7 @@ Status TableMetadataBuilder::Impl::RemoveProperties(
 Status TableMetadataBuilder::Impl::SetCurrentSchema(int32_t schema_id) {
   if (schema_id == kLastAdded) {
     if (!last_added_schema_id_.has_value()) {
-      return InvalidArgument("Cannot set last added schema: no schema has been added");
+      return ValidationFailed("Cannot set last added schema: no schema has been added");
     }
     return SetCurrentSchema(last_added_schema_id_.value());
   }
@@ -686,57 +686,55 @@ Status TableMetadataBuilder::Impl::SetCurrentSchema(int32_t schema_id) {
   return {};
 }
 
-Status TableMetadataBuilder::Impl::RemoveSchemas(const std::vector<int32_t>& schema_ids) {
-  std::unordered_set<int32_t> schema_ids_to_remove(schema_ids.begin(), schema_ids.end());
+Status TableMetadataBuilder::Impl::RemoveSchemas(
+    const std::unordered_set<int32_t>& schema_ids) {
   auto current_schema_id = metadata_.current_schema_id.value_or(Schema::kInitialSchemaId);
-  if (!schema_ids_to_remove.contains(current_schema_id)) {
-    return InvalidArgument("Cannot remove current schema: {}", current_schema_id);
-  }
+  ICEBERG_PRECHECK(!schema_ids.contains(current_schema_id),
+                   "Cannot remove current schema: {}", current_schema_id);
 
-  if (!schema_ids_to_remove.empty()) {
+  if (!schema_ids.empty()) {
     metadata_.schemas = metadata_.schemas | std::views::filter([&](const auto& schema) {
-                          return !schema_ids_to_remove.contains(
+                          return !schema_ids.contains(
                               schema->schema_id().value_or(Schema::kInitialSchemaId));
                         }) |
                         std::ranges::to<std::vector<std::shared_ptr<Schema>>>();
-    changes_.push_back(
-        std::make_unique<table::RemoveSchemas>(std::move(schema_ids_to_remove)));
+    changes_.push_back(std::make_unique<table::RemoveSchemas>(schema_ids));
   }
 
   return {};
 }
 
 Result<int32_t> TableMetadataBuilder::Impl::AddSchema(const Schema& schema,
                                                       int32_t new_last_column_id) {
-  if (new_last_column_id < metadata_.last_column_id) {
-    return InvalidArgument("Invalid last column ID: {} < {} (previous last column ID)",
-                           new_last_column_id, metadata_.last_column_id);
-  }
+  ICEBERG_PRECHECK(new_last_column_id >= metadata_.last_column_id,
+                   "Invalid last column ID: {} < {} (previous last column ID)",
+                   new_last_column_id, metadata_.last_column_id);
 
   ICEBERG_RETURN_UNEXPECTED(schema.Validate(metadata_.format_version));
 
   auto new_schema_id = ReuseOrCreateNewSchemaId(schema);
-  if (schemas_by_id_.find(new_schema_id) != schemas_by_id_.end()) {
+  if (schemas_by_id_.find(new_schema_id) != schemas_by_id_.end() &&
+      new_last_column_id == metadata_.last_column_id) {
     // 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 =
         last_added_schema_id_.has_value() &&
-        std::ranges::find_if(changes_, [new_schema_id](const auto& change) {
+        std::ranges::any_of(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;
-        }) != changes_.cend();
+        });
     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);
+  auto new_schema =
+      std::make_shared<Schema>(schema.fields() | std::ranges::to<std::vector>(),
+                               new_schema_id, schema.IdentifierFieldIds());
 
   metadata_.schemas.push_back(new_schema);
   schemas_by_id_.emplace(new_schema_id, new_schema);
@@ -934,8 +932,7 @@ TableMetadataBuilder& TableMetadataBuilder::AddSchema(
     std::shared_ptr<Schema> const& schema) {
   ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto highest_field_id, schema->HighestFieldId());
   auto new_last_column_id = std::max(impl_->metadata().last_column_id, highest_field_id);
-  ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto schema_id,
-                                   impl_->AddSchema(*schema, new_last_column_id));
+  ICEBERG_BUILDER_RETURN_IF_ERROR(impl_->AddSchema(*schema, new_last_column_id));
   return *this;
 }
 
@@ -959,7 +956,7 @@ TableMetadataBuilder& TableMetadataBuilder::RemovePartitionSpecs(
 }
 
 TableMetadataBuilder& TableMetadataBuilder::RemoveSchemas(
-    const std::vector<int32_t>& schema_ids) {
+    const std::unordered_set<int32_t>& schema_ids) {
   ICEBERG_BUILDER_RETURN_IF_ERROR(impl_->RemoveSchemas(schema_ids));
   return *this;
 }

src/iceberg/table_metadata.h

@@ -292,7 +292,7 @@ class ICEBERG_EXPORT TableMetadataBuilder : public ErrorCollector {
   ///
   /// \param schema_ids The IDs of schemas to remove
   /// \return Reference to this builder for method chaining
-  TableMetadataBuilder& RemoveSchemas(const std::vector<int32_t>& schema_ids);
+  TableMetadataBuilder& RemoveSchemas(const std::unordered_set<int32_t>& schema_ids);
 
   /// \brief Set the default sort order for the table
   ///

src/iceberg/table_update.cc

@@ -52,7 +52,7 @@ void UpgradeFormatVersion::GenerateRequirements(TableUpdateContext& context) con
 // AddSchema
 
 void AddSchema::ApplyTo(TableMetadataBuilder& builder) const {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+  builder.AddSchema(schema_);
 }
 
 void AddSchema::GenerateRequirements(TableUpdateContext& context) const {
@@ -62,7 +62,7 @@ void AddSchema::GenerateRequirements(TableUpdateContext& context) const {
 // SetCurrentSchema
 
 void SetCurrentSchema::ApplyTo(TableMetadataBuilder& builder) const {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+  builder.SetCurrentSchema(schema_id_);
 }
 
 void SetCurrentSchema::GenerateRequirements(TableUpdateContext& context) const {
@@ -103,7 +103,7 @@ void RemovePartitionSpecs::GenerateRequirements(TableUpdateContext& context) con
 // RemoveSchemas
 
 void RemoveSchemas::ApplyTo(TableMetadataBuilder& builder) const {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+  builder.RemoveSchemas(schema_ids_);
 }
 
 void RemoveSchemas::GenerateRequirements(TableUpdateContext& context) const {