refactor: decouple PartitionSpec from Schema

Author: HeartLinked

src/iceberg/json_internal.cc

@@ -519,8 +519,7 @@ Result<std::unique_ptr<PartitionField>> PartitionFieldFromJson(
                                           std::move(transform));
 }
 
-Result<std::unique_ptr<PartitionSpec>> PartitionSpecFromJson(
-    const std::shared_ptr<Schema>& schema, const nlohmann::json& json) {
+Result<std::unique_ptr<PartitionSpec>> PartitionSpecFromJson(const nlohmann::json& json) {
   ICEBERG_ASSIGN_OR_RAISE(auto spec_id, GetJsonValue<int32_t>(json, kSpecId));
   ICEBERG_ASSIGN_OR_RAISE(auto fields, GetJsonValue<nlohmann::json>(json, kFields));
 
@@ -529,7 +528,7 @@ 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));
+  return std::make_unique<PartitionSpec>(spec_id, std::move(partition_fields));
 }
 
 Result<std::unique_ptr<SnapshotRef>> SnapshotRefFromJson(const nlohmann::json& json) {
@@ -853,11 +852,9 @@ Result<std::shared_ptr<Schema>> ParseSchemas(
 ///
 /// \param[in] json The JSON object to parse.
 /// \param[in] format_version The format version of the table.
-/// \param[in] current_schema The current schema.
 /// \param[out] default_spec_id The default partition spec ID.
 /// \param[out] partition_specs The list of partition specs.
 Status ParsePartitionSpecs(const nlohmann::json& json, int8_t format_version,
-                           const std::shared_ptr<Schema>& current_schema,
                            int32_t& default_spec_id,
                            std::vector<std::shared_ptr<PartitionSpec>>& partition_specs) {
   if (json.contains(kPartitionSpecs)) {
@@ -870,8 +867,7 @@ Status ParsePartitionSpecs(const nlohmann::json& json, int8_t format_version,
     ICEBERG_ASSIGN_OR_RAISE(default_spec_id, GetJsonValue<int32_t>(json, kDefaultSpecId));
 
     for (const auto& spec_json : spec_array) {
-      ICEBERG_ASSIGN_OR_RAISE(auto spec,
-                              PartitionSpecFromJson(current_schema, spec_json));
+      ICEBERG_ASSIGN_OR_RAISE(auto spec, PartitionSpecFromJson(spec_json));
       partition_specs.push_back(std::move(spec));
     }
   } else {
@@ -902,8 +898,8 @@ 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));
+    auto spec =
+        std::make_unique<PartitionSpec>(PartitionSpec::kInitialSpecId, std::move(fields));
     default_spec_id = spec->spec_id();
     partition_specs.push_back(std::move(spec));
   }
@@ -977,9 +973,9 @@ Result<std::unique_ptr<TableMetadata>> TableMetadataFromJson(const nlohmann::jso
       ParseSchemas(json, table_metadata->format_version,
                    table_metadata->current_schema_id, table_metadata->schemas));
 
-  ICEBERG_RETURN_UNEXPECTED(ParsePartitionSpecs(
-      json, table_metadata->format_version, current_schema,
-      table_metadata->default_spec_id, table_metadata->partition_specs));
+  ICEBERG_RETURN_UNEXPECTED(ParsePartitionSpecs(json, table_metadata->format_version,
+                                                table_metadata->default_spec_id,
+                                                table_metadata->partition_specs));
 
   if (json.contains(kLastPartitionId)) {
     ICEBERG_ASSIGN_OR_RAISE(table_metadata->last_partition_id,

src/iceberg/json_internal.h

@@ -177,7 +177,7 @@ ICEBERG_EXPORT Result<std::string> ToJsonString(const PartitionSpec& partition_s
 /// \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.
 ICEBERG_EXPORT Result<std::unique_ptr<PartitionSpec>> PartitionSpecFromJson(
-    const std::shared_ptr<Schema>& schema, const nlohmann::json& json);
+    const nlohmann::json& json);
 
 /// \brief Serializes a `SnapshotRef` object to JSON.
 ///

src/iceberg/manifest_adapter.cc

@@ -152,7 +152,11 @@ Result<std::shared_ptr<StructType>> ManifestEntryAdapter::GetManifestEntryType()
   if (partition_spec_ == nullptr) [[unlikely]] {
     return ManifestEntry::TypeFromPartitionType(nullptr);
   }
-  ICEBERG_ASSIGN_OR_RAISE(auto partition_type, partition_spec_->PartitionType());
+  std::shared_ptr<StructType> partition_type = nullptr;
+  if (table_schema_ != nullptr) {
+    ICEBERG_ASSIGN_OR_RAISE(partition_type,
+                            partition_spec_->PartitionType(table_schema_));
+  }
   return ManifestEntry::TypeFromPartitionType(std::move(partition_type));
 }
 

src/iceberg/manifest_adapter.h

@@ -61,8 +61,10 @@ class ICEBERG_EXPORT ManifestAdapter {
 /// Implemented by different versions with version-specific schemas.
 class ICEBERG_EXPORT ManifestEntryAdapter : public ManifestAdapter {
  public:
-  explicit ManifestEntryAdapter(std::shared_ptr<PartitionSpec> partition_spec)
-      : partition_spec_(std::move(partition_spec)) {}
+  ManifestEntryAdapter(std::shared_ptr<PartitionSpec> partition_spec,
+                       std::shared_ptr<Schema> table_schema)
+      : partition_spec_(std::move(partition_spec)),
+        table_schema_(std::move(table_schema)) {}
   ~ManifestEntryAdapter() override;
 
   virtual Status Append(const ManifestEntry& entry) = 0;
@@ -96,6 +98,7 @@ class ICEBERG_EXPORT ManifestEntryAdapter : public ManifestAdapter {
 
  protected:
   std::shared_ptr<PartitionSpec> partition_spec_;
+  std::shared_ptr<Schema> table_schema_;
   std::shared_ptr<Schema> manifest_schema_;
 };
 

src/iceberg/manifest_entry.cc

@@ -38,7 +38,7 @@ bool ManifestEntry::operator==(const ManifestEntry& other) const {
 
 std::shared_ptr<StructType> DataFile::Type(std::shared_ptr<StructType> partition_type) {
   if (!partition_type) {
-    partition_type = PartitionSpec::Unpartitioned()->schema();
+    partition_type = std::make_shared<StructType>(std::vector<SchemaField>{});
   }
   return std::make_shared<StructType>(std::vector<SchemaField>{
       kContent,

src/iceberg/manifest_writer.cc

@@ -68,9 +68,10 @@ 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> table_schema) {
+  auto adapter = std::make_unique<ManifestEntryAdapterV1>(
+      snapshot_id, std::move(partition_spec), std::move(table_schema));
   ICEBERG_RETURN_UNEXPECTED(adapter->Init());
   ICEBERG_RETURN_UNEXPECTED(adapter->StartAppending());
 
@@ -83,9 +84,10 @@ 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) {
-  auto adapter =
-      std::make_unique<ManifestEntryAdapterV2>(snapshot_id, std::move(partition_spec));
+    std::shared_ptr<FileIO> file_io, std::shared_ptr<PartitionSpec> partition_spec,
+    std::shared_ptr<Schema> table_schema) {
+  auto adapter = std::make_unique<ManifestEntryAdapterV2>(
+      snapshot_id, std::move(partition_spec), std::move(table_schema));
   ICEBERG_RETURN_UNEXPECTED(adapter->Init());
   ICEBERG_RETURN_UNEXPECTED(adapter->StartAppending());
 
@@ -99,9 +101,9 @@ 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) {
-  auto adapter = std::make_unique<ManifestEntryAdapterV3>(snapshot_id, first_row_id,
-                                                          std::move(partition_spec));
+    std::shared_ptr<PartitionSpec> partition_spec, std::shared_ptr<Schema> table_schema) {
+  auto adapter = std::make_unique<ManifestEntryAdapterV3>(
+      snapshot_id, first_row_id, std::move(partition_spec), std::move(table_schema));
   ICEBERG_RETURN_UNEXPECTED(adapter->Init());
   ICEBERG_RETURN_UNEXPECTED(adapter->StartAppending());
 

src/iceberg/manifest_writer.h

@@ -59,32 +59,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 table_schema Schema containing the source fields referenced by partition
+  /// spec.
   /// \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> table_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 table_schema Schema containing the source fields referenced by partition
+  /// spec.
   /// \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);
+      std::shared_ptr<FileIO> file_io, std::shared_ptr<PartitionSpec> partition_spec,
+      std::shared_ptr<Schema> table_schema);
 
   /// \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 table_schema Schema containing the source fields referenced by partition
+  /// spec.
   /// \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);
+      std::shared_ptr<PartitionSpec> partition_spec,
+      std::shared_ptr<Schema> table_schema);
 
  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,19 +48,17 @@ 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::shared_ptr<StructType>> PartitionSpec::PartitionType(
+    std::shared_ptr<Schema> schema) {
   if (fields_.empty()) {
     return nullptr;
   }
@@ -75,7 +73,7 @@ Result<std::shared_ptr<StructType>> PartitionSpec::PartitionType() {
   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

src/iceberg/partition_spec.h

@@ -23,6 +23,7 @@
 /// Partition specs for Iceberg tables.
 
 #include <cstdint>
+#include <memory>
 #include <mutex>
 #include <optional>
 #include <span>
@@ -32,6 +33,7 @@
 #include "iceberg/iceberg_export.h"
 #include "iceberg/partition_field.h"
 #include "iceberg/result.h"
+#include "iceberg/schema.h"
 #include "iceberg/util/formattable.h"
 
 namespace iceberg {
@@ -56,24 +58,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();
+  Result<std::shared_ptr<StructType>> PartitionType(std::shared_ptr<Schema> schema);
 
   std::string ToString() const override;
 
@@ -87,7 +85,6 @@ 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_;

src/iceberg/table_scan.cc

@@ -269,12 +269,17 @@ 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 table_schema, context_.table_metadata->Schema());
+  ICEBERG_ASSIGN_OR_RAISE(auto partition_type,
+                          partition_spec->PartitionType(table_schema));
 
   for (const auto& manifest_file : manifest_files) {
     ICEBERG_ASSIGN_OR_RAISE(
         auto manifest_reader,
-        ManifestReader::Make(manifest_file, file_io_, partition_schema));
+        ManifestReader::Make(manifest_file, file_io_, partition_type));
     ICEBERG_ASSIGN_OR_RAISE(auto manifests, manifest_reader->Entries());
 
     // TODO(gty404): filter manifests using partition spec and filter expression