feat: add PartitionSpec::Make and PartitionSpec::Validate

Author: HeartLinked

src/iceberg/partition_spec.cc

@@ -20,15 +20,19 @@
 #include "iceberg/partition_spec.h"
 
 #include <algorithm>
+#include <cstdint>
 #include <format>
 #include <memory>
 #include <ranges>
+#include <unordered_map>
 
+#include "iceberg/result.h"
 #include "iceberg/schema.h"
 #include "iceberg/schema_field.h"
 #include "iceberg/transform.h"
 #include "iceberg/util/formatter.h"  // IWYU pragma: keep
 #include "iceberg/util/macros.h"
+#include "iceberg/util/type_util.h"
 
 namespace iceberg {
 
@@ -104,4 +108,68 @@ bool PartitionSpec::Equals(const PartitionSpec& other) const {
   return spec_id_ == other.spec_id_ && fields_ == other.fields_;
 }
 
+Status PartitionSpec::Validate(const Schema& schema, bool allowMissingFields) const {
+  std::unordered_map<int32_t, int32_t> parents = indexParents(schema);
+  for (const auto& partition_field : fields_) {
+    ICEBERG_ASSIGN_OR_RAISE(auto source_field,
+                            schema.FindFieldById(partition_field.source_id()));
+    // In the case the underlying field is dropped, we cannot check if they are compatible
+    if (allowMissingFields && !source_field.has_value()) {
+      continue;
+    }
+    auto field_transform = partition_field.transform();
+
+    // In the case of a Version 1 partition-spec field gets deleted, it is replaced with a
+    // void transform, see: https://iceberg.apache.org/spec/#partition-transforms. We
+    // don't care about the source type since a VoidTransform is always compatible and
+    // skip the checks
+    if (field_transform->transform_type() != TransformType::kVoid) {
+      if (!source_field.has_value()) {
+        return InvalidArgument("Cannot find source column for partition field: {}",
+                               partition_field);
+      }
+      const auto& source_type = source_field.value().get().type();
+      if (!partition_field.transform()->CanTransform(*source_type)) {
+        return InvalidArgument("Invalid source type {} for transform {}",
+                               source_type->ToString(),
+                               partition_field.transform()->ToString());
+      }
+
+      // The only valid parent types for a PartitionField are StructTypes. This must be
+      // checked recursively.
+      auto parent_id_iter = parents.find(partition_field.source_id());
+      while (parent_id_iter != parents.end()) {
+        int32_t parent_id = parent_id_iter->second;
+        ICEBERG_ASSIGN_OR_RAISE(auto parent_field, schema.FindFieldById(parent_id));
+        if (!parent_field.has_value()) {
+          return InvalidArgument("Cannot find parent field with ID: {}", parent_id);
+        }
+        const auto& parent_type = parent_field.value().get().type();
+        if (parent_type->type_id() != TypeId::kStruct) {
+          return InvalidArgument("Invalid partition field parent type: {}",
+                                 parent_type->ToString());
+        }
+        parent_id_iter = parents.find(parent_id);
+      }
+    }
+  }
+  return {};
+}
+
+Result<std::unique_ptr<PartitionSpec>> PartitionSpec::Make(
+    const Schema& schema, int32_t spec_id, std::vector<PartitionField> fields,
+    std::optional<int32_t> last_assigned_field_id, bool allowMissingFields) {
+  auto partition_spec =
+      std::make_unique<PartitionSpec>(spec_id, std::move(fields), last_assigned_field_id);
+  ICEBERG_RETURN_UNEXPECTED(partition_spec->Validate(schema, allowMissingFields));
+  return partition_spec;
+}
+
+Result<std::unique_ptr<PartitionSpec>> PartitionSpec::Make(
+    int32_t spec_id, std::vector<PartitionField> fields,
+    std::optional<int32_t> last_assigned_field_id) {
+  return std::make_unique<PartitionSpec>(spec_id, std::move(fields),
+                                         last_assigned_field_id);
+}
+
 }  // namespace iceberg

