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, 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,19 @@ 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 Full 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, int32_t schema_id = kInitialSchemaId,
+      const std::vector<std::string>& identifier_field_names = {});
 
   /// \brief Get the schema ID.
   ///
@@ -78,6 +90,12 @@ class ICEBERG_EXPORT Schema : public StructType {
   Result<std::optional<std::reference_wrapper<const SchemaField>>> FindFieldById(
       int32_t field_id) const;
 
+  /// \brief Returns the full column name for the given id.
+  ///
+  /// \param field_id The id of the field to get the full name for.
+  /// \return The full 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,45 @@ class ICEBERG_EXPORT Schema : public StructType {
   Result<std::unique_ptr<Schema>> Project(
       const std::unordered_set<int32_t>& field_ids) const;
 
+  const std::vector<int32_t>& IdentifierFieldIds() const;
+  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 full 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 full 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

src/iceberg/test/assign_id_visitor_test.cc

@@ -0,0 +1,188 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <memory>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include "iceberg/schema.h"
+#include "iceberg/schema_field.h"
+#include "iceberg/test/matchers.h"
+#include "iceberg/type.h"
+#include "iceberg/util/type_util.h"
+
+namespace iceberg {
+
+namespace {
+
+Schema CreateFlatSchema() {
+  return Schema({
+      SchemaField::MakeRequired(/*field_id=*/10, "id", iceberg::int64()),
+      SchemaField::MakeOptional(/*field_id=*/20, "name", iceberg::string()),
+      SchemaField::MakeOptional(/*field_id=*/30, "age", iceberg::int32()),
+      SchemaField::MakeRequired(/*field_id=*/40, "data", iceberg::float64()),
+  });
+}
+
+std::shared_ptr<Type> CreateListOfStruct() {
+  return std::make_shared<ListType>(SchemaField::MakeOptional(
+      /*field_id=*/101, "element",
+      std::make_shared<StructType>(std::vector<SchemaField>{
+          SchemaField::MakeOptional(/*field_id=*/102, "x", iceberg::int32()),
+          SchemaField::MakeRequired(/*field_id=*/103, "y", iceberg::string()),
+      })));
+}
+
+std::shared_ptr<Type> CreateMapWithStructValue() {
+  return std::make_shared<MapType>(
+      SchemaField::MakeRequired(/*field_id=*/201, "key", iceberg::string()),
+      SchemaField::MakeRequired(
+          /*field_id=*/202, "value",
+          std::make_shared<StructType>(std::vector<SchemaField>{
+              SchemaField::MakeRequired(/*field_id=*/203, "id", iceberg::int64()),
+              SchemaField::MakeOptional(/*field_id=*/204, "name", iceberg::string()),
+          })));
+}
+
+std::shared_ptr<Type> CreateNestedStruct() {
+  return std::make_shared<StructType>(std::vector<SchemaField>{
+      SchemaField::MakeRequired(/*field_id=*/301, "outer_id", iceberg::int64()),
+      SchemaField::MakeRequired(
+          /*field_id=*/302, "nested",
+          std::make_shared<StructType>(std::vector<SchemaField>{
+              SchemaField::MakeOptional(/*field_id=*/303, "inner_id", iceberg::int32()),
+              SchemaField::MakeRequired(/*field_id=*/304, "inner_name",
+                                        iceberg::string()),
+          })),
+  });
+}
+
+Schema CreateNestedSchema(std::vector<int32_t> identifier_field_ids = {}) {
+  return Schema(
+      {
+          SchemaField::MakeRequired(/*field_id=*/10, "id", iceberg::int64()),
+          SchemaField::MakeOptional(/*field_id=*/20, "list", CreateListOfStruct()),
+          SchemaField::MakeOptional(/*field_id=*/30, "map", CreateMapWithStructValue()),
+          SchemaField::MakeRequired(/*field_id=*/40, "struct", CreateNestedStruct()),
+      },
+      Schema::kInitialSchemaId, std::move(identifier_field_ids));
+}
+
+}  // namespace
+
+TEST(AssignFreshIdVisitorTest, FlatSchema) {
+  Schema schema = CreateFlatSchema();
+
+  std::atomic<int32_t> id = 0;
+  auto next_id = [&id]() { return ++id; };
+  ICEBERG_UNWRAP_OR_FAIL(auto fresh_schema,
+                         AssignFreshIds(Schema::kInitialSchemaId, schema, next_id));
+
+  ASSERT_EQ(fresh_schema->fields().size(), schema.fields().size());
+  EXPECT_EQ(Schema(
+                {
+                    SchemaField::MakeRequired(/*field_id=*/1, "id", iceberg::int64()),
+                    SchemaField::MakeOptional(/*field_id=*/2, "name", iceberg::string()),
+                    SchemaField::MakeOptional(/*field_id=*/3, "age", iceberg::int32()),
+                    SchemaField::MakeRequired(/*field_id=*/4, "data", iceberg::float64()),
+                },
+                Schema::kInitialSchemaId),
+            *fresh_schema);
+}
+
+TEST(AssignFreshIdVisitorTest, NestedSchema) {
+  Schema schema = CreateNestedSchema();
+  std::atomic<int32_t> id = 0;
+  auto next_id = [&id]() { return ++id; };
+  ICEBERG_UNWRAP_OR_FAIL(auto fresh_schema,
+                         AssignFreshIds(Schema::kInitialSchemaId, schema, next_id));
+
+  ASSERT_EQ(4, fresh_schema->fields().size());
+  for (int32_t i = 0; i < fresh_schema->fields().size(); ++i) {
+    EXPECT_EQ(i + 1, fresh_schema->fields()[i].field_id());
+  }
+
+  auto list_field = fresh_schema->fields()[1];
+  auto list_type = std::dynamic_pointer_cast<ListType>(list_field.type());
+  ASSERT_TRUE(list_type);
+  auto list_element_field = list_type->fields()[0];
+  EXPECT_EQ(5, list_element_field.field_id());
+  auto list_element_type =
+      std::dynamic_pointer_cast<StructType>(list_element_field.type());
+  ASSERT_TRUE(list_element_type);
+  EXPECT_EQ(StructType(std::vector<SchemaField>{
+                SchemaField::MakeOptional(/*field_id=*/6, "x", iceberg::int32()),
+                SchemaField::MakeRequired(/*field_id=*/7, "y", iceberg::string()),
+            }),
+            *list_element_type);
+
+  auto map_field = fresh_schema->fields()[2];
+  auto map_type = std::dynamic_pointer_cast<MapType>(map_field.type());
+  ASSERT_TRUE(map_type);
+  EXPECT_EQ(8, map_type->fields()[0].field_id());
+  auto map_value_field = map_type->fields()[1];
+  EXPECT_EQ(9, map_value_field.field_id());
+  auto map_value_type = std::dynamic_pointer_cast<StructType>(map_value_field.type());
+  ASSERT_TRUE(map_value_type);
+  EXPECT_EQ(StructType(std::vector<SchemaField>{
+                SchemaField::MakeRequired(/*field_id=*/10, "id", iceberg::int64()),
+                SchemaField::MakeOptional(/*field_id=*/11, "name", iceberg::string()),
+            }),
+            *map_value_type);
+
+  auto struct_field = fresh_schema->fields()[3];
+  auto struct_type = std::dynamic_pointer_cast<StructType>(struct_field.type());
+  ASSERT_TRUE(struct_type);
+
+  auto expect_nested_struct_type = std::make_shared<StructType>(std::vector<SchemaField>{
+      SchemaField::MakeOptional(/*field_id=*/14, "inner_id", iceberg::int32()),
+      SchemaField::MakeRequired(/*field_id=*/15, "inner_name", iceberg::string()),
+  });
+  EXPECT_EQ(StructType(std::vector<SchemaField>{
+                SchemaField::MakeRequired(/*field_id=*/12, "outer_id", iceberg::int64()),
+                SchemaField::MakeRequired(
+                    /*field_id=*/13, "nested", expect_nested_struct_type)}),
+            *struct_type);
+
+  auto nested_struct_field = struct_type->fields()[1];
+  auto nested_struct_type =
+      std::dynamic_pointer_cast<StructType>(nested_struct_field.type());
+  ASSERT_TRUE(nested_struct_type);
+  EXPECT_EQ(*expect_nested_struct_type, *nested_struct_type);
+}
+
+TEST(AssignFreshIdVisitorTest, RefreshIdentifierId) {
+  std::atomic<int32_t> id = 0;
+  auto next_id = [&id]() { return ++id; };
+
+  Schema invalid_schema = CreateNestedSchema({10, 400});
+  // Invalid identified field id
+  auto result = AssignFreshIds(Schema::kInitialSchemaId, invalid_schema, next_id);
+  EXPECT_THAT(result, IsError(ErrorKind::kInvalidSchema));
+  EXPECT_THAT(result, HasErrorMessage("Cannot find"));
+
+  id = 0;
+  Schema schema = CreateNestedSchema({10, 301});
+  ICEBERG_UNWRAP_OR_FAIL(auto fresh_schema,
+                         AssignFreshIds(Schema::kInitialSchemaId, schema, next_id));
+  EXPECT_THAT(fresh_schema->IdentifierFieldIds(), testing::ElementsAre(1, 12));
+}
+
+}  // namespace iceberg