fix: align manifest partition reads with reassigned IDs

Use unknown types for missing partition source fields when building manifest
partition schemas. Add Java-style schema ID reassignment maps and use them to
map raw manifest partition field IDs back to their original IDs.

Also add transform result type support so unbound partition specs can derive
partition field types consistently with Java.

Author: wgtmac

src/iceberg/manifest/manifest_reader.cc

@@ -861,7 +861,7 @@ Result<std::vector<ManifestEntry>> ManifestReaderImpl::LiveEntries() {
 }
 
 Result<std::vector<ManifestEntry>> ManifestReaderImpl::ReadEntries(bool only_live) {
-  ICEBERG_ASSIGN_OR_RAISE(auto partition_type, spec_->PartitionType(*schema_));
+  ICEBERG_ASSIGN_OR_RAISE(auto partition_type, spec_->RawPartitionType(*schema_));
   auto data_file_schema = DataFile::Type(std::move(partition_type))->ToSchema();
 
   std::shared_ptr<Schema> projected_data_file_schema;

src/iceberg/partition_spec.cc

@@ -74,24 +74,16 @@ Result<std::unique_ptr<StructType>> PartitionSpec::PartitionType(
 
   std::vector<SchemaField> partition_fields;
   for (const auto& partition_field : fields_) {
-    // Get the source field from the original schema by source_id
     ICEBERG_ASSIGN_OR_RAISE(auto source_field,
                             schema.FindFieldById(partition_field.source_id()));
-    if (!source_field.has_value()) {
-      // TODO(xiao.dong) when source field is missing,
-      // should return an error or just use UNKNOWN type
-      return InvalidSchema("Cannot find source field for partition field:{}",
-                           partition_field.field_id());
+    std::shared_ptr<Type> result_type;
+    if (source_field.has_value()) {
+      auto source_field_type = source_field.value().get().type();
+      result_type = partition_field.transform()->ResultType(std::move(source_field_type));
+    } else {
+      result_type = unknown();
     }
-    auto source_field_type = source_field.value().get().type();
-    // Bind the transform to the source field type to get the result type
-    ICEBERG_ASSIGN_OR_RAISE(auto transform_function,
-                            partition_field.transform()->Bind(source_field_type));
-
-    auto result_type = transform_function->ResultType();
 
-    // Create the partition field with the transform result type
-    // Partition fields are always optional (can be null)
     partition_fields.emplace_back(partition_field.field_id(),
                                   std::string(partition_field.name()),
                                   std::move(result_type),
@@ -101,6 +93,29 @@ Result<std::unique_ptr<StructType>> PartitionSpec::PartitionType(
   return std::make_unique<StructType>(std::move(partition_fields));
 }
 
+Result<std::unique_ptr<StructType>> PartitionSpec::RawPartitionType(
+    const Schema& schema) const {
+  const auto& ids_to_original = schema.IdsToOriginal();
+  if (ids_to_original.empty()) {
+    return PartitionType(schema);
+  }
+
+  ICEBERG_ASSIGN_OR_RAISE(auto partition_type, PartitionType(schema));
+  std::vector<SchemaField> raw_partition_fields;
+  raw_partition_fields.reserve(partition_type->fields().size());
+  for (const auto& field : partition_type->fields()) {
+    auto original_id = ids_to_original.find(field.field_id());
+    if (original_id == ids_to_original.end()) {
+      return InvalidSchema("Cannot find original field ID for reassigned field ID: {}",
+                           field.field_id());
+    }
+    raw_partition_fields.emplace_back(original_id->second, std::string(field.name()),
+                                      field.type(), field.optional(),
+                                      std::string(field.doc()));
+  }
+  return std::make_unique<StructType>(std::move(raw_partition_fields));
+}
+
 Result<std::string> PartitionSpec::PartitionPath(const PartitionValues& data) const {
   ICEBERG_PRECHECK(fields_.size() == data.num_fields(),
                    "Partition spec and data mismatch, expected field num {}, got {}",

src/iceberg/partition_spec.h

@@ -64,6 +64,9 @@ class ICEBERG_EXPORT PartitionSpec : public util::Formattable {
   /// \brief Get the partition type binding to the input schema.
   Result<std::unique_ptr<StructType>> PartitionType(const Schema& schema) const;
 
+  /// \brief Get the partition type as physically written in manifest files.
+  Result<std::unique_ptr<StructType>> RawPartitionType(const Schema& schema) const;
+
   /// \brief Get the partition path for the given partition data.
   Result<std::string> PartitionPath(const PartitionValues& data) const;
 

src/iceberg/schema.cc

@@ -35,8 +35,103 @@
 
 namespace iceberg {
 
+struct SchemaReassignIdContext {
+  Schema::IdMap ids_to_reassigned;
+  Schema::IdMap ids_to_original;
+};
+
 namespace {
 
+const Schema::IdMap& EmptyIdMap() {
+  static const Schema::IdMap kEmpty;
+  return kEmpty;
+}
+
+void RecordIdReassignment(int32_t old_id, int32_t new_id,
+                          Schema::IdMap& ids_to_reassigned,
+                          Schema::IdMap& ids_to_original) {
+  if (new_id != old_id) {
+    ids_to_reassigned[old_id] = new_id;
+    ids_to_original[new_id] = old_id;
+  }
+}
+
+SchemaField ReassignField(const SchemaField& field, int32_t new_id,
+                          const Schema::GetId& get_id, Schema::IdMap& ids_to_reassigned,
+                          Schema::IdMap& ids_to_original);
+
+std::shared_ptr<Type> ReassignTypeIds(const std::shared_ptr<Type>& type,
+                                      const Schema::GetId& get_id,
+                                      Schema::IdMap& ids_to_reassigned,
+                                      Schema::IdMap& ids_to_original) {
+  switch (type->type_id()) {
+    case TypeId::kStruct: {
+      const auto& struct_type = static_cast<const StructType&>(*type);
+      const auto& fields = struct_type.fields();
+      std::vector<int32_t> new_ids;
+      new_ids.reserve(fields.size());
+      for (const auto& field : fields) {
+        const auto new_id = get_id(field.field_id());
+        RecordIdReassignment(field.field_id(), new_id, ids_to_reassigned,
+                             ids_to_original);
+        new_ids.push_back(new_id);
+      }
+
+      std::vector<SchemaField> reassigned_fields;
+      reassigned_fields.reserve(fields.size());
+      for (size_t i = 0; i < fields.size(); ++i) {
+        reassigned_fields.emplace_back(ReassignField(fields[i], new_ids[i], get_id,
+                                                     ids_to_reassigned, ids_to_original));
+      }
+      return std::make_shared<StructType>(std::move(reassigned_fields));
+    }
+    case TypeId::kList: {
+      const auto& list_type = static_cast<const ListType&>(*type);
+      const auto& element = list_type.element();
+      const auto new_id = get_id(element.field_id());
+      RecordIdReassignment(element.field_id(), new_id, ids_to_reassigned,
+                           ids_to_original);
+      return std::make_shared<ListType>(
+          ReassignField(element, new_id, get_id, ids_to_reassigned, ids_to_original));
+    }
+    case TypeId::kMap: {
+      const auto& map_type = static_cast<const MapType&>(*type);
+      const auto& key = map_type.key();
+      const auto& value = map_type.value();
+      const auto new_key_id = get_id(key.field_id());
+      const auto new_value_id = get_id(value.field_id());
+      RecordIdReassignment(key.field_id(), new_key_id, ids_to_reassigned,
+                           ids_to_original);
+      RecordIdReassignment(value.field_id(), new_value_id, ids_to_reassigned,
+                           ids_to_original);
+      return std::make_shared<MapType>(
+          ReassignField(key, new_key_id, get_id, ids_to_reassigned, ids_to_original),
+          ReassignField(value, new_value_id, get_id, ids_to_reassigned, ids_to_original));
+    }
+    default:
+      return type;
+  }
+}
+
+SchemaField ReassignField(const SchemaField& field, int32_t new_id,
+                          const Schema::GetId& get_id, Schema::IdMap& ids_to_reassigned,
+                          Schema::IdMap& ids_to_original) {
+  return {new_id, std::string(field.name()),
+          ReassignTypeIds(field.type(), get_id, ids_to_reassigned, ids_to_original),
+          field.optional(), std::string(field.doc())};
+}
+
+std::vector<SchemaField> ReassignIds(std::vector<SchemaField> fields,
+                                     const Schema::GetId& get_id,
+                                     SchemaReassignIdContext& reassign_id_context) {
+  auto reassigned_type = ReassignTypeIds(std::make_shared<StructType>(std::move(fields)),
+                                         get_id, reassign_id_context.ids_to_reassigned,
+                                         reassign_id_context.ids_to_original);
+  const auto& reassigned_fields =
+      internal::checked_cast<const StructType&>(*reassigned_type).fields();
+  return {reassigned_fields.begin(), reassigned_fields.end()};
+}
+
 Status ValidateFieldNullability(const Type& type) {
   auto validate_field = [&](const SchemaField& field) -> Status {
     ICEBERG_PRECHECK(field.optional() || field.type()->type_id() != TypeId::kUnknown,
@@ -73,17 +168,23 @@ Status ValidateFieldNullability(const Type& type) {
 
 }  // namespace
 
-Schema::Schema(std::vector<SchemaField> fields, int32_t schema_id)
+Schema::Schema(std::vector<SchemaField> fields, int32_t schema_id, GetId get_id)
     : StructType(std::move(fields)),
       schema_id_(schema_id),
-      cache_(std::make_unique<SchemaCache>(this)) {}
+      cache_(std::make_unique<SchemaCache>(this)) {
+  if (get_id) {
+    reassign_id_context_ = std::make_unique<SchemaReassignIdContext>();
+    fields_ = ReassignIds(std::move(fields_), get_id, *reassign_id_context_);
+  }
+}
 
 Schema::~Schema() = default;
 
 Result<std::unique_ptr<Schema>> Schema::Make(std::vector<SchemaField> fields,
                                              int32_t schema_id,
-                                             std::vector<int32_t> identifier_field_ids) {
-  auto schema = std::make_unique<Schema>(std::move(fields), schema_id);
+                                             std::vector<int32_t> identifier_field_ids,
+                                             GetId get_id) {
+  auto schema = std::make_unique<Schema>(std::move(fields), schema_id, std::move(get_id));
 
   if (!identifier_field_ids.empty()) {
     auto id_to_parent = IndexParents(*schema);
@@ -99,8 +200,8 @@ Result<std::unique_ptr<Schema>> Schema::Make(std::vector<SchemaField> fields,
 
 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);
+    const std::vector<std::string>& identifier_field_names, GetId get_id) {
+  auto schema = std::make_unique<Schema>(std::move(fields), schema_id, std::move(get_id));
 
   std::vector<int32_t> fresh_identifier_ids;
   for (const auto& name : identifier_field_names) {
@@ -181,6 +282,14 @@ const std::shared_ptr<Schema>& Schema::EmptySchema() {
 
 int32_t Schema::schema_id() const { return schema_id_; }
 
+const Schema::IdMap& Schema::IdsToReassigned() const {
+  return reassign_id_context_ ? reassign_id_context_->ids_to_reassigned : EmptyIdMap();
+}
+
+const Schema::IdMap& Schema::IdsToOriginal() const {
+  return reassign_id_context_ ? reassign_id_context_->ids_to_original : EmptyIdMap();
+}
+
 std::string Schema::ToString() const {
   std::string repr = "schema<";
   for (const auto& field : fields_) {

src/iceberg/schema.h

@@ -24,6 +24,7 @@
 /// and any utility functions.  See iceberg/type.h and iceberg/field.h as well.
 
 #include <cstdint>
+#include <functional>
 #include <optional>
 #include <string>
 #include <unordered_map>
@@ -40,6 +41,7 @@
 namespace iceberg {
 
 class SchemaCache;
+struct SchemaReassignIdContext;
 
 /// \brief A schema for a Table.
 ///
@@ -55,7 +57,14 @@ class ICEBERG_EXPORT Schema : public StructType {
   /// \brief Special value to select all columns from manifest files.
   static constexpr std::string_view kAllColumns = "*";
 
-  explicit Schema(std::vector<SchemaField> fields, int32_t schema_id = kInitialSchemaId);
+  /// \brief Maps an original field ID to its reassigned ID.
+  ///
+  /// The mapping is total: return the original ID when no reassignment is needed.
+  using GetId = std::function<int32_t(int32_t)>;
+  using IdMap = std::unordered_map<int32_t, int32_t>;
+
+  explicit Schema(std::vector<SchemaField> fields, int32_t schema_id = kInitialSchemaId,
+                  GetId get_id = {});
 
   ~Schema() override;
 
@@ -64,21 +73,24 @@ class ICEBERG_EXPORT Schema : public StructType {
   /// \param fields The fields that make up the schema.
   /// \param schema_id The unique identifier for this schema (default:kInitialSchemaId).
   /// \param identifier_field_ids Field IDs that uniquely identify rows in the table.
+  /// \param get_id Function mapping each original field ID to its reassigned ID.
   /// \return A new Schema instance or Status if failed.
   static Result<std::unique_ptr<Schema>> Make(std::vector<SchemaField> fields,
                                               int32_t schema_id,
-                                              std::vector<int32_t> identifier_field_ids);
+                                              std::vector<int32_t> identifier_field_ids,
+                                              GetId get_id = {});
 
   /// \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.
+  /// \param get_id Function mapping each original field ID to its reassigned ID.
   /// \return A new Schema instance or Status if failed.
   static Result<std::unique_ptr<Schema>> Make(
       std::vector<SchemaField> fields, int32_t schema_id,
-      const std::vector<std::string>& identifier_field_names);
+      const std::vector<std::string>& identifier_field_names, GetId get_id = {});
 
   /// \brief Validate that the identifier field with the given ID is valid for the schema
   ///
@@ -166,6 +178,12 @@ class ICEBERG_EXPORT Schema : public StructType {
   /// \brief Return the field IDs of the identifier fields.
   const std::vector<int32_t>& IdentifierFieldIds() const;
 
+  /// \brief Return a map of original field IDs to reassigned field IDs.
+  const IdMap& IdsToReassigned() const;
+
+  /// \brief Return a map of reassigned field IDs to original field IDs.
+  const IdMap& IdsToOriginal() const;
+
   /// \brief Return the canonical field names of the identifier fields.
   Result<std::vector<std::string>> IdentifierFieldNames() const;
 
@@ -196,6 +214,7 @@ class ICEBERG_EXPORT Schema : public StructType {
   const int32_t schema_id_;
   // Field IDs that uniquely identify rows in the table.
   std::vector<int32_t> identifier_field_ids_;
+  std::unique_ptr<SchemaReassignIdContext> reassign_id_context_;
   // Cache for schema mappings to facilitate fast lookups.
   std::unique_ptr<SchemaCache> cache_;
 };

src/iceberg/test/assign_id_visitor_test.cc

@@ -106,6 +106,8 @@ TEST(AssignFreshIdVisitorTest, FlatSchema) {
                 },
                 Schema::kInitialSchemaId),
             *fresh_schema);
+  EXPECT_TRUE(fresh_schema->IdsToReassigned().empty());
+  EXPECT_TRUE(fresh_schema->IdsToOriginal().empty());
 }
 
 TEST(AssignFreshIdVisitorTest, NestedSchema) {
@@ -169,6 +171,55 @@ TEST(AssignFreshIdVisitorTest, NestedSchema) {
   EXPECT_EQ(*expect_nested_struct_type, *nested_struct_type);
 }
 
+TEST(AssignFreshIdVisitorTest, GetIdMaps) {
+  ICEBERG_UNWRAP_OR_FAIL(auto schema, CreateNestedSchema());
+  std::vector<SchemaField> fields(schema->fields().begin(), schema->fields().end());
+  auto reassign_id = [](int32_t old_id) { return old_id + 1000; };
+
+  Schema reassigned_schema(std::move(fields), Schema::kInitialSchemaId, reassign_id);
+
+  EXPECT_EQ(reassigned_schema.fields()[0].field_id(), 1010);
+  EXPECT_EQ(reassigned_schema.fields()[1].field_id(), 1020);
+  auto list_type =
+      std::dynamic_pointer_cast<ListType>(reassigned_schema.fields()[1].type());
+  ASSERT_TRUE(list_type);
+  EXPECT_EQ(list_type->element().field_id(), 1101);
+
+  EXPECT_EQ(reassigned_schema.IdsToReassigned().size(), 15U);
+  EXPECT_THAT(
+      reassigned_schema.IdsToReassigned(),
+      testing::UnorderedElementsAre(
+          testing::Pair(10, 1010), testing::Pair(20, 1020), testing::Pair(30, 1030),
+          testing::Pair(40, 1040), testing::Pair(101, 1101), testing::Pair(102, 1102),
+          testing::Pair(103, 1103), testing::Pair(201, 1201), testing::Pair(202, 1202),
+          testing::Pair(203, 1203), testing::Pair(204, 1204), testing::Pair(301, 1301),
+          testing::Pair(302, 1302), testing::Pair(303, 1303), testing::Pair(304, 1304)));
+  EXPECT_THAT(
+      reassigned_schema.IdsToOriginal(),
+      testing::UnorderedElementsAre(
+          testing::Pair(1010, 10), testing::Pair(1020, 20), testing::Pair(1030, 30),
+          testing::Pair(1040, 40), testing::Pair(1101, 101), testing::Pair(1102, 102),
+          testing::Pair(1103, 103), testing::Pair(1201, 201), testing::Pair(1202, 202),
+          testing::Pair(1203, 203), testing::Pair(1204, 204), testing::Pair(1301, 301),
+          testing::Pair(1302, 302), testing::Pair(1303, 303), testing::Pair(1304, 304)));
+}
+
+TEST(AssignFreshIdVisitorTest, GetIdIdentifierNames) {
+  ICEBERG_UNWRAP_OR_FAIL(auto schema, CreateNestedSchema());
+  std::vector<SchemaField> fields(schema->fields().begin(), schema->fields().end());
+  auto reassign_id = [](int32_t old_id) { return old_id + 1000; };
+
+  ICEBERG_UNWRAP_OR_FAIL(
+      auto reassigned_schema,
+      Schema::Make(std::move(fields), Schema::kInitialSchemaId,
+                   std::vector<std::string>{"id", "struct.outer_id"}, reassign_id));
+
+  EXPECT_THAT(reassigned_schema->IdentifierFieldIds(), testing::ElementsAre(1010, 1301));
+  ICEBERG_UNWRAP_OR_FAIL(auto identifier_field_names,
+                         reassigned_schema->IdentifierFieldNames());
+  EXPECT_THAT(identifier_field_names, testing::ElementsAre("id", "struct.outer_id"));
+}
+
 TEST(AssignFreshIdVisitorTest, RefreshIdentifierId) {
   int32_t id = 0;
   auto next_id = [&id]() { return ++id; };