feat: assign fresh field ids for new schema

Author: WZhuo

src/iceberg/schema.cc

@@ -33,8 +33,28 @@
 
 namespace iceberg {
 
-Schema::Schema(std::vector<SchemaField> fields, std::optional<int32_t> schema_id)
-    : StructType(std::move(fields)), schema_id_(schema_id) {}
+Schema::Schema(std::vector<SchemaField> fields, std::optional<int32_t> schema_id,
+               std::vector<int32_t> identifier_field_ids)
+    : StructType(std::move(fields)),
+      schema_id_(schema_id),
+      identifier_field_ids_(std::move(identifier_field_ids)) {}
+
+Result<std::unique_ptr<Schema>> Schema::Make(
+    std::vector<SchemaField> fields, std::optional<int32_t> schema_id,
+    const std::vector<std::string>& identifier_field_names) {
+  auto schema = std::make_unique<Schema>(std::move(fields), schema_id);
+
+  std::vector<int32_t> fresh_identifier_ids;
+  for (const auto& name : identifier_field_names) {
+    ICEBERG_ASSIGN_OR_RAISE(auto field, schema->FindFieldByName(name));
+    if (!field) {
+      return InvalidSchema("Cannot find identifier field: {}", name);
+    }
+    fresh_identifier_ids.push_back(field.value().get().field_id());
+  }
+  schema->identifier_field_ids_ = std::move(fresh_identifier_ids);
+  return schema;
+}
 
 std::optional<int32_t> Schema::schema_id() const { return schema_id_; }
 
@@ -48,15 +68,16 @@ std::string Schema::ToString() const {
 }
 
 bool Schema::Equals(const Schema& other) const {
-  return schema_id_ == other.schema_id_ && fields_ == other.fields_;
+  return schema_id_ == other.schema_id_ && fields_ == other.fields_ &&
+         identifier_field_ids_ == other.identifier_field_ids_;
 }
 
 Result<std::optional<std::reference_wrapper<const SchemaField>>> Schema::FindFieldByName(
     std::string_view name, bool case_sensitive) const {
   if (case_sensitive) {
-    ICEBERG_ASSIGN_OR_RAISE(auto name_to_id, name_to_id_.Get(*this));
-    auto it = name_to_id.get().find(name);
-    if (it == name_to_id.get().end()) {
+    ICEBERG_ASSIGN_OR_RAISE(auto name_id_map, name_id_map_.Get(*this));
+    auto it = name_id_map.get().name_to_id.find(name);
+    if (it == name_id_map.get().name_to_id.end()) {
       return std::nullopt;
     };
     return FindFieldById(it->second);
@@ -77,21 +98,22 @@ Schema::InitIdToFieldMap(const Schema& self) {
   return id_to_field;
 }
 
-Result<std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>>>
-Schema::InitNameToIdMap(const Schema& self) {
-  std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>> name_to_id;
-  NameToIdVisitor visitor(name_to_id, /*case_sensitive=*/true);
+Result<Schema::NameIdMap> Schema::InitNameIdMap(const Schema& self) {
+  NameIdMap name_id_map;
+  NameToIdVisitor visitor(name_id_map.name_to_id, &name_id_map.id_to_name,
+                          /*case_sensitive=*/true);
   ICEBERG_RETURN_UNEXPECTED(
       VisitTypeInline(self, &visitor, /*path=*/"", /*short_path=*/""));
   visitor.Finish();
-  return name_to_id;
+  return name_id_map;
 }
 
 Result<std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>>>
 Schema::InitLowerCaseNameToIdMap(const Schema& self) {
   std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>>
       lowercase_name_to_id;
-  NameToIdVisitor visitor(lowercase_name_to_id, /*case_sensitive=*/false);
+  NameToIdVisitor visitor(lowercase_name_to_id, /*id_to_name=*/nullptr,
+                          /*case_sensitive=*/false);
   ICEBERG_RETURN_UNEXPECTED(
       VisitTypeInline(self, &visitor, /*path=*/"", /*short_path=*/""));
   visitor.Finish();
@@ -108,6 +130,16 @@ Result<std::optional<std::reference_wrapper<const SchemaField>>> Schema::FindFie
   return it->second;
 }
 
+Result<std::optional<std::string_view>> Schema::FindColumnNameById(
+    int32_t field_id) const {
+  ICEBERG_ASSIGN_OR_RAISE(auto name_id_map, name_id_map_.Get(*this));
+  auto it = name_id_map.get().id_to_name.find(field_id);
+  if (it == name_id_map.get().id_to_name.end()) {
+    return std::nullopt;
+  }
+  return it->second;
+}
+
 Result<std::unordered_map<int32_t, std::vector<size_t>>> Schema::InitIdToPositionPath(
     const Schema& self) {
   PositionPathVisitor visitor;
@@ -179,4 +211,21 @@ Result<std::unique_ptr<Schema>> Schema::Project(
                         std::nullopt);
 }
 
+const std::vector<int32_t>& Schema::IdentifierFieldIds() const {
+  return identifier_field_ids_;
+}
+
+Result<std::vector<std::string>> Schema::IdentifierFieldNames() const {
+  std::vector<std::string> names;
+  names.reserve(identifier_field_ids_.size());
+  for (auto id : identifier_field_ids_) {
+    ICEBERG_ASSIGN_OR_RAISE(auto name, FindColumnNameById(id));
+    if (!name.has_value()) {
+      return InvalidSchema("Cannot find the field of the specified field id: {}", id);
+    }
+    names.emplace_back(name.value());
+  }
+  return names;
+}
+
 }  // namespace iceberg

src/iceberg/schema.h

@@ -49,7 +49,18 @@ class ICEBERG_EXPORT Schema : public StructType {
   static constexpr int32_t kInvalidColumnId = -1;
 
   explicit Schema(std::vector<SchemaField> fields,
-                  std::optional<int32_t> schema_id = std::nullopt);
+                  std::optional<int32_t> schema_id = std::nullopt,
+                  std::vector<int32_t> identifier_field_ids = {});
+
+  /// \brief Create a schema.
+  ///
+  /// \param fields The fields that make up the schema.
+  /// \param schema_id The unique identifier for this schema (default: kInitialSchemaId).
+  /// \param identifier_field_names Canonical names of fields that uniquely identify rows
+  /// in the table (default: empty). \return A new Schema instance or Status if failed.
+  static Result<std::unique_ptr<Schema>> Make(
+      std::vector<SchemaField> fields, std::optional<int32_t> schema_id = std::nullopt,
+      const std::vector<std::string>& identifier_field_names = {});
 
   /// \brief Get the schema ID.
   ///
@@ -78,6 +89,13 @@ class ICEBERG_EXPORT Schema : public StructType {
   Result<std::optional<std::reference_wrapper<const SchemaField>>> FindFieldById(
       int32_t field_id) const;
 
+  /// \brief Returns the canonical field name for the given id.
+  ///
+  /// \param field_id The id of the field to get the canonical name for.
+  /// \return The canocinal column name of the field with the given id, or std::nullopt if
+  /// not found.
+  Result<std::optional<std::string_view>> FindColumnNameById(int32_t field_id) const;
+
   /// \brief Get the accessor to access the field by field id.
   ///
   /// \param field_id The id of the field to get the accessor for.
@@ -103,26 +121,48 @@ class ICEBERG_EXPORT Schema : public StructType {
   Result<std::unique_ptr<Schema>> Project(
       const std::unordered_set<int32_t>& field_ids) const;
 
+  /// \brief Return the field IDs of the identifier fields.
+  const std::vector<int32_t>& IdentifierFieldIds() const;
+
+  /// \brief Return the canonical field names of the identifier fields.
+  Result<std::vector<std::string>> IdentifierFieldNames() const;
+
   friend bool operator==(const Schema& lhs, const Schema& rhs) { return lhs.Equals(rhs); }
 
  private:
   /// \brief Compare two schemas for equality.
   bool Equals(const Schema& other) const;
 
+  struct NameIdMap {
+    /// \brief Mapping from canonical field name to ID
+    ///
+    /// \note Short names for maps and lists are included for any name that does not
+    /// conflict with a canonical name. For example, a list, 'l', of structs with field
+    /// 'x' will produce short name 'l.x' in addition to canonical name 'l.element.x'.
+    std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>> name_to_id;
+
+    /// \brief Mapping from field ID to canonical name
+    ///
+    /// \note Canonical names, but not short names are set, for example
+    /// 'list.element.field' instead of 'list.field'.
+    std::unordered_map<int32_t, std::string> id_to_name;
+  };
+
   static Result<std::unordered_map<int32_t, std::reference_wrapper<const SchemaField>>>
   InitIdToFieldMap(const Schema&);
-  static Result<std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>>>
-  InitNameToIdMap(const Schema&);
+  static Result<NameIdMap> InitNameIdMap(const Schema&);
   static Result<std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>>>
   InitLowerCaseNameToIdMap(const Schema&);
   static Result<std::unordered_map<int32_t, std::vector<size_t>>> InitIdToPositionPath(
       const Schema&);
 
   const std::optional<int32_t> schema_id_;
+  /// Field IDs that uniquely identify rows in the table.
+  std::vector<int32_t> identifier_field_ids_;
   /// Mapping from field id to field.
   Lazy<InitIdToFieldMap> id_to_field_;
   /// Mapping from field name to field id.
-  Lazy<InitNameToIdMap> name_to_id_;
+  Lazy<InitNameIdMap> name_id_map_;
   /// Mapping from lowercased field name to field id
   Lazy<InitLowerCaseNameToIdMap> lowercase_name_to_id_;
   /// Mapping from field id to (nested) position path to access the field.

src/iceberg/test/CMakeLists.txt

@@ -54,6 +54,7 @@ endfunction()
 
 add_iceberg_test(schema_test
                  SOURCES
+                 assign_id_visitor_test.cc
                  name_mapping_test.cc
                  partition_field_test.cc
                  partition_spec_test.cc