diff --git a/src/iceberg/json_serde.cc b/src/iceberg/json_serde.cc index c1c3ffb73..eec3f55bd 100644 --- a/src/iceberg/json_serde.cc +++ b/src/iceberg/json_serde.cc @@ -1013,15 +1013,22 @@ nlohmann::json ToJson(const PartitionField& partition_field) { return json; } -nlohmann::json ToJson(const PartitionSpec& partition_spec) { - nlohmann::json json; - json[kSpecId] = partition_spec.spec_id(); - +nlohmann::json PartitionFieldsToJson(const PartitionSpec& partition_spec) { nlohmann::json fields_json = nlohmann::json::array(); for (const auto& field : partition_spec.fields()) { fields_json.push_back(ToJson(field)); } - json[kFields] = fields_json; + return fields_json; +} + +Result PartitionFieldsToJsonString(const PartitionSpec& partition_spec) { + return ToJsonString(PartitionFieldsToJson(partition_spec)); +} + +nlohmann::json ToJson(const PartitionSpec& partition_spec) { + nlohmann::json json; + json[kSpecId] = partition_spec.spec_id(); + json[kFields] = PartitionFieldsToJson(partition_spec); return json; } diff --git a/src/iceberg/json_serde_internal.h b/src/iceberg/json_serde_internal.h index 2465df140..ea71dfdda 100644 --- a/src/iceberg/json_serde_internal.h +++ b/src/iceberg/json_serde_internal.h @@ -158,6 +158,25 @@ ICEBERG_EXPORT nlohmann::json ToJson(const PartitionField& partition_field); ICEBERG_EXPORT Result> PartitionFieldFromJson( const nlohmann::json& json, bool allow_field_id_missing = false); +/// \brief Serializes only the fields of a `PartitionSpec` to a JSON array. +/// +/// Unlike `ToJson(PartitionSpec)`, the result does not include the spec ID. This is +/// the representation required for the `partition-spec` key-value metadata of manifest +/// files. Each `PartitionField` is serialized as described in the +/// `ToJson(PartitionField)` function. +/// +/// \param partition_spec The `PartitionSpec` whose fields are to be serialized. +/// \return A JSON array of the `PartitionField` objects of the `PartitionSpec`. +ICEBERG_EXPORT nlohmann::json PartitionFieldsToJson(const PartitionSpec& partition_spec); + +/// \brief Serializes only the fields of a `PartitionSpec` to a JSON array string. +/// +/// \param partition_spec The `PartitionSpec` whose fields are to be serialized. +/// \return A JSON array string of the `PartitionField` objects of the +/// `PartitionSpec`, as described in `PartitionFieldsToJson(PartitionSpec)`. +ICEBERG_EXPORT Result PartitionFieldsToJsonString( + const PartitionSpec& partition_spec); + /// \brief Serializes a `PartitionSpec` object to JSON. /// /// This function converts a `PartitionSpec` object into a JSON representation. diff --git a/src/iceberg/manifest/v1_metadata.cc b/src/iceberg/manifest/v1_metadata.cc index cb3b895b8..23416f64a 100644 --- a/src/iceberg/manifest/v1_metadata.cc +++ b/src/iceberg/manifest/v1_metadata.cc @@ -76,7 +76,9 @@ std::shared_ptr ManifestEntryAdapterV1::DataFileSchema( Status ManifestEntryAdapterV1::Init() { ICEBERG_ASSIGN_OR_RAISE(metadata_["schema"], ToJsonString(*current_schema_)) - ICEBERG_ASSIGN_OR_RAISE(metadata_["partition-spec"], ToJsonString(*partition_spec_)); + metadata_["schema-id"] = std::to_string(current_schema_->schema_id()); + ICEBERG_ASSIGN_OR_RAISE(metadata_["partition-spec"], + PartitionFieldsToJsonString(*partition_spec_)); metadata_["partition-spec-id"] = std::to_string(partition_spec_->spec_id()); metadata_["format-version"] = "1"; diff --git a/src/iceberg/manifest/v2_metadata.cc b/src/iceberg/manifest/v2_metadata.cc index 2e1e9e7c5..6c7a5c760 100644 --- a/src/iceberg/manifest/v2_metadata.cc +++ b/src/iceberg/manifest/v2_metadata.cc @@ -75,7 +75,9 @@ std::shared_ptr ManifestEntryAdapterV2::DataFileType( Status ManifestEntryAdapterV2::Init() { ICEBERG_ASSIGN_OR_RAISE(metadata_["schema"], ToJsonString(*current_schema_)) - ICEBERG_ASSIGN_OR_RAISE(metadata_["partition-spec"], ToJsonString(*partition_spec_)); + metadata_["schema-id"] = std::to_string(current_schema_->schema_id()); + ICEBERG_ASSIGN_OR_RAISE(metadata_["partition-spec"], + PartitionFieldsToJsonString(*partition_spec_)); metadata_["partition-spec-id"] = std::to_string(partition_spec_->spec_id()); metadata_["format-version"] = "2"; metadata_["content"] = content_ == ManifestContent::kData ? "data" : "delete"; diff --git a/src/iceberg/manifest/v3_metadata.cc b/src/iceberg/manifest/v3_metadata.cc index ad577b240..396da7acb 100644 --- a/src/iceberg/manifest/v3_metadata.cc +++ b/src/iceberg/manifest/v3_metadata.cc @@ -86,7 +86,9 @@ std::shared_ptr ManifestEntryAdapterV3::DataFileType( Status ManifestEntryAdapterV3::Init() { ICEBERG_ASSIGN_OR_RAISE(metadata_["schema"], ToJsonString(*current_schema_)) - ICEBERG_ASSIGN_OR_RAISE(metadata_["partition-spec"], ToJsonString(*partition_spec_)); + metadata_["schema-id"] = std::to_string(current_schema_->schema_id()); + ICEBERG_ASSIGN_OR_RAISE(metadata_["partition-spec"], + PartitionFieldsToJsonString(*partition_spec_)); metadata_["partition-spec-id"] = std::to_string(partition_spec_->spec_id()); metadata_["format-version"] = "3"; metadata_["content"] = content_ == ManifestContent::kData ? "data" : "delete"; diff --git a/src/iceberg/test/json_serde_test.cc b/src/iceberg/test/json_serde_test.cc index 51bc2af11..db6bbafbf 100644 --- a/src/iceberg/test/json_serde_test.cc +++ b/src/iceberg/test/json_serde_test.cc @@ -366,6 +366,27 @@ TEST(JsonInternalTest, PartitionSpec) { EXPECT_EQ(*spec, *parsed_spec_result.value()); } +TEST(JsonInternalTest, PartitionFieldsToJson) { + ICEBERG_UNWRAP_OR_FAIL( + auto spec, + PartitionSpec::Make(1, {PartitionField(3, 101, "region", Transform::Identity()), + PartitionField(5, 102, "ts_day", Transform::Day())})); + nlohmann::json expected_fields = R"([ + {"source-id": 3, "field-id": 101, "transform": "identity", "name": "region"}, + {"source-id": 5, "field-id": 102, "transform": "day", "name": "ts_day"} + ])"_json; + + // Only the fields array, as used for the manifest `partition-spec` metadata. + EXPECT_EQ(PartitionFieldsToJson(*spec), expected_fields); + EXPECT_EQ(PartitionFieldsToJson(*PartitionSpec::Unpartitioned()), + nlohmann::json::array()); + + // The general serializer still produces the full spec object. + EXPECT_EQ(ToJson(*spec), nlohmann::json({{"spec-id", 1}, {"fields", expected_fields}})); + EXPECT_EQ(ToJson(*PartitionSpec::Unpartitioned()), + R"({"spec-id": 0, "fields": []})"_json); +} + TEST(JsonInternalTest, SortOrderFromJson) { auto identity_transform = Transform::Identity(); SortField st1(5, identity_transform, SortDirection::kAscending, NullOrder::kFirst); diff --git a/src/iceberg/test/manifest_writer_versions_test.cc b/src/iceberg/test/manifest_writer_versions_test.cc index 24669d5b5..038753f28 100644 --- a/src/iceberg/test/manifest_writer_versions_test.cc +++ b/src/iceberg/test/manifest_writer_versions_test.cc @@ -24,11 +24,13 @@ #include #include +#include #include "iceberg/arrow/arrow_io_util.h" #include "iceberg/avro/avro_register.h" #include "iceberg/constants.h" #include "iceberg/file_format.h" +#include "iceberg/file_reader.h" #include "iceberg/manifest/manifest_entry.h" #include "iceberg/manifest/manifest_list.h" #include "iceberg/manifest/manifest_reader.h" @@ -624,4 +626,88 @@ TEST_F(ManifestWriterVersionsTest, TestV3ManifestRewriteWithInheritance) { DataFile::Content::kData, kFirstRowId); } +class ManifestWriterHeaderTest : public ManifestWriterVersionsTest, + public ::testing::WithParamInterface { + protected: + /// \brief Write a manifest with a single data file and return the key-value + /// metadata stored in the header of the written Avro file. + std::unordered_map WriteManifestAndReadHeader( + std::shared_ptr spec, std::shared_ptr schema, + PartitionValues partition) { + auto data_file = CreateDataFile(kFirstRowId); + data_file->partition = std::move(partition); + + const std::string manifest_path = CreateManifestPath(); + auto writer_result = ManifestWriter::MakeWriter( + GetParam(), kSnapshotId, manifest_path, file_io_, std::move(spec), + std::move(schema), ManifestContent::kData, kFirstRowId); + EXPECT_THAT(writer_result, IsOk()); + auto writer = std::move(writer_result.value()); + EXPECT_THAT(writer->WriteAddedEntry(std::move(data_file)), IsOk()); + EXPECT_THAT(writer->Close(), IsOk()); + + auto manifest_result = writer->ToManifestFile(); + EXPECT_THAT(manifest_result, IsOk()); + + auto reader_result = ReaderFactoryRegistry::Open( + FileFormatType::kAvro, + {.path = manifest_result->manifest_path, + .length = static_cast(manifest_result->manifest_length), + .io = file_io_, + .projection = std::make_shared(std::vector{})}); + EXPECT_THAT(reader_result, IsOk()); + auto metadata_result = reader_result.value()->Metadata(); + EXPECT_THAT(metadata_result, IsOk()); + return std::move(metadata_result.value()); + } +}; + +TEST_P(ManifestWriterHeaderTest, UnpartitionedManifestHeader) { + auto metadata = WriteManifestAndReadHeader(PartitionSpec::Unpartitioned(), schema_, + PartitionValues{}); + + // The header holds only the partition fields array, not the spec object. + ASSERT_TRUE(metadata.contains("partition-spec")); + EXPECT_EQ(metadata.at("partition-spec"), "[]"); + EXPECT_EQ(metadata.at("partition-spec-id"), "0"); + + // The schema ID is a separate header, not only embedded in the schema JSON. It is + // optional in v1 and required in v2 and later; it is always written. + ASSERT_TRUE(metadata.contains("schema-id")); + EXPECT_EQ(metadata.at("schema-id"), std::to_string(Schema::kInitialSchemaId)); + EXPECT_EQ(metadata.at("format-version"), std::to_string(GetParam())); +} + +TEST_P(ManifestWriterHeaderTest, PartitionedManifestHeader) { + constexpr int32_t kSchemaId = 7; + constexpr int32_t kSpecId = 3; + auto schema = std::make_shared( + std::vector(schema_->fields().begin(), schema_->fields().end()), + kSchemaId); + ICEBERG_UNWRAP_OR_FAIL( + std::shared_ptr spec, + PartitionSpec::Make(kSpecId, + {PartitionField(3, 1000, "category", Transform::Identity()), + PartitionField(2, 1001, "timestamp_hour", Transform::Hour()), + PartitionField(1, 1002, "id_bucket", Transform::Bucket(16))})); + + auto metadata = WriteManifestAndReadHeader(spec, schema, kPartition); + + ASSERT_TRUE(metadata.contains("partition-spec")); + auto partition_fields = nlohmann::json::parse(metadata.at("partition-spec")); + EXPECT_EQ(partition_fields, R"([ + {"name": "category", "transform": "identity", "source-id": 3, "field-id": 1000}, + {"name": "timestamp_hour", "transform": "hour", "source-id": 2, "field-id": 1001}, + {"name": "id_bucket", "transform": "bucket[16]", "source-id": 1, "field-id": 1002} + ])"_json); + EXPECT_EQ(metadata.at("partition-spec-id"), std::to_string(kSpecId)); + + ASSERT_TRUE(metadata.contains("schema-id")); + EXPECT_EQ(metadata.at("schema-id"), std::to_string(kSchemaId)); + EXPECT_EQ(nlohmann::json::parse(metadata.at("schema"))["schema-id"], kSchemaId); +} + +INSTANTIATE_TEST_SUITE_P(ManifestWriterHeaderVersions, ManifestWriterHeaderTest, + ::testing::Values(1, 2, 3)); + } // namespace iceberg