refactor: decouple PartitionSpec from Schema

Author: HeartLinked

src/iceberg/json_internal.cc

@@ -529,7 +529,9 @@ Result<std::unique_ptr<PartitionSpec>> PartitionSpecFromJson(
     ICEBERG_ASSIGN_OR_RAISE(auto partition_field, PartitionFieldFromJson(field_json));
     partition_fields.push_back(std::move(*partition_field));
   }
-  return std::make_unique<PartitionSpec>(schema, spec_id, std::move(partition_fields));
+  // TODO(Li Feiyang):use a new PartitionSpec::Make to find the source field of each
+  // partition field from schema and then verify it
+  return std::make_unique<PartitionSpec>(spec_id, std::move(partition_fields));
 }
 
 Result<std::unique_ptr<SnapshotRef>> SnapshotRefFromJson(const nlohmann::json& json) {
@@ -902,8 +904,10 @@ Status ParsePartitionSpecs(const nlohmann::json& json, int8_t format_version,
                           std::move(field->transform()));
     }
 
-    auto spec = std::make_unique<PartitionSpec>(
-        current_schema, PartitionSpec::kInitialSpecId, std::move(fields));
+    // TODO(Li Feiyang):use a new PartitionSpec::Make to find the source field of each
+    // partition field from schema and then verify it
+    auto spec =
+        std::make_unique<PartitionSpec>(PartitionSpec::kInitialSpecId, std::move(fields));
     default_spec_id = spec->spec_id();
     partition_specs.push_back(std::move(spec));
   }

src/iceberg/json_internal.h

