iceberg/
partitioning.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Partition type utilities for Iceberg tables.
19
20use std::cmp::Reverse;
21use std::collections::{HashMap, HashSet};
22
23use crate::spec::{
24    NestedField, NestedFieldRef, PartitionField, PartitionSpec, Schema, StructType, Transform, Type,
25};
26use crate::{Error, ErrorKind, Result};
27
28/// Computes the unified partition type across all partition specs in the table.
29///
30/// This is equivalent to Java's `Partitioning.partitionType(table)`. The result is a
31/// StructType containing all partition fields ever used across all specs, enabling correct
32/// representation of the `_partition` metadata column when partition evolution has occurred.
33///
34/// Matches Java's `buildPartitionProjectionType` behavior:
35/// - Specs are sorted by spec_id in descending order (newer specs first), so newer field
36///   names take precedence when deduplicating by field_id.
37/// - Unknown transforms cause an error.
38/// - Fields whose source column was dropped from the schema are skipped.
39/// - Two specs defining the same field_id must be compatible (same source, compatible
40///   transforms); V1 tables do not guarantee field ids are unique across specs.
41/// - When a newer spec marks a field as Void (dropped) but an older spec has it with a
42///   real transform, the older spec's type is preserved while the newer spec's name is kept.
43/// - Fields are deduplicated by field_id; each unique field_id appears exactly once.
44/// - Output fields are sorted by field_id ascending.
45///
46/// # Arguments
47/// * `partition_specs` - Iterator over all partition specs in the table
48/// * `schema` - The current table schema (needed to determine result types of transforms)
49pub fn compute_unified_partition_type<'a>(
50    partition_specs: impl Iterator<Item = &'a PartitionSpec>,
51    schema: &Schema,
52) -> Result<StructType> {
53    let mut specs: Vec<&PartitionSpec> = partition_specs.collect();
54    specs.sort_by_key(|s| Reverse(s.spec_id()));
55
56    let active_field_ids = all_active_field_ids(specs.iter().copied(), schema);
57
58    let mut field_map: HashMap<i32, &PartitionField> = HashMap::new();
59    let mut type_map: HashMap<i32, Type> = HashMap::new();
60    let mut name_map: HashMap<i32, String> = HashMap::new();
61
62    for spec in &specs {
63        for field in spec.fields() {
64            let field_id = field.field_id;
65
66            // Reject unknown transforms up front: we cannot determine their result type,
67            // so we cannot build a partition column for them. This check must precede the
68            // active_field_ids filter below, otherwise an unknown transform could be
69            // silently skipped.
70            if matches!(field.transform, Transform::Unknown) {
71                return Err(Error::new(
72                    ErrorKind::DataInvalid,
73                    format!(
74                        "Partition field '{}' uses an unknown transform whose result type \
75                         cannot be determined",
76                        field.name
77                    ),
78                ));
79            }
80
81            if !active_field_ids.contains(&field_id) {
82                continue;
83            }
84
85            let source_field = match schema.field_by_id(field.source_id) {
86                Some(f) => f,
87                None => continue,
88            };
89
90            match field_map.get(&field_id) {
91                None => {
92                    let res_type = field.transform.result_type(&source_field.field_type)?;
93                    field_map.insert(field_id, field);
94                    type_map.insert(field_id, res_type);
95                    name_map.insert(field_id, field.name.clone());
96                }
97                Some(existing) => {
98                    // V1 tables do not guarantee field ids are unique across specs, so two
99                    // specs may define the same field id. They must be compatible.
100                    if !equivalent_ignoring_names(field, existing) {
101                        return Err(Error::new(
102                            ErrorKind::DataInvalid,
103                            format!(
104                                "Conflicting partition fields for field id {field_id}: \
105                                 '{}' and '{}'",
106                                field.name, existing.name
107                            ),
108                        ));
109                    }
110
111                    // Use the correct type for dropped partitions in v1 tables: if the
112                    // newer spec voided the field but an older spec has a real transform,
113                    // keep the older spec's type.
114                    if is_void_transform(existing) && !is_void_transform(field) {
115                        let res_type = field.transform.result_type(&source_field.field_type)?;
116                        field_map.insert(field_id, field);
117                        type_map.insert(field_id, res_type);
118                    }
119                }
120            }
121        }
122    }
123
124    let mut field_ids: Vec<i32> = field_map.keys().copied().collect();
125    field_ids.sort();
126
127    let struct_fields = field_ids
128        .into_iter()
129        .map(|fid| -> Result<NestedFieldRef> {
130            let name = name_map.get(&fid).ok_or_else(|| {
131                Error::new(
132                    ErrorKind::Unexpected,
133                    format!("Missing name for partition field {fid}"),
134                )
135            })?;
136            let ty = type_map.remove(&fid).ok_or_else(|| {
137                Error::new(
138                    ErrorKind::Unexpected,
139                    format!("Missing type for partition field {fid}"),
140                )
141            })?;
142            Ok(NestedField::optional(fid, name, ty).into())
143        })
144        .collect::<Result<Vec<_>>>()?;
145
146    Ok(StructType::new(struct_fields))
147}
148
149fn is_void_transform(field: &PartitionField) -> bool {
150    matches!(field.transform, Transform::Void)
151}
152
153/// Two partition fields with the same field id are compatible if they share the same
154/// source id and have compatible transforms. Matches Java's
155/// `Partitioning.equivalentIgnoringNames`.
156fn equivalent_ignoring_names(field: &PartitionField, other: &PartitionField) -> bool {
157    field.field_id == other.field_id
158        && field.source_id == other.source_id
159        && compatible_transforms(&field.transform, &other.transform)
160}
161
162/// Transforms are compatible if they are equal, or if either is Void (a dropped field).
163/// Matches Java's `Partitioning.compatibleTransforms`.
164fn compatible_transforms(t1: &Transform, t2: &Transform) -> bool {
165    t1 == t2 || matches!(t1, Transform::Void) || matches!(t2, Transform::Void)
166}
167
168fn all_active_field_ids<'a>(
169    partition_specs: impl Iterator<Item = &'a PartitionSpec>,
170    schema: &Schema,
171) -> HashSet<i32> {
172    partition_specs
173        .flat_map(|spec| spec.fields().iter())
174        .filter(|field| schema.field_by_id(field.source_id).is_some())
175        .map(|field| field.field_id)
176        .collect()
177}
178
179#[cfg(test)]
180mod tests {
181    use std::sync::Arc;
182
183    use super::*;
184    use crate::spec::{NestedField, PrimitiveType, Transform, Type, UnboundPartitionSpec};
185
186    fn test_schema() -> Schema {
187        Schema::builder()
188            .with_fields(vec![
189                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
190                NestedField::required(2, "data", Type::Primitive(PrimitiveType::String)).into(),
191                NestedField::required(3, "ts", Type::Primitive(PrimitiveType::Timestamp)).into(),
192                NestedField::required(4, "category", Type::Primitive(PrimitiveType::String)).into(),
193            ])
194            .build()
195            .unwrap()
196    }
197
198    fn build_spec(
199        schema: &Schema,
200        spec_id: i32,
201        fields: Vec<(i32, &str, Transform)>,
202    ) -> PartitionSpec {
203        let mut builder = UnboundPartitionSpec::builder().with_spec_id(spec_id);
204        for (source_id, name, transform) in fields {
205            builder = builder
206                .add_partition_field(source_id, name, transform)
207                .unwrap();
208        }
209        builder.build().bind(schema.clone()).unwrap()
210    }
211
212    #[test]
213    fn test_single_spec_identity() {
214        let schema = test_schema();
215        let spec = build_spec(&schema, 0, vec![(4, "category", Transform::Identity)]);
216
217        let result = compute_unified_partition_type([&spec].into_iter(), &schema).unwrap();
218        assert_eq!(result.fields().len(), 1);
219        assert_eq!(result.fields()[0].name, "category");
220        assert_eq!(
221            *result.fields()[0].field_type,
222            Type::Primitive(PrimitiveType::String)
223        );
224    }
225
226    #[test]
227    fn test_single_spec_with_year_transform() {
228        let schema = test_schema();
229        let spec = build_spec(&schema, 0, vec![(3, "ts_year", Transform::Year)]);
230
231        let result = compute_unified_partition_type([&spec].into_iter(), &schema).unwrap();
232        assert_eq!(result.fields().len(), 1);
233        assert_eq!(result.fields()[0].name, "ts_year");
234        assert_eq!(
235            *result.fields()[0].field_type,
236            Type::Primitive(PrimitiveType::Int)
237        );
238    }
239
240    #[test]
241    fn test_unpartitioned() {
242        let schema = test_schema();
243        let spec = PartitionSpec::unpartition_spec();
244        let result = compute_unified_partition_type([&spec].into_iter(), &schema).unwrap();
245        assert!(result.fields().is_empty());
246    }
247
248    #[test]
249    fn test_multiple_fields_sorted_by_id() {
250        let schema = test_schema();
251        let spec = build_spec(&schema, 0, vec![
252            (3, "ts_year", Transform::Year),
253            (4, "category", Transform::Identity),
254        ]);
255
256        let result = compute_unified_partition_type([&spec].into_iter(), &schema).unwrap();
257        assert_eq!(result.fields().len(), 2);
258        assert!(result.fields()[0].id < result.fields()[1].id);
259    }
260
261    #[test]
262    fn test_newer_name_takes_precedence() {
263        let schema = test_schema();
264
265        // Spec 0: old name
266        let spec_v0 = PartitionSpec::builder(Arc::new(schema.clone()))
267            .with_spec_id(0)
268            .add_unbound_field(crate::spec::UnboundPartitionField {
269                source_id: 4,
270                field_id: Some(1000),
271                name: "cat_old".to_string(),
272                transform: Transform::Identity,
273            })
274            .unwrap()
275            .build()
276            .unwrap();
277
278        // Spec 1: newer name, same field_id
279        let spec_v1 = PartitionSpec::builder(Arc::new(schema.clone()))
280            .with_spec_id(1)
281            .add_unbound_field(crate::spec::UnboundPartitionField {
282                source_id: 4,
283                field_id: Some(1000),
284                name: "cat_new".to_string(),
285                transform: Transform::Identity,
286            })
287            .unwrap()
288            .build()
289            .unwrap();
290
291        let result =
292            compute_unified_partition_type([&spec_v0, &spec_v1].into_iter(), &schema).unwrap();
293        assert_eq!(result.fields().len(), 1);
294        assert_eq!(result.fields()[0].name, "cat_new");
295    }
296
297    #[test]
298    fn test_void_replaced_by_older_non_void() {
299        let schema = test_schema();
300
301        // Spec 0 (older): category partitioned by identity
302        let spec_v0 = PartitionSpec::builder(Arc::new(schema.clone()))
303            .with_spec_id(0)
304            .add_unbound_field(crate::spec::UnboundPartitionField {
305                source_id: 4,
306                field_id: Some(1000),
307                name: "category".to_string(),
308                transform: Transform::Identity,
309            })
310            .unwrap()
311            .build()
312            .unwrap();
313
314        // Spec 1 (newer): same field_id voided (partition dropped)
315        let spec_v1 = PartitionSpec::builder(Arc::new(schema.clone()))
316            .with_spec_id(1)
317            .add_unbound_field(crate::spec::UnboundPartitionField {
318                source_id: 4,
319                field_id: Some(1000),
320                name: "category_v2".to_string(),
321                transform: Transform::Void,
322            })
323            .unwrap()
324            .build()
325            .unwrap();
326
327        let result =
328            compute_unified_partition_type([&spec_v0, &spec_v1].into_iter(), &schema).unwrap();
329
330        assert_eq!(result.fields().len(), 1);
331        // Name from newer spec
332        assert_eq!(result.fields()[0].name, "category_v2");
333        // Type from older non-void spec
334        assert_eq!(
335            *result.fields()[0].field_type,
336            Type::Primitive(PrimitiveType::String)
337        );
338    }
339
340    #[test]
341    fn test_dropped_source_column_skipped() {
342        // Schema without field 4 (category was dropped)
343        let schema = Schema::builder()
344            .with_fields(vec![
345                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
346                NestedField::required(2, "data", Type::Primitive(PrimitiveType::String)).into(),
347            ])
348            .build()
349            .unwrap();
350
351        // Spec references source_id=4 which no longer exists in the schema.
352        // Deserialize directly since the builder rejects unknown source columns.
353        let spec = serde_json::from_value::<PartitionSpec>(serde_json::json!({
354            "spec-id": 0,
355            "fields": [{
356                "source-id": 4,
357                "field-id": 1000,
358                "name": "category",
359                "transform": "identity"
360            }]
361        }))
362        .unwrap();
363
364        let result = compute_unified_partition_type([&spec].into_iter(), &schema).unwrap();
365        assert!(result.fields().is_empty());
366    }
367
368    #[test]
369    fn test_evolution_adds_new_field() {
370        let schema = test_schema();
371
372        // Spec 0: partition by category
373        let spec_v0 = build_spec(&schema, 0, vec![(4, "category", Transform::Identity)]);
374
375        // Spec 1: partition by category + ts_year
376        let spec_v1 = PartitionSpec::builder(Arc::new(schema.clone()))
377            .with_spec_id(1)
378            .add_unbound_field(crate::spec::UnboundPartitionField {
379                source_id: 4,
380                field_id: Some(spec_v0.fields()[0].field_id),
381                name: "category".to_string(),
382                transform: Transform::Identity,
383            })
384            .unwrap()
385            .add_partition_field("ts", "ts_year", Transform::Year)
386            .unwrap()
387            .build()
388            .unwrap();
389
390        let result =
391            compute_unified_partition_type([&spec_v0, &spec_v1].into_iter(), &schema).unwrap();
392        assert_eq!(result.fields().len(), 2);
393    }
394
395    #[test]
396    fn test_unknown_transform_errors() {
397        let schema = test_schema();
398
399        // A spec using an unknown transform. Deserialize directly since the builder
400        // validates transforms.
401        let spec = serde_json::from_value::<PartitionSpec>(serde_json::json!({
402            "spec-id": 0,
403            "fields": [{
404                "source-id": 4,
405                "field-id": 1000,
406                "name": "category",
407                "transform": "unknown"
408            }]
409        }))
410        .unwrap();
411
412        let err = compute_unified_partition_type([&spec].into_iter(), &schema).unwrap_err();
413        assert_eq!(err.kind(), ErrorKind::DataInvalid);
414    }
415
416    #[test]
417    fn test_conflicting_partition_fields_error() {
418        let schema = test_schema();
419
420        // Spec 0: field id 1000 -> source 4 (category), identity
421        let spec_v0 = serde_json::from_value::<PartitionSpec>(serde_json::json!({
422            "spec-id": 0,
423            "fields": [{
424                "source-id": 4,
425                "field-id": 1000,
426                "name": "category",
427                "transform": "identity"
428            }]
429        }))
430        .unwrap();
431
432        // Spec 1: field id 1000 reused for a different source (ts) and transform (year).
433        // This conflicts with spec 0 and must be rejected.
434        let spec_v1 = serde_json::from_value::<PartitionSpec>(serde_json::json!({
435            "spec-id": 1,
436            "fields": [{
437                "source-id": 3,
438                "field-id": 1000,
439                "name": "ts_year",
440                "transform": "year"
441            }]
442        }))
443        .unwrap();
444
445        let err =
446            compute_unified_partition_type([&spec_v0, &spec_v1].into_iter(), &schema).unwrap_err();
447        assert_eq!(err.kind(), ErrorKind::DataInvalid);
448    }
449}
iceberg/partitioning.rs

iceberg/
partitioning.rs
