feat:implement AddSortOrder for table metadata

Author: HeartLinked

src/iceberg/table_metadata.cc

@@ -21,8 +21,11 @@
 
 #include <algorithm>
 #include <chrono>
+#include <cstdint>
 #include <format>
+#include <optional>
 #include <string>
+#include <unordered_map>
 
 #include <nlohmann/json.hpp>
 
@@ -216,6 +219,9 @@ struct TableMetadataBuilder::Impl {
 
   // Change tracking
   std::vector<std::unique_ptr<TableUpdate>> changes;
+  std::optional<int32_t> last_added_schema_id;
+  std::optional<int32_t> last_added_order_id;
+  std::optional<int32_t> last_added_spec_id;
 
   // Error collection (since methods return *this and cannot throw)
   std::vector<Error> errors;
@@ -224,6 +230,11 @@ struct TableMetadataBuilder::Impl {
   std::optional<std::string> metadata_location;
   std::optional<std::string> previous_metadata_location;
 
+  // indexes for convenience
+  std::unordered_map<int32_t, std::shared_ptr<Schema>> schemas_by_id;
+  std::unordered_map<int32_t, std::shared_ptr<PartitionSpec>> specs_by_id;
+  std::unordered_map<int32_t, std::shared_ptr<SortOrder>> sort_orders_by_id;
+
   // Constructor for new table
   explicit Impl(int8_t format_version) : base(nullptr), metadata{} {
     metadata.format_version = format_version;
@@ -239,7 +250,38 @@ struct TableMetadataBuilder::Impl {
 
   // Constructor from existing metadata
   explicit Impl(const TableMetadata* base_metadata)
-      : base(base_metadata), metadata(*base_metadata) {}
+      : base(base_metadata), metadata(*base_metadata) {
+    // Initialize index maps from base metadata
+    for (const auto& schema : metadata.schemas) {
+      if (schema->schema_id().has_value()) {
+        schemas_by_id.emplace(schema->schema_id().value(), schema);
+      }
+    }
+
+    for (const auto& spec : metadata.partition_specs) {
+      specs_by_id.emplace(spec->spec_id(), spec);
+    }
+
+    for (const auto& order : metadata.sort_orders) {
+      sort_orders_by_id.emplace(order->order_id(), order);
+    }
+  }
+
+  int32_t reuseOrCreateNewSortOrderId(const SortOrder& new_order) {
+    if (new_order.is_unsorted()) {
+      return SortOrder::kUnsortedOrderId;
+    }
+    // determine the next order id
+    int32_t new_order_id = SortOrder::kInitialSortOrderId;
+    for (const auto& order : metadata.sort_orders) {
+      if (order->SameOrder(new_order)) {
+        return order->order_id();
+      } else if (new_order_id <= order->order_id()) {
+        new_order_id = order->order_id() + 1;
+      }
+    }
+    return new_order_id;
+  }
 };
 
 TableMetadataBuilder::TableMetadataBuilder(int8_t format_version)
@@ -361,7 +403,68 @@ TableMetadataBuilder& TableMetadataBuilder::SetDefaultSortOrder(int32_t order_id
 
 TableMetadataBuilder& TableMetadataBuilder::AddSortOrder(
     std::shared_ptr<SortOrder> order) {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+  int32_t new_order_id = impl_->reuseOrCreateNewSortOrderId(*order);
+
+  if (impl_->sort_orders_by_id.find(new_order_id) != impl_->sort_orders_by_id.end()) {
+    // update last_added_order_id if the order was added in this set of changes (since it
+    // is now the last)
+    bool is_new_order = false;
+    for (const auto& change : impl_->changes) {
+      auto* add_sort_order = dynamic_cast<table::AddSortOrder*>(change.get());
+      if (add_sort_order && add_sort_order->sort_order()->order_id() == new_order_id) {
+        is_new_order = true;
+        break;
+      }
+    }
+    impl_->last_added_order_id =
+        is_new_order ? std::make_optional(new_order_id) : std::nullopt;
+    return *this;
+  }
+
+  // Get current schema for validation
+  auto schema_result = impl_->metadata.Schema();
+  if (!schema_result) {
+    impl_->errors.emplace_back(
+        ErrorKind::kInvalidArgument,
+        std::format("Cannot find current schema: {}", schema_result.error().message));
+    return *this;
+  }
+
+  auto schema = schema_result.value();
+
+  // Validate the sort order against the schema
+  auto validate_status = order->Validate(*schema);
+  if (!validate_status) {
+    impl_->errors.emplace_back(
+        ErrorKind::kInvalidArgument,
+        std::format("Sort order validation failed: {}", validate_status.error().message));
+    return *this;
+  }
+
+  std::shared_ptr<SortOrder> new_order;
+  if (order->is_unsorted()) {
+    new_order = SortOrder::Unsorted();
+  } else {
+    // TODO(Li Feiyang): rebuild the sort order using new column ids (freshSortOrder)
+    // For now, create a new sort order with the provided fields
+    auto sort_order_result = SortOrder::Make(
+        *schema, new_order_id,
+        std::vector<SortField>(order->fields().begin(), order->fields().end()));
+    if (!sort_order_result) {
+      impl_->errors.emplace_back(ErrorKind::kInvalidArgument,
+                                 std::format("Failed to create sort order: {}",
+                                             sort_order_result.error().message));
+      return *this;
+    }
+    new_order = std::move(sort_order_result.value());
+  }
+
+  impl_->metadata.sort_orders.push_back(new_order);
+  impl_->sort_orders_by_id.emplace(new_order_id, new_order);
+
+  impl_->changes.push_back(std::make_unique<table::AddSortOrder>(new_order));
+  impl_->last_added_order_id = new_order_id;
+  return *this;
 }
 
 TableMetadataBuilder& TableMetadataBuilder::AddSnapshot(

src/iceberg/table_update.cc

@@ -124,11 +124,12 @@ Status RemoveSchemas::GenerateRequirements(TableUpdateContext& context) const {
 // AddSortOrder
 
 void AddSortOrder::ApplyTo(TableMetadataBuilder& builder) const {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+  builder.AddSortOrder(sort_order_);
 }
 
 Status AddSortOrder::GenerateRequirements(TableUpdateContext& context) const {
-  return NotImplemented("AddTableSortOrder::GenerateRequirements not implemented");
+  // AddSortOrder doesn't generate any requirements
+  return {};
 }
 
 // SetDefaultSortOrder

src/iceberg/test/table_metadata_builder_test.cc

@@ -19,20 +19,33 @@
 
 #include <memory>
 #include <string>
+#include <vector>
 
 #include <gtest/gtest.h>
 
 #include "iceberg/partition_spec.h"
+#include "iceberg/schema.h"
 #include "iceberg/snapshot.h"
+#include "iceberg/sort_field.h"
 #include "iceberg/sort_order.h"
 #include "iceberg/table_metadata.h"
 #include "iceberg/table_update.h"
 #include "iceberg/test/matchers.h"
+#include "iceberg/transform.h"
+#include "iceberg/type.h"
 
 namespace iceberg {
 
 namespace {
 
+// Helper function to create a simple schema for testing
+std::shared_ptr<Schema> CreateTestSchema() {
+  auto field1 = SchemaField::MakeRequired(1, "id", int32());
+  auto field2 = SchemaField::MakeRequired(2, "data", string());
+  auto field3 = SchemaField::MakeRequired(3, "ts", timestamp());
+  return std::make_shared<Schema>(std::vector<SchemaField>{field1, field2, field3}, 0);
+}
+
 // Helper function to create base metadata for tests
 std::unique_ptr<TableMetadata> CreateBaseMetadata() {
   auto metadata = std::make_unique<TableMetadata>();
@@ -41,11 +54,14 @@ std::unique_ptr<TableMetadata> CreateBaseMetadata() {
   metadata->location = "s3://bucket/test";
   metadata->last_sequence_number = 0;
   metadata->last_updated_ms = TimePointMs{std::chrono::milliseconds(1000)};
-  metadata->last_column_id = 0;
+  metadata->last_column_id = 3;
+  metadata->current_schema_id = 0;
+  metadata->schemas.push_back(CreateTestSchema());
   metadata->default_spec_id = PartitionSpec::kInitialSpecId;
   metadata->last_partition_id = 0;
   metadata->current_snapshot_id = Snapshot::kInvalidSnapshotId;
   metadata->default_sort_order_id = SortOrder::kInitialSortOrderId;
+  metadata->sort_orders.push_back(SortOrder::Unsorted());
   metadata->next_row_id = TableMetadata::kInitialRowId;
   return metadata;
 }
@@ -124,17 +140,121 @@ TEST(TableMetadataBuilderTest, AssignUUID) {
   EXPECT_EQ(metadata->table_uuid, "TEST-UUID-ABCD");  // Original case preserved
 }
 
+// Test AddSortOrder
+TEST(TableMetadataBuilderTest, AddSortOrderBasic) {
+  auto base = CreateBaseMetadata();
+  auto builder = TableMetadataBuilder::BuildFrom(base.get());
+  auto schema = CreateTestSchema();
+
+  // 1. Add unsorted - should reuse existing unsorted order
+  builder->AddSortOrder(SortOrder::Unsorted());
+  ICEBERG_UNWRAP_OR_FAIL(auto metadata, builder->Build());
+  ASSERT_EQ(metadata->sort_orders.size(), 1);
+  EXPECT_TRUE(metadata->sort_orders[0]->is_unsorted());
+
+  // 2. Add basic sort order
+  builder = TableMetadataBuilder::BuildFrom(base.get());
+  SortField field1(1, Transform::Identity(), SortDirection::kAscending,
+                   NullOrder::kFirst);
+  ICEBERG_UNWRAP_OR_FAIL(auto order1,
+                         SortOrder::Make(*schema, 1, std::vector<SortField>{field1}));
+  builder->AddSortOrder(std::move(order1));
+  ICEBERG_UNWRAP_OR_FAIL(metadata, builder->Build());
+  ASSERT_EQ(metadata->sort_orders.size(), 2);
+  EXPECT_EQ(metadata->sort_orders[1]->order_id(), 1);
+
+  // 3. Add duplicate - should be idempotent
+  builder = TableMetadataBuilder::BuildFrom(base.get());
+  ICEBERG_UNWRAP_OR_FAIL(auto order2,
+                         SortOrder::Make(*schema, 1, std::vector<SortField>{field1}));
+  ICEBERG_UNWRAP_OR_FAIL(auto order3,
+                         SortOrder::Make(*schema, 1, std::vector<SortField>{field1}));
+  builder->AddSortOrder(std::move(order2));
+  builder->AddSortOrder(std::move(order3));  // Duplicate
+  ICEBERG_UNWRAP_OR_FAIL(metadata, builder->Build());
+  ASSERT_EQ(metadata->sort_orders.size(), 2);  // Only one added
+
+  // 4. Add multiple different orders + verify ID reassignment
+  builder = TableMetadataBuilder::BuildFrom(base.get());
+  SortField field2(2, Transform::Identity(), SortDirection::kDescending,
+                   NullOrder::kLast);
+  // User provides ID=99, Builder should reassign to ID=1
+  ICEBERG_UNWRAP_OR_FAIL(auto order4,
+                         SortOrder::Make(*schema, 99, std::vector<SortField>{field1}));
+  ICEBERG_UNWRAP_OR_FAIL(
+      auto order5, SortOrder::Make(*schema, 2, std::vector<SortField>{field1, field2}));
+  builder->AddSortOrder(std::move(order4));
+  builder->AddSortOrder(std::move(order5));
+  ICEBERG_UNWRAP_OR_FAIL(metadata, builder->Build());
+  ASSERT_EQ(metadata->sort_orders.size(), 3);
+  EXPECT_EQ(metadata->sort_orders[1]->order_id(), 1);  // Reassigned from 99
+  EXPECT_EQ(metadata->sort_orders[2]->order_id(), 2);
+}
+
+TEST(TableMetadataBuilderTest, AddSortOrderInvalid) {
+  auto base = CreateBaseMetadata();
+  auto schema = CreateTestSchema();
+
+  // 1. Invalid field ID
+  auto builder = TableMetadataBuilder::BuildFrom(base.get());
+  SortField invalid_field(999, Transform::Identity(), SortDirection::kAscending,
+                          NullOrder::kFirst);
+  ICEBERG_UNWRAP_OR_FAIL(auto order1,
+                         SortOrder::Make(1, std::vector<SortField>{invalid_field}));
+  builder->AddSortOrder(std::move(order1));
+  ASSERT_THAT(builder->Build(), IsError(ErrorKind::kCommitFailed));
+  ASSERT_THAT(builder->Build(),
+              HasErrorMessage("Cannot find source column for sort field"));
+
+  // 2. Invalid transform (Day transform on string type)
+  builder = TableMetadataBuilder::BuildFrom(base.get());
+  SortField invalid_transform(2, Transform::Day(), SortDirection::kAscending,
+                              NullOrder::kFirst);
+  ICEBERG_UNWRAP_OR_FAIL(auto order2,
+                         SortOrder::Make(1, std::vector<SortField>{invalid_transform}));
+  builder->AddSortOrder(std::move(order2));
+  ASSERT_THAT(builder->Build(), IsError(ErrorKind::kCommitFailed));
+  ASSERT_THAT(builder->Build(), HasErrorMessage("Invalid source type"));
+
+  // 3. Without schema
+  builder = TableMetadataBuilder::BuildFromEmpty(2);
+  builder->AssignUUID("test-uuid");
+  SortField field(1, Transform::Identity(), SortDirection::kAscending, NullOrder::kFirst);
+  ICEBERG_UNWRAP_OR_FAIL(auto order3,
+                         SortOrder::Make(*schema, 1, std::vector<SortField>{field}));
+  builder->AddSortOrder(std::move(order3));
+  ASSERT_THAT(builder->Build(), IsError(ErrorKind::kCommitFailed));
+  ASSERT_THAT(builder->Build(), HasErrorMessage("Cannot find current schema"));
+}
+
 // Test applying TableUpdate to builder
 TEST(TableMetadataBuilderTest, ApplyUpdate) {
+  // Use base metadata with schema for all update tests
+  auto base = CreateBaseMetadata();
+  auto builder = TableMetadataBuilder::BuildFrom(base.get());
+
   // Apply AssignUUID update
-  auto builder = TableMetadataBuilder::BuildFromEmpty(2);
-  table::AssignUUID update("apply-uuid");
-  update.ApplyTo(*builder);
-  // TODO(Li Feiyang): Add more update and `apply` once other build methods are
-  // implemented
+  table::AssignUUID uuid_update("apply-uuid");
+  uuid_update.ApplyTo(*builder);
+
+  // Apply AddSortOrder update
+  auto schema = CreateTestSchema();
+  SortField sort_field(1, Transform::Identity(), SortDirection::kAscending,
+                       NullOrder::kFirst);
+  ICEBERG_UNWRAP_OR_FAIL(auto sort_order_unique,
+                         SortOrder::Make(*schema, 1, std::vector<SortField>{sort_field}));
+  auto sort_order = std::shared_ptr<SortOrder>(std::move(sort_order_unique));
+
+  table::AddSortOrder sort_order_update(sort_order);
+  sort_order_update.ApplyTo(*builder);
+
+  // TODO(Li Feiyang): Apply more build methods once they are implemented
 
+  // Verify all updates were applied
   ICEBERG_UNWRAP_OR_FAIL(auto metadata, builder->Build());
   EXPECT_EQ(metadata->table_uuid, "apply-uuid");
+  ASSERT_EQ(metadata->sort_orders.size(), 2);
+  EXPECT_EQ(metadata->sort_orders[1]->order_id(), 1);
 }
 
 }  // namespace iceberg