fix: review comments

Author: zhjwpku

src/iceberg/update/update_partition_spec.cc

@@ -20,6 +20,7 @@
 #include "iceberg/update/update_partition_spec.h"
 
 #include <format>
+#include <memory>
 
 #include "iceberg/expression/term.h"
 #include "iceberg/partition_field.h"
@@ -45,24 +46,19 @@ Result<std::shared_ptr<UpdatePartitionSpec>> UpdatePartitionSpec::Make(
 
 UpdatePartitionSpec::UpdatePartitionSpec(std::shared_ptr<Transaction> transaction)
     : PendingUpdate(std::move(transaction)) {
-  const TableMetadata* base_metadata = transaction_->base();
-  if (base_metadata == nullptr) [[unlikely]] {
-    AddError(ErrorKind::kInvalidArgument,
-             "Base table metadata is required to construct UpdatePartitionSpec");
-    return;
-  }
-  format_version_ = base_metadata->format_version;
+  const TableMetadata& base_metadata = transaction_->current();
+  format_version_ = base_metadata.format_version;
 
   // Get the current/default partition spec
-  auto spec_result = base_metadata->PartitionSpec();
+  auto spec_result = base_metadata.PartitionSpec();
   if (!spec_result.has_value()) {
     AddError(spec_result.error());
     return;
   }
   spec_ = std::move(spec_result.value());
 
   // Get the current schema
-  auto schema_result = base_metadata->Schema();
+  auto schema_result = base_metadata.Schema();
   if (!schema_result.has_value()) {
     AddError(schema_result.error());
     return;
@@ -112,39 +108,43 @@ UpdatePartitionSpec& UpdatePartitionSpec::AddField(const std::string& source_nam
   return AddFieldInternal(nullptr, source_id, Transform::Identity());
 }
 
-UpdatePartitionSpec& UpdatePartitionSpec::AddField(std::shared_ptr<NamedReference> term,
-                                                   std::string part_name) {
-  // Bind the term to get the source field
-  ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto bound_ref, term->Bind(*schema_, case_sensitive_));
-  int32_t source_id = bound_ref->field().field_id();
-
-  // Reference terms use identity transform
-  return AddFieldInternal(part_name.empty() ? nullptr : &part_name, source_id,
-                          Transform::Identity());
-}
+UpdatePartitionSpec& UpdatePartitionSpec::AddField(const std::shared_ptr<Term>& term,
+                                                   const std::string& part_name) {
+  ICEBERG_BUILDER_CHECK(term->is_unbound(), "Cannot add bound term to partition spec");
+  // Bind the term to get the source id
+  if (term->kind() == Term::Kind::kReference) {
+    const auto& ref = std::dynamic_pointer_cast<NamedReference>(term);
+    ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto bound_ref,
+                                     ref->Bind(*schema_, case_sensitive_));
+    int32_t source_id = bound_ref->field().field_id();
+    return AddFieldInternal(part_name.empty() ? nullptr : &part_name, source_id,
+                            Transform::Identity());
+  } else if (term->kind() == Term::Kind::kTransform) {
+    const auto& unbound_transform = std::dynamic_pointer_cast<UnboundTransform>(term);
+    ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto bound_transform,
+                                     unbound_transform->Bind(*schema_, case_sensitive_));
+    int32_t source_id = bound_transform->reference()->field().field_id();
+    return AddFieldInternal(part_name.empty() ? nullptr : &part_name, source_id,
+                            bound_transform->transform());
+  }
 
-UpdatePartitionSpec& UpdatePartitionSpec::AddField(std::shared_ptr<UnboundTransform> term,
-                                                   std::string part_name) {
-  // Bind the term to get the source field and transform
-  ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto bound_transform,
-                                   term->Bind(*schema_, case_sensitive_));
-  int32_t source_id = bound_transform->reference()->field().field_id();
-  return AddFieldInternal(part_name.empty() ? nullptr : &part_name, source_id,
-                          bound_transform->transform());
+  ICEBERG_BUILDER_CHECK(false, "Cannot add {} term to partition spec", term->ToString());
+  std::unreachable();
 }
 
 UpdatePartitionSpec& UpdatePartitionSpec::AddFieldInternal(
-    const std::string* name, int32_t source_id, std::shared_ptr<Transform> transform) {
-  // Cache transform string to avoid repeated ToString() calls
-  const std::string transform_str = transform->ToString();
-  TransformKey validation_key{source_id, transform_str};
-
+    const std::string* name, int32_t source_id,
+    const std::shared_ptr<Transform>& transform) {
   // Check for duplicate name in added fields
   if (name != nullptr) {
     ICEBERG_BUILDER_CHECK(!added_field_names_.contains(*name),
                           "Cannot add duplicate partition field: {}", *name);
   }
 
+  // Cache transform string to avoid repeated ToString() calls
+  const std::string transform_str = transform->ToString();
+  TransformKey validation_key{source_id, transform_str};
+
   // Check if this field already exists in the current spec
   auto existing_it = transform_to_field_.find(validation_key);
   if (existing_it != transform_to_field_.end()) {
@@ -157,18 +157,18 @@ UpdatePartitionSpec& UpdatePartitionSpec::AddFieldInternal(
 
     ICEBERG_BUILDER_CHECK(
         is_deleted,
-        "Cannot add duplicate partition field {} for source {} with transform {}, "
+        "Cannot add duplicate partition field '{}' for source {} with transform {}, "
         "conflicts with {}",
-        name ? *name : "unknown", source_id, transform_str, existing->ToString());
+        name ? *name : "", source_id, transform_str, existing->ToString());
   }
 
   // Check if already being added
   auto added_it = transform_to_added_field_.find(validation_key);
   ICEBERG_BUILDER_CHECK(
       added_it == transform_to_added_field_.end(),
-      "Cannot add duplicate partition field {} for source {} with transform {}, "
+      "Cannot add duplicate partition field '{}' for source {} with transform {}, "
       "already added: {}",
-      name ? *name : "unknown", source_id, transform_str, added_it->second);
+      name ? *name : "", source_id, transform_str, added_it->second);
 
   // Create or recycle the partition field
   PartitionField new_field = RecycleOrCreatePartitionField(source_id, transform, name);
