iceberg/spec/

partition.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

/*!
 * Partitioning
 */
use std::sync::Arc;

use itertools::Itertools;
use serde::{Deserialize, Serialize};
use typed_builder::TypedBuilder;

use super::transform::Transform;
use super::{NestedField, PrimitiveType, Schema, SchemaRef, StructType};
use crate::Result;
use crate::error::invalid_data;
use crate::spec::Struct;

pub(crate) const UNPARTITIONED_LAST_ASSIGNED_ID: i32 = 999;
pub(crate) const DEFAULT_PARTITION_SPEC_ID: i32 = 0;

/// Partition fields capture the transform from table data to partition values.
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, TypedBuilder)]
#[serde(rename_all = "kebab-case")]
pub struct PartitionField {
    /// A source column id from the table’s schema
    pub source_id: i32,
    /// A partition field id that is used to identify a partition field and is unique within a partition spec.
    /// In v2 table metadata, it is unique across all partition specs.
    pub field_id: i32,
    /// A partition name.
    pub name: String,
    /// A transform that is applied to the source column to produce a partition value.
    pub transform: Transform,
}

impl PartitionField {
    /// To unbound partition field
    pub fn into_unbound(self) -> UnboundPartitionField {
        self.into()
    }
}

/// Reference to [`PartitionSpec`].
pub type PartitionSpecRef = Arc<PartitionSpec>;
/// Partition spec that defines how to produce a tuple of partition values from a record.
///
/// A [`PartitionSpec`] is originally obtained by binding an [`UnboundPartitionSpec`] to a schema and is
/// only guaranteed to be valid for that schema. The main difference between [`PartitionSpec`] and
/// [`UnboundPartitionSpec`] is that the former has field ids assigned,
/// while field ids are optional for [`UnboundPartitionSpec`].
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
#[serde(rename_all = "kebab-case")]
pub struct PartitionSpec {
    /// Identifier for PartitionSpec
    spec_id: i32,
    /// Details of the partition spec
    fields: Vec<PartitionField>,
}

impl PartitionSpec {
    /// Create a new partition spec builder with the given schema.
    pub fn builder(schema: impl Into<SchemaRef>) -> PartitionSpecBuilder {
        PartitionSpecBuilder::new(schema)
    }

    /// Fields of the partition spec
    pub fn fields(&self) -> &[PartitionField] {
        &self.fields
    }

    /// Spec id of the partition spec
    pub fn spec_id(&self) -> i32 {
        self.spec_id
    }

    /// Get a new unpartitioned partition spec
    pub fn unpartition_spec() -> Self {
        Self {
            spec_id: DEFAULT_PARTITION_SPEC_ID,
            fields: vec![],
        }
    }

    /// Returns if the partition spec is unpartitioned.
    ///
    /// A [`PartitionSpec`] is unpartitioned if it has no fields or all fields are [`Transform::Void`] transform.
    pub fn is_unpartitioned(&self) -> bool {
        self.fields.is_empty() || self.fields.iter().all(|f| f.transform == Transform::Void)
    }

    /// Returns the partition type of this partition spec.
    /// If a source column is absent, preserves fixed transform result types and uses
    /// unknown for result types that depend on the source type.
    pub fn partition_type(&self, schema: &Schema) -> Result<StructType> {
        let mut struct_fields = Vec::with_capacity(self.fields.len());
        for partition_field in &self.fields {
            let res_type = match schema.field_by_id(partition_field.source_id) {
                Some(field) => partition_field.transform.result_type(&field.field_type)?,
                // Historical specs may reference dropped source columns. Retain every
                // field's position and any result type that is independent of its source.
                None => match partition_field.transform {
                    Transform::Bucket(_) | Transform::Year | Transform::Month | Transform::Hour => {
                        PrimitiveType::Int.into()
                    }
                    Transform::Day => PrimitiveType::Date.into(),
                    Transform::Unknown => PrimitiveType::String.into(),
                    Transform::Identity | Transform::Truncate(_) | Transform::Void => {
                        PrimitiveType::Unknown.into()
                    }
                },
            };
            struct_fields.push(
                NestedField::optional(partition_field.field_id, &partition_field.name, res_type)
                    .into(),
            );
        }
        Ok(StructType::new(struct_fields))
    }

    /// Convert to unbound partition spec
    pub fn into_unbound(self) -> UnboundPartitionSpec {
        self.into()
    }

    /// Change the spec id of the partition spec
    pub fn with_spec_id(self, spec_id: i32) -> Self {
        Self { spec_id, ..self }
    }

    /// Check if this partition spec has sequential partition ids.
    /// Sequential ids start from 1000 and increment by 1 for each field.
    /// This is required for spec version 1
    pub fn has_sequential_ids(&self) -> bool {
        has_sequential_ids(self.fields.iter().map(|f| f.field_id))
    }

    /// Get the highest field id in the partition spec.
    pub fn highest_field_id(&self) -> Option<i32> {
        self.fields.iter().map(|f| f.field_id).max()
    }

    /// Check if this partition spec is compatible with another partition spec.
    ///
    /// Returns true if the partition spec is equal to the other spec with partition field ids ignored and
    /// spec_id ignored. The following must be identical:
    /// * The number of fields
    /// * Field order
    /// * Field names
    /// * Source column ids
    /// * Transforms
    pub fn is_compatible_with(&self, other: &PartitionSpec) -> bool {
        if self.fields.len() != other.fields.len() {
            return false;
        }

        for (this_field, other_field) in self.fields.iter().zip(other.fields.iter()) {
            if this_field.source_id != other_field.source_id
                || this_field.name != other_field.name
                || this_field.transform != other_field.transform
            {
                return false;
            }
        }

        true
    }

    /// Returns partition path string containing partition type and partition
    /// value as key-value pairs.
    pub fn partition_to_path(&self, data: &Struct, schema: SchemaRef) -> String {
        let partition_type = self.partition_type(&schema).unwrap();
        let field_types = partition_type.fields();

        self.fields
            .iter()
            .enumerate()
            .map(|(i, field)| {
                let value = data[i].as_ref();
                form_urlencoded::Serializer::new(String::new())
                    .append_pair(
                        &field.name,
                        &field
                            .transform
                            .to_human_string(&field_types[i].field_type, value),
                    )
                    .finish()
            })
            .join("/")
    }
}

/// A partition key represents a specific partition in a table, containing the partition spec,
/// schema, and the actual partition values.
#[derive(Clone, Debug)]
pub struct PartitionKey {
    /// The partition spec that contains the partition fields.
    spec: PartitionSpec,
    /// The schema to which the partition spec is bound.
    schema: SchemaRef,
    /// Partition fields' values in struct.
    data: Struct,
}

impl PartitionKey {
    /// Creates a new partition key with the given spec, schema, and data.
    pub fn new(spec: PartitionSpec, schema: SchemaRef, data: Struct) -> Self {
        Self { spec, schema, data }
    }

    /// Creates a new partition key from another partition key, with a new data field.
    pub fn copy_with_data(&self, data: Struct) -> Self {
        Self {
            spec: self.spec.clone(),
            schema: self.schema.clone(),
            data,
        }
    }

    /// Generates a partition path based on the partition values.
    pub fn to_path(&self) -> String {
        self.spec.partition_to_path(&self.data, self.schema.clone())
    }

    /// Returns `true` if the partition key is absent (`None`)
    /// or represents an unpartitioned spec.
    pub fn is_effectively_none(partition_key: Option<&PartitionKey>) -> bool {
        match partition_key {
            None => true,
            Some(pk) => pk.spec.is_unpartitioned(),
        }
    }

    /// Returns the associated [`PartitionSpec`].
    pub fn spec(&self) -> &PartitionSpec {
        &self.spec
    }

    /// Returns the associated [`SchemaRef`].
    pub fn schema(&self) -> &SchemaRef {
        &self.schema
    }

    /// Returns the associated [`Struct`].
    pub fn data(&self) -> &Struct {
        &self.data
    }
}

