Add _partition as well

Author: Fokko

crates/iceberg/src/metadata_columns.rs

@@ -41,6 +41,9 @@ pub const RESERVED_FIELD_ID_DELETED: i32 = i32::MAX - 3;
 /// Reserved field ID for the spec ID (_spec_id) column per Iceberg spec
 pub const RESERVED_FIELD_ID_SPEC_ID: i32 = i32::MAX - 4;
 
+/// Reserved field ID for the partition (_partition) column per Iceberg spec
+pub const RESERVED_FIELD_ID_PARTITION: i32 = i32::MAX - 5;
+
 /// Reserved field ID for the file path in position delete files
 pub const RESERVED_FIELD_ID_DELETE_FILE_PATH: i32 = i32::MAX - 101;
 
@@ -74,6 +77,9 @@ pub const RESERVED_COL_NAME_DELETED: &str = "_deleted";
 /// Reserved column name for the spec ID metadata column
 pub const RESERVED_COL_NAME_SPEC_ID: &str = "_spec_id";
 
+/// Reserved column name for the partition metadata column
+pub const RESERVED_COL_NAME_PARTITION: &str = "_partition";
+
 /// Reserved column name for the file path in position delete files
 pub const RESERVED_COL_NAME_DELETE_FILE_PATH: &str = "file_path";
 
@@ -328,8 +334,57 @@ pub fn last_updated_sequence_number_field() -> &'static NestedFieldRef {
     &LAST_UPDATED_SEQUENCE_NUMBER_FIELD
 }
 
+/// Creates the Iceberg field definition for the _partition metadata column.
+///
+/// The _partition field is a struct whose fields depend on the partition spec.
+/// This function creates the field dynamically with the provided partition fields.
+///
+/// # Arguments
+/// * `partition_fields` - The fields that make up the partition struct
+///
+/// # Returns
+/// A new _partition field definition as an Iceberg NestedField
+///
+/// # Example
+/// ```
+/// use std::sync::Arc;
+///
+/// use iceberg::metadata_columns::partition_field;
+/// use iceberg::spec::{NestedField, PrimitiveType, Type};
+///
+/// let fields = vec![
+///     Arc::new(NestedField::required(
+///         1,
+///         "year",
+///         Type::Primitive(PrimitiveType::Int),
+///     )),
+///     Arc::new(NestedField::required(
+///         2,
+///         "month",
+///         Type::Primitive(PrimitiveType::Int),
+///     )),
+/// ];
+/// let partition_field = partition_field(fields);
+/// ```
+pub fn partition_field(partition_fields: Vec<NestedFieldRef>) -> NestedFieldRef {
+    use crate::spec::StructType;
+
+    Arc::new(
+        NestedField::required(
+            RESERVED_FIELD_ID_PARTITION,
+            RESERVED_COL_NAME_PARTITION,
+            Type::Struct(StructType::new(partition_fields)),
+        )
+        .with_doc("Partition to which a row belongs"),
+    )
+}
+
 /// Returns the Iceberg field definition for a metadata field ID.
 ///
+/// Note: This function does not support `_partition` (field ID `i32::MAX - 5`) because
+/// it's a struct field that requires dynamic partition fields. Use `partition_field()`
+/// instead to create the `_partition` field with the appropriate partition fields.
+///
 /// # Arguments
 /// * `field_id` - The metadata field ID
 ///
@@ -341,15 +396,17 @@ pub fn get_metadata_field(field_id: i32) -> Result<&'static NestedFieldRef> {
         RESERVED_FIELD_ID_POS => Ok(pos_field()),
         RESERVED_FIELD_ID_DELETED => Ok(deleted_field()),
         RESERVED_FIELD_ID_SPEC_ID => Ok(spec_id_field()),
+        RESERVED_FIELD_ID_PARTITION => Err(Error::new(
+            ErrorKind::Unexpected,
+            "The _partition field must be created using partition_field() with appropriate partition fields",
+        )),
         RESERVED_FIELD_ID_DELETE_FILE_PATH => Ok(delete_file_path_field()),
         RESERVED_FIELD_ID_DELETE_FILE_POS => Ok(delete_file_pos_field()),
         RESERVED_FIELD_ID_CHANGE_TYPE => Ok(change_type_field()),
         RESERVED_FIELD_ID_CHANGE_ORDINAL => Ok(change_ordinal_field()),
         RESERVED_FIELD_ID_COMMIT_SNAPSHOT_ID => Ok(commit_snapshot_id_field()),
         RESERVED_FIELD_ID_ROW_ID => Ok(row_id_field()),
