vendor/iceberg-arrow59/src/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.
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_unbound_field(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.
pub fn add_unbound_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)?;
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 the transformation is identity
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 == Transform::Identity {
if schema_collision.id == field.source_id()? {
Ok(())
} else {
Err(invalid_data!(
"Cannot create identity 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 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()).into(),
NestedField::optional(1001, "id_bucket", PrimitiveType::Int.into()).into(),
NestedField::optional(1002, "id_truncate", PrimitiveType::String.into()).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()).into(),
NestedField::optional(1001, "id_bucket", PrimitiveType::Int.into()).into(),
NestedField::optional(1002, "id_truncate", PrimitiveType::Unknown.into()).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()).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()).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"));
}
#[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() {
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_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}"
);
}
}