/// Reference to [`UnboundPartitionSpec`].
pub type UnboundPartitionSpecRef = Arc<UnboundPartitionSpec>;
/// Unbound partition field can be built without a schema and later bound to a schema.
///
/// The fields are private so that an instance is always well formed: in particular
/// `source_ids` holds at least one id, which [`UnboundPartitionField::builder`] checks when the
/// field is built.
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, TypedBuilder)]
#[serde(
    try_from = "self::_serde::UnboundPartitionFieldSerde",
    into = "self::_serde::UnboundPartitionFieldSerde"
)]
#[builder(build_method(into = Result<UnboundPartitionField>))]
pub struct UnboundPartitionField {
    /// The source column ids from the table’s schema. A single-argument transform reads one
    /// id; a v3 multi-argument transform reads several.
    source_ids: Vec<i32>,
    /// A partition field id that is used to identify a partition field and is unique within a partition spec.
    /// In v2 table metadata, it is unique across all partition specs.
    #[builder(default, setter(strip_option(fallback = field_id_opt)))]
    field_id: Option<i32>,
    /// A partition name.
    #[builder(setter(into))]
    name: String,
    /// A transform that is applied to the source column to produce a partition value.
    transform: Transform,
}

impl UnboundPartitionField {
    /// The single source column id this field reads.
    ///
    /// Returns an error for a multi-argument field, which reads several columns and therefore
    /// has no single source id. Use [`Self::source_ids`] to handle both shapes.
    pub fn source_id(&self) -> Result<i32> {
        match self.source_ids.as_slice() {
            [source_id] => Ok(*source_id),
            source_ids => Err(invalid_data!(
                "Partition field '{}' reads {} source columns and has no single source id",
                self.name,
                source_ids.len()
            )),
        }
    }

    /// The source column ids this field reads, in order. Never empty.
    pub fn source_ids(&self) -> &[i32] {
        &self.source_ids
    }

    /// The partition field id, when one was assigned.
    pub fn field_id(&self) -> Option<i32> {
        self.field_id
    }

    /// The partition name.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// The transform applied to the source columns to produce a partition value.
    pub fn transform(&self) -> Transform {
        self.transform
    }

    /// Return this field with the given partition field id assigned.
    pub fn with_field_id(self, field_id: i32) -> Self {
        Self {
            field_id: Some(field_id),
            ..self
        }
    }

    fn validate(&self) -> Result<()> {
        if self.source_ids.is_empty() {
            return Err(invalid_data!("Empty source-ids is not allowed"));
        }
        Ok(())
    }
}

impl From<UnboundPartitionField> for Result<UnboundPartitionField> {
    fn from(field: UnboundPartitionField) -> Self {
        field.validate()?;
        Ok(field)
    }
}

mod _serde {
    use serde::{Deserialize, Serialize};

    use super::UnboundPartitionField;
    use crate::Error;
    use crate::error::invalid_data;
    use crate::spec::Transform;

    /// Per the spec a single-argument field carries `source-id` and a multi-argument field
    /// carries `source-ids`. Either spelling is read, but not both at once; the one that
    /// matches the field is written.
    #[derive(Serialize, Deserialize)]
    #[serde(rename_all = "kebab-case")]
    pub(super) struct UnboundPartitionFieldSerde {
        #[serde(default, skip_serializing_if = "Option::is_none")]
        source_id: Option<i32>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        source_ids: Option<Vec<i32>>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        field_id: Option<i32>,
        name: String,
        transform: Transform,
    }

    impl TryFrom<UnboundPartitionFieldSerde> for UnboundPartitionField {
        type Error = Error;

        fn try_from(value: UnboundPartitionFieldSerde) -> Result<Self, Error> {
            let source_ids = match (value.source_id, value.source_ids) {
                (Some(source_id), None) => vec![source_id],
                (None, Some(source_ids)) => source_ids,
                (Some(_), Some(_)) => {
                    return Err(invalid_data!(
                        "source-id and source-ids are mutually exclusive"
                    ));
                }
                (None, None) => {
                    return Err(invalid_data!(
                        "Either `source-id` or `source-ids` must be present"
                    ));
                }
            };

            UnboundPartitionField::builder()
                .source_ids(source_ids)
                .field_id_opt(value.field_id)
                .name(value.name)
                .transform(value.transform)
                .build()
        }
    }

    impl From<UnboundPartitionField> for UnboundPartitionFieldSerde {
        fn from(value: UnboundPartitionField) -> Self {
            let (source_id, source_ids) = match value.source_ids.as_slice() {
                [source_id] => (Some(*source_id), None),
                _ => (None, Some(value.source_ids)),
            };
            Self {
                source_id,
                source_ids,
                field_id: value.field_id,
                name: value.name,
                transform: value.transform,
            }
        }
    }
}

/// Unbound partition spec can be built without a schema and later bound to a schema.
/// They are used to transport schema information as part of the REST specification.
/// The main difference to [`PartitionSpec`] is that the field ids are optional.
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, Default)]
#[serde(rename_all = "kebab-case")]
pub struct UnboundPartitionSpec {
    /// Identifier for PartitionSpec
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) spec_id: Option<i32>,
    /// Details of the partition spec
    pub(crate) fields: Vec<UnboundPartitionField>,
}

impl UnboundPartitionSpec {
    /// Create unbound partition spec builder
    pub fn builder() -> UnboundPartitionSpecBuilder {
        UnboundPartitionSpecBuilder::default()
    }

    /// Bind this unbound partition spec to a schema.
    ///
    /// Unlike [`PartitionSpecBuilder::add_unbound_field`], this accepts redundant fields, such
    /// as two time granularities of one column, because spec evolution produces them.
    pub fn bind(self, schema: impl Into<SchemaRef>) -> Result<PartitionSpec> {
        PartitionSpecBuilder::new_from_unbound(self, schema)?.build()
    }

    /// Spec id of the partition spec
    pub fn spec_id(&self) -> Option<i32> {
        self.spec_id
    }

    /// Fields of the partition spec
    pub fn fields(&self) -> &[UnboundPartitionField] {
        &self.fields
    }

    /// Change the spec id of the partition spec
    pub fn with_spec_id(self, spec_id: i32) -> Self {
        Self {
            spec_id: Some(spec_id),
            ..self
        }
    }
}

fn has_sequential_ids(field_ids: impl Iterator<Item = i32>) -> bool {
    for (index, field_id) in field_ids.enumerate() {
        let expected_id = (UNPARTITIONED_LAST_ASSIGNED_ID as i64)
            .checked_add(1)
            .and_then(|id| id.checked_add(index as i64))
            .unwrap_or(i64::MAX);

        if field_id as i64 != expected_id {
            return false;
        }
    }

    true
}

impl From<PartitionField> for UnboundPartitionField {
    fn from(field: PartitionField) -> Self {
        UnboundPartitionField {
            source_ids: vec![field.source_id],
            field_id: Some(field.field_id),
            name: field.name,
            transform: field.transform,
        }
    }
}

impl From<PartitionSpec> for UnboundPartitionSpec {
    fn from(spec: PartitionSpec) -> Self {
        UnboundPartitionSpec {
            spec_id: Some(spec.spec_id),
            fields: spec.fields.into_iter().map(Into::into).collect(),
        }
    }
}

/// Create a new UnboundPartitionSpec
#[derive(Debug, Default)]
pub struct UnboundPartitionSpecBuilder {
    spec_id: Option<i32>,
    fields: Vec<UnboundPartitionField>,
}

impl UnboundPartitionSpecBuilder {
    /// Create a new partition spec builder with the given schema.
    pub fn new() -> Self {
        Self {
            spec_id: None,
            fields: vec![],
        }
    }

    /// Set the spec id for the partition spec.
    pub fn with_spec_id(mut self, spec_id: i32) -> Self {
        self.spec_id = Some(spec_id);
        self
    }