-        RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER => {
-            Ok(last_updated_sequence_number_field())
-        }
+        RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER => Ok(last_updated_sequence_number_field()),
         _ if is_metadata_field(field_id) => {
             // Future metadata fields can be added here
             Err(Error::new(
@@ -379,6 +436,7 @@ pub fn get_metadata_field_id(column_name: &str) -> Result<i32> {
         RESERVED_COL_NAME_POS => Ok(RESERVED_FIELD_ID_POS),
         RESERVED_COL_NAME_DELETED => Ok(RESERVED_FIELD_ID_DELETED),
         RESERVED_COL_NAME_SPEC_ID => Ok(RESERVED_FIELD_ID_SPEC_ID),
+        RESERVED_COL_NAME_PARTITION => Ok(RESERVED_FIELD_ID_PARTITION),
         RESERVED_COL_NAME_DELETE_FILE_PATH => Ok(RESERVED_FIELD_ID_DELETE_FILE_PATH),
         RESERVED_COL_NAME_DELETE_FILE_POS => Ok(RESERVED_FIELD_ID_DELETE_FILE_POS),
         RESERVED_COL_NAME_CHANGE_TYPE => Ok(RESERVED_FIELD_ID_CHANGE_TYPE),
@@ -409,6 +467,7 @@ pub fn is_metadata_field(field_id: i32) -> bool {
             | RESERVED_FIELD_ID_POS
             | RESERVED_FIELD_ID_DELETED
             | RESERVED_FIELD_ID_SPEC_ID
+            | RESERVED_FIELD_ID_PARTITION
             | RESERVED_FIELD_ID_DELETE_FILE_PATH
             | RESERVED_FIELD_ID_DELETE_FILE_POS
             | RESERVED_FIELD_ID_CHANGE_TYPE
@@ -429,3 +488,85 @@ pub fn is_metadata_field(field_id: i32) -> bool {
 pub fn is_metadata_column_name(column_name: &str) -> bool {
     get_metadata_field_id(column_name).is_ok()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::spec::PrimitiveType;
+
+    #[test]
+    fn test_partition_field_creation() {
+        // Create partition fields for a hypothetical year/month partition
+        let partition_fields = vec![
+            Arc::new(NestedField::required(
+                1000,
+                "year",
+                Type::Primitive(PrimitiveType::Int),
+            )),
+            Arc::new(NestedField::required(
+                1001,
+                "month",
+                Type::Primitive(PrimitiveType::Int),
+            )),
+        ];
+
+        // Create the _partition metadata field
+        let partition = partition_field(partition_fields);
+
+        // Verify field properties
+        assert_eq!(partition.id, RESERVED_FIELD_ID_PARTITION);
+        assert_eq!(partition.name, RESERVED_COL_NAME_PARTITION);
+        assert!(partition.required);
+
+        // Verify it's a struct type with correct fields
+        if let Type::Struct(struct_type) = partition.field_type.as_ref() {
+            assert_eq!(struct_type.fields().len(), 2);
+            assert_eq!(struct_type.fields()[0].name, "year");
+            assert_eq!(struct_type.fields()[1].name, "month");
+        } else {
+            panic!("Expected struct type for _partition field");
+        }
+    }
+
+    #[test]
+    fn test_partition_field_id_recognized() {
+        assert!(is_metadata_field(RESERVED_FIELD_ID_PARTITION));
+    }
+
+    #[test]
+    fn test_partition_field_name_recognized() {
+        assert_eq!(
+            get_metadata_field_id(RESERVED_COL_NAME_PARTITION).unwrap(),
+            RESERVED_FIELD_ID_PARTITION
+        );
+    }
+
+    #[test]
+    fn test_get_metadata_field_returns_error_for_partition() {
+        // partition field requires dynamic creation, so get_metadata_field should return an error
+        let result = get_metadata_field(RESERVED_FIELD_ID_PARTITION);
+        assert!(result.is_err());
+        assert!(
+            result
+                .unwrap_err()
+                .to_string()
+                .contains("partition_field()")
+        );
+    }
+
+    #[test]
+    fn test_all_metadata_field_ids() {
+        // Test that all non-partition metadata fields can be retrieved
+        assert!(get_metadata_field(RESERVED_FIELD_ID_FILE).is_ok());
+        assert!(get_metadata_field(RESERVED_FIELD_ID_POS).is_ok());
+        assert!(get_metadata_field(RESERVED_FIELD_ID_DELETED).is_ok());
+        assert!(get_metadata_field(RESERVED_FIELD_ID_SPEC_ID).is_ok());
+        assert!(get_metadata_field(RESERVED_FIELD_ID_DELETE_FILE_PATH).is_ok());
+        assert!(get_metadata_field(RESERVED_FIELD_ID_DELETE_FILE_POS).is_ok());
+        assert!(get_metadata_field(RESERVED_FIELD_ID_CHANGE_TYPE).is_ok());
+        assert!(get_metadata_field(RESERVED_FIELD_ID_CHANGE_ORDINAL).is_ok());
+        assert!(get_metadata_field(RESERVED_FIELD_ID_COMMIT_SNAPSHOT_ID).is_ok());
+        assert!(get_metadata_field(RESERVED_FIELD_ID_ROW_ID).is_ok());
+        assert!(get_metadata_field(RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER).is_ok());
+    }
+}