resolve comments

Author: gty404

src/iceberg/partition_spec.cc

@@ -131,6 +131,8 @@ bool PartitionSpec::Equals(const PartitionSpec& other) const {
 }
 
 Status PartitionSpec::Validate(const Schema& schema, bool allow_missing_fields) const {
+  ICEBERG_RETURN_UNEXPECTED(ValidatePartitionName(schema, *this));
+
   std::unordered_map<int32_t, int32_t> parents = IndexParents(schema);
   for (const auto& partition_field : fields_) {
     ICEBERG_ASSIGN_OR_RAISE(auto source_field,
@@ -177,9 +179,10 @@ Status PartitionSpec::Validate(const Schema& schema, bool allow_missing_fields)
   return {};
 }
 
-Status PartitionSpec::ValidatePartitionName(const Schema& schema) const {
+Status PartitionSpec::ValidatePartitionName(const Schema& schema,
+                                            const PartitionSpec& spec) {
   std::unordered_set<std::string> partition_names;
-  for (const auto& partition_field : fields_) {
+  for (const auto& partition_field : spec.fields()) {
     auto name = std::string(partition_field.name());
     ICEBERG_PRECHECK(!name.empty(), "Cannot use empty partition name: {}", name);
 
@@ -190,9 +193,10 @@ Status PartitionSpec::ValidatePartitionName(const Schema& schema) const {
 
     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 (transform_type == TransformType::kIdentity ||
+        transform_type == TransformType::kVoid) {
+      // for identity/nulls 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(

src/iceberg/partition_spec.h

@@ -79,14 +79,14 @@ class ICEBERG_EXPORT PartitionSpec : public util::Formattable {
 
   /// \brief Validates the partition spec against a schema.
   /// \param schema The schema to validate against.
-  /// \param allowMissingFields Whether to skip validation for partition fields whose
+  /// \param allow_missing_fields Whether to skip validation for partition fields whose
   /// source columns have been dropped from the schema.
   /// \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;
+  static Status ValidatePartitionName(const Schema& schema, const PartitionSpec& spec);
 
   /// \brief Get the partition fields by source ID.
   /// \param source_id The id of the source field.

src/iceberg/table_metadata.cc

@@ -457,20 +457,6 @@ class TableMetadataBuilder::Impl {
   /// \return The ID to use for this schema (reused if exists, new otherwise
   int32_t ReuseOrCreateNewSchemaId(const Schema& new_schema) const;
 
-  /// \brief Update partition spec to use a new schema
-  /// \param schema The new schema to bind to
-  /// \param partition_spec The partition spec to update
-  /// \return The updated partition spec bound to the new schema
-  static Result<std::shared_ptr<PartitionSpec>> UpdateSpecSchema(
-      const Schema& schema, const PartitionSpec& partition_spec);
-
-  /// \brief Update sort order to use a new schema
-  /// \param schema The new schema to bind to
-  /// \param sort_order The sort order to update
-  /// \return The updated sort order bound to the new schema
-  static Result<std::unique_ptr<SortOrder>> UpdateSortOrderSchema(
-      const Schema& schema, const SortOrder& sort_order);
-
  private:
   // Base metadata (nullptr for new tables)
   const TableMetadata* base_;
@@ -542,7 +528,7 @@ Status TableMetadataBuilder::Impl::UpgradeFormatVersion(int8_t new_format_versio
 }
 
 Status TableMetadataBuilder::Impl::SetDefaultSortOrder(int32_t order_id) {
-  if (order_id == -1) {
+  if (order_id == kLastAdded) {
     if (!last_added_order_id_.has_value()) {
       return InvalidArgument(
           "Cannot set last added sort order: no sort order has been added");
@@ -607,7 +593,7 @@ Result<int32_t> TableMetadataBuilder::Impl::AddSortOrder(const SortOrder& order)
 }
 
 Status TableMetadataBuilder::Impl::SetDefaultPartitionSpec(int32_t spec_id) {
-  if (spec_id == -1) {
+  if (spec_id == kLastAdded) {
     if (!last_added_spec_id_.has_value()) {
       return ValidationFailed(
           "Cannot set last added partition spec: no partition spec has been added");
@@ -656,8 +642,7 @@ Result<int32_t> TableMetadataBuilder::Impl::AddPartitionSpec(const PartitionSpec
 
   ICEBERG_ASSIGN_OR_RAISE(
       std::shared_ptr<PartitionSpec> new_spec,
-      PartitionSpec::Make(new_spec_id, std::vector<PartitionField>(spec.fields().begin(),
-                                                                   spec.fields().end())));
+      PartitionSpec::Make(new_spec_id, spec.fields() | std::ranges::to<std::vector>()));
   metadata_.last_partition_id =
       std::max(metadata_.last_partition_id, new_spec->last_assigned_field_id());
   metadata_.partition_specs.push_back(new_spec);
@@ -718,30 +703,41 @@ Status TableMetadataBuilder::Impl::SetCurrentSchema(int32_t schema_id) {
   }
 
   auto it = schemas_by_id_.find(schema_id);
-  if (it == schemas_by_id_.end()) {
-    return InvalidArgument("Cannot set current schema to unknown schema: {}", schema_id);
-  }
+  ICEBERG_PRECHECK(it != 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 : metadata_.partition_specs) {
-    ICEBERG_ASSIGN_OR_RAISE(auto updated_spec, UpdateSpecSchema(*schema, *spec));
+  for (const auto& partition_spec : metadata_.partition_specs) {
+    ICEBERG_ASSIGN_OR_RAISE(
+        auto updated_spec,
+        PartitionSpec::Make(partition_spec->spec_id(),
+                            partition_spec->fields() | std::ranges::to<std::vector>()));
+
+    ICEBERG_RETURN_UNEXPECTED(
+        PartitionSpec::ValidatePartitionName(*schema, *updated_spec));
+
     updated_specs.push_back(std::move(updated_spec));
   }
   metadata_.partition_specs = std::move(updated_specs);
+
   specs_by_id_.clear();
   for (const auto& spec : metadata_.partition_specs) {
     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 : metadata_.sort_orders) {
-    ICEBERG_ASSIGN_OR_RAISE(auto updated_order, UpdateSortOrderSchema(*schema, *order));
+  for (const auto& sort_order : metadata_.sort_orders) {
+    ICEBERG_ASSIGN_OR_RAISE(
+        auto updated_order,
+        SortOrder::Make(sort_order->order_id(),
+                        sort_order->fields() | std::ranges::to<std::vector>()));
     updated_orders.push_back(std::move(updated_order));
   }
   metadata_.sort_orders = std::move(updated_orders);
+
   sort_orders_by_id_.clear();
   for (const auto& order : metadata_.sort_orders) {
     sort_orders_by_id_.emplace(order->order_id(), order);
@@ -768,8 +764,8 @@ Status TableMetadataBuilder::Impl::RemoveSchemas(
 
   if (!schema_ids.empty()) {
     metadata_.schemas = metadata_.schemas | std::views::filter([&](const auto& schema) {
-                          return !schema_ids.contains(
-                              schema->schema_id().value_or(Schema::kInitialSchemaId));
+                          return schema->schema_id().has_value() &&
+                                 !schema_ids.contains(schema->schema_id().value());
                         }) |
                         std::ranges::to<std::vector<std::shared_ptr<Schema>>>();
     changes_.push_back(std::make_unique<table::RemoveSchemas>(schema_ids));
@@ -787,8 +783,8 @@ Result<int32_t> TableMetadataBuilder::Impl::AddSchema(const Schema& schema,
   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() &&
-      new_last_column_id == metadata_.last_column_id) {
+  auto schema_found = schemas_by_id_.contains(new_schema_id);
+  if (schema_found && 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 =
@@ -797,24 +793,27 @@ Result<int32_t> TableMetadataBuilder::Impl::AddSchema(const Schema& schema,
           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;
+          auto* add_schema = internal::checked_cast<table::AddSchema*>(change.get());
+          return add_schema->schema()->schema_id() == std::make_optional(new_schema_id);
         });
     last_added_schema_id_ =
         is_new_schema ? std::make_optional(new_schema_id) : std::nullopt;
     return new_schema_id;
   }
 
+  metadata_.last_column_id = new_last_column_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);
+  if (!schema_found) {
+    metadata_.schemas.push_back(new_schema);
+    schemas_by_id_.emplace(new_schema_id, new_schema);
+  }
 
   changes_.push_back(std::make_unique<table::AddSchema>(new_schema, new_last_column_id));
-  metadata_.last_column_id = new_last_column_id;
+
   last_added_schema_id_ = new_schema_id;
 
   return new_schema_id;
@@ -834,14 +833,16 @@ Result<std::unique_ptr<TableMetadata>> TableMetadataBuilder::Impl::Build() {
   auto schema_it = schemas_by_id_.find(current_schema_id);
   ICEBERG_PRECHECK(schema_it != schemas_by_id_.end(),
                    "Current schema ID {} not found in schemas", current_schema_id);
-  const auto& current_schema = schema_it->second;
   {
+    const auto& current_schema = schema_it->second;
+
     auto spec_it = specs_by_id_.find(metadata_.default_spec_id);
-    // FIXME(GuoTao.yu): Default spec must exist after we support update partition spec
-    if (spec_it != specs_by_id_.end()) {
-      ICEBERG_RETURN_UNEXPECTED(
-          spec_it->second->Validate(*current_schema, /*allow_missing_fields=*/false));
-    }
+    ICEBERG_PRECHECK(spec_it != specs_by_id_.end(),
+                     "Default spec ID {} not found in partition specs",
+                     metadata_.default_spec_id);
+    ICEBERG_RETURN_UNEXPECTED(
+        spec_it->second->Validate(*current_schema, /*allow_missing_fields=*/false));
+
     auto sort_order_it = sort_orders_by_id_.find(metadata_.default_sort_order_id);
     ICEBERG_PRECHECK(sort_order_it != sort_orders_by_id_.end(),
                      "Default sort order ID {} not found in sort orders",
@@ -913,42 +914,6 @@ int32_t TableMetadataBuilder::Impl::ReuseOrCreateNewSchemaId(
   return new_schema_id;
 }
 
-Result<std::shared_ptr<PartitionSpec>> TableMetadataBuilder::Impl::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::Impl::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(int8_t format_version)
     : impl_(std::make_unique<Impl>(format_version)) {}
 

src/iceberg/table_metadata.h

@@ -249,7 +249,7 @@ class ICEBERG_EXPORT TableMetadataBuilder : public ErrorCollector {
   /// \param schema The schema to set as current
   /// \param new_last_column_id The highest column ID in the schema
   /// \return Reference to this builder for method chaining
-  TableMetadataBuilder& SetCurrentSchema(std::shared_ptr<Schema> const& schema,
+  TableMetadataBuilder& SetCurrentSchema(const std::shared_ptr<Schema>& schema,
                                          int32_t new_last_column_id);
 
   /// \brief Set the current schema by schema ID
@@ -262,7 +262,7 @@ class ICEBERG_EXPORT TableMetadataBuilder : public ErrorCollector {
   ///
   /// \param schema The schema to add
   /// \return Reference to this builder for method chaining
-  TableMetadataBuilder& AddSchema(std::shared_ptr<Schema> const& schema);
+  TableMetadataBuilder& AddSchema(const std::shared_ptr<Schema>& schema);
 
   /// \brief Set the default partition spec for the table
   ///

src/iceberg/test/metadata_io_test.cc

@@ -30,6 +30,7 @@
 #include "iceberg/arrow/arrow_fs_file_io_internal.h"
 #include "iceberg/file_io.h"
 #include "iceberg/json_internal.h"
+#include "iceberg/partition_spec.h"
 #include "iceberg/result.h"
 #include "iceberg/schema.h"
 #include "iceberg/snapshot.h"
@@ -65,6 +66,7 @@ class MetadataIOTest : public TempFileTestBase {
                          .last_sequence_number = 0,
                          .schemas = {schema},
                          .current_schema_id = 1,
+                         .partition_specs = {PartitionSpec::Unpartitioned()},
                          .default_spec_id = 0,
                          .last_partition_id = 0,
                          .properties = TableProperties::FromMap({{"key", "value"}}),

src/iceberg/test/table_metadata_builder_test.cc

@@ -81,6 +81,7 @@ TEST(TableMetadataBuilderTest, BuildFromEmpty) {
   auto schema = CreateTestSchema();
   builder->SetCurrentSchema(schema, schema->HighestFieldId().value());
   builder->SetDefaultSortOrder(SortOrder::Unsorted());
+  builder->SetDefaultPartitionSpec(PartitionSpec::Unpartitioned());
   builder->AssignUUID("new-uuid-5678");
 
   ICEBERG_UNWRAP_OR_FAIL(auto metadata, builder->Build());
@@ -160,6 +161,7 @@ TEST(TableMetadataBuilderTest, AssignUUID) {
   auto schema = CreateTestSchema();
   builder->SetCurrentSchema(schema, schema->HighestFieldId().value());
   builder->SetDefaultSortOrder(SortOrder::Unsorted());
+  builder->SetDefaultPartitionSpec(PartitionSpec::Unpartitioned());
   builder->AssignUUID("new-uuid-5678");
   ICEBERG_UNWRAP_OR_FAIL(auto metadata, builder->Build());
   EXPECT_EQ(metadata->table_uuid, "new-uuid-5678");
@@ -187,6 +189,7 @@ TEST(TableMetadataBuilderTest, AssignUUID) {
   builder = TableMetadataBuilder::BuildFromEmpty(2);
   builder->SetCurrentSchema(schema, schema->HighestFieldId().value());
   builder->SetDefaultSortOrder(SortOrder::Unsorted());
+  builder->SetDefaultPartitionSpec(PartitionSpec::Unpartitioned());
   builder->AssignUUID();
   ICEBERG_UNWRAP_OR_FAIL(metadata, builder->Build());
   EXPECT_FALSE(metadata->table_uuid.empty());
@@ -238,6 +241,7 @@ TEST(TableMetadataBuilderTest, UpgradeFormatVersion) {
   auto schema = CreateTestSchema();
   builder->SetCurrentSchema(schema, schema->HighestFieldId().value());
   builder->SetDefaultSortOrder(SortOrder::Unsorted());
+  builder->SetDefaultPartitionSpec(PartitionSpec::Unpartitioned());
   builder->UpgradeFormatVersion(2);
 
   ICEBERG_UNWRAP_OR_FAIL(auto metadata, builder->Build());
@@ -438,17 +442,6 @@ TEST(TableMetadataBuilderTest, AddSchemaBasic) {
   EXPECT_EQ(metadata->last_column_id, 6);
 }
 
-TEST(TableMetadataBuilderTest, AddSchemaInvalid) {
-  auto base = CreateBaseMetadata();
-
-  auto builder = TableMetadataBuilder::BuildFrom(base.get());
-  auto field_low_id = SchemaField::MakeRequired(1, "low_id", int32());
-  auto schema_low_id =
-      std::make_shared<Schema>(std::vector<SchemaField>{field_low_id}, 1);
-  // Manually try to set a lower last_column_id via SetCurrentSchema
-  // This is tested indirectly through AddSchemaInternal validation
-}
-
 // Test SetCurrentSchema
 TEST(TableMetadataBuilderTest, SetCurrentSchemaBasic) {
   auto base = CreateBaseMetadata();