    /// Add a new partition field to the partition spec from an unbound partition field.
    pub fn add_partition_field(mut self, field: UnboundPartitionField) -> Result<Self> {
        self.check_name_set_and_unique(&field.name)?;
        self.check_for_redundant_partitions(&field.source_ids, &field.transform)?;
        if let Some(partition_field_id) = field.field_id {
            self.check_partition_id_unique(partition_field_id)?;
        }
        self.fields.push(field);
        Ok(self)
    }

    /// Add multiple partition fields to the partition spec.
    pub fn add_partition_fields(
        self,
        fields: impl IntoIterator<Item = UnboundPartitionField>,
    ) -> Result<Self> {
        let mut builder = self;
        for field in fields {
            builder = builder.add_partition_field(field)?;
        }
        Ok(builder)
    }

    /// Build the unbound partition spec.
    pub fn build(self) -> UnboundPartitionSpec {
        UnboundPartitionSpec {
            spec_id: self.spec_id,
            fields: self.fields,
        }
    }
}

/// Create valid partition specs for a given schema.
#[derive(Debug)]
pub struct PartitionSpecBuilder {
    spec_id: Option<i32>,
    last_assigned_field_id: i32,
    fields: Vec<UnboundPartitionField>,
    schema: SchemaRef,
}

impl PartitionSpecBuilder {
    /// Create a new partition spec builder with the given schema.
    pub fn new(schema: impl Into<SchemaRef>) -> Self {
        Self {
            spec_id: None,
            fields: vec![],
            last_assigned_field_id: UNPARTITIONED_LAST_ASSIGNED_ID,
            schema: schema.into(),
        }
    }

    /// Create a new partition spec builder from an existing unbound partition spec.
    pub fn new_from_unbound(
        unbound: UnboundPartitionSpec,
        schema: impl Into<SchemaRef>,
    ) -> Result<Self> {
        let mut builder =
            Self::new(schema).with_spec_id(unbound.spec_id.unwrap_or(DEFAULT_PARTITION_SPEC_ID));

        for field in unbound.fields {
            builder = builder.add_field_allowing_redundancy(field)?;
        }
        Ok(builder)
    }

    /// Set the last assigned field id for the partition spec.
    ///
    /// Set this field when a new partition spec is created for an existing TableMetaData.
    /// As `field_id` must be unique in V2 metadata, this should be set to
    /// the highest field id used previously.
    pub fn with_last_assigned_field_id(mut self, last_assigned_field_id: i32) -> Self {
        self.last_assigned_field_id = last_assigned_field_id;
        self
    }

    /// Set the spec id for the partition spec.
    pub fn with_spec_id(mut self, spec_id: i32) -> Self {
        self.spec_id = Some(spec_id);
        self
    }

    /// Add a new partition field to the partition spec.
    pub fn add_partition_field(
        self,
        source_name: impl AsRef<str>,
        target_name: impl Into<String>,
        transform: Transform,
    ) -> Result<Self> {
        let source_id = self
            .schema
            .field_by_name(source_name.as_ref())
            .ok_or_else(|| {
                invalid_data!(
                    "Cannot find source column with name: {} in schema",
                    source_name.as_ref()
                )
            })?
            .id;
        let field = UnboundPartitionField {
            source_ids: vec![source_id],
            field_id: None,
            name: target_name.into(),
            transform,
        };

        self.add_unbound_field(field)
    }

    /// Add a new partition field to the partition spec.
    ///
    /// If partition field id is set, it is used as the field id.
    /// Otherwise, a new `field_id` is assigned.
    ///
    /// Rejects a field that is redundant with an existing one, such as a second time
    /// granularity of the same column. [`UnboundPartitionSpec::bind`] accepts such fields.
    pub fn add_unbound_field(self, field: UnboundPartitionField) -> Result<Self> {
        self.check_name_set_and_unique(&field.name)?;
        self.check_for_redundant_partitions(&field.source_ids, &field.transform)?;
        self.add_field_allowing_redundancy(field)
    }

    /// Same as [`Self::add_unbound_field`] without the redundancy check, as Java's
    /// `PartitionSpec.Builder.add` does. Used to bind persisted specs, where evolution can leave
    /// several time granularities or void fields for one column.
    fn add_field_allowing_redundancy(mut self, field: UnboundPartitionField) -> Result<Self> {
        self.check_name_set_and_unique(&field.name)?;
        Self::check_name_does_not_collide_with_schema(&field, &self.schema)?;
        Self::check_transform_compatibility(&field, &self.schema)?;
        if let Some(partition_field_id) = field.field_id {
            self.check_partition_id_unique(partition_field_id)?;
        }

        // Non-fallible from here
        self.fields.push(field);
        Ok(self)
    }

    /// Wrapper around `with_unbound_fields` to add multiple partition fields.
    pub fn add_unbound_fields(
        self,
        fields: impl IntoIterator<Item = UnboundPartitionField>,
    ) -> Result<Self> {
        let mut builder = self;
        for field in fields {
            builder = builder.add_unbound_field(field)?;
        }
        Ok(builder)
    }

    /// Build a bound partition spec with the given schema.
    pub fn build(self) -> Result<PartitionSpec> {
        let fields = Self::set_field_ids(self.fields, self.last_assigned_field_id)?;
        Ok(PartitionSpec {
            spec_id: self.spec_id.unwrap_or(DEFAULT_PARTITION_SPEC_ID),
            fields,
        })
    }

    fn set_field_ids(
        fields: Vec<UnboundPartitionField>,
        last_assigned_field_id: i32,
    ) -> Result<Vec<PartitionField>> {
        let mut last_assigned_field_id = last_assigned_field_id;
        // Already assigned partition ids. If we see one of these during iteration,
        // we skip it.
        let assigned_ids = fields
            .iter()
            .filter_map(|f| f.field_id)
            .collect::<std::collections::HashSet<_>>();

        fn _check_add_1(prev: i32) -> Result<i32> {
            prev.checked_add(1)
                .ok_or_else(|| invalid_data!("Cannot assign more partition ids. Overflow."))
        }

        let mut bound_fields = Vec::with_capacity(fields.len());
        for field in fields.into_iter() {
            let partition_field_id = if let Some(partition_field_id) = field.field_id {
                last_assigned_field_id = std::cmp::max(last_assigned_field_id, partition_field_id);
                partition_field_id
            } else {
                last_assigned_field_id = _check_add_1(last_assigned_field_id)?;
                while assigned_ids.contains(&last_assigned_field_id) {
                    last_assigned_field_id = _check_add_1(last_assigned_field_id)?;
                }
                last_assigned_field_id
            };

            bound_fields.push(PartitionField {
                source_id: field.source_id()?,
                field_id: partition_field_id,
                name: field.name,
                transform: field.transform,
            })
        }

        Ok(bound_fields)
    }

    /// Ensure that the partition name is unique among columns in the schema.
    /// Duplicate names are allowed if:
    /// 1. The column is sourced from the column with the same name.
    /// 2. AND [`Transform::may_share_name_with_source`] holds for the transform.
    fn check_name_does_not_collide_with_schema(
        field: &UnboundPartitionField,
        schema: &Schema,
    ) -> Result<()> {
        match schema.field_by_name(field.name.as_str()) {
            Some(schema_collision) => {
                if field.transform.may_share_name_with_source() {
                    if schema_collision.id == field.source_id()? {
                        Ok(())
                    } else {
                        Err(invalid_data!(
                            "Cannot create identity or void partition sourced from different field in schema. Field name '{}' has id `{}` in schema but partition source id is `{}`",
                            field.name,
                            schema_collision.id,
                            field.source_id()?
                        ))
                    }
                } else {
                    Err(invalid_data!(
                        "Cannot create partition with name: '{}' that conflicts with schema field and is not an identity or void transform.",
                        field.name
                    ))
                }
            }
            None => Ok(()),
        }
    }