@@ -173,6 +173,7 @@ ICEBERG_EXPORT Result<std::string> ToJsonString(const PartitionSpec& partition_s
 /// objects. Each `PartitionField` will be parsed using the `PartitionFieldFromJson`
 /// function.
 ///
+/// \param schema The current schema.
 /// \param json The JSON object representing a `PartitionSpec`.
 /// \return An `expected` value containing either a `PartitionSpec` object or an error. If
 /// the JSON is malformed or missing expected fields, an error will be returned.

src/iceberg/manifest_adapter.cc

@@ -21,6 +21,8 @@
 
 #include <nanoarrow/nanoarrow.h>
 
+#include <utility>
+
 #include "iceberg/arrow/nanoarrow_status_internal.h"
 #include "iceberg/manifest_entry.h"
 #include "iceberg/manifest_list.h"
@@ -139,9 +141,9 @@ Result<ArrowArray*> ManifestAdapter::FinishAppending() {
   return &array_;
 }
 
-ManifestEntryAdapter::ManifestEntryAdapter(std::shared_ptr<PartitionSpec> partition_spec,
+ManifestEntryAdapter::ManifestEntryAdapter(std::shared_ptr<PartitionSpec> partition_spec, std::shared_ptr<Schema> current_schema,
                                            ManifestContent content)
-    : partition_spec_(std::move(partition_spec)), content_(content) {
+    : partition_spec_(std::move(partition_spec)), current_schema_(std::move(current_schema)), content_(content) {
   if (!partition_spec_) {
     partition_spec_ = PartitionSpec::Unpartitioned();
   }

src/iceberg/manifest_adapter.h

@@ -61,8 +61,8 @@ class ICEBERG_EXPORT ManifestAdapter {
 /// Implemented by different versions with version-specific schemas.
 class ICEBERG_EXPORT ManifestEntryAdapter : public ManifestAdapter {
  public:
-  ManifestEntryAdapter(std::shared_ptr<PartitionSpec> partition_spec,
-                       ManifestContent content);
+  ManifestEntryAdapter(std::shared_ptr<PartitionSpec> partition_spec, std::shared_ptr<Schema> current_schema,
+     ManifestContent content);
 
   ~ManifestEntryAdapter() override;
 
@@ -92,6 +92,7 @@ class ICEBERG_EXPORT ManifestEntryAdapter : public ManifestAdapter {
 
  protected:
   std::shared_ptr<PartitionSpec> partition_spec_;
+  std::shared_ptr<Schema> current_schema_;
   std::shared_ptr<Schema> manifest_schema_;
   const ManifestContent content_;
 };

src/iceberg/manifest_writer.cc

@@ -70,9 +70,14 @@ Result<std::unique_ptr<Writer>> OpenFileWriter(
 
 Result<std::unique_ptr<ManifestWriter>> ManifestWriter::MakeV1Writer(
     std::optional<int64_t> snapshot_id, std::string_view manifest_location,
-    std::shared_ptr<FileIO> file_io, std::shared_ptr<PartitionSpec> partition_spec) {
-  auto adapter =
-      std::make_unique<ManifestEntryAdapterV1>(snapshot_id, std::move(partition_spec));
+    std::shared_ptr<FileIO> file_io, std::shared_ptr<PartitionSpec> partition_spec,
+    std::shared_ptr<Schema> current_schema) {
+  if (manifest_location.empty() || !file_io || !partition_spec || !current_schema) {
+    return InvalidArgument("Invalid arguments to create V1 ManifestWriter");
+  }
+
+  auto adapter = std::make_unique<ManifestEntryAdapterV1>(
+      snapshot_id, std::move(partition_spec), std::move(current_schema));
   ICEBERG_RETURN_UNEXPECTED(adapter->Init());
   ICEBERG_RETURN_UNEXPECTED(adapter->StartAppending());
 
@@ -85,10 +90,12 @@ Result<std::unique_ptr<ManifestWriter>> ManifestWriter::MakeV1Writer(
 
 Result<std::unique_ptr<ManifestWriter>> ManifestWriter::MakeV2Writer(
     std::optional<int64_t> snapshot_id, std::string_view manifest_location,
-    std::shared_ptr<FileIO> file_io, std::shared_ptr<PartitionSpec> partition_spec,
-    ManifestContent content) {
+    std::shared_ptr<FileIO> file_io, std::shared_ptr<PartitionSpec> partition_spec, std::shared_ptr<Schema> current_schema, ManifestContent content) {
+  if (manifest_location.empty() || !file_io || !partition_spec || !current_schema) {
+    return InvalidArgument("Invalid arguments to create V2 ManifestWriter");
+  }
   auto adapter = std::make_unique<ManifestEntryAdapterV2>(
-      snapshot_id, std::move(partition_spec), content);
+      snapshot_id, std::move(partition_spec), std::move(current_schema), content);
   ICEBERG_RETURN_UNEXPECTED(adapter->Init());
   ICEBERG_RETURN_UNEXPECTED(adapter->StartAppending());
 
@@ -102,9 +109,13 @@ Result<std::unique_ptr<ManifestWriter>> ManifestWriter::MakeV2Writer(
 Result<std::unique_ptr<ManifestWriter>> ManifestWriter::MakeV3Writer(
     std::optional<int64_t> snapshot_id, std::optional<int64_t> first_row_id,
     std::string_view manifest_location, std::shared_ptr<FileIO> file_io,
-    std::shared_ptr<PartitionSpec> partition_spec, ManifestContent content) {
+    std::shared_ptr<PartitionSpec> partition_spec, std::shared_ptr<Schema> current_schema, ManifestContent content) {
   auto adapter = std::make_unique<ManifestEntryAdapterV3>(
-      snapshot_id, first_row_id, std::move(partition_spec), content);
+      snapshot_id, first_row_id, std::move(partition_spec), current_schema, content);
+  if (manifest_location.empty() || !file_io || !partition_spec || !current_schema) {
+    return InvalidArgument("Invalid arguments to create V3 ManifestWriter");
+  }
+
   ICEBERG_RETURN_UNEXPECTED(adapter->Init());
   ICEBERG_RETURN_UNEXPECTED(adapter->StartAppending());
 

src/iceberg/manifest_writer.h

@@ -62,35 +62,41 @@ class ICEBERG_EXPORT ManifestWriter {
   /// \param manifest_location Path to the manifest file.
   /// \param file_io File IO implementation to use.
   /// \param partition_spec Partition spec for the manifest.
+  /// \param current_schema Current table schema.
   /// \return A Result containing the writer or an error.
   static Result<std::unique_ptr<ManifestWriter>> MakeV1Writer(
       std::optional<int64_t> snapshot_id, std::string_view manifest_location,
-      std::shared_ptr<FileIO> file_io, std::shared_ptr<PartitionSpec> partition_spec);
+      std::shared_ptr<FileIO> file_io, std::shared_ptr<PartitionSpec> partition_spec,
+      std::shared_ptr<Schema> current_schema);
 
   /// \brief Creates a writer for a manifest file.
   /// \param snapshot_id ID of the snapshot.
   /// \param manifest_location Path to the manifest file.
   /// \param file_io File IO implementation to use.
   /// \param partition_spec Partition spec for the manifest.
+  /// \param current_schema Schema containing the source fields referenced by partition
+  /// spec.
   /// \param content Content of the manifest.
   /// \return A Result containing the writer or an error.
   static Result<std::unique_ptr<ManifestWriter>> MakeV2Writer(
       std::optional<int64_t> snapshot_id, std::string_view manifest_location,
       std::shared_ptr<FileIO> file_io, std::shared_ptr<PartitionSpec> partition_spec,
-      ManifestContent content);
+      std::shared_ptr<Schema> current_schema, ManifestContent content);
 
   /// \brief Creates a writer for a manifest file.
   /// \param snapshot_id ID of the snapshot.
   /// \param first_row_id First row ID of the snapshot.
   /// \param manifest_location Path to the manifest file.
   /// \param file_io File IO implementation to use.
   /// \param partition_spec Partition spec for the manifest.
+  /// \param current_schema Schema containing the source fields referenced by partition
+  /// spec.
   /// \param content Content of the manifest.
   /// \return A Result containing the writer or an error.
   static Result<std::unique_ptr<ManifestWriter>> MakeV3Writer(
       std::optional<int64_t> snapshot_id, std::optional<int64_t> first_row_id,
       std::string_view manifest_location, std::shared_ptr<FileIO> file_io,
-      std::shared_ptr<PartitionSpec> partition_spec, ManifestContent content);
+      std::shared_ptr<PartitionSpec> partition_spec, std::shared_ptr<Schema> current_schema, ManifestContent content);
 
  private:
   static constexpr int64_t kBatchSize = 1024;

src/iceberg/partition_spec.cc

@@ -21,6 +21,7 @@
 
 #include <algorithm>
 #include <format>
+#include <memory>
 #include <ranges>
 
 #include "iceberg/schema.h"
@@ -31,10 +32,9 @@
 
 namespace iceberg {
 
-PartitionSpec::PartitionSpec(std::shared_ptr<Schema> schema, int32_t spec_id,
-                             std::vector<PartitionField> fields,
+PartitionSpec::PartitionSpec(int32_t spec_id, std::vector<PartitionField> fields,
                              std::optional<int32_t> last_assigned_field_id)
-    : schema_(std::move(schema)), spec_id_(spec_id), fields_(std::move(fields)) {
+    : spec_id_(spec_id), fields_(std::move(fields)) {
   if (last_assigned_field_id) {
     last_assigned_field_id_ = last_assigned_field_id.value();
   } else if (fields_.empty()) {
@@ -48,34 +48,25 @@ PartitionSpec::PartitionSpec(std::shared_ptr<Schema> schema, int32_t spec_id,
 
 const std::shared_ptr<PartitionSpec>& PartitionSpec::Unpartitioned() {
   static const std::shared_ptr<PartitionSpec> unpartitioned =
-      std::make_shared<PartitionSpec>(
-          /*schema=*/std::make_shared<Schema>(std::vector<SchemaField>{}), kInitialSpecId,
-          std::vector<PartitionField>{}, kLegacyPartitionDataIdStart - 1);
+      std::make_shared<PartitionSpec>(kInitialSpecId, std::vector<PartitionField>{},
+                                      kLegacyPartitionDataIdStart - 1);
   return unpartitioned;
 }
 
-const std::shared_ptr<Schema>& PartitionSpec::schema() const { return schema_; }
-
 int32_t PartitionSpec::spec_id() const { return spec_id_; }
 
 std::span<const PartitionField> PartitionSpec::fields() const { return fields_; }
 
-Result<std::shared_ptr<StructType>> PartitionSpec::PartitionType() {
+Result<std::unique_ptr<StructType>> PartitionSpec::PartitionType(const Schema& schema) {
   if (fields_.empty()) {
     return nullptr;
   }
-  {
-    std::scoped_lock<std::mutex> lock(mutex_);
-    if (partition_type_ != nullptr) {
-      return partition_type_;
-    }
-  }
 
   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()));
+                            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
@@ -97,11 +88,7 @@ Result<std::shared_ptr<StructType>> PartitionSpec::PartitionType() {
                                   /*optional=*/true);
   }
 
-  std::scoped_lock<std::mutex> lock(mutex_);
-  if (partition_type_ == nullptr) {
-    partition_type_ = std::make_shared<StructType>(std::move(partition_fields));
-  }
-  return partition_type_;
+  return std::make_unique<StructType>(std::move(partition_fields));
 }
 
 std::string PartitionSpec::ToString() const {

src/iceberg/partition_spec.h

@@ -23,7 +23,7 @@
 /// Partition specs for Iceberg tables.
 
 #include <cstdint>
-#include <mutex>
+#include <memory>
 #include <optional>
 #include <span>
 #include <string>
@@ -32,6 +32,7 @@
 #include "iceberg/iceberg_export.h"
 #include "iceberg/partition_field.h"
 #include "iceberg/result.h"
+#include "iceberg/type_fwd.h"
 #include "iceberg/util/formattable.h"
 
 namespace iceberg {
@@ -56,24 +57,20 @@ class ICEBERG_EXPORT PartitionSpec : public util::Formattable {
   /// \param fields The partition fields.
   /// \param last_assigned_field_id The last assigned field ID. If not provided, it will
   /// be calculated from the fields.
-  PartitionSpec(std::shared_ptr<Schema> schema, int32_t spec_id,
-                std::vector<PartitionField> fields,
+  PartitionSpec(int32_t spec_id, std::vector<PartitionField> fields,
                 std::optional<int32_t> last_assigned_field_id = std::nullopt);
 
   /// \brief Get an unsorted partition spec singleton.
   static const std::shared_ptr<PartitionSpec>& Unpartitioned();
 
-  /// \brief Get the table schema
-  const std::shared_ptr<Schema>& schema() const;
-
   /// \brief Get the spec ID.
   int32_t spec_id() const;
 
   /// \brief Get a list view of the partition fields.
   std::span<const PartitionField> fields() const;
 
-  /// \brief Get the partition type.
-  Result<std::shared_ptr<StructType>> PartitionType();
+  /// \brief Get the partition type binding to the input schema.
+  Result<std::unique_ptr<StructType>> PartitionType(const Schema&);
 
   std::string ToString() const override;
 
@@ -87,14 +84,9 @@ class ICEBERG_EXPORT PartitionSpec : public util::Formattable {
   /// \brief Compare two partition specs for equality.
   bool Equals(const PartitionSpec& other) const;
 
-  std::shared_ptr<Schema> schema_;
   const int32_t spec_id_;
   std::vector<PartitionField> fields_;
   int32_t last_assigned_field_id_;
-
-  // FIXME: use similar lazy initialization pattern as in StructType
-  std::mutex mutex_;
-  std::shared_ptr<StructType> partition_type_;
 };
 
 }  // namespace iceberg

src/iceberg/table_scan.cc

@@ -31,6 +31,7 @@
 #include "iceberg/schema_field.h"
 #include "iceberg/snapshot.h"
 #include "iceberg/table_metadata.h"
+#include "iceberg/type.h"
 #include "iceberg/util/macros.h"
 
 namespace iceberg {
@@ -269,7 +270,12 @@ Result<std::vector<std::shared_ptr<FileScanTask>>> DataTableScan::PlanFiles() co
 
   std::vector<std::shared_ptr<FileScanTask>> tasks;
   ICEBERG_ASSIGN_OR_RAISE(auto partition_spec, context_.table_metadata->PartitionSpec());
-  auto partition_schema = partition_spec->schema();
+  // auto partition_schema = context_.table_metadata->Schema().value();
+
+  // Get the table schema and partition type
+  ICEBERG_ASSIGN_OR_RAISE(auto current_schema, context_.table_metadata->Schema());
+  ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr<StructType> partition_schema,
+                          partition_spec->PartitionType(*current_schema));
 
   for (const auto& manifest_file : manifest_files) {
     ICEBERG_ASSIGN_OR_RAISE(

src/iceberg/test/json_internal_test.cc

@@ -144,11 +144,9 @@ TEST(JsonPartitionTest, PartitionSpec) {
       std::vector<SchemaField>{SchemaField(3, "region", iceberg::string(), false),
                                SchemaField(5, "ts", iceberg::int64(), false)},
       /*schema_id=*/100);
-
   auto identity_transform = Transform::Identity();
-  PartitionSpec spec(schema, 1,
-                     {PartitionField(3, 101, "region", identity_transform),
-                      PartitionField(5, 102, "ts", identity_transform)});
+  PartitionSpec spec(1, {PartitionField(3, 101, "region", identity_transform),
+                         PartitionField(5, 102, "ts", identity_transform)});
   auto json = ToJson(spec);
   nlohmann::json expected_json = R"({"spec-id": 1,
                                      "fields": [