@@ -244,27 +244,32 @@ UpdatePartitionSpec& UpdatePartitionSpec::RemoveField(const std::string& name) {
   return *this;
 }
 
-UpdatePartitionSpec& UpdatePartitionSpec::RemoveField(
-    std::shared_ptr<NamedReference> term) {
-  // Bind the term to get the source field
-  ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto bound_ref, term->Bind(*schema_, case_sensitive_));
-  int32_t source_id = bound_ref->field().field_id();
-
-  // Reference terms use identity transform
-  TransformKey key{source_id, Transform::Identity()->ToString()};
-  return RemoveFieldByTransform(key, term->ToString());
-}
-
-UpdatePartitionSpec& UpdatePartitionSpec::RemoveField(
-    std::shared_ptr<UnboundTransform> term) {
-  // Bind the term to get the source field and transform
-  ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto bound_transform,
-                                   term->Bind(*schema_, case_sensitive_));
-  int32_t source_id = bound_transform->reference()->field().field_id();
-  auto transform = bound_transform->transform();
+UpdatePartitionSpec& UpdatePartitionSpec::RemoveField(const std::shared_ptr<Term>& term) {
+  ICEBERG_BUILDER_CHECK(term->is_unbound(),
+                        "Cannot remove bound term from partition spec");
+  // Bind the term to get the source id
+  if (term->kind() == Term::Kind::kReference) {
+    const auto& ref = std::dynamic_pointer_cast<NamedReference>(term);
+    ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto bound_ref,
+                                     ref->Bind(*schema_, case_sensitive_));
+    int32_t source_id = bound_ref->field().field_id();
+    // Reference terms use identity transform
+    TransformKey key{source_id, Transform::Identity()->ToString()};
+    return RemoveFieldByTransform(key, term->ToString());
+  } else if (term->kind() == Term::Kind::kTransform) {
+    const auto& unbound_transform = std::dynamic_pointer_cast<UnboundTransform>(term);
+    ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto bound_transform,
+                                     unbound_transform->Bind(*schema_, case_sensitive_));
+    int32_t source_id = bound_transform->reference()->field().field_id();
+    auto transform = bound_transform->transform();
+
+    TransformKey key{source_id, transform->ToString()};
+    return RemoveFieldByTransform(key, term->ToString());
+  }
 