    /// Ensure that the transformation of the field is compatible with type of the field
    /// in the schema. Implicitly also checks if the source field exists in the schema.
    fn check_transform_compatibility(field: &UnboundPartitionField, schema: &Schema) -> Result<()> {
        // A multi-argument field is rejected here on purpose: a bound `PartitionField` has
        // no place for more than one source id yet.
        let source_id = field.source_id()?;
        let schema_field = schema.field_by_id(source_id).ok_or_else(|| {
            invalid_data!("Cannot find partition source field with id `{source_id}` in schema")
        })?;

        if field.transform != Transform::Void {
            if !schema_field.field_type.is_primitive() {
                return Err(invalid_data!(
                    "Cannot partition by non-primitive source field: '{}'.",
                    schema_field.field_type
                ));
            }

            if field
                .transform
                .result_type(&schema_field.field_type)
                .is_err()
            {
                return Err(invalid_data!(
                    "Invalid source type: '{}' for transform: '{}'.",
                    schema_field.field_type,
                    field.transform.dedup_name()
                ));
            }
        }

        Ok(())
    }
}

/// Contains checks that are common to both PartitionSpecBuilder and UnboundPartitionSpecBuilder
trait CorePartitionSpecValidator {
    /// Ensure that the partition name is unique among the partition fields and is not empty.
    fn check_name_set_and_unique(&self, name: &str) -> Result<()> {
        if name.is_empty() {
            return Err(invalid_data!("Cannot use empty partition name"));
        }

        if self.fields().iter().any(|f| f.name == name) {
            return Err(invalid_data!(
                "Cannot use partition name more than once: {name}"
            ));
        }
        Ok(())
    }

    /// For a single source-column transformations must be unique.
    fn check_for_redundant_partitions(
        &self,
        source_ids: &[i32],
        transform: &Transform,
    ) -> Result<()> {
        let collision = self.fields().iter().find(|f| {
            f.source_ids == source_ids && f.transform.dedup_name() == transform.dedup_name()
        });

        if let Some(collision) = collision {
            Err(invalid_data!(
                "Cannot add redundant partition with source ids `{:?}` and transform `{}`. A partition with the same source ids and transform already exists with name `{}`",
                source_ids,
                transform.dedup_name(),
                collision.name
            ))
        } else {
            Ok(())
        }
    }

    /// Check field / partition_id unique within the partition spec if set
    fn check_partition_id_unique(&self, field_id: i32) -> Result<()> {
        if self.fields().iter().any(|f| f.field_id == Some(field_id)) {
            return Err(invalid_data!(
                "Cannot use field id more than once in one PartitionSpec: {field_id}"
            ));
        }

        Ok(())
    }

    fn fields(&self) -> &Vec<UnboundPartitionField>;
}

impl CorePartitionSpecValidator for PartitionSpecBuilder {
    fn fields(&self) -> &Vec<UnboundPartitionField> {
        &self.fields
    }
}

impl CorePartitionSpecValidator for UnboundPartitionSpecBuilder {
    fn fields(&self) -> &Vec<UnboundPartitionField> {
        &self.fields
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ErrorKind;
    use crate::spec::{Literal, PrimitiveType, Type};

    #[test]
    fn test_partition_spec() {
        let spec = r#"
        {
        "spec-id": 1,
        "fields": [ {
            "source-id": 4,
            "field-id": 1000,
            "name": "ts_day",
            "transform": "day"
            }, {
            "source-id": 1,
            "field-id": 1001,
            "name": "id_bucket",
            "transform": "bucket[16]"
            }, {
            "source-id": 2,
            "field-id": 1002,
            "name": "id_truncate",
            "transform": "truncate[4]"
            } ]
        }
        "#;

        let partition_spec: PartitionSpec = serde_json::from_str(spec).unwrap();
        assert_eq!(4, partition_spec.fields[0].source_id);
        assert_eq!(1000, partition_spec.fields[0].field_id);
        assert_eq!("ts_day", partition_spec.fields[0].name);
        assert_eq!(Transform::Day, partition_spec.fields[0].transform);

        assert_eq!(1, partition_spec.fields[1].source_id);
        assert_eq!(1001, partition_spec.fields[1].field_id);
        assert_eq!("id_bucket", partition_spec.fields[1].name);
        assert_eq!(Transform::Bucket(16), partition_spec.fields[1].transform);

        assert_eq!(2, partition_spec.fields[2].source_id);
        assert_eq!(1002, partition_spec.fields[2].field_id);
        assert_eq!("id_truncate", partition_spec.fields[2].name);
        assert_eq!(Transform::Truncate(4), partition_spec.fields[2].transform);
    }

    #[test]
    fn test_is_unpartitioned() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();
        let partition_spec = PartitionSpec::builder(schema.clone())
            .with_spec_id(1)
            .build()
            .unwrap();
        assert!(
            partition_spec.is_unpartitioned(),
            "Empty partition spec should be unpartitioned"
        );

        let partition_spec = PartitionSpec::builder(schema.clone())
            .add_unbound_fields(vec![
                UnboundPartitionField::builder()
                    .source_ids(vec![1])
                    .name("id".to_string())
                    .transform(Transform::Identity)
                    .build()
                    .unwrap(),
                UnboundPartitionField::builder()
                    .source_ids(vec![2])
                    .name("name_string".to_string())
                    .transform(Transform::Void)
                    .build()
                    .unwrap(),
            ])
            .unwrap()
            .with_spec_id(1)
            .build()
            .unwrap();
        assert!(
            !partition_spec.is_unpartitioned(),
            "Partition spec with one non void transform should not be unpartitioned"
        );

        let partition_spec = PartitionSpec::builder(schema.clone())
            .with_spec_id(1)
            .add_unbound_fields(vec![
                UnboundPartitionField::builder()
                    .source_ids(vec![1])
                    .name("id_void".to_string())
                    .transform(Transform::Void)
                    .build()
                    .unwrap(),
                UnboundPartitionField::builder()
                    .source_ids(vec![2])
                    .name("name_void".to_string())
                    .transform(Transform::Void)
                    .build()
                    .unwrap(),
            ])
            .unwrap()
            .build()
            .unwrap();
        assert!(
            partition_spec.is_unpartitioned(),
            "Partition spec with all void field should be unpartitioned"
        );
    }

    #[test]
    fn test_unbound_partition_spec() {
        let spec = r#"
		{
		"spec-id": 1,
		"fields": [ {
			"source-id": 4,
			"field-id": 1000,
			"name": "ts_day",
			"transform": "day"
			}, {
			"source-id": 1,
			"field-id": 1001,
			"name": "id_bucket",
			"transform": "bucket[16]"
			}, {
			"source-id": 2,
			"field-id": 1002,
			"name": "id_truncate",
			"transform": "truncate[4]"
			} ]
		}
		"#;

        let partition_spec: UnboundPartitionSpec = serde_json::from_str(spec).unwrap();
        assert_eq!(Some(1), partition_spec.spec_id);

        assert_eq!(4, partition_spec.fields[0].source_id().unwrap());
        assert_eq!(Some(1000), partition_spec.fields[0].field_id);
        assert_eq!("ts_day", partition_spec.fields[0].name);
        assert_eq!(Transform::Day, partition_spec.fields[0].transform);

        assert_eq!(1, partition_spec.fields[1].source_id().unwrap());
        assert_eq!(Some(1001), partition_spec.fields[1].field_id);
        assert_eq!("id_bucket", partition_spec.fields[1].name);
        assert_eq!(Transform::Bucket(16), partition_spec.fields[1].transform);

        assert_eq!(2, partition_spec.fields[2].source_id().unwrap());
        assert_eq!(Some(1002), partition_spec.fields[2].field_id);
        assert_eq!("id_truncate", partition_spec.fields[2].name);
        assert_eq!(Transform::Truncate(4), partition_spec.fields[2].transform);

        let spec = r#"
		{
		"fields": [ {
			"source-id": 4,
			"name": "ts_day",
			"transform": "day"
			} ]
		}
		"#;
        let partition_spec: UnboundPartitionSpec = serde_json::from_str(spec).unwrap();
        assert_eq!(None, partition_spec.spec_id);

        assert_eq!(4, partition_spec.fields[0].source_id().unwrap());
        assert_eq!(None, partition_spec.fields[0].field_id);
        assert_eq!("ts_day", partition_spec.fields[0].name);
        assert_eq!(Transform::Day, partition_spec.fields[0].transform);
    }