src/iceberg/partition_spec.h

@@ -80,6 +80,38 @@ class ICEBERG_EXPORT PartitionSpec : public util::Formattable {
     return lhs.Equals(rhs);
   }
 
+  /// \brief Validates the partition spec against a schema.
+  /// \param schema The schema to validate against.
+  /// \param allowMissingFields Whether to skip validation for partition fields whose
+  /// source columns have been dropped from the schema.
+  /// \return Error status if the partition spec is invalid.
+  Status Validate(const Schema& schema, bool allowMissingFields) const;
+
+  /// \brief Create a PartitionSpec binding to a schema.
+  /// \param schema The schema to bind the partition spec to.
+  /// \param spec_id The spec ID.
+  /// \param fields The partition fields.
+  /// \param last_assigned_field_id The last assigned field ID assigned to ensure new
+  /// fields get unique IDs.
+  /// \param allowMissingFields Whether to skip validation for partition fields whose
+  /// source columns have been dropped from the schema.
+  /// \return A Result containing the partition spec or an error.
+  static Result<std::unique_ptr<PartitionSpec>> Make(
+      const Schema& schema, int32_t spec_id, std::vector<PartitionField> fields,
+      std::optional<int32_t> last_assigned_field_id = std::nullopt,
+      bool allowMissingFields = false);
+
+  /// \brief Create a PartitionSpec without binding to a schema.
+  /// \param spec_id The spec ID.
+  /// \param fields The partition fields.
+  /// \param last_assigned_field_id The last assigned field ID assigned to ensure new
+  /// fields get unique IDs.
+  /// \return A Result containing the partition spec or an error.
+  /// \note This method does not check whether the sort fields are valid for any schema.
+  static Result<std::unique_ptr<PartitionSpec>> Make(
+      int32_t spec_id, std::vector<PartitionField> fields,
+      std::optional<int32_t> last_assigned_field_id = std::nullopt);
+
  private:
   /// \brief Compare two partition specs for equality.
   bool Equals(const PartitionSpec& other) const;

src/iceberg/sort_order.cc