-  TransformKey key{source_id, transform->ToString()};
-  return RemoveFieldByTransform(key, term->ToString());
+  ICEBERG_BUILDER_CHECK(false, "Cannot remove {} term from partition spec",
+                        term->ToString());
+  std::unreachable();
 }
 
 UpdatePartitionSpec& UpdatePartitionSpec::RemoveFieldByTransform(
@@ -461,25 +466,20 @@ UpdatePartitionSpec::IndexSpecByTransform(const PartitionSpec& spec) {
 }
 
 void UpdatePartitionSpec::BuildHistoricalFieldsIndex() {
-  const TableMetadata* base_metadata = transaction_->base();
-  if (base_metadata == nullptr) {
-    return;
-  }
+  const TableMetadata& base_metadata = transaction_->current();
 
   // Count total fields across all specs to reserve capacity
   size_t total_fields = 0;
-  for (const auto& partition_spec : base_metadata->partition_specs) {
+  for (const auto& partition_spec : base_metadata.partition_specs) {
     total_fields += partition_spec->fields().size();
   }
   historical_fields_.reserve(total_fields);
 
   // Index all fields from all historical partition specs
   // Later specs override earlier ones for the same (source_id, transform) key
-  for (const auto& partition_spec : base_metadata->partition_specs) {
+  for (const auto& partition_spec : base_metadata.partition_specs) {
     for (const auto& field : partition_spec->fields()) {
       TransformKey key{field.source_id(), field.transform()->ToString()};
-      // Use emplace to only insert if key doesn't exist, preserving first occurrence
-      // This ensures we get the earliest field ID for recycling
       historical_fields_.emplace(key, field);
     }
   }

src/iceberg/update/update_partition_spec.h

@@ -28,9 +28,9 @@
 #include <unordered_set>
 #include <vector>
 
+#include "iceberg/expression/term.h"
 #include "iceberg/iceberg_export.h"
 #include "iceberg/result.h"
-#include "iceberg/table_identifier.h"
 #include "iceberg/type_fwd.h"
 #include "iceberg/update/pending_update.h"
 
@@ -61,43 +61,27 @@ class ICEBERG_EXPORT UpdatePartitionSpec : public PendingUpdate {
 
   /// \brief Add a new partition field with a custom name.
   ///
-  /// \param term The named reference representing the source column.
+  /// \param term The term representing the source column, should be unbound.
   /// \param part_name Name for the partition field.
   /// \return Reference to this for method chaining.
-  UpdatePartitionSpec& AddField(std::shared_ptr<NamedReference> term,
-                                std::string part_name = "");
-
-  /// \brief Add a new partition field with a custom name from an unbound transform.
-  ///
-  /// \param term The unbound transform representing the partition transform.
-  /// \param part_name Name for the partition field.
-  /// \return Reference to this for method chaining.
-  UpdatePartitionSpec& AddField(std::shared_ptr<UnboundTransform> term,
-                                std::string part_name = "");
+  UpdatePartitionSpec& AddField(const std::shared_ptr<Term>& term,
+                                const std::string& part_name = "");
 
   /// \brief Remove a partition field by name.
   ///
   /// \param name Name of the partition field to remove.
   /// \return Reference to this for method chaining.
   UpdatePartitionSpec& RemoveField(const std::string& name);
 
-  /// \brief Remove a partition field by its transform term.
+  /// \brief Remove a partition field by its source term.
   ///
   /// The partition field with the same transform and source reference will be removed.
   /// If the term is a reference and does not have a transform, the identity transform
   /// is used.
   ///
-  /// \param term The named reference representing the partition transform to remove.
-  /// \return Reference to this for method chaining.
-  UpdatePartitionSpec& RemoveField(std::shared_ptr<NamedReference> term);
-
-  /// \brief Remove a partition field by its transform term.
-  ///
-  /// The partition field with the same transform and source reference will be removed.
-  ///
-  /// \param term The unbound transform representing the partition transform to remove.
+  /// \param term The term representing the source column, should be unbound.
   /// \return Reference to this for method chaining.
-  UpdatePartitionSpec& RemoveField(std::shared_ptr<UnboundTransform> term);
+  UpdatePartitionSpec& RemoveField(const std::shared_ptr<Term>& term);
 
   /// \brief Rename a field in the partition spec.
   ///
@@ -133,17 +117,27 @@ class ICEBERG_EXPORT UpdatePartitionSpec : public PendingUpdate {
   /// \brief Assign a new partition field ID.
   int32_t AssignFieldId();
 
-  /// \brief Recycle or create a partition field.
   ///
   /// In V2, searches for a similar partition field in historical specs.
   /// If not found or in V1, creates a new PartitionField.
+
+  /// \brief Recycle or create a partition field.
+  ///
+  /// In V2 it searches for a similar partition field in historical partition specs. Tries
+  /// to match on source field ID, transform type and target name (optional). If not found
+  /// or in V1 cases it creates a new PartitionField.
+  ///
+  /// \param source_id The source field ID.
+  /// \param transform The transform function.
+  /// \param name The target partition field name, if specified.
+  /// \return The recycled or newly created partition field.
   PartitionField RecycleOrCreatePartitionField(int32_t source_id,
                                                std::shared_ptr<Transform> transform,
                                                const std::string* name);
 
   /// \brief Internal implementation of AddField with resolved source ID and transform.
   UpdatePartitionSpec& AddFieldInternal(const std::string* name, int32_t source_id,
-                                        std::shared_ptr<Transform> transform);
+                                        const std::shared_ptr<Transform>& transform);
 
   /// \brief Generate a partition field name from the source and transform.
   std::string GeneratePartitionName(int32_t source_id,