    #[test]
    fn test_unbound_partition_spec_serialization_skips_none_fields() {
        let spec = UnboundPartitionSpec::builder()
            .add_partition_field(
                UnboundPartitionField::builder()
                    .source_ids(vec![4])
                    .name("ts_day")
                    .transform(Transform::Day)
                    .build()
                    .unwrap(),
            )
            .unwrap()
            .build();

        let value = serde_json::to_value(&spec).unwrap();
        let object = value.as_object().unwrap();
        assert!(!object.contains_key("spec-id"));
        let field = object["fields"][0].as_object().unwrap();
        assert!(!field.contains_key("field-id"));

        let value = serde_json::to_value(spec.with_spec_id(1)).unwrap();
        let object = value.as_object().unwrap();
        assert_eq!(Some(&serde_json::json!(1)), object.get("spec-id"));

        // Explicit nulls must still deserialize to `None` for backwards
        // compatibility.
        let spec: UnboundPartitionSpec = serde_json::from_str(
            r#"{
                "spec-id": null,
                "fields": [
                    {"source-id": 4, "name": "ts_day", "transform": "day", "field-id": null}
                ]
            }"#,
        )
        .unwrap();
        assert_eq!(None, spec.spec_id);
        assert_eq!(None, spec.fields[0].field_id);
    }

    #[test]
    fn test_new_unpartition() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();
        let partition_spec = PartitionSpec::builder(schema.clone())
            .with_spec_id(0)
            .build()
            .unwrap();
        let partition_type = partition_spec.partition_type(&schema).unwrap();
        assert_eq!(0, partition_type.fields().len());

        let unpartition_spec = PartitionSpec::unpartition_spec();
        assert_eq!(partition_spec, unpartition_spec);
    }

    #[test]
    fn test_partition_type() {
        let spec = r#"
            {
            "spec-id": 1,
            "fields": [ {
                "source-id": 4,
                "field-id": 1000,
                "name": "ts_day",
                "transform": "day"
                }, {
                "source-id": 1,
                "field-id": 1001,
                "name": "id_bucket",
                "transform": "bucket[16]"
                }, {
                "source-id": 2,
                "field-id": 1002,
                "name": "id_truncate",
                "transform": "truncate[4]"
                } ]
            }
            "#;

        let partition_spec: PartitionSpec = serde_json::from_str(spec).unwrap();
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
                NestedField::required(3, "ts", Type::Primitive(PrimitiveType::Timestamp)).into(),
                NestedField::required(4, "ts_day", Type::Primitive(PrimitiveType::Timestamp))
                    .into(),
                NestedField::required(5, "id_bucket", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(6, "id_truncate", Type::Primitive(PrimitiveType::Int)).into(),
            ])
            .build()
            .unwrap();

        let partition_type = partition_spec.partition_type(&schema).unwrap();
        assert_eq!(3, partition_type.fields().len());
        assert_eq!(
            *partition_type.fields()[0],
            NestedField::optional(
                partition_spec.fields[0].field_id,
                &partition_spec.fields[0].name,
                Type::Primitive(PrimitiveType::Date)
            )
        );
        assert_eq!(
            *partition_type.fields()[1],
            NestedField::optional(
                partition_spec.fields[1].field_id,
                &partition_spec.fields[1].name,
                Type::Primitive(PrimitiveType::Int)
            )
        );
        assert_eq!(
            *partition_type.fields()[2],
            NestedField::optional(
                partition_spec.fields[2].field_id,
                &partition_spec.fields[2].name,
                Type::Primitive(PrimitiveType::String)
            )
        );
    }

    #[test]
    fn test_partition_empty() {
        let spec = r#"
            {
            "spec-id": 1,
            "fields": []
            }
            "#;

        let partition_spec: PartitionSpec = serde_json::from_str(spec).unwrap();
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
                NestedField::required(3, "ts", Type::Primitive(PrimitiveType::Timestamp)).into(),
                NestedField::required(4, "ts_day", Type::Primitive(PrimitiveType::Timestamp))
                    .into(),
                NestedField::required(5, "id_bucket", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(6, "id_truncate", Type::Primitive(PrimitiveType::Int)).into(),
            ])
            .build()
            .unwrap();

        let partition_type = partition_spec.partition_type(&schema).unwrap();
        assert_eq!(0, partition_type.fields().len());
    }

    #[test]
    fn test_partition_type_with_dropped_source_column() {
        let spec = r#"
        {
        "spec-id": 1,
        "fields": [ {
            "source-id": 4,
            "field-id": 1000,
            "name": "ts_day",
            "transform": "day"
            }, {
            "source-id": 1,
            "field-id": 1001,
            "name": "id_bucket",
            "transform": "bucket[16]"
            }, {
            "source-id": 2,
            "field-id": 1002,
            "name": "id_truncate",
            "transform": "truncate[4]"
            } ]
        }
        "#;

        let partition_spec: PartitionSpec = serde_json::from_str(spec).unwrap();
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();

        assert_eq!(
            partition_spec.partition_type(&schema).unwrap(),
            StructType::new(vec![
                NestedField::optional(1000, "ts_day", PrimitiveType::Date).into(),
                NestedField::optional(1001, "id_bucket", PrimitiveType::Int).into(),
                NestedField::optional(1002, "id_truncate", PrimitiveType::String).into(),
            ])
        );

        // Dropping all sources must retain every partition field in its original position.
        let empty_schema = Schema::builder().build().unwrap();
        assert_eq!(
            partition_spec.partition_type(&empty_schema).unwrap(),
            StructType::new(vec![
                NestedField::optional(1000, "ts_day", PrimitiveType::Date).into(),
                NestedField::optional(1001, "id_bucket", PrimitiveType::Int).into(),
                NestedField::optional(1002, "id_truncate", PrimitiveType::Unknown).into(),
            ])
        );

        // Missing sources must not suppress validation of transforms on remaining sources.
        let incompatible_schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", PrimitiveType::Boolean).into(),
            ])
            .build()
            .unwrap();
        let err = partition_spec
            .partition_type(&incompatible_schema)
            .unwrap_err();
        assert_eq!(err.kind(), ErrorKind::DataInvalid);
        assert_eq!(
            err.message(),
            "boolean is not a valid input type of bucket transform"
        );
    }

    #[test]
    fn test_partition_type_without_source_type() {
        let schema = Schema::builder().build().unwrap();
        for (transform, expected_type) in [
            (Transform::Identity, PrimitiveType::Unknown),
            (Transform::Truncate(4), PrimitiveType::Unknown),
            (Transform::Void, PrimitiveType::Unknown),
            (Transform::Bucket(16), PrimitiveType::Int),
            (Transform::Year, PrimitiveType::Int),
            (Transform::Month, PrimitiveType::Int),
            (Transform::Day, PrimitiveType::Date),
            (Transform::Hour, PrimitiveType::Int),
            (Transform::Unknown, PrimitiveType::String),
        ] {
            let spec = PartitionSpec {
                spec_id: 0,
                fields: vec![PartitionField {
                    source_id: 1,
                    field_id: 1000,
                    name: "partition".to_string(),
                    transform,
                }],
            };
            assert_eq!(
                spec.partition_type(&schema).unwrap(),
                StructType::new(vec![
                    NestedField::optional(1000, "partition", expected_type).into(),
                ])
            );
        }
    }

    #[test]
    fn test_builder_disallow_duplicate_names() {
        UnboundPartitionSpec::builder()
            .add_partition_field(
                UnboundPartitionField::builder()
                    .source_ids(vec![1])
                    .name("ts_day")
                    .transform(Transform::Day)
                    .build()
                    .unwrap(),
            )
            .unwrap()
            .add_partition_field(
                UnboundPartitionField::builder()
                    .source_ids(vec![2])
                    .name("ts_day")
                    .transform(Transform::Day)
                    .build()
                    .unwrap(),
            )
            .unwrap_err();
    }

    #[test]
    fn test_builder_disallow_duplicate_field_ids() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();
        PartitionSpec::builder(schema.clone())
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: Some(1000),
                name: "id".to_string(),
                transform: Transform::Identity,
            })
            .unwrap()
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![2],
                field_id: Some(1000),
                name: "id_bucket".to_string(),
                transform: Transform::Bucket(16),
            })
            .unwrap_err();
    }

    #[test]
    fn test_builder_auto_assign_field_ids() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
                NestedField::required(3, "ts", Type::Primitive(PrimitiveType::Timestamp)).into(),
            ])
            .build()
            .unwrap();
        let spec = PartitionSpec::builder(schema.clone())
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                name: "id".to_string(),
                transform: Transform::Identity,
                field_id: Some(1012),
            })
            .unwrap()
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![2],
                name: "name_void".to_string(),
                transform: Transform::Void,
                field_id: None,
            })
            .unwrap()
            // Should keep its ID even if its lower
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![3],
                name: "year".to_string(),
                transform: Transform::Year,
                field_id: Some(1),
            })
            .unwrap()
            .build()
            .unwrap();

        assert_eq!(1012, spec.fields[0].field_id);
        assert_eq!(1013, spec.fields[1].field_id);
        assert_eq!(1, spec.fields[2].field_id);
    }

    #[test]
    fn test_builder_valid_schema() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();

        PartitionSpec::builder(schema.clone())
            .with_spec_id(1)
            .build()
            .unwrap();

        let spec = PartitionSpec::builder(schema.clone())
            .with_spec_id(1)
            .add_partition_field("id", "id_bucket[16]", Transform::Bucket(16))
            .unwrap()
            .build()
            .unwrap();

        assert_eq!(spec, PartitionSpec {
            spec_id: 1,
            fields: vec![PartitionField {
                source_id: 1,
                field_id: 1000,
                name: "id_bucket[16]".to_string(),
                transform: Transform::Bucket(16),
            }],
        });
        assert_eq!(
            spec.partition_type(&schema).unwrap(),
            StructType::new(vec![
                NestedField::optional(1000, "id_bucket[16]", Type::Primitive(PrimitiveType::Int))
                    .into()
            ])
        )
    }

    #[test]
    fn test_collision_with_schema_name() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
            ])
            .build()
            .unwrap();

        PartitionSpec::builder(schema.clone())
            .with_spec_id(1)
            .build()
            .unwrap();

        let err = PartitionSpec::builder(schema)
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: None,
                name: "id".to_string(),
                transform: Transform::Bucket(16),
            })
            .unwrap_err();
        assert!(err.message().contains("conflicts with schema"))
    }

    #[test]
    fn test_builder_collision_is_ok_for_identity_transforms() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "number", Type::Primitive(PrimitiveType::Int)).into(),
            ])
            .build()
            .unwrap();

        PartitionSpec::builder(schema.clone())
            .with_spec_id(1)
            .build()
            .unwrap();

        PartitionSpec::builder(schema.clone())
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: None,
                name: "id".to_string(),
                transform: Transform::Identity,
            })
            .unwrap()
            .build()
            .unwrap();

        // Not OK for different source id
        PartitionSpec::builder(schema)
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![2],
                field_id: None,
                name: "id".to_string(),
                transform: Transform::Identity,
            })
            .unwrap_err();
    }

    #[test]
    fn test_builder_all_source_ids_must_exist() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
                NestedField::required(3, "ts", Type::Primitive(PrimitiveType::Timestamp)).into(),
            ])
            .build()
            .unwrap();

        // Valid
        PartitionSpec::builder(schema.clone())
            .with_spec_id(1)
            .add_unbound_fields(vec![
                UnboundPartitionField {
                    source_ids: vec![1],
                    field_id: None,
                    name: "id_bucket".to_string(),
                    transform: Transform::Bucket(16),
                },
                UnboundPartitionField {
                    source_ids: vec![2],
                    field_id: None,
                    name: "name".to_string(),
                    transform: Transform::Identity,
                },
            ])
            .unwrap()
            .build()
            .unwrap();

        // Invalid
        PartitionSpec::builder(schema)
            .with_spec_id(1)
            .add_unbound_fields(vec![
                UnboundPartitionField {
                    source_ids: vec![1],
                    field_id: None,
                    name: "id_bucket".to_string(),
                    transform: Transform::Bucket(16),
                },
                UnboundPartitionField {
                    source_ids: vec![4],
                    field_id: None,
                    name: "name".to_string(),
                    transform: Transform::Identity,
                },
            ])
            .unwrap_err();
    }

    #[test]
    fn test_builder_disallows_variant_source() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::optional(2, "v", Type::Variant(crate::spec::VariantType)).into(),
            ])
            .build()
            .unwrap();

        let err = PartitionSpec::builder(schema)
            .with_spec_id(1)
            .add_unbound_fields(vec![UnboundPartitionField {
                source_ids: vec![2],
                field_id: None,
                name: "v_part".to_string(),
                transform: Transform::Identity,
            }])
            .expect_err("variant must not be allowed as a partition source");

        assert_eq!(
            err.message(),
            "Cannot partition by non-primitive source field: 'variant'."
        );
    }

    #[test]
    fn test_builder_disallows_redundant() {
        let err = UnboundPartitionSpec::builder()
            .with_spec_id(1)
            .add_partition_field(
                UnboundPartitionField::builder()
                    .source_ids(vec![1])
                    .name("id_bucket[16]")
                    .transform(Transform::Bucket(16))
                    .build()
                    .unwrap(),
            )
            .unwrap()
            .add_partition_field(
                UnboundPartitionField::builder()
                    .source_ids(vec![1])
                    .name("id_bucket_with_other_name")
                    .transform(Transform::Bucket(16))
                    .build()
                    .unwrap(),
            )
            .unwrap_err();
        assert!(err.message().contains("redundant partition"));
    }

    fn ts_schema() -> Schema {
        Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "ts", Type::Primitive(PrimitiveType::Timestamptz)).into(),
                NestedField::required(2, "id", Type::Primitive(PrimitiveType::Long)).into(),
            ])
            .build()
            .unwrap()
    }

    fn ts_field(field_id: i32, name: &str, transform: Transform) -> UnboundPartitionField {
        UnboundPartitionField {
            source_ids: vec![1],
            field_id: Some(field_id),
            name: name.to_string(),
            transform,
        }
    }

    #[test]
    fn test_add_unbound_field_rejects_second_time_granularity() {
        let err = PartitionSpec::builder(ts_schema())
            .add_unbound_field(ts_field(1000, "ts_day", Transform::Day))
            .unwrap()
            .add_unbound_field(ts_field(1001, "ts_hour", Transform::Hour))
            .unwrap_err();
        assert!(err.message().contains("redundant partition"));
    }

    #[test]
    fn test_bind_accepts_evolved_spec_fields() {
        let fields = vec![
            ts_field(1000, "ts_day", Transform::Day),
            ts_field(1001, "ts_hour", Transform::Hour),
            ts_field(1002, "ts", Transform::Void),
            ts_field(1003, "ts_null", Transform::Void),
        ];
        let unbound = UnboundPartitionSpec {
            spec_id: Some(1),
            fields: fields.clone(),
        };

        let spec = unbound.bind(ts_schema()).unwrap();

        let expected = PartitionSpec {
            spec_id: 1,
            fields: fields
                .into_iter()
                .map(|f| PartitionField {
                    source_id: 1,
                    field_id: f.field_id.unwrap(),
                    name: f.name,
                    transform: f.transform,
                })
                .collect(),
        };
        assert_eq!(spec, expected);
    }

    #[test]
    fn test_bind_rejects_void_field_named_after_other_column() {
        let unbound = UnboundPartitionSpec {
            spec_id: Some(1),
            fields: vec![ts_field(1000, "id", Transform::Void)],
        };

        let err = unbound.bind(ts_schema()).unwrap_err();

        assert_eq!(err.kind(), ErrorKind::DataInvalid);
        assert!(err.message().contains("Field name 'id'"));
    }

    #[test]
    fn test_bind_accepts_identical_fields_with_different_names() {
        let bucket_field = |field_id: i32, name: &str| UnboundPartitionField {
            source_ids: vec![2],
            field_id: Some(field_id),
            name: name.to_string(),
            transform: Transform::Bucket(16),
        };
        let unbound = UnboundPartitionSpec {
            spec_id: Some(1),
            fields: vec![
                bucket_field(1000, "id_bucket"),
                bucket_field(1001, "id_bucket_2"),
            ],
        };

        let spec = unbound.bind(ts_schema()).unwrap();

        let expected = PartitionSpec {
            spec_id: 1,
            fields: vec![
                PartitionField {
                    source_id: 2,
                    field_id: 1000,
                    name: "id_bucket".to_string(),
                    transform: Transform::Bucket(16),
                },
                PartitionField {
                    source_id: 2,
                    field_id: 1001,
                    name: "id_bucket_2".to_string(),
                    transform: Transform::Bucket(16),
                },
            ],
        };
        assert_eq!(spec, expected);
    }

    #[test]
    fn test_bind_rejects_repeated_partition_name() {
        let unbound = UnboundPartitionSpec {
            spec_id: Some(1),
            fields: vec![
                ts_field(1000, "ts_part", Transform::Day),
                ts_field(1001, "ts_part", Transform::Hour),
            ],
        };

        let err = unbound.bind(ts_schema()).unwrap_err();

        assert!(
            err.message()
                .contains("Cannot use partition name more than once")
        );
    }

    #[test]
    fn test_add_unbound_field_checks_name_before_redundancy() {
        let err = PartitionSpec::builder(ts_schema())
            .add_unbound_field(ts_field(1000, "ts_part", Transform::Day))
            .unwrap()
            .add_unbound_field(ts_field(1001, "ts_part", Transform::Hour))
            .unwrap_err();

        assert!(
            err.message()
                .contains("Cannot use partition name more than once")
        );
    }

    #[test]
    fn test_builder_incompatible_transforms_disallowed() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
            ])
            .build()
            .unwrap();

        PartitionSpec::builder(schema)
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: None,
                name: "id_year".to_string(),
                transform: Transform::Year,
            })
            .unwrap_err();
    }

    #[test]
    fn test_build_unbound_specs_without_partition_id() {
        let spec = UnboundPartitionSpec::builder()
            .with_spec_id(1)
            .add_partition_fields(vec![UnboundPartitionField {
                source_ids: vec![1],
                field_id: None,
                name: "id_bucket[16]".to_string(),
                transform: Transform::Bucket(16),
            }])
            .unwrap()
            .build();

        assert_eq!(spec, UnboundPartitionSpec {
            spec_id: Some(1),
            fields: vec![UnboundPartitionField {
                source_ids: vec![1],
                field_id: None,
                name: "id_bucket[16]".to_string(),
                transform: Transform::Bucket(16),
            }]
        });
    }

    #[test]
    fn test_is_compatible_with() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();

        let partition_spec_1 = PartitionSpec::builder(schema.clone())
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: None,
                name: "id_bucket".to_string(),
                transform: Transform::Bucket(16),
            })
            .unwrap()
            .build()
            .unwrap();

        let partition_spec_2 = PartitionSpec::builder(schema)
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: None,
                name: "id_bucket".to_string(),
                transform: Transform::Bucket(16),
            })
            .unwrap()
            .build()
            .unwrap();

        assert!(partition_spec_1.is_compatible_with(&partition_spec_2));
    }

    #[test]
    fn test_not_compatible_with_transform_different() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
            ])
            .build()
            .unwrap();

        let partition_spec_1 = PartitionSpec::builder(schema.clone())
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: None,
                name: "id_bucket".to_string(),
                transform: Transform::Bucket(16),
            })
            .unwrap()
            .build()
            .unwrap();

        let partition_spec_2 = PartitionSpec::builder(schema)
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: None,
                name: "id_bucket".to_string(),
                transform: Transform::Bucket(32),
            })
            .unwrap()
            .build()
            .unwrap();

        assert!(!partition_spec_1.is_compatible_with(&partition_spec_2));
    }

    #[test]
    fn test_not_compatible_with_source_id_different() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();

        let partition_spec_1 = PartitionSpec::builder(schema.clone())
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: None,
                name: "id_bucket".to_string(),
                transform: Transform::Bucket(16),
            })
            .unwrap()
            .build()
            .unwrap();

        let partition_spec_2 = PartitionSpec::builder(schema)
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![2],
                field_id: None,
                name: "id_bucket".to_string(),
                transform: Transform::Bucket(16),
            })
            .unwrap()
            .build()
            .unwrap();

        assert!(!partition_spec_1.is_compatible_with(&partition_spec_2));
    }

    #[test]
    fn test_not_compatible_with_order_different() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();

        let partition_spec_1 = PartitionSpec::builder(schema.clone())
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: None,
                name: "id_bucket".to_string(),
                transform: Transform::Bucket(16),
            })
            .unwrap()
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![2],
                field_id: None,
                name: "name".to_string(),
                transform: Transform::Identity,
            })
            .unwrap()
            .build()
            .unwrap();

        let partition_spec_2 = PartitionSpec::builder(schema)
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![2],
                field_id: None,
                name: "name".to_string(),
                transform: Transform::Identity,
            })
            .unwrap()
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: None,
                name: "id_bucket".to_string(),
                transform: Transform::Bucket(16),
            })
            .unwrap()
            .build()
            .unwrap();

        assert!(!partition_spec_1.is_compatible_with(&partition_spec_2));
    }

    #[test]
    fn test_highest_field_id_unpartitioned() {
        let spec = PartitionSpec::builder(Schema::builder().with_fields(vec![]).build().unwrap())
            .with_spec_id(1)
            .build()
            .unwrap();

        assert!(spec.highest_field_id().is_none());
    }

    #[test]
    fn test_highest_field_id() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();

        let spec = PartitionSpec::builder(schema)
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: Some(1001),
                name: "id".to_string(),
                transform: Transform::Identity,
            })
            .unwrap()
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![2],
                field_id: Some(1000),
                name: "name".to_string(),
                transform: Transform::Identity,
            })
            .unwrap()
            .build()
            .unwrap();

        assert_eq!(Some(1001), spec.highest_field_id());
    }

    #[test]
    fn test_has_sequential_ids() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();

        let spec = PartitionSpec::builder(schema)
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: Some(1000),
                name: "id".to_string(),
                transform: Transform::Identity,
            })
            .unwrap()
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![2],
                field_id: Some(1001),
                name: "name".to_string(),
                transform: Transform::Identity,
            })
            .unwrap()
            .build()
            .unwrap();

        assert_eq!(1000, spec.fields[0].field_id);
        assert_eq!(1001, spec.fields[1].field_id);
        assert!(spec.has_sequential_ids());
    }

    #[test]
    fn test_sequential_ids_must_start_at_1000() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();

        let spec = PartitionSpec::builder(schema)
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: Some(999),
                name: "id".to_string(),
                transform: Transform::Identity,
            })
            .unwrap()
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![2],
                field_id: Some(1000),
                name: "name".to_string(),
                transform: Transform::Identity,
            })
            .unwrap()
            .build()
            .unwrap();

        assert_eq!(999, spec.fields[0].field_id);
        assert_eq!(1000, spec.fields[1].field_id);
        assert!(!spec.has_sequential_ids());
    }

    #[test]
    fn test_sequential_ids_must_have_no_gaps() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();

        let spec = PartitionSpec::builder(schema)
            .with_spec_id(1)
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![1],
                field_id: Some(1000),
                name: "id".to_string(),
                transform: Transform::Identity,
            })
            .unwrap()
            .add_unbound_field(UnboundPartitionField {
                source_ids: vec![2],
                field_id: Some(1002),
                name: "name".to_string(),
                transform: Transform::Identity,
            })
            .unwrap()
            .build()
            .unwrap();

        assert_eq!(1000, spec.fields[0].field_id);
        assert_eq!(1002, spec.fields[1].field_id);
        assert!(!spec.has_sequential_ids());
    }

    #[test]
    fn test_partition_to_path_with_dropped_source_column() {
        for (source_type, transform, value, expected) in [
            (
                PrimitiveType::Int,
                Transform::Identity,
                Literal::int(1),
                "p=1",
            ),
            (
                PrimitiveType::Timestamp,
                Transform::Identity,
                Literal::long(1),
                "p=1",
            ),
            (
                PrimitiveType::Int,
                Transform::Truncate(10),
                Literal::int(10),
                "p=10",
            ),
        ] {
            let schema = Schema::builder()
                .with_fields(vec![
                    NestedField::required(1, "kept", PrimitiveType::Int).into(),
                    NestedField::required(2, "dropped", source_type).into(),
                ])
                .build()
                .unwrap();
            let spec = PartitionSpec::builder(schema)
                .add_partition_field("dropped", "p", transform)
                .unwrap()
                .build()
                .unwrap();
            let current_schema = Schema::builder()
                .with_fields(vec![
                    NestedField::required(1, "kept", PrimitiveType::Int).into(),
                ])
                .build()
                .unwrap();
            assert_eq!(
                spec.partition_to_path(&Struct::from_iter([Some(value)]), current_schema.into()),
                expected
            );
        }
    }

    #[test]
    fn test_partition_to_path() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
                NestedField::required(3, "timestamp", Type::Primitive(PrimitiveType::Timestamp))
                    .into(),
                NestedField::required(4, "empty", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();

        let spec = PartitionSpec::builder(schema.clone())
            .add_partition_field("id", "id", Transform::Identity)
            .unwrap()
            .add_partition_field("name", "name", Transform::Identity)
            .unwrap()
            .add_partition_field("timestamp", "ts_hour", Transform::Hour)
            .unwrap()
            .add_partition_field("empty", "empty_void", Transform::Void)
            .unwrap()
            .build()
            .unwrap();

        let data = Struct::from_iter([
            Some(Literal::int(42)),
            Some(Literal::string("alice")),
            Some(Literal::int(1000)),
            Some(Literal::string("empty")),
        ]);

        assert_eq!(
            spec.partition_to_path(&data, schema.into()),
            "id=42/name=alice/ts_hour=1970-02-11-16/empty_void=null"
        );
    }

    #[test]
    fn test_partition_to_path_float_matches_java() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "d", Type::Primitive(PrimitiveType::Double)).into(),
                NestedField::required(2, "f", Type::Primitive(PrimitiveType::Float)).into(),
            ])
            .build()
            .unwrap();

        let spec = PartitionSpec::builder(schema.clone())
            .add_partition_field("d", "d", Transform::Identity)
            .unwrap()
            .add_partition_field("f", "f", Transform::Identity)
            .unwrap()
            .build()
            .unwrap();

        let data = Struct::from_iter([Some(Literal::double(1.0)), Some(Literal::float(1.0e7_f32))]);

        assert_eq!(
            spec.partition_to_path(&data, schema.into()),
            "d=1.0/f=1.0E7"
        );
    }

    #[test]
    fn test_partition_to_path_escaped_strings() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "\"esc\"#1", Type::Primitive(PrimitiveType::String))
                    .into(),
                NestedField::required(2, "data", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();

        let spec = PartitionSpec::builder(schema.clone())
            .add_partition_field("\"esc\"#1", "\"esc\"#1", Transform::Identity)
            .unwrap()
            .build()
            .unwrap();

        let data = Struct::from_iter([
            Some(Literal::string("a/b/c/d")),
            Some(Literal::string("val#1")),
        ]);

        assert_eq!(
            spec.partition_to_path(&data, schema.into()),
            "%22esc%22%231=a%2Fb%2Fc%2Fd"
        );
    }

    #[test]
    fn test_partition_to_path_escaped_field_name() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "\"esc\"#1", Type::Primitive(PrimitiveType::String))
                    .into(),
                NestedField::required(2, "data", Type::Primitive(PrimitiveType::String)).into(),
            ])
            .build()
            .unwrap();

        let spec = PartitionSpec::builder(schema.clone())
            .add_partition_field("data", "data", Transform::Identity)
            .unwrap()
            .add_partition_field("data", "data_truc_10", Transform::Truncate(10))
            .unwrap()
            .build()
            .unwrap();

        let data = Struct::from_iter([
            Some(Literal::string("a/b/c/d")),
            Some(Literal::string("a/b/c/d")),
        ]);

        assert_eq!(
            spec.partition_to_path(&data, schema.into()),
            "data=a%2Fb%2Fc%2Fd/data_truc_10=a%2Fb%2Fc%2Fd"
        );
    }

    #[test]
    fn test_unbound_partition_field_reads_source_ids_only() {
        let field: UnboundPartitionField = serde_json::from_str(
            r#"{"source-ids": [1, 2], "name": "m", "transform": "bucket[4]"}"#,
        )
        .unwrap();

        assert_eq!([1, 2], field.source_ids());
        // a multi-argument field has no single source id
        assert!(
            field
                .source_id()
                .unwrap_err()
                .to_string()
                .contains("has no single source id")
        );

        let serialized = serde_json::to_value(&field).unwrap();
        assert_eq!(
            Some(&serde_json::json!([1, 2])),
            serialized.get("source-ids")
        );
        assert!(serialized.get("source-id").is_none());
    }

    #[test]
    fn test_unbound_partition_field_single_source_id_round_trip() {
        let field: UnboundPartitionField =
            serde_json::from_str(r#"{"source-id": 1, "name": "m", "transform": "identity"}"#)
                .unwrap();

        assert_eq!(1, field.source_id().unwrap());
        assert_eq!([1], field.source_ids());

        // a single-argument field keeps writing source-id
        let serialized = serde_json::to_value(&field).unwrap();
        assert_eq!(Some(&serde_json::json!(1)), serialized.get("source-id"));
        assert!(serialized.get("source-ids").is_none());
    }

    #[test]
    fn test_unbound_partition_field_rejects_malformed_source_ids() {
        for (input, expected) in [
            (
                r#"{"source-ids": [], "name": "m", "transform": "identity"}"#,
                "Empty source-ids is not allowed",
            ),
            (
                r#"{"name": "m", "transform": "identity"}"#,
                "Either `source-id` or `source-ids` must be present",
            ),
            (
                r#"{"source-id": 1, "source-ids": [1], "name": "m", "transform": "identity"}"#,
                "mutually exclusive",
            ),
        ] {
            let err = serde_json::from_str::<UnboundPartitionField>(input).unwrap_err();
            assert!(
                err.to_string().contains(expected),
                "unexpected error for {input}: {err}"
            );
        }
    }

    #[test]
    fn test_binding_a_multi_argument_field_fails_loudly() {
        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "a", Type::Primitive(PrimitiveType::Int)).into(),
                NestedField::required(2, "b", Type::Primitive(PrimitiveType::Int)).into(),
            ])
            .build()
            .unwrap();

        let spec: UnboundPartitionSpec = serde_json::from_str(
            r#"{"spec-id": 1, "fields": [{"source-ids": [1, 2], "field-id": 1000, "name": "m", "transform": "bucket[4]"}]}"#,
        )
        .unwrap();

        // A bound PartitionField still carries a single source id, so binding has to refuse
        // rather than quietly keep the first one.
        let err = spec.bind(schema).unwrap_err();
        assert!(
            err.to_string().contains("has no single source id"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn test_unbound_partition_field_builder_rejects_empty_source_ids() {
        let err = UnboundPartitionField::builder()
            .source_ids(vec![])
            .name("m")
            .transform(Transform::Identity)
            .build()
            .unwrap_err();
        assert!(
            err.to_string().contains("Empty source-ids is not allowed"),
            "unexpected error: {err}"
        );
    }
}