@@ -90,7 +90,7 @@ bool SortOrder::Equals(const SortOrder& other) const {
 Status SortOrder::Validate(const Schema& schema) const {
   for (const auto& field : fields_) {
     ICEBERG_ASSIGN_OR_RAISE(auto schema_field, schema.FindFieldById(field.source_id()));
-    if (!schema_field.has_value() || schema_field == std::nullopt) {
+    if (!schema_field.has_value()) {
       return InvalidArgument("Cannot find source column for sort field: {}", field);
     }
 

src/iceberg/sort_order.h

@@ -88,7 +88,6 @@ class ICEBERG_EXPORT SortOrder : public util::Formattable {
   /// \param fields The sort fields.
   /// \return A Result containing the SortOrder or an error.
   /// \note This method does not check whether the sort fields are valid for any schema.
-  /// Use IsBoundToSchema to check if the sort order is valid for a given schema.
   static Result<std::unique_ptr<SortOrder>> Make(int32_t sort_id,
                                                  std::vector<SortField> fields);
 

src/iceberg/test/partition_spec_test.cc

@@ -152,4 +152,174 @@ TEST(PartitionSpecTest, PartitionTypeTest) {
   EXPECT_EQ(pt_field3, partition_type.value()->fields()[2]);
 }
 
+TEST(PartitionSpecTest, MissingSourceColumn) {
+  SchemaField field1(1, "id", int64(), false);
+  SchemaField field2(2, "ts", timestamp(), false);
+  SchemaField field3(6, "s", string(), false);
+  Schema schema({field1, field2, field3}, Schema::kInitialSchemaId);
+
+  // Try to create partition field with non-existent source column ID 999
+  PartitionField pt_field_invalid(999, 1000, "missing_partition", Transform::Identity());
+
+  auto result = PartitionSpec::Make(schema, 1, {pt_field_invalid});
+  EXPECT_FALSE(result.has_value());
+  EXPECT_THAT(result.error().message,
+              ::testing::HasSubstr("Cannot find source column for partition field"));
+}
+
+TEST(PartitionSpecTest, InvalidTransformForType) {
+  // Test Day transform on string type (should fail)
+  SchemaField field_string(6, "s", string(), false);
+  Schema schema_string({field_string}, Schema::kInitialSchemaId);
+
+  PartitionField pt_field_invalid(6, 1005, "s_day", Transform::Day());
+
+  auto result_string = PartitionSpec::Make(schema_string, 1, {pt_field_invalid});
+  EXPECT_FALSE(result_string.has_value());
+  EXPECT_THAT(result_string.error().message, ::testing::HasSubstr("Invalid source type"));
+}
+
+TEST(PartitionSpecTest, SourceIdNotFound) {
+  SchemaField field1(1, "id", int64(), false);
+  SchemaField field2(2, "ts", timestamp(), false);
+  Schema schema({field1, field2}, Schema::kInitialSchemaId);
+
+  // Try to create partition field with source ID 99 which doesn't exist
+  PartitionField pt_field_invalid(99, 1000, "Test", Transform::Identity());
+
+  auto result = PartitionSpec::Make(schema, 1, {pt_field_invalid});
+  EXPECT_FALSE(result.has_value());
+  EXPECT_THAT(result.error().message,
+              ::testing::HasSubstr("Cannot find source column for partition field"));
+}
+
+TEST(PartitionSpecTest, PartitionFieldInStruct) {
+  SchemaField field1(1, "id", int64(), false);
+  SchemaField field2(2, "ts", timestamp(), false);
+  Schema base_schema({field1, field2}, Schema::kInitialSchemaId);
+
+  // Create a struct that contains the base schema fields
+  auto struct_type =
+      std::make_shared<StructType>(std::vector<SchemaField>{field1, field2});
+  SchemaField outer_struct(11, "MyStruct", struct_type, false);
+
+  Schema schema({outer_struct}, Schema::kInitialSchemaId);
+
+  // Partition on a field within the struct (id field with ID 1)
+  PartitionField pt_field(1, 1000, "id_partition", Transform::Identity());
+
+  auto result = PartitionSpec::Make(schema, 1, {pt_field});
+  EXPECT_TRUE(result.has_value()) << result.error().message;
+}
+
+TEST(PartitionSpecTest, PartitionFieldInStructInStruct) {
+  SchemaField field1(1, "id", int64(), false);
+  SchemaField field2(2, "ts", timestamp(), false);
+
+  // Create inner struct
+  auto inner_struct =
+      std::make_shared<StructType>(std::vector<SchemaField>{field1, field2});
+  SchemaField inner_field(11, "Inner", inner_struct, false);
+
+  // Create outer struct containing inner struct
+  auto outer_struct = std::make_shared<StructType>(std::vector<SchemaField>{inner_field});
+  SchemaField outer_field(12, "Outer", outer_struct, true);
+
+  Schema schema({outer_field}, Schema::kInitialSchemaId);
+
+  // Partition on a field deep within nested structs (id field with ID 1)
+  PartitionField pt_field(1, 1000, "id_partition", Transform::Identity());
+
+  auto result = PartitionSpec::Make(schema, 1, {pt_field});
+  EXPECT_TRUE(result.has_value()) << result.error().message;
+}
+
+TEST(PartitionSpecTest, PartitionFieldInList) {
+  auto list_type = std::make_shared<ListType>(1, int32(), /*element_required=*/false);
+  SchemaField list_field(2, "MyList", list_type, false);
+
+  Schema schema({list_field}, Schema::kInitialSchemaId);
+
+  // Try to partition on the list element field (field ID 1 is the element)
+  PartitionField pt_field(1, 1000, "element_partition", Transform::Identity());
+
+  auto result = PartitionSpec::Make(schema, 1, {pt_field});
+  EXPECT_FALSE(result.has_value());
+  EXPECT_THAT(result.error().message,
+              ::testing::HasSubstr("Invalid partition field parent"));
+}
+
+TEST(PartitionSpecTest, PartitionFieldInStructInList) {
+  auto struct_in_list = std::make_shared<StructType>(
+      std::vector<SchemaField>{SchemaField(1, "Foo", int32(), true)});
+  auto list_type = std::make_shared<ListType>(2, struct_in_list,
+                                              /*element_required=*/false);
+  SchemaField list_field(3, "MyList", list_type, false);
+
+  Schema schema({list_field}, Schema::kInitialSchemaId);
+
+  // Try to partition on a field within a struct inside a list (Foo field with ID 1)
+  PartitionField pt_field(1, 1000, "foo_partition", Transform::Identity());
+
+  auto result = PartitionSpec::Make(schema, 1, {pt_field});
+  EXPECT_FALSE(result.has_value());
+  EXPECT_THAT(result.error().message,
+              ::testing::HasSubstr("Invalid partition field parent"));
+}
+
+TEST(PartitionSpecTest, PartitionFieldInMap) {
+  SchemaField key_field(1, "key", int32(), false);
+  SchemaField value_field(2, "value", int32(), false);
+  auto map_type = std::make_shared<MapType>(key_field, value_field);
+  SchemaField map_field(3, "MyMap", map_type, false);
+
+  Schema schema({map_field}, Schema::kInitialSchemaId);
+
+  // Try to partition on the map key field (field ID 1 is the key)
+  PartitionField pt_field_key(1, 1000, "key_partition", Transform::Identity());
+
+  auto result_key = PartitionSpec::Make(schema, 1, {pt_field_key});
+  EXPECT_FALSE(result_key.has_value());
+  EXPECT_THAT(result_key.error().message,
+              ::testing::HasSubstr("Invalid partition field parent"));
+
+  // Try to partition on the map value field (field ID 2 is the value)
+  PartitionField pt_field_value(2, 1001, "value_partition", Transform::Identity());
+
+  auto result_value = PartitionSpec::Make(schema, 1, {pt_field_value});
+  EXPECT_FALSE(result_value.has_value());
+  EXPECT_THAT(result_value.error().message,
+              ::testing::HasSubstr("Invalid partition field parent"));
+}
+
+TEST(PartitionSpecTest, PartitionFieldInStructInMap) {
+  auto struct_key = std::make_shared<StructType>(
+      std::vector<SchemaField>{SchemaField(1, "Foo", int32(), true)});
+  auto struct_value = std::make_shared<StructType>(
+      std::vector<SchemaField>{SchemaField(2, "Bar", int32(), true)});
+
+  SchemaField key_field(3, "key", struct_key, false);
+  SchemaField value_field(4, "value", struct_value, false);
+  auto map_type = std::make_shared<MapType>(key_field, value_field);
+  SchemaField map_field(5, "MyMap", map_type, false);
+
+  Schema schema({map_field}, Schema::kInitialSchemaId);
+
+  // Try to partition on a field within the key struct (Foo field with ID 1)
+  PartitionField pt_field_key(1, 1000, "foo_partition", Transform::Identity());
+
+  auto result_key = PartitionSpec::Make(schema, 1, {pt_field_key});
+  EXPECT_FALSE(result_key.has_value());
+  EXPECT_THAT(result_key.error().message,
+              ::testing::HasSubstr("Invalid partition field parent"));
+
+  // Try to partition on a field within the value struct (Bar field with ID 2)
+  PartitionField pt_field_value(2, 1001, "bar_partition", Transform::Identity());
+
+  auto result_value = PartitionSpec::Make(schema, 1, {pt_field_value});
+  EXPECT_FALSE(result_value.has_value());
+  EXPECT_THAT(result_value.error().message,
+              ::testing::HasSubstr("Invalid partition field parent"));
+}
+
 }  // namespace iceberg