From 06b4a8590ddaf93a46c1ef80f01858840accaae8 Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Thu, 16 Jul 2026 13:39:07 -0700 Subject: [PATCH 1/6] feat: add support for _partition metadata column --- crates/iceberg/public-api.txt | 3 + crates/iceberg/src/arrow/mod.rs | 1 + crates/iceberg/src/arrow/reader/pipeline.rs | 25 +- crates/iceberg/src/arrow/reader/row_filter.rs | 2 + .../src/arrow/record_batch_transformer.rs | 581 +++++++++++++++--- crates/iceberg/src/arrow/value.rs | 74 +-- crates/iceberg/src/lib.rs | 1 + crates/iceberg/src/partitioning.rs | 101 +++ crates/iceberg/src/scan/context.rs | 10 +- crates/iceberg/src/scan/mod.rs | 20 +- crates/iceberg/src/scan/task.rs | 18 +- 11 files changed, 692 insertions(+), 144 deletions(-) create mode 100644 crates/iceberg/src/partitioning.rs diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 610675c83d..8e90d08d34 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -1158,6 +1158,8 @@ pub fn iceberg::metadata_columns::partition_field(partition_fields: alloc::vec:: pub fn iceberg::metadata_columns::pos_field() -> &'static iceberg::spec::NestedFieldRef pub fn iceberg::metadata_columns::row_id_field() -> &'static iceberg::spec::NestedFieldRef pub fn iceberg::metadata_columns::spec_id_field() -> &'static iceberg::spec::NestedFieldRef +pub mod iceberg::partitioning +pub fn iceberg::partitioning::compute_unified_partition_type<'a>(partition_specs: impl core::iter::traits::iterator::Iterator, schema: &iceberg::spec::Schema) -> iceberg::Result pub mod iceberg::puffin pub enum iceberg::puffin::CompressionCodec pub iceberg::puffin::CompressionCodec::Gzip(u8) @@ -1274,6 +1276,7 @@ pub iceberg::scan::FileScanTask::project_field_ids: alloc::vec::Vec pub iceberg::scan::FileScanTask::record_count: core::option::Option pub iceberg::scan::FileScanTask::schema: iceberg::spec::SchemaRef pub iceberg::scan::FileScanTask::start: u64 +pub iceberg::scan::FileScanTask::unified_partition_type: core::option::Option> impl iceberg::scan::FileScanTask pub fn iceberg::scan::FileScanTask::data_file_path(&self) -> &str pub fn iceberg::scan::FileScanTask::predicate(&self) -> core::option::Option<&iceberg::expr::BoundPredicate> diff --git a/crates/iceberg/src/arrow/mod.rs b/crates/iceberg/src/arrow/mod.rs index 4e0286cbde..1473809436 100644 --- a/crates/iceberg/src/arrow/mod.rs +++ b/crates/iceberg/src/arrow/mod.rs @@ -32,6 +32,7 @@ mod reader; /// RecordBatch projection utilities pub mod record_batch_projector; pub(crate) mod record_batch_transformer; +pub(crate) use record_batch_transformer::build_partition_column_constant; mod scan_metrics; mod value; diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index 579f047266..9d05139c43 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -34,6 +34,7 @@ use super::{ ArrowFileReader, ArrowReader, ParquetReadOptions, add_fallback_field_ids_to_arrow_schema, apply_name_mapping_to_arrow_schema, }; +use crate::arrow::build_partition_column_constant; use crate::arrow::caching_delete_file_loader::CachingDeleteFileLoader; use crate::arrow::int96::coerce_int96_timestamps; use crate::arrow::record_batch_transformer::RecordBatchTransformerBuilder; @@ -42,11 +43,11 @@ use crate::encryption::StandardKeyMetadata; use crate::error::Result; use crate::io::{FileIO, FileMetadata, FileRead}; use crate::metadata_columns::{ - RESERVED_COL_NAME_POS, RESERVED_FIELD_ID_FILE, RESERVED_FIELD_ID_POS, - RESERVED_FIELD_ID_SPEC_ID, is_metadata_field, + RESERVED_COL_NAME_POS, RESERVED_FIELD_ID_FILE, RESERVED_FIELD_ID_PARTITION, + RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_SPEC_ID, is_metadata_field, }; use crate::scan::{ArrowRecordBatchStream, FileScanTask, FileScanTaskStream}; -use crate::spec::Datum; +use crate::spec::{Datum, PartitionSpec, Struct}; use crate::{Error, ErrorKind}; impl ArrowReader { @@ -311,6 +312,24 @@ impl FileScanTaskReader { record_batch_transformer_builder.with_virtual_field(RESERVED_FIELD_ID_POS); } + // Add the _partition metadata struct column if it's in the projected fields. + // Computed lazily here at read time from the unified partition type + task's spec + data. + if task + .project_field_ids() + .contains(&RESERVED_FIELD_ID_PARTITION) + && let Some(unified_type) = &task.unified_partition_type + { + let (spec, partition_data) = match (&task.partition_spec, &task.partition) { + (Some(spec), Some(data)) => (spec.clone(), data.clone()), + // Unpartitioned table or missing spec: build_partition_column_constant + // handles the empty unified_type case with an early return. + _ => (Arc::new(PartitionSpec::unpartition_spec()), Struct::empty()), + }; + let constant = build_partition_column_constant(unified_type, &spec, &partition_data)?; + record_batch_transformer_builder = + record_batch_transformer_builder.with_partition_column_precomputed(constant); + } + let mut record_batch_transformer = record_batch_transformer_builder.build(); if let Some(batch_size) = self.batch_size { diff --git a/crates/iceberg/src/arrow/reader/row_filter.rs b/crates/iceberg/src/arrow/reader/row_filter.rs index a94c159d1d..d82b9a0aff 100644 --- a/crates/iceberg/src/arrow/reader/row_filter.rs +++ b/crates/iceberg/src/arrow/reader/row_filter.rs @@ -1147,6 +1147,7 @@ mod tests { partition: None, partition_spec: None, name_mapping: None, + unified_partition_type: None, case_sensitive: false, key_metadata: None, }; @@ -1246,6 +1247,7 @@ mod tests { partition: None, partition_spec: None, name_mapping: None, + unified_partition_type: None, case_sensitive: false, key_metadata: None, }; diff --git a/crates/iceberg/src/arrow/record_batch_transformer.rs b/crates/iceberg/src/arrow/record_batch_transformer.rs index ec91eb7b3c..61f3a105ad 100644 --- a/crates/iceberg/src/arrow/record_batch_transformer.rs +++ b/crates/iceberg/src/arrow/record_batch_transformer.rs @@ -20,21 +20,39 @@ use std::sync::Arc; use arrow_array::{ Array as ArrowArray, ArrayRef, Int32Array, RecordBatch, RecordBatchOptions, RunArray, + StructArray, }; use arrow_cast::cast; use arrow_schema::{ - DataType, Field, FieldRef, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, SchemaRef, + DataType, Field, FieldRef, Fields, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, + SchemaRef, }; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; use crate::arrow::value::{create_primitive_array_repeated, create_primitive_array_single_element}; use crate::arrow::{datum_to_arrow_type_with_ree, schema_to_arrow_schema, type_to_arrow_type}; -use crate::metadata_columns::get_metadata_field; +use crate::metadata_columns::{ + RESERVED_COL_NAME_PARTITION, RESERVED_FIELD_ID_PARTITION, get_metadata_field, +}; use crate::spec::{ - Datum, Literal, PartitionSpec, PrimitiveLiteral, Schema as IcebergSchema, Struct, Transform, + Datum, Literal, PartitionSpec, PrimitiveLiteral, Schema as IcebergSchema, Struct, StructType, + Transform, }; use crate::{Error, ErrorKind, Result}; +/// Create an Arrow Field with PARQUET_FIELD_ID_META_KEY metadata attached. +fn field_with_id( + name: impl Into, + data_type: DataType, + nullable: bool, + field_id: i32, +) -> Field { + Field::new(name, data_type, nullable).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + field_id.to_string(), + )])) +} + /// Build a map of field ID to constant value (as Datum) for identity-partitioned fields. /// /// Implements Iceberg spec "Column Projection" rule #1: use partition metadata constants @@ -140,6 +158,13 @@ pub(crate) enum ColumnSource { target_type: DataType, value: Option, }, + + // A struct column where each child is a constant primitive value. + // Used for the _partition metadata column. + AddStructConstant { + fields: Fields, + child_values: Vec>, + }, // The iceberg spec refers to other permissible schema evolution actions // (see https://iceberg.apache.org/spec/#schema-evolution): // renaming fields, deleting fields and reordering fields. @@ -186,16 +211,82 @@ enum SchemaComparison { /// Builder for RecordBatchTransformer to improve ergonomics when constructing with optional parameters. /// -/// Constant fields are pre-computed for both virtual/metadata fields (like _file) and -/// identity-partitioned fields to avoid duplicate work during batch processing. +/// All per-file constants (scalar metadata like `_file`, identity partition values, +/// and the `_partition` struct) are stored in a single `metadata_columns` map keyed by +/// field_id. This unified representation (via [`MetadataColumnSource`]) means the +/// transformer handles all constant columns through one code path. #[derive(Debug)] pub(crate) struct RecordBatchTransformerBuilder { snapshot_schema: Arc, projected_iceberg_field_ids: Vec, - constant_fields: HashMap, + metadata_columns: HashMap, virtual_fields: HashSet, } +/// Unified source for metadata/virtual column constants. +/// +/// Both scalar metadata columns (like `_file`) and the struct `_partition` column are +/// per-file constants synthesized during scan. This enum unifies both representations +/// so the transformer can handle them through a single `metadata_columns` map keyed +/// by field_id, replacing the previous split between `constant_fields: HashMap` +/// and `partition_column: Option`. +/// +/// # Design for #2699 +/// +/// When `_pos` is implemented (via arrow-rs RowNumber — a real source column passed through, +/// not synthesized), it will NOT use this enum. `_pos` is not a constant and belongs in the +/// virtual_fields / pass-through axis. This enum is exclusively for constant-per-file values. +/// +/// The `ColumnSource` enum's `Add` and `AddStructConstant` variants map 1:1 to this enum's +/// variants during transform generation. #2699 may collapse `ColumnSource::Add` and +/// `ColumnSource::AddStructConstant` into a single `ColumnSource::AddMetadata { source: MetadataColumnSource }` +/// variant, but that refactor can happen independently. +#[derive(Debug, Clone, PartialEq)] +pub enum MetadataColumnSource { + /// A scalar constant (e.g., `_file` path, `_spec_id`, identity partition values). + /// The Datum carries both the Iceberg type and the value. + Scalar(Datum), + /// A struct constant (the `_partition` column). Each child is a primitive constant + /// or null (for partition evolution gaps). + Struct(PartitionColumnConstant), +} + +/// Pre-computed data for the _partition struct constant. +#[derive(Debug, Clone, PartialEq)] +pub struct PartitionColumnConstant { + fields: Fields, + child_values: Vec>, +} + +impl PartitionColumnConstant { + /// Create a new PartitionColumnConstant, validating that fields and child_values have + /// the same length. + pub fn new(fields: Fields, child_values: Vec>) -> Result { + if fields.len() != child_values.len() { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "PartitionColumnConstant: fields length ({}) != child_values length ({})", + fields.len(), + child_values.len() + ), + )); + } + Ok(Self { + fields, + child_values, + }) + } + + pub(crate) fn fields(&self) -> &Fields { + &self.fields + } + + pub(crate) fn child_values(&self) -> &[Option] { + &self.child_values + } +} + impl RecordBatchTransformerBuilder { pub(crate) fn new( snapshot_schema: Arc, @@ -204,19 +295,16 @@ impl RecordBatchTransformerBuilder { Self { snapshot_schema, projected_iceberg_field_ids: projected_iceberg_field_ids.to_vec(), - constant_fields: HashMap::new(), + metadata_columns: HashMap::new(), virtual_fields: HashSet::new(), } } - /// Add a constant value for a specific field ID. + /// Add a scalar constant value for a specific field ID. /// This is used for virtual/metadata fields like _file that have constant values per batch. - /// - /// # Arguments - /// * `field_id` - The field ID to associate with the constant - /// * `datum` - The constant value (with type) for this field pub(crate) fn with_constant(mut self, field_id: i32, datum: Datum) -> Self { - self.constant_fields.insert(field_id, datum); + self.metadata_columns + .insert(field_id, MetadataColumnSource::Scalar(datum)); self } @@ -224,19 +312,18 @@ impl RecordBatchTransformerBuilder { /// /// Both partition_spec and partition_data must be provided together since the spec defines /// which fields are identity-partitioned, and the data provides their constant values. - /// This method computes the partition constants and merges them into constant_fields. + /// This method computes the partition constants and merges them into metadata_columns. pub(crate) fn with_partition( mut self, partition_spec: Arc, partition_data: Struct, ) -> Result { - // Compute partition constants for identity-transformed fields (already returns Datum) let partition_constants = constants_map(&partition_spec, &partition_data, &self.snapshot_schema)?; - // Add partition constants to constant_fields for (field_id, datum) in partition_constants { - self.constant_fields.insert(field_id, datum); + self.metadata_columns + .insert(field_id, MetadataColumnSource::Scalar(datum)); } Ok(self) @@ -249,11 +336,23 @@ impl RecordBatchTransformerBuilder { self } + /// Set a pre-computed _partition column constant directly. + pub(crate) fn with_partition_column_precomputed( + mut self, + partition_column: PartitionColumnConstant, + ) -> Self { + self.metadata_columns.insert( + RESERVED_FIELD_ID_PARTITION, + MetadataColumnSource::Struct(partition_column), + ); + self + } + pub(crate) fn build(self) -> RecordBatchTransformer { RecordBatchTransformer { snapshot_schema: self.snapshot_schema, projected_iceberg_field_ids: self.projected_iceberg_field_ids, - constant_fields: self.constant_fields, + metadata_columns: self.metadata_columns, virtual_fields: self.virtual_fields, batch_transform: None, } @@ -294,10 +393,8 @@ impl RecordBatchTransformerBuilder { pub(crate) struct RecordBatchTransformer { snapshot_schema: Arc, projected_iceberg_field_ids: Vec, - // Pre-computed constant field information: field_id -> Datum - // Includes both virtual/metadata fields (like _file) and identity-partitioned fields - // Datum holds both the Iceberg type and the value - constant_fields: HashMap, + // Unified map of all per-file constant columns (metadata, identity partition, _partition struct). + metadata_columns: HashMap, // Field IDs whose data is delivered by the source batch as an arrow-rs virtual column // (e.g. _pos via RowNumber). These fields bypass the snapshot-schema lookup and the @@ -344,7 +441,7 @@ impl RecordBatchTransformer { record_batch.schema_ref(), self.snapshot_schema.as_ref(), &self.projected_iceberg_field_ids, - &self.constant_fields, + &self.metadata_columns, &self.virtual_fields, )?); @@ -364,7 +461,7 @@ impl RecordBatchTransformer { source_schema: &ArrowSchemaRef, snapshot_schema: &IcebergSchema, projected_iceberg_field_ids: &[i32], - constant_fields: &HashMap, + metadata_columns: &HashMap, virtual_fields: &HashSet, ) -> Result { let mapped_unprojected_arrow_schema = Arc::new(schema_to_arrow_schema(snapshot_schema)?); @@ -376,38 +473,54 @@ impl RecordBatchTransformer { let fields: Result> = projected_iceberg_field_ids .iter() .map(|field_id| { - // Metadata/virtual fields (like _file, _spec_id) don't exist in the table - // schema, so build their Arrow field from the metadata column definition and - // the pre-computed constant's type. - if let Some(datum) = constant_fields.get(field_id) - && let Ok(iceberg_field) = get_metadata_field(*field_id) - { - let arrow_type = datum_to_arrow_type_with_ree(datum); - let arrow_field = - Field::new(&iceberg_field.name, arrow_type, !iceberg_field.required) - .with_metadata(HashMap::from([( - PARQUET_FIELD_ID_META_KEY.to_string(), - iceberg_field.id.to_string(), - )])); - Ok(Arc::new(arrow_field)) - } else if virtual_fields.contains(field_id) { + match metadata_columns.get(field_id) { + Some(MetadataColumnSource::Struct(pc)) => { + let struct_type = DataType::Struct(pc.fields().clone()); + let nullable = pc.fields().is_empty(); + let arrow_field = field_with_id( + RESERVED_COL_NAME_PARTITION, + struct_type, + nullable, + *field_id, + ); + return Ok(Arc::new(arrow_field)); + } + Some(MetadataColumnSource::Scalar(datum)) => { + if let Ok(iceberg_field) = get_metadata_field(*field_id) { + let arrow_type = datum_to_arrow_type_with_ree(datum); + let arrow_field = field_with_id( + &iceberg_field.name, + arrow_type, + !iceberg_field.required, + iceberg_field.id, + ); + return Ok(Arc::new(arrow_field)); + } + // Identity partition constant -- fall through to use the + // mapped schema field (read from file if present). + } + None => {} + } + + if virtual_fields.contains(field_id) { let virtual_field = get_metadata_field(*field_id)?; let arrow_type = type_to_arrow_type(&virtual_field.field_type)?; - let arrow_field = - Field::new(&virtual_field.name, arrow_type, !virtual_field.required) - .with_metadata(HashMap::from([( - PARQUET_FIELD_ID_META_KEY.to_string(), - virtual_field.id.to_string(), - )])); - Ok(Arc::new(arrow_field)) - } else { - // Regular fields and identity-partitioned constant fields both exist in the - // table schema, so use the mapped Arrow field as-is. - Ok(field_id_to_mapped_schema_map - .get(field_id) - .ok_or(Error::new(ErrorKind::Unexpected, "field not found"))? - .0 - .clone()) + let arrow_field = field_with_id( + &virtual_field.name, + arrow_type, + !virtual_field.required, + virtual_field.id, + ); + return Ok(Arc::new(arrow_field)); + } + + // Regular fields and identity-partitioned constant fields both exist in the + // table schema, so use the mapped Arrow field as-is. + Ok(field_id_to_mapped_schema_map + .get(field_id) + .ok_or(Error::new(ErrorKind::Unexpected, "field not found"))? + .0 + .clone()) } }) .collect(); @@ -423,7 +536,7 @@ impl RecordBatchTransformer { snapshot_schema, projected_iceberg_field_ids, field_id_to_mapped_schema_map, - constant_fields, + metadata_columns, virtual_fields, )?, target_schema, @@ -492,7 +605,7 @@ impl RecordBatchTransformer { snapshot_schema: &IcebergSchema, projected_iceberg_field_ids: &[i32], field_id_to_mapped_schema_map: HashMap, - constant_fields: &HashMap, + metadata_columns: &HashMap, virtual_fields: &HashSet, ) -> Result> { let field_id_to_source_schema_map = @@ -501,39 +614,42 @@ impl RecordBatchTransformer { projected_iceberg_field_ids .iter() .map(|field_id| { - // Check if this is a constant field (metadata/virtual or identity-partitioned). - // - // Metadata/virtual fields (like _file, _spec_id) never exist in the data file, - // so they always use their pre-computed constant value. - // - // For identity-partitioned fields, the Iceberg spec's "Column Projection" rules - // only apply to "field ids which are not present in a data file". When the column - // IS present in the Parquet file, it must be read from the file; the partition - // metadata constant is only a fallback for when the column is absent (e.g. add_files). - if let Some(datum) = constant_fields.get(field_id) { - let is_metadata_field = get_metadata_field(*field_id).is_ok(); - let present_in_file = field_id_to_source_schema_map.contains_key(field_id); - - if is_metadata_field || !present_in_file { - let arrow_type = if is_metadata_field { - datum_to_arrow_type_with_ree(datum) - } else { - field_id_to_mapped_schema_map - .get(field_id) - .ok_or(Error::new( - ErrorKind::Unexpected, - "could not find field in schema", - ))? - .0 - .data_type() - .clone() - }; - - return Ok(ColumnSource::Add { - value: Some(datum.literal().clone()), - target_type: arrow_type, + match metadata_columns.get(field_id) { + Some(MetadataColumnSource::Struct(pc)) => { + return Ok(ColumnSource::AddStructConstant { + fields: pc.fields().clone(), + child_values: pc.child_values().to_vec(), }); } + Some(MetadataColumnSource::Scalar(datum)) => { + let is_metadata = get_metadata_field(*field_id).is_ok(); + let present_in_file = + field_id_to_source_schema_map.contains_key(field_id); + + if is_metadata || !present_in_file { + let arrow_type = if is_metadata { + datum_to_arrow_type_with_ree(datum) + } else { + field_id_to_mapped_schema_map + .get(field_id) + .ok_or(Error::new( + ErrorKind::Unexpected, + "could not find field in schema", + ))? + .0 + .data_type() + .clone() + }; + + return Ok(ColumnSource::Add { + value: Some(datum.literal().clone()), + target_type: arrow_type, + }); + } + // Identity partition field present in the file -- fall through + // to read from the file instead of using the constant. + } + None => {} } if virtual_fields.contains(field_id) { @@ -682,6 +798,11 @@ impl RecordBatchTransformer { ColumnSource::Add { target_type, value } => { Self::create_column(target_type, value, num_rows)? } + + ColumnSource::AddStructConstant { + fields, + child_values, + } => Self::create_struct_column(fields, child_values, num_rows)?, }) }) .collect() @@ -723,6 +844,89 @@ impl RecordBatchTransformer { create_primitive_array_repeated(target_type, prim_lit, num_rows) } } + + fn create_struct_column( + fields: &Fields, + child_values: &[Option], + num_rows: usize, + ) -> Result { + if fields.is_empty() { + let nulls = arrow_buffer::NullBuffer::new_null(num_rows); + return Ok(Arc::new(StructArray::new_empty_fields( + num_rows, + Some(nulls), + ))); + } + + let child_arrays: Vec = fields + .iter() + .zip(child_values.iter()) + .map(|(field, value)| { + create_primitive_array_repeated(field.data_type(), value, num_rows) + }) + .collect::>()?; + + Ok(Arc::new(StructArray::try_new( + fields.clone(), + child_arrays, + None, + )?)) + } +} + +/// Builds a [`PartitionColumnConstant`] from the unified partition type and a file's +/// partition spec/data. +/// +/// If `unified_partition_type` has no fields (unpartitioned table), returns an empty constant +/// that renders as a null struct column. +/// +/// For each field in the unified partition type: +/// - If it corresponds to a field in this file's partition spec, use the value from partition_data +/// - Otherwise (partition evolution), use null +pub fn build_partition_column_constant( + unified_partition_type: &StructType, + partition_spec: &PartitionSpec, + partition_data: &Struct, +) -> Result { + // Unpartitioned table: empty struct rendered as null + if unified_partition_type.fields().is_empty() { + return PartitionColumnConstant::new(Fields::empty(), vec![]); + } + + use crate::arrow::type_to_arrow_type; + + let spec_fields = partition_spec.fields(); + + let mut arrow_fields = Vec::with_capacity(unified_partition_type.fields().len()); + let mut child_values = Vec::with_capacity(unified_partition_type.fields().len()); + + for unified_field in unified_partition_type.fields() { + let arrow_type = type_to_arrow_type(&unified_field.field_type)?; + // Don't attach PARQUET_FIELD_ID_META_KEY metadata to child fields -- + // Spark's output schema for _partition doesn't include it, and Arrow's + // Field::eq checks metadata, causing the schema adapter to insert an + // unnecessary cast operation per batch. + // Always nullable: partition evolution means any field may be absent + // for files written under a different spec. + let arrow_field = Field::new(&unified_field.name, arrow_type, true); + arrow_fields.push(Arc::new(arrow_field)); + + // Find matching field in this file's partition spec by field_id + let value = spec_fields + .iter() + .position(|f| f.field_id == unified_field.id) + .and_then(|pos| { + // Get the value from partition_data at this position + match &partition_data[pos] { + Some(Literal::Primitive(prim)) => Some(prim.clone()), + _ => None, + } + }); + + child_values.push(value); + } + + PartitionColumnConstant::new(Fields::from(arrow_fields), child_values) } #[cfg(test)] @@ -737,6 +941,7 @@ mod test { use arrow_schema::{DataType, Field, Schema as ArrowSchema}; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + use crate::arrow::build_partition_column_constant; use crate::arrow::record_batch_transformer::{ RecordBatchTransformer, RecordBatchTransformerBuilder, }; @@ -2008,4 +2213,206 @@ mod test { assert_eq!(id_col.len(), 3); assert_eq!(id_col.values(), &[100, 200, 300]); } + + #[test] + fn partition_column_struct_constant() { + use arrow_array::StructArray; + + use crate::metadata_columns::RESERVED_FIELD_ID_PARTITION; + use crate::spec::Transform; + + let snapshot_schema = Arc::new( + Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(2, "name", Type::Primitive(PrimitiveType::String)).into(), + ]) + .build() + .unwrap(), + ); + + // Partition spec: identity(id) + let partition_spec = Arc::new( + crate::spec::PartitionSpec::builder(snapshot_schema.clone()) + .with_spec_id(0) + .add_partition_field("id", "id", Transform::Identity) + .unwrap() + .build() + .unwrap(), + ); + + // Unified partition type: just one field (same as the spec's type) + let unified_partition_type = partition_spec.partition_type(&snapshot_schema).unwrap(); + + // Partition data: id=42 + let partition_data = Struct::from_iter(vec![Some(Literal::int(42))]); + + // Parquet file has both columns + let parquet_schema = Arc::new(ArrowSchema::new(vec![ + simple_field("id", DataType::Int32, false, "1"), + simple_field("name", DataType::Utf8, true, "2"), + ])); + + // Project id, name, and _partition + let projected_field_ids = [1, 2, RESERVED_FIELD_ID_PARTITION]; + + let partition_column = build_partition_column_constant( + &unified_partition_type, + &partition_spec, + &partition_data, + ) + .unwrap(); + let mut transformer = + RecordBatchTransformerBuilder::new(snapshot_schema, &projected_field_ids) + .with_partition_column_precomputed(partition_column) + .build(); + + let parquet_batch = RecordBatch::try_new(parquet_schema, vec![ + Arc::new(Int32Array::from(vec![100, 200, 300])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ]) + .unwrap(); + + let result = transformer.process_record_batch(parquet_batch).unwrap(); + + assert_eq!(result.num_columns(), 3); + assert_eq!(result.num_rows(), 3); + + // id column from file + let id_col = result + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(id_col.values(), &[100, 200, 300]); + + // name column from file + let name_col = result + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(name_col.value(0), "a"); + + // _partition struct column + let partition_col = result + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(partition_col.num_columns(), 1); + assert_eq!(partition_col.len(), 3); + + let inner = partition_col + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(inner.value(0), 42); + assert_eq!(inner.value(1), 42); + assert_eq!(inner.value(2), 42); + } + + #[test] + fn partition_column_with_evolution() { + use arrow_array::StructArray; + + use crate::metadata_columns::RESERVED_FIELD_ID_PARTITION; + use crate::partitioning::compute_unified_partition_type; + use crate::spec::Transform; + + // Schema with two fields that could be partition sources + let snapshot_schema = Arc::new( + Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::required(1, "year", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(2, "month", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(3, "data", Type::Primitive(PrimitiveType::String)).into(), + ]) + .build() + .unwrap(), + ); + + // Old spec: partition by year only + let spec_v0 = crate::spec::PartitionSpec::builder(snapshot_schema.clone()) + .with_spec_id(0) + .add_partition_field("year", "year", Transform::Identity) + .unwrap() + .build() + .unwrap(); + + // New spec: partition by year and month + let spec_v1 = crate::spec::PartitionSpec::builder(snapshot_schema.clone()) + .with_spec_id(1) + .add_partition_field("year", "year", Transform::Identity) + .unwrap() + .add_partition_field("month", "month", Transform::Identity) + .unwrap() + .build() + .unwrap(); + + // Unified type includes both year and month + let unified_partition_type = + compute_unified_partition_type([&spec_v0, &spec_v1].into_iter(), &snapshot_schema) + .unwrap(); + + assert_eq!(unified_partition_type.fields().len(), 2); + + // File written with spec_v0 (only has year=2023) + let partition_data = Struct::from_iter(vec![Some(Literal::int(2023))]); + + let parquet_schema = Arc::new(ArrowSchema::new(vec![simple_field( + "data", + DataType::Utf8, + true, + "3", + )])); + + let projected_field_ids = [3, RESERVED_FIELD_ID_PARTITION]; + + let partition_column = + build_partition_column_constant(&unified_partition_type, &spec_v0, &partition_data) + .unwrap(); + let mut transformer = + RecordBatchTransformerBuilder::new(snapshot_schema, &projected_field_ids) + .with_partition_column_precomputed(partition_column) + .build(); + + let parquet_batch = + RecordBatch::try_new(parquet_schema, vec![Arc::new(StringArray::from(vec![ + "hello", "world", + ]))]) + .unwrap(); + + let result = transformer.process_record_batch(parquet_batch).unwrap(); + + assert_eq!(result.num_columns(), 2); + assert_eq!(result.num_rows(), 2); + + // _partition struct has 2 fields: year (present) and month (null for this old spec file) + let partition_col = result + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(partition_col.num_columns(), 2); + + let year_col = partition_col + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(year_col.value(0), 2023); + assert_eq!(year_col.value(1), 2023); + + let month_col = partition_col + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(month_col.is_null(0)); + assert!(month_col.is_null(1)); + } } diff --git a/crates/iceberg/src/arrow/value.rs b/crates/iceberg/src/arrow/value.rs index 34f6ce6e94..d22e565a4a 100644 --- a/crates/iceberg/src/arrow/value.rs +++ b/crates/iceberg/src/arrow/value.rs @@ -21,7 +21,7 @@ use arrow_array::{ Array, ArrayRef, BinaryArray, BooleanArray, Date32Array, Decimal128Array, FixedSizeBinaryArray, FixedSizeListArray, Float32Array, Float64Array, Int32Array, Int64Array, LargeBinaryArray, LargeListArray, LargeStringArray, ListArray, MapArray, StringArray, StructArray, - Time64MicrosecondArray, TimestampMicrosecondArray, TimestampNanosecondArray, + Time64MicrosecondArray, TimestampMicrosecondArray, TimestampNanosecondArray, new_null_array, }; use arrow_buffer::NullBuffer; use arrow_schema::{DataType, FieldRef, TimeUnit}; @@ -820,37 +820,22 @@ pub(crate) fn create_primitive_array_repeated( num_rows: usize, ) -> Result { Ok(match (data_type, prim_lit) { + // --- Primitive Some arms --- (DataType::Boolean, Some(PrimitiveLiteral::Boolean(value))) => { Arc::new(BooleanArray::from(vec![*value; num_rows])) } - (DataType::Boolean, None) => { - let vals: Vec> = vec![None; num_rows]; - Arc::new(BooleanArray::from(vals)) - } (DataType::Int32, Some(PrimitiveLiteral::Int(value))) => { Arc::new(Int32Array::from(vec![*value; num_rows])) } - (DataType::Int32, None) => { - let vals: Vec> = vec![None; num_rows]; - Arc::new(Int32Array::from(vals)) - } (DataType::Date32, Some(PrimitiveLiteral::Int(value))) => { Arc::new(Date32Array::from(vec![*value; num_rows])) } - (DataType::Date32, None) => { - let vals: Vec> = vec![None; num_rows]; - Arc::new(Date32Array::from(vals)) - } (DataType::Int64, Some(PrimitiveLiteral::Int(value))) => { Arc::new(Int64Array::from(vec![i64::from(*value); num_rows])) } (DataType::Int64, Some(PrimitiveLiteral::Long(value))) => { Arc::new(Int64Array::from(vec![*value; num_rows])) } - (DataType::Int64, None) => { - let vals: Vec> = vec![None; num_rows]; - Arc::new(Int64Array::from(vals)) - } ( DataType::Timestamp(TimeUnit::Microsecond, timezone), Some(PrimitiveLiteral::Long(value)), @@ -862,15 +847,6 @@ pub(crate) fn create_primitive_array_repeated( Arc::new(array) } } - (DataType::Timestamp(TimeUnit::Microsecond, timezone), None) => { - let vals: Vec> = vec![None; num_rows]; - let array = TimestampMicrosecondArray::from(vals); - if let Some(timezone) = timezone { - Arc::new(array.with_timezone(timezone.clone())) - } else { - Arc::new(array) - } - } ( DataType::Timestamp(TimeUnit::Nanosecond, timezone), Some(PrimitiveLiteral::Long(value)), @@ -882,42 +858,32 @@ pub(crate) fn create_primitive_array_repeated( Arc::new(array) } } - (DataType::Timestamp(TimeUnit::Nanosecond, timezone), None) => { - let vals: Vec> = vec![None; num_rows]; - let array = TimestampNanosecondArray::from(vals); - if let Some(timezone) = timezone { - Arc::new(array.with_timezone(timezone.clone())) - } else { - Arc::new(array) - } - } (DataType::Float32, Some(PrimitiveLiteral::Float(value))) => { Arc::new(Float32Array::from(vec![value.0; num_rows])) } - (DataType::Float32, None) => { - let vals: Vec> = vec![None; num_rows]; - Arc::new(Float32Array::from(vals)) - } (DataType::Float64, Some(PrimitiveLiteral::Double(value))) => { Arc::new(Float64Array::from(vec![value.0; num_rows])) } - (DataType::Float64, None) => { - let vals: Vec> = vec![None; num_rows]; - Arc::new(Float64Array::from(vals)) - } (DataType::Utf8, Some(PrimitiveLiteral::String(value))) => { Arc::new(StringArray::from(vec![value.clone(); num_rows])) } - (DataType::Utf8, None) => { - let vals: Vec> = vec![None; num_rows]; - Arc::new(StringArray::from(vals)) - } (DataType::Binary, Some(PrimitiveLiteral::Binary(value))) => { Arc::new(BinaryArray::from_vec(vec![value; num_rows])) } - (DataType::Binary, None) => { - let vals: Vec> = vec![None; num_rows]; - Arc::new(BinaryArray::from_opt_vec(vals)) + (DataType::LargeBinary, Some(PrimitiveLiteral::Binary(value))) => { + Arc::new(LargeBinaryArray::from_vec(vec![value; num_rows])) + } + (DataType::FixedSizeBinary(len), Some(PrimitiveLiteral::Binary(value))) => { + let repeated: Vec<&[u8]> = vec![value.as_slice(); num_rows]; + Arc::new(FixedSizeBinaryArray::try_from_iter(repeated.into_iter()).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("Failed to create FixedSizeBinary({len}) array: {e}"), + ) + })?) + } + (DataType::Time64(TimeUnit::Microsecond), Some(PrimitiveLiteral::Long(value))) => { + Arc::new(Time64MicrosecondArray::from(vec![*value; num_rows])) } (DataType::Decimal128(precision, scale), Some(PrimitiveLiteral::Int128(value))) => { Arc::new( @@ -947,6 +913,8 @@ pub(crate) fn create_primitive_array_repeated( })?, ) } + + // --- Special-case None arms --- (DataType::Decimal128(precision, scale), None) => { let vals: Vec> = vec![None; num_rows]; Arc::new( @@ -963,7 +931,7 @@ pub(crate) fn create_primitive_array_repeated( ) } (DataType::Struct(fields), None) => { - // Create a StructArray filled with nulls + // Create a StructArray filled with nulls, recursively creating null children let null_arrays: Vec = fields .iter() .map(|field| create_primitive_array_repeated(field.data_type(), &None, num_rows)) @@ -976,6 +944,10 @@ pub(crate) fn create_primitive_array_repeated( )) } (DataType::Null, _) => Arc::new(arrow_array::NullArray::new(num_rows)), + + // --- Catch-all null arm: use arrow-rs new_null_array for any remaining DataType --- + (dt, None) => new_null_array(dt, num_rows), + (dt, _) => { return Err(Error::new( ErrorKind::Unexpected, diff --git a/crates/iceberg/src/lib.rs b/crates/iceberg/src/lib.rs index 4e346460f5..301992d15e 100644 --- a/crates/iceberg/src/lib.rs +++ b/crates/iceberg/src/lib.rs @@ -100,6 +100,7 @@ pub mod writer; mod delete_vector; pub mod metadata_columns; +pub mod partitioning; pub mod puffin; /// Utility functions and modules. pub mod util; diff --git a/crates/iceberg/src/partitioning.rs b/crates/iceberg/src/partitioning.rs new file mode 100644 index 0000000000..87fa830a07 --- /dev/null +++ b/crates/iceberg/src/partitioning.rs @@ -0,0 +1,101 @@ +// 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. + +//! Partition type utilities for Iceberg tables. + +use std::cmp::Reverse; +use std::collections::HashSet; + +use crate::spec::{NestedField, NestedFieldRef, PartitionSpec, Schema, StructType, Transform}; +use crate::{Error, ErrorKind, Result}; + +/// Computes the unified partition type across all partition specs in the table. +/// +/// This is equivalent to Java's `Partitioning.partitionType(table)`. The result is a +/// StructType containing all partition fields ever used across all specs, enabling correct +/// representation of the `_partition` metadata column when partition evolution has occurred. +/// +/// Matches Java's behavior: +/// - Specs are sorted by spec_id in descending order (newer specs first), so newer field +/// names take precedence when deduplicating by field_id. +/// - Void and unknown transform fields are skipped. +/// - Fields are deduplicated by field_id — each unique field_id appears exactly once. +/// +/// # Arguments +/// * `partition_specs` - Iterator over all partition specs in the table +/// * `schema` - The current table schema (needed to determine result types of transforms) +pub fn compute_unified_partition_type<'a>( + partition_specs: impl Iterator, + schema: &Schema, +) -> Result { + let mut seen_field_ids = HashSet::new(); + let mut struct_fields: Vec = Vec::new(); + + // Sort specs by spec_id descending (newer first) to match Java's behavior: + // newer field names take precedence when deduplicating by field_id. + let mut specs: Vec<&PartitionSpec> = partition_specs.collect(); + specs.sort_by_key(|s| Reverse(s.spec_id())); + + for spec in specs { + for field in spec.fields() { + if seen_field_ids.contains(&field.field_id) { + continue; + } + + // Claim this field_id for the current (newest) spec, even if it's void. + // This ensures that if a newer spec marks a field Void, an older spec + // won't re-add it with a stale name. + seen_field_ids.insert(field.field_id); + + // Skip void transforms (dropped partition columns) + if matches!(field.transform, Transform::Void) { + continue; + } + + // Reject unknown transforms — the table uses a spec feature that + // this version of iceberg-rust doesn't support. Matching Java's + // behavior where getResultType() throws for unknown transforms. + if matches!(field.transform, Transform::Unknown) { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + format!( + "Partition field '{}' uses an unknown transform that is not \ + supported by this version of iceberg-rust", + field.name + ), + )); + } + + // Skip fields whose source column was dropped (partition evolution + // allows dropping columns; Java's Partitioning.partitionType skips these). + let source_field = match schema.field_by_id(field.source_id) { + Some(f) => f, + None => continue, + }; + + let res_type = field.transform.result_type(&source_field.field_type)?; + let nested = NestedField::optional(field.field_id, &field.name, res_type).into(); + struct_fields.push(nested); + } + } + + // Sort by field ID ascending to match Java's buildPartitionProjectionType + // which sorts by fieldMap.keySet().stream().sorted(Comparator.naturalOrder()). + struct_fields.sort_by_key(|f| f.id); + + Ok(StructType::new(struct_fields)) +} diff --git a/crates/iceberg/src/scan/context.rs b/crates/iceberg/src/scan/context.rs index 9181d02727..176bcba459 100644 --- a/crates/iceberg/src/scan/context.rs +++ b/crates/iceberg/src/scan/context.rs @@ -29,7 +29,7 @@ use crate::scan::{ }; use crate::spec::{ ManifestContentType, ManifestEntryRef, ManifestFile, ManifestList, NameMapping, - PartitionSpecRef, SchemaRef, SnapshotRef, TableMetadataRef, + PartitionSpecRef, SchemaRef, SnapshotRef, StructType, TableMetadataRef, }; use crate::{Error, ErrorKind, Result}; @@ -49,6 +49,7 @@ pub(crate) struct ManifestFileContext { name_mapping: Option>, case_sensitive: bool, partition_spec: Option, + unified_partition_type: Option>, } /// Wraps a [`ManifestEntryRef`] alongside the objects that are needed @@ -65,6 +66,7 @@ pub(crate) struct ManifestEntryContext { pub name_mapping: Option>, pub case_sensitive: bool, pub partition_spec: Option, + pub unified_partition_type: Option>, } impl ManifestFileContext { @@ -83,6 +85,7 @@ impl ManifestFileContext { name_mapping, case_sensitive, partition_spec, + unified_partition_type, } = self; let manifest = object_cache.get_manifest(&manifest_file).await?; @@ -100,6 +103,7 @@ impl ManifestFileContext { name_mapping: name_mapping.clone(), case_sensitive, partition_spec: partition_spec.clone(), + unified_partition_type: unified_partition_type.clone(), }; sender @@ -141,6 +145,7 @@ impl ManifestEntryContext { .with_partition(Some(self.manifest_entry.data_file.partition.clone())) .with_partition_spec(self.partition_spec.clone()) .with_name_mapping(self.name_mapping) + .with_unified_partition_type(self.unified_partition_type.clone()) .with_case_sensitive(self.case_sensitive) .with_key_metadata(self.manifest_entry.data_file.key_metadata().map(Box::from)) .build()) @@ -165,6 +170,8 @@ pub(crate) struct PlanContext { pub partition_filter_cache: Arc, pub manifest_evaluator_cache: Arc, pub expression_evaluator_cache: Arc, + + pub unified_partition_type: Option>, } impl PlanContext { @@ -294,6 +301,7 @@ impl PlanContext { .table_metadata .partition_spec_by_id(manifest_file.partition_spec_id) .cloned(), + unified_partition_type: self.unified_partition_type.clone(), } } } diff --git a/crates/iceberg/src/scan/mod.rs b/crates/iceberg/src/scan/mod.rs index 19f46c2cb1..e3f19620b2 100644 --- a/crates/iceberg/src/scan/mod.rs +++ b/crates/iceberg/src/scan/mod.rs @@ -37,7 +37,10 @@ use crate::delete_file_index::DeleteFileIndex; use crate::expr::visitors::inclusive_metrics_evaluator::InclusiveMetricsEvaluator; use crate::expr::{Bind, BoundPredicate, Predicate}; use crate::io::FileIO; -use crate::metadata_columns::{get_metadata_field_id, is_metadata_column_name}; +use crate::metadata_columns::{ + RESERVED_FIELD_ID_PARTITION, get_metadata_field_id, is_metadata_column_name, +}; +use crate::partitioning::compute_unified_partition_type; use crate::runtime::Runtime; use crate::spec::{DEFAULT_SCHEMA_NAME_MAPPING, DataContentType, NameMapping, SnapshotRef}; use crate::table::Table; @@ -300,6 +303,20 @@ impl<'a> TableScanBuilder<'a> { .transpose()? .map(Arc::new); + // Compute unified partition type if _partition is projected + let unified_partition_type = if field_ids.contains(&RESERVED_FIELD_ID_PARTITION) { + let upt = compute_unified_partition_type( + self.table + .metadata() + .partition_specs_iter() + .map(|s| s.as_ref()), + &schema, + )?; + Some(Arc::new(upt)) + } else { + None + }; + let plan_context = PlanContext { snapshot, table_metadata: self.table.metadata_ref(), @@ -313,6 +330,7 @@ impl<'a> TableScanBuilder<'a> { partition_filter_cache: Arc::new(PartitionFilterCache::new()), manifest_evaluator_cache: Arc::new(ManifestEvaluatorCache::new()), expression_evaluator_cache: Arc::new(ExpressionEvaluatorCache::new()), + unified_partition_type, }; Ok(TableScan { diff --git a/crates/iceberg/src/scan/task.rs b/crates/iceberg/src/scan/task.rs index faeac51be9..24a6036870 100644 --- a/crates/iceberg/src/scan/task.rs +++ b/crates/iceberg/src/scan/task.rs @@ -25,7 +25,7 @@ use crate::Result; use crate::expr::BoundPredicate; use crate::spec::{ DataContentType, DataFileFormat, ManifestEntryRef, NameMapping, PartitionSpec, Schema, - SchemaRef, Struct, + SchemaRef, Struct, StructType, }; /// A stream of [`FileScanTask`]. @@ -116,6 +116,22 @@ pub struct FileScanTask { #[builder(default)] pub name_mapping: Option>, + /// The unified partition type across all specs in the table. + /// When `RESERVED_FIELD_ID_PARTITION` is in the projected field IDs, the reader + /// uses this type along with the task's partition_spec and partition data to + /// materialize the `_partition` struct column at read time. + /// + /// This is a table-level value (same for all tasks in a scan), stored per-task + /// so that readers are self-contained without needing back-pointers to table + /// metadata. The cost is one Arc clone per task. + /// Serde: not yet implemented (same pattern as partition, partition_spec, name_mapping). + #[serde(default)] + #[serde(skip_serializing_if = "Option::is_none")] + #[serde(serialize_with = "serialize_not_implemented")] + #[serde(deserialize_with = "deserialize_not_implemented")] + #[builder(default)] + pub unified_partition_type: Option>, + /// Whether this scan task should treat column names as case-sensitive when binding predicates. pub case_sensitive: bool, From 804cdfcab1a01ef4cdcc65449aca324b10f143b5 Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Mon, 20 Jul 2026 15:52:32 -0700 Subject: [PATCH 2/6] fix after rebase --- crates/iceberg/src/arrow/record_batch_transformer.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/iceberg/src/arrow/record_batch_transformer.rs b/crates/iceberg/src/arrow/record_batch_transformer.rs index 61f3a105ad..3af22be638 100644 --- a/crates/iceberg/src/arrow/record_batch_transformer.rs +++ b/crates/iceberg/src/arrow/record_batch_transformer.rs @@ -521,7 +521,6 @@ impl RecordBatchTransformer { .ok_or(Error::new(ErrorKind::Unexpected, "field not found"))? .0 .clone()) - } }) .collect(); From 2ddb7b4e61ad7371143b0c6c086a88be39b2d6c0 Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Tue, 21 Jul 2026 10:42:43 -0700 Subject: [PATCH 3/6] fix public api --- crates/iceberg/public-api.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 8e90d08d34..8a62277295 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -1291,7 +1291,7 @@ impl core::fmt::Debug for iceberg::scan::FileScanTask pub fn iceberg::scan::FileScanTask::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl core::marker::StructuralPartialEq for iceberg::scan::FileScanTask impl iceberg::scan::FileScanTask -pub fn iceberg::scan::FileScanTask::builder() -> FileScanTaskBuilder<((), (), (), (), (), (), (), (), (), (), (), (), (), (), ())> +pub fn iceberg::scan::FileScanTask::builder() -> FileScanTaskBuilder<((), (), (), (), (), (), (), (), (), (), (), (), (), (), (), ())> impl serde_core::ser::Serialize for iceberg::scan::FileScanTask pub fn iceberg::scan::FileScanTask::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg::scan::FileScanTask From eebccd66b3143d54054974a7061645ea302bfba3 Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Wed, 22 Jul 2026 12:52:11 -0700 Subject: [PATCH 4/6] address review comments --- .../src/arrow/record_batch_transformer.rs | 106 +++++++++--------- crates/iceberg/src/scan/mod.rs | 4 +- 2 files changed, 56 insertions(+), 54 deletions(-) diff --git a/crates/iceberg/src/arrow/record_batch_transformer.rs b/crates/iceberg/src/arrow/record_batch_transformer.rs index 3af22be638..33ee9087d6 100644 --- a/crates/iceberg/src/arrow/record_batch_transformer.rs +++ b/crates/iceberg/src/arrow/record_batch_transformer.rs @@ -212,61 +212,47 @@ enum SchemaComparison { /// Builder for RecordBatchTransformer to improve ergonomics when constructing with optional parameters. /// /// All per-file constants (scalar metadata like `_file`, identity partition values, -/// and the `_partition` struct) are stored in a single `metadata_columns` map keyed by -/// field_id. This unified representation (via [`MetadataColumnSource`]) means the +/// and the `_partition` struct) are stored in a single `constant_fields` map keyed by +/// field_id. This unified representation (via [`ColumnConstant`]) means the /// transformer handles all constant columns through one code path. #[derive(Debug)] pub(crate) struct RecordBatchTransformerBuilder { snapshot_schema: Arc, projected_iceberg_field_ids: Vec, - metadata_columns: HashMap, + constant_fields: HashMap, virtual_fields: HashSet, } -/// Unified source for metadata/virtual column constants. +/// A per-file constant value for a column. /// -/// Both scalar metadata columns (like `_file`) and the struct `_partition` column are -/// per-file constants synthesized during scan. This enum unifies both representations -/// so the transformer can handle them through a single `metadata_columns` map keyed -/// by field_id, replacing the previous split between `constant_fields: HashMap` -/// and `partition_column: Option`. -/// -/// # Design for #2699 -/// -/// When `_pos` is implemented (via arrow-rs RowNumber — a real source column passed through, -/// not synthesized), it will NOT use this enum. `_pos` is not a constant and belongs in the -/// virtual_fields / pass-through axis. This enum is exclusively for constant-per-file values. -/// -/// The `ColumnSource` enum's `Add` and `AddStructConstant` variants map 1:1 to this enum's -/// variants during transform generation. #2699 may collapse `ColumnSource::Add` and -/// `ColumnSource::AddStructConstant` into a single `ColumnSource::AddMetadata { source: MetadataColumnSource }` -/// variant, but that refactor can happen independently. +/// Covers both scalar constants (metadata columns like `_file` and `_spec_id`, +/// as well as identity partition source fields) and the struct `_partition` column. #[derive(Debug, Clone, PartialEq)] -pub enum MetadataColumnSource { +pub enum ColumnConstant { /// A scalar constant (e.g., `_file` path, `_spec_id`, identity partition values). /// The Datum carries both the Iceberg type and the value. Scalar(Datum), /// A struct constant (the `_partition` column). Each child is a primitive constant /// or null (for partition evolution gaps). - Struct(PartitionColumnConstant), + Struct(StructConstant), } -/// Pre-computed data for the _partition struct constant. +/// Pre-computed data for a struct constant column. #[derive(Debug, Clone, PartialEq)] -pub struct PartitionColumnConstant { +pub struct StructConstant { fields: Fields, child_values: Vec>, } -impl PartitionColumnConstant { - /// Create a new PartitionColumnConstant, validating that fields and child_values have +impl StructConstant { + /// Create a new StructConstant, validating that fields and child_values have /// the same length. pub fn new(fields: Fields, child_values: Vec>) -> Result { if fields.len() != child_values.len() { return Err(Error::new( ErrorKind::DataInvalid, format!( - "PartitionColumnConstant: fields length ({}) != child_values length ({})", + "StructConstant: fields length ({}) != child_values length ({})", fields.len(), child_values.len() ), @@ -295,7 +281,7 @@ impl RecordBatchTransformerBuilder { Self { snapshot_schema, projected_iceberg_field_ids: projected_iceberg_field_ids.to_vec(), - metadata_columns: HashMap::new(), + constant_fields: HashMap::new(), virtual_fields: HashSet::new(), } } @@ -303,8 +289,8 @@ impl RecordBatchTransformerBuilder { /// Add a scalar constant value for a specific field ID. /// This is used for virtual/metadata fields like _file that have constant values per batch. pub(crate) fn with_constant(mut self, field_id: i32, datum: Datum) -> Self { - self.metadata_columns - .insert(field_id, MetadataColumnSource::Scalar(datum)); + self.constant_fields + .insert(field_id, ColumnConstant::Scalar(datum)); self } @@ -312,7 +298,7 @@ impl RecordBatchTransformerBuilder { /// /// Both partition_spec and partition_data must be provided together since the spec defines /// which fields are identity-partitioned, and the data provides their constant values. - /// This method computes the partition constants and merges them into metadata_columns. + /// This method computes the partition constants and merges them into constant_fields. pub(crate) fn with_partition( mut self, partition_spec: Arc, @@ -322,8 +308,8 @@ impl RecordBatchTransformerBuilder { constants_map(&partition_spec, &partition_data, &self.snapshot_schema)?; for (field_id, datum) in partition_constants { - self.metadata_columns - .insert(field_id, MetadataColumnSource::Scalar(datum)); + self.constant_fields + .insert(field_id, ColumnConstant::Scalar(datum)); } Ok(self) @@ -339,11 +325,11 @@ impl RecordBatchTransformerBuilder { /// Set a pre-computed _partition column constant directly. pub(crate) fn with_partition_column_precomputed( mut self, - partition_column: PartitionColumnConstant, + partition_column: StructConstant, ) -> Self { - self.metadata_columns.insert( + self.constant_fields.insert( RESERVED_FIELD_ID_PARTITION, - MetadataColumnSource::Struct(partition_column), + ColumnConstant::Struct(partition_column), ); self } @@ -352,7 +338,7 @@ impl RecordBatchTransformerBuilder { RecordBatchTransformer { snapshot_schema: self.snapshot_schema, projected_iceberg_field_ids: self.projected_iceberg_field_ids, - metadata_columns: self.metadata_columns, + constant_fields: self.constant_fields, virtual_fields: self.virtual_fields, batch_transform: None, } @@ -394,7 +380,7 @@ pub(crate) struct RecordBatchTransformer { snapshot_schema: Arc, projected_iceberg_field_ids: Vec, // Unified map of all per-file constant columns (metadata, identity partition, _partition struct). - metadata_columns: HashMap, + constant_fields: HashMap, // Field IDs whose data is delivered by the source batch as an arrow-rs virtual column // (e.g. _pos via RowNumber). These fields bypass the snapshot-schema lookup and the @@ -441,7 +427,7 @@ impl RecordBatchTransformer { record_batch.schema_ref(), self.snapshot_schema.as_ref(), &self.projected_iceberg_field_ids, - &self.metadata_columns, + &self.constant_fields, &self.virtual_fields, )?); @@ -461,7 +447,7 @@ impl RecordBatchTransformer { source_schema: &ArrowSchemaRef, snapshot_schema: &IcebergSchema, projected_iceberg_field_ids: &[i32], - metadata_columns: &HashMap, + constant_fields: &HashMap, virtual_fields: &HashSet, ) -> Result { let mapped_unprojected_arrow_schema = Arc::new(schema_to_arrow_schema(snapshot_schema)?); @@ -473,8 +459,10 @@ impl RecordBatchTransformer { let fields: Result> = projected_iceberg_field_ids .iter() .map(|field_id| { - match metadata_columns.get(field_id) { - Some(MetadataColumnSource::Struct(pc)) => { + match constant_fields.get(field_id) { + Some(ColumnConstant::Struct(pc)) + if *field_id == RESERVED_FIELD_ID_PARTITION => + { let struct_type = DataType::Struct(pc.fields().clone()); let nullable = pc.fields().is_empty(); let arrow_field = field_with_id( @@ -485,7 +473,13 @@ impl RecordBatchTransformer { ); return Ok(Arc::new(arrow_field)); } - Some(MetadataColumnSource::Scalar(datum)) => { + Some(ColumnConstant::Struct(_)) => { + return Err(Error::new( + ErrorKind::Unexpected, + format!("Unexpected struct constant for field id {field_id}"), + )); + } + Some(ColumnConstant::Scalar(datum)) => { if let Ok(iceberg_field) = get_metadata_field(*field_id) { let arrow_type = datum_to_arrow_type_with_ree(datum); let arrow_field = field_with_id( @@ -535,7 +529,7 @@ impl RecordBatchTransformer { snapshot_schema, projected_iceberg_field_ids, field_id_to_mapped_schema_map, - metadata_columns, + constant_fields, virtual_fields, )?, target_schema, @@ -604,7 +598,7 @@ impl RecordBatchTransformer { snapshot_schema: &IcebergSchema, projected_iceberg_field_ids: &[i32], field_id_to_mapped_schema_map: HashMap, - metadata_columns: &HashMap, + constant_fields: &HashMap, virtual_fields: &HashSet, ) -> Result> { let field_id_to_source_schema_map = @@ -613,14 +607,22 @@ impl RecordBatchTransformer { projected_iceberg_field_ids .iter() .map(|field_id| { - match metadata_columns.get(field_id) { - Some(MetadataColumnSource::Struct(pc)) => { + match constant_fields.get(field_id) { + Some(ColumnConstant::Struct(pc)) + if *field_id == RESERVED_FIELD_ID_PARTITION => + { return Ok(ColumnSource::AddStructConstant { fields: pc.fields().clone(), child_values: pc.child_values().to_vec(), }); } - Some(MetadataColumnSource::Scalar(datum)) => { + Some(ColumnConstant::Struct(_)) => { + return Err(Error::new( + ErrorKind::Unexpected, + format!("Unexpected struct constant for field id {field_id}"), + )); + } + Some(ColumnConstant::Scalar(datum)) => { let is_metadata = get_metadata_field(*field_id).is_ok(); let present_in_file = field_id_to_source_schema_map.contains_key(field_id); @@ -873,7 +875,7 @@ impl RecordBatchTransformer { } } -/// Builds a [`PartitionColumnConstant`] from the unified partition type and a file's +/// Builds a [`StructConstant`] from the unified partition type and a file's /// partition spec/data. /// /// If `unified_partition_type` has no fields (unpartitioned table), returns an empty constant @@ -886,10 +888,10 @@ pub fn build_partition_column_constant( unified_partition_type: &StructType, partition_spec: &PartitionSpec, partition_data: &Struct, -) -> Result { +) -> Result { // Unpartitioned table: empty struct rendered as null if unified_partition_type.fields().is_empty() { - return PartitionColumnConstant::new(Fields::empty(), vec![]); + return StructConstant::new(Fields::empty(), vec![]); } use crate::arrow::type_to_arrow_type; @@ -925,7 +927,7 @@ pub fn build_partition_column_constant( child_values.push(value); } - PartitionColumnConstant::new(Fields::from(arrow_fields), child_values) + StructConstant::new(Fields::from(arrow_fields), child_values) } #[cfg(test)] diff --git a/crates/iceberg/src/scan/mod.rs b/crates/iceberg/src/scan/mod.rs index e3f19620b2..a1c85c7117 100644 --- a/crates/iceberg/src/scan/mod.rs +++ b/crates/iceberg/src/scan/mod.rs @@ -305,14 +305,14 @@ impl<'a> TableScanBuilder<'a> { // Compute unified partition type if _partition is projected let unified_partition_type = if field_ids.contains(&RESERVED_FIELD_ID_PARTITION) { - let upt = compute_unified_partition_type( + let partition_type = compute_unified_partition_type( self.table .metadata() .partition_specs_iter() .map(|s| s.as_ref()), &schema, )?; - Some(Arc::new(upt)) + Some(Arc::new(partition_type)) } else { None }; From 46bc0fb4e63bb0ab568829155f1db493b60ffcbe Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Sun, 26 Jul 2026 14:57:03 -0700 Subject: [PATCH 5/6] address review comments --- crates/iceberg/src/arrow/mod.rs | 2 +- crates/iceberg/src/arrow/reader/pipeline.rs | 8 +- .../src/arrow/record_batch_transformer.rs | 141 ++++---- crates/iceberg/src/partitioning.rs | 310 ++++++++++++++++-- 4 files changed, 345 insertions(+), 116 deletions(-) diff --git a/crates/iceberg/src/arrow/mod.rs b/crates/iceberg/src/arrow/mod.rs index 1473809436..28cf6e94a4 100644 --- a/crates/iceberg/src/arrow/mod.rs +++ b/crates/iceberg/src/arrow/mod.rs @@ -32,7 +32,7 @@ mod reader; /// RecordBatch projection utilities pub mod record_batch_projector; pub(crate) mod record_batch_transformer; -pub(crate) use record_batch_transformer::build_partition_column_constant; +pub(crate) use record_batch_transformer::build_partition_constant; mod scan_metrics; mod value; diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index 9d05139c43..5c0ea6b72f 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -34,7 +34,7 @@ use super::{ ArrowFileReader, ArrowReader, ParquetReadOptions, add_fallback_field_ids_to_arrow_schema, apply_name_mapping_to_arrow_schema, }; -use crate::arrow::build_partition_column_constant; +use crate::arrow::build_partition_constant; use crate::arrow::caching_delete_file_loader::CachingDeleteFileLoader; use crate::arrow::int96::coerce_int96_timestamps; use crate::arrow::record_batch_transformer::RecordBatchTransformerBuilder; @@ -321,13 +321,13 @@ impl FileScanTaskReader { { let (spec, partition_data) = match (&task.partition_spec, &task.partition) { (Some(spec), Some(data)) => (spec.clone(), data.clone()), - // Unpartitioned table or missing spec: build_partition_column_constant + // Unpartitioned table or missing spec: build_partition_constant // handles the empty unified_type case with an early return. _ => (Arc::new(PartitionSpec::unpartition_spec()), Struct::empty()), }; - let constant = build_partition_column_constant(unified_type, &spec, &partition_data)?; + let constant = build_partition_constant(unified_type, &spec, &partition_data)?; record_batch_transformer_builder = - record_batch_transformer_builder.with_partition_column_precomputed(constant); + record_batch_transformer_builder.with_partition_constant(constant); } let mut record_batch_transformer = record_batch_transformer_builder.build(); diff --git a/crates/iceberg/src/arrow/record_batch_transformer.rs b/crates/iceberg/src/arrow/record_batch_transformer.rs index 33ee9087d6..c9334f38e7 100644 --- a/crates/iceberg/src/arrow/record_batch_transformer.rs +++ b/crates/iceberg/src/arrow/record_batch_transformer.rs @@ -228,7 +228,7 @@ pub(crate) struct RecordBatchTransformerBuilder { /// Covers both scalar constants (metadata columns like `_file` and `_spec_id`, /// as well as identity partition source fields) and the struct `_partition` column. #[derive(Debug, Clone, PartialEq)] -pub enum ColumnConstant { +pub(crate) enum ColumnConstant { /// A scalar constant (e.g., `_file` path, `_spec_id`, identity partition values). /// The Datum carries both the Iceberg type and the value. Scalar(Datum), @@ -239,7 +239,7 @@ pub enum ColumnConstant { /// Pre-computed data for a struct constant column. #[derive(Debug, Clone, PartialEq)] -pub struct StructConstant { +pub(crate) struct StructConstant { fields: Fields, child_values: Vec>, } @@ -247,7 +247,7 @@ pub struct StructConstant { impl StructConstant { /// Create a new StructConstant, validating that fields and child_values have /// the same length. - pub fn new(fields: Fields, child_values: Vec>) -> Result { + pub(crate) fn new(fields: Fields, child_values: Vec>) -> Result { if fields.len() != child_values.len() { return Err(Error::new( ErrorKind::DataInvalid, @@ -323,10 +323,7 @@ impl RecordBatchTransformerBuilder { } /// Set a pre-computed _partition column constant directly. - pub(crate) fn with_partition_column_precomputed( - mut self, - partition_column: StructConstant, - ) -> Self { + pub(crate) fn with_partition_constant(mut self, partition_column: StructConstant) -> Self { self.constant_fields.insert( RESERVED_FIELD_ID_PARTITION, ColumnConstant::Struct(partition_column), @@ -476,7 +473,11 @@ impl RecordBatchTransformer { Some(ColumnConstant::Struct(_)) => { return Err(Error::new( ErrorKind::Unexpected, - format!("Unexpected struct constant for field id {field_id}"), + format!( + "Struct column constants are only supported for the `_partition` \ + metadata column (field id {RESERVED_FIELD_ID_PARTITION}), but \ + one was set for field id {field_id}" + ), )); } Some(ColumnConstant::Scalar(datum)) => { @@ -619,7 +620,11 @@ impl RecordBatchTransformer { Some(ColumnConstant::Struct(_)) => { return Err(Error::new( ErrorKind::Unexpected, - format!("Unexpected struct constant for field id {field_id}"), + format!( + "Struct column constants are only supported for the `_partition` \ + metadata column (field id {RESERVED_FIELD_ID_PARTITION}), but \ + one was set for field id {field_id}" + ), )); } Some(ColumnConstant::Scalar(datum)) => { @@ -884,7 +889,7 @@ impl RecordBatchTransformer { /// For each field in the unified partition type: /// - If it corresponds to a field in this file's partition spec, use the value from partition_data /// - Otherwise (partition evolution), use null -pub fn build_partition_column_constant( +pub(crate) fn build_partition_constant( unified_partition_type: &StructType, partition_spec: &PartitionSpec, partition_data: &Struct, @@ -940,9 +945,9 @@ mod test { StringArray, }; use arrow_schema::{DataType, Field, Schema as ArrowSchema}; - use parquet::arrow::PARQUET_FIELD_ID_META_KEY; - use crate::arrow::build_partition_column_constant; + use super::field_with_id; + use crate::arrow::build_partition_constant; use crate::arrow::record_batch_transformer::{ RecordBatchTransformer, RecordBatchTransformerBuilder, }; @@ -1072,8 +1077,8 @@ mod test { .build(); let file_schema = Arc::new(ArrowSchema::new(vec![ - simple_field("id", DataType::Int32, false, "1"), - simple_field("name", DataType::Utf8, true, "2"), + field_with_id("id", DataType::Int32, false, 1), + field_with_id("name", DataType::Utf8, true, 2), ])); let file_batch = RecordBatch::try_new(file_schema, vec![ @@ -1140,11 +1145,11 @@ mod test { RecordBatchTransformerBuilder::new(snapshot_schema, &projected_iceberg_field_ids) .build(); - let file_schema = Arc::new(ArrowSchema::new(vec![simple_field( + let file_schema = Arc::new(ArrowSchema::new(vec![field_with_id( "id", DataType::Int32, false, - "1", + 1, )])); let file_batch = RecordBatch::try_new(file_schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]) @@ -1194,8 +1199,8 @@ mod test { .build(); let file_schema = Arc::new(ArrowSchema::new(vec![ - simple_field("id", DataType::Int32, false, "1"), - simple_field("data", DataType::Utf8, false, "2"), + field_with_id("id", DataType::Int32, false, 1), + field_with_id("data", DataType::Utf8, false, 2), ])); let file_batch = RecordBatch::try_new(file_schema, vec![ @@ -1315,38 +1320,30 @@ mod test { fn arrow_schema_already_same_as_target() -> Arc { Arc::new(ArrowSchema::new(vec![ - simple_field("a", DataType::Utf8, true, "10"), - simple_field("b", DataType::Int64, false, "11"), - simple_field("c", DataType::Float64, false, "12"), - simple_field("e", DataType::Utf8, true, "14"), - simple_field("f", DataType::Utf8, false, "15"), + field_with_id("a", DataType::Utf8, true, 10), + field_with_id("b", DataType::Int64, false, 11), + field_with_id("c", DataType::Float64, false, 12), + field_with_id("e", DataType::Utf8, true, 14), + field_with_id("f", DataType::Utf8, false, 15), ])) } fn arrow_schema_promotion_addition_and_renaming_required() -> Arc { Arc::new(ArrowSchema::new(vec![ - simple_field("b", DataType::Int32, false, "11"), - simple_field("c", DataType::Float32, false, "12"), - simple_field("d", DataType::Int32, false, "13"), - simple_field("e_old", DataType::Utf8, true, "14"), + field_with_id("b", DataType::Int32, false, 11), + field_with_id("c", DataType::Float32, false, 12), + field_with_id("d", DataType::Int32, false, 13), + field_with_id("e_old", DataType::Utf8, true, 14), ])) } fn arrow_schema_no_promotion_addition_or_renaming_required() -> Arc { Arc::new(ArrowSchema::new(vec![ - simple_field("d", DataType::Int32, false, "13"), - simple_field("e", DataType::Utf8, true, "14"), + field_with_id("d", DataType::Int32, false, 13), + field_with_id("e", DataType::Utf8, true, 14), ])) } - /// Create a simple arrow field with metadata. - fn simple_field(name: &str, ty: DataType, nullable: bool, value: &str) -> Field { - Field::new(name, ty, nullable).with_metadata(HashMap::from([( - PARQUET_FIELD_ID_META_KEY.to_string(), - value.to_string(), - )])) - } - /// Test for add_files with Parquet files that have NO field IDs (Hive tables). /// /// This reproduces the scenario from Iceberg spec where: @@ -1526,8 +1523,8 @@ mod test { // Parquet file contains both id and name columns let parquet_schema = Arc::new(ArrowSchema::new(vec![ - simple_field("id", DataType::Int32, false, "1"), - simple_field("name", DataType::Utf8, true, "2"), + field_with_id("id", DataType::Int32, false, 1), + field_with_id("name", DataType::Utf8, true, 2), ])); let projected_field_ids = [1, 2]; // id, name @@ -1647,8 +1644,8 @@ mod test { // Parquet file contains only id and name (dept is in partition path) let parquet_schema = Arc::new(ArrowSchema::new(vec![ - simple_field("id", DataType::Int32, false, "1"), - simple_field("name", DataType::Utf8, true, "3"), + field_with_id("id", DataType::Int32, false, 1), + field_with_id("name", DataType::Utf8, true, 3), ])); let projected_field_ids = [1, 2, 3]; // id, dept, name @@ -1737,8 +1734,8 @@ mod test { let partition_data = Struct::from_iter(vec![Some(Literal::string("engineering"))]); let parquet_schema = Arc::new(ArrowSchema::new(vec![ - simple_field("id", DataType::Int32, false, "1"), - simple_field("name", DataType::Utf8, true, "3"), + field_with_id("id", DataType::Int32, false, 1), + field_with_id("name", DataType::Utf8, true, 3), ])); let projected_field_ids = [1, 3]; // id, name -- dept is gone @@ -1833,8 +1830,8 @@ mod test { // Parquet file has OLD column name "id" but SAME field_id=1 // Field-ID-based mapping should find this despite name mismatch let parquet_schema = Arc::new(ArrowSchema::new(vec![ - simple_field("id", DataType::Int32, false, "1"), - simple_field("name", DataType::Utf8, true, "2"), + field_with_id("id", DataType::Int32, false, 1), + field_with_id("name", DataType::Utf8, true, 2), ])); let projected_field_ids = [1, 2]; // row_id (field_id=1), name (field_id=2) @@ -1938,8 +1935,8 @@ mod test { // Has id (field_id=1) and data (field_id=3, assigned by ArrowReader via name mapping) // Missing: dept (in partition), category (has default), notes (no default) let parquet_schema = Arc::new(ArrowSchema::new(vec![ - simple_field("id", DataType::Int32, false, "1"), - simple_field("data", DataType::Utf8, false, "3"), + field_with_id("id", DataType::Int32, false, 1), + field_with_id("data", DataType::Utf8, false, 3), ])); let projected_field_ids = [1, 2, 3, 4, 5]; // id, dept, data, category, notes @@ -2030,11 +2027,11 @@ mod test { // Partition has null value for the data column let partition_data = Struct::from_iter(vec![None]); - let file_schema = Arc::new(ArrowSchema::new(vec![simple_field( + let file_schema = Arc::new(ArrowSchema::new(vec![field_with_id( "id", DataType::Int32, true, - "1", + 1, )])); let projected_field_ids = [1, 2]; @@ -2085,14 +2082,9 @@ mod test { // Simulate what arrow-rs's virtual-column reader produces: file columns // followed by the _pos Int64 column carrying absolute file row indices. let parquet_schema = Arc::new(ArrowSchema::new(vec![ - simple_field("id", DataType::Int32, false, "1"), - simple_field("name", DataType::Utf8, true, "2"), - simple_field( - "_pos", - DataType::Int64, - false, - &RESERVED_FIELD_ID_POS.to_string(), - ), + field_with_id("id", DataType::Int32, false, 1), + field_with_id("name", DataType::Utf8, true, 2), + field_with_id("_pos", DataType::Int64, false, RESERVED_FIELD_ID_POS), ])); let projected_field_ids = [1, 2, RESERVED_FIELD_ID_POS]; @@ -2160,14 +2152,9 @@ mod test { ); let parquet_schema = Arc::new(ArrowSchema::new(vec![ - simple_field("id_long", DataType::Int64, false, "1"), - simple_field("name", DataType::Utf8, true, "2"), - simple_field( - "_pos", - DataType::Int64, - false, - &RESERVED_FIELD_ID_POS.to_string(), - ), + field_with_id("id_long", DataType::Int64, false, 1), + field_with_id("name", DataType::Utf8, true, 2), + field_with_id("_pos", DataType::Int64, false, RESERVED_FIELD_ID_POS), ])); let projected_field_ids = [RESERVED_FIELD_ID_POS, 2, 1]; @@ -2251,22 +2238,19 @@ mod test { // Parquet file has both columns let parquet_schema = Arc::new(ArrowSchema::new(vec![ - simple_field("id", DataType::Int32, false, "1"), - simple_field("name", DataType::Utf8, true, "2"), + field_with_id("id", DataType::Int32, false, 1), + field_with_id("name", DataType::Utf8, true, 2), ])); // Project id, name, and _partition let projected_field_ids = [1, 2, RESERVED_FIELD_ID_PARTITION]; - let partition_column = build_partition_column_constant( - &unified_partition_type, - &partition_spec, - &partition_data, - ) - .unwrap(); + let partition_column = + build_partition_constant(&unified_partition_type, &partition_spec, &partition_data) + .unwrap(); let mut transformer = RecordBatchTransformerBuilder::new(snapshot_schema, &projected_field_ids) - .with_partition_column_precomputed(partition_column) + .with_partition_constant(partition_column) .build(); let parquet_batch = RecordBatch::try_new(parquet_schema, vec![ @@ -2364,21 +2348,20 @@ mod test { // File written with spec_v0 (only has year=2023) let partition_data = Struct::from_iter(vec![Some(Literal::int(2023))]); - let parquet_schema = Arc::new(ArrowSchema::new(vec![simple_field( + let parquet_schema = Arc::new(ArrowSchema::new(vec![field_with_id( "data", DataType::Utf8, true, - "3", + 3, )])); let projected_field_ids = [3, RESERVED_FIELD_ID_PARTITION]; let partition_column = - build_partition_column_constant(&unified_partition_type, &spec_v0, &partition_data) - .unwrap(); + build_partition_constant(&unified_partition_type, &spec_v0, &partition_data).unwrap(); let mut transformer = RecordBatchTransformerBuilder::new(snapshot_schema, &projected_field_ids) - .with_partition_column_precomputed(partition_column) + .with_partition_constant(partition_column) .build(); let parquet_batch = diff --git a/crates/iceberg/src/partitioning.rs b/crates/iceberg/src/partitioning.rs index 87fa830a07..b0c6781057 100644 --- a/crates/iceberg/src/partitioning.rs +++ b/crates/iceberg/src/partitioning.rs @@ -18,9 +18,11 @@ //! Partition type utilities for Iceberg tables. use std::cmp::Reverse; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; -use crate::spec::{NestedField, NestedFieldRef, PartitionSpec, Schema, StructType, Transform}; +use crate::spec::{ + NestedField, NestedFieldRef, PartitionSpec, Schema, StructType, Transform, Type, +}; use crate::{Error, ErrorKind, Result}; /// Computes the unified partition type across all partition specs in the table. @@ -29,11 +31,15 @@ use crate::{Error, ErrorKind, Result}; /// StructType containing all partition fields ever used across all specs, enabling correct /// representation of the `_partition` metadata column when partition evolution has occurred. /// -/// Matches Java's behavior: +/// Matches Java's `buildPartitionProjectionType` behavior: /// - Specs are sorted by spec_id in descending order (newer specs first), so newer field /// names take precedence when deduplicating by field_id. -/// - Void and unknown transform fields are skipped. -/// - Fields are deduplicated by field_id — each unique field_id appears exactly once. +/// - Unknown transforms cause an error. +/// - Fields whose source column was dropped from the schema are skipped. +/// - When a newer spec marks a field as Void (dropped) but an older spec has it with a +/// real transform, the older spec's type is preserved while the newer spec's name is kept. +/// - Fields are deduplicated by field_id; each unique field_id appears exactly once. +/// - Output fields are sorted by field_id ascending. /// /// # Arguments /// * `partition_specs` - Iterator over all partition specs in the table @@ -42,33 +48,23 @@ pub fn compute_unified_partition_type<'a>( partition_specs: impl Iterator, schema: &Schema, ) -> Result { - let mut seen_field_ids = HashSet::new(); - let mut struct_fields: Vec = Vec::new(); - - // Sort specs by spec_id descending (newer first) to match Java's behavior: - // newer field names take precedence when deduplicating by field_id. let mut specs: Vec<&PartitionSpec> = partition_specs.collect(); specs.sort_by_key(|s| Reverse(s.spec_id())); - for spec in specs { - for field in spec.fields() { - if seen_field_ids.contains(&field.field_id) { - continue; - } + let active_field_ids = all_active_field_ids(specs.iter().copied(), schema); - // Claim this field_id for the current (newest) spec, even if it's void. - // This ensures that if a newer spec marks a field Void, an older spec - // won't re-add it with a stale name. - seen_field_ids.insert(field.field_id); + let mut field_map: HashMap = HashMap::new(); + let mut type_map: HashMap = HashMap::new(); + let mut name_map: HashMap = HashMap::new(); - // Skip void transforms (dropped partition columns) - if matches!(field.transform, Transform::Void) { + for spec in &specs { + for field in spec.fields() { + let field_id = field.field_id; + + if !active_field_ids.contains(&field_id) { continue; } - // Reject unknown transforms — the table uses a spec feature that - // this version of iceberg-rust doesn't support. Matching Java's - // behavior where getResultType() throws for unknown transforms. if matches!(field.transform, Transform::Unknown) { return Err(Error::new( ErrorKind::FeatureUnsupported, @@ -80,22 +76,272 @@ pub fn compute_unified_partition_type<'a>( )); } - // Skip fields whose source column was dropped (partition evolution - // allows dropping columns; Java's Partitioning.partitionType skips these). let source_field = match schema.field_by_id(field.source_id) { Some(f) => f, None => continue, }; - let res_type = field.transform.result_type(&source_field.field_type)?; - let nested = NestedField::optional(field.field_id, &field.name, res_type).into(); - struct_fields.push(nested); + match field_map.get(&field_id) { + None => { + let res_type = field.transform.result_type(&source_field.field_type)?; + field_map.insert(field_id, field); + type_map.insert(field_id, res_type); + name_map.insert(field_id, field.name.clone()); + } + Some(existing) => { + if is_void_transform(existing) && !is_void_transform(field) { + let res_type = field.transform.result_type(&source_field.field_type)?; + field_map.insert(field_id, field); + type_map.insert(field_id, res_type); + } + } + } } } - // Sort by field ID ascending to match Java's buildPartitionProjectionType - // which sorts by fieldMap.keySet().stream().sorted(Comparator.naturalOrder()). - struct_fields.sort_by_key(|f| f.id); + let mut field_ids: Vec = field_map.keys().copied().collect(); + field_ids.sort(); + + let struct_fields: Vec = field_ids + .into_iter() + .map(|fid| { + let name = &name_map[&fid]; + let ty = type_map.remove(&fid).unwrap(); + NestedField::optional(fid, name, ty).into() + }) + .collect(); Ok(StructType::new(struct_fields)) } + +fn is_void_transform(field: &crate::spec::PartitionField) -> bool { + matches!(field.transform, Transform::Void) +} + +fn all_active_field_ids<'a>( + partition_specs: impl Iterator, + schema: &Schema, +) -> HashSet { + partition_specs + .flat_map(|spec| spec.fields().iter()) + .filter(|field| schema.field_by_id(field.source_id).is_some()) + .map(|field| field.field_id) + .collect() +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use crate::spec::{NestedField, PrimitiveType, Transform, Type, UnboundPartitionSpec}; + + fn test_schema() -> Schema { + Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(2, "data", Type::Primitive(PrimitiveType::String)).into(), + NestedField::required(3, "ts", Type::Primitive(PrimitiveType::Timestamp)).into(), + NestedField::required(4, "category", Type::Primitive(PrimitiveType::String)).into(), + ]) + .build() + .unwrap() + } + + fn build_spec( + schema: &Schema, + spec_id: i32, + fields: Vec<(i32, &str, Transform)>, + ) -> PartitionSpec { + let mut builder = UnboundPartitionSpec::builder().with_spec_id(spec_id); + for (source_id, name, transform) in fields { + builder = builder + .add_partition_field(source_id, name, transform) + .unwrap(); + } + builder.build().bind(schema.clone()).unwrap() + } + + #[test] + fn test_single_spec_identity() { + let schema = test_schema(); + let spec = build_spec(&schema, 0, vec![(4, "category", Transform::Identity)]); + + let result = compute_unified_partition_type([&spec].into_iter(), &schema).unwrap(); + assert_eq!(result.fields().len(), 1); + assert_eq!(result.fields()[0].name, "category"); + assert_eq!( + *result.fields()[0].field_type, + Type::Primitive(PrimitiveType::String) + ); + } + + #[test] + fn test_single_spec_with_year_transform() { + let schema = test_schema(); + let spec = build_spec(&schema, 0, vec![(3, "ts_year", Transform::Year)]); + + let result = compute_unified_partition_type([&spec].into_iter(), &schema).unwrap(); + assert_eq!(result.fields().len(), 1); + assert_eq!(result.fields()[0].name, "ts_year"); + assert_eq!( + *result.fields()[0].field_type, + Type::Primitive(PrimitiveType::Int) + ); + } + + #[test] + fn test_unpartitioned() { + let schema = test_schema(); + let spec = PartitionSpec::unpartition_spec(); + let result = compute_unified_partition_type([&spec].into_iter(), &schema).unwrap(); + assert!(result.fields().is_empty()); + } + + #[test] + fn test_multiple_fields_sorted_by_id() { + let schema = test_schema(); + let spec = build_spec(&schema, 0, vec![ + (3, "ts_year", Transform::Year), + (4, "category", Transform::Identity), + ]); + + let result = compute_unified_partition_type([&spec].into_iter(), &schema).unwrap(); + assert_eq!(result.fields().len(), 2); + assert!(result.fields()[0].id < result.fields()[1].id); + } + + #[test] + fn test_newer_name_takes_precedence() { + let schema = test_schema(); + + // Spec 0: old name + let spec_v0 = PartitionSpec::builder(Arc::new(schema.clone())) + .with_spec_id(0) + .add_unbound_field(crate::spec::UnboundPartitionField { + source_id: 4, + field_id: Some(1000), + name: "cat_old".to_string(), + transform: Transform::Identity, + }) + .unwrap() + .build() + .unwrap(); + + // Spec 1: newer name, same field_id + let spec_v1 = PartitionSpec::builder(Arc::new(schema.clone())) + .with_spec_id(1) + .add_unbound_field(crate::spec::UnboundPartitionField { + source_id: 4, + field_id: Some(1000), + name: "cat_new".to_string(), + transform: Transform::Identity, + }) + .unwrap() + .build() + .unwrap(); + + let result = + compute_unified_partition_type([&spec_v0, &spec_v1].into_iter(), &schema).unwrap(); + assert_eq!(result.fields().len(), 1); + assert_eq!(result.fields()[0].name, "cat_new"); + } + + #[test] + fn test_void_replaced_by_older_non_void() { + let schema = test_schema(); + + // Spec 0 (older): category partitioned by identity + let spec_v0 = PartitionSpec::builder(Arc::new(schema.clone())) + .with_spec_id(0) + .add_unbound_field(crate::spec::UnboundPartitionField { + source_id: 4, + field_id: Some(1000), + name: "category".to_string(), + transform: Transform::Identity, + }) + .unwrap() + .build() + .unwrap(); + + // Spec 1 (newer): same field_id voided (partition dropped) + let spec_v1 = PartitionSpec::builder(Arc::new(schema.clone())) + .with_spec_id(1) + .add_unbound_field(crate::spec::UnboundPartitionField { + source_id: 4, + field_id: Some(1000), + name: "category_v2".to_string(), + transform: Transform::Void, + }) + .unwrap() + .build() + .unwrap(); + + let result = + compute_unified_partition_type([&spec_v0, &spec_v1].into_iter(), &schema).unwrap(); + + assert_eq!(result.fields().len(), 1); + // Name from newer spec + assert_eq!(result.fields()[0].name, "category_v2"); + // Type from older non-void spec + assert_eq!( + *result.fields()[0].field_type, + Type::Primitive(PrimitiveType::String) + ); + } + + #[test] + fn test_dropped_source_column_skipped() { + // Schema without field 4 (category was dropped) + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(2, "data", Type::Primitive(PrimitiveType::String)).into(), + ]) + .build() + .unwrap(); + + // Spec references source_id=4 which no longer exists in the schema. + // Deserialize directly since the builder rejects unknown source columns. + let spec = serde_json::from_value::(serde_json::json!({ + "spec-id": 0, + "fields": [{ + "source-id": 4, + "field-id": 1000, + "name": "category", + "transform": "identity" + }] + })) + .unwrap(); + + let result = compute_unified_partition_type([&spec].into_iter(), &schema).unwrap(); + assert!(result.fields().is_empty()); + } + + #[test] + fn test_evolution_adds_new_field() { + let schema = test_schema(); + + // Spec 0: partition by category + let spec_v0 = build_spec(&schema, 0, vec![(4, "category", Transform::Identity)]); + + // Spec 1: partition by category + ts_year + let spec_v1 = PartitionSpec::builder(Arc::new(schema.clone())) + .with_spec_id(1) + .add_unbound_field(crate::spec::UnboundPartitionField { + source_id: 4, + field_id: Some(spec_v0.fields()[0].field_id), + name: "category".to_string(), + transform: Transform::Identity, + }) + .unwrap() + .add_partition_field("ts", "ts_year", Transform::Year) + .unwrap() + .build() + .unwrap(); + + let result = + compute_unified_partition_type([&spec_v0, &spec_v1].into_iter(), &schema).unwrap(); + assert_eq!(result.fields().len(), 2); + } +} From 65322f1d3b7db527eae89acd8f1eec7b1fccb249 Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Wed, 29 Jul 2026 10:45:19 -0700 Subject: [PATCH 6/6] address review comments --- crates/iceberg/src/arrow/reader/pipeline.rs | 17 ++- .../src/arrow/record_batch_transformer.rs | 8 ++ crates/iceberg/src/partitioning.rs | 134 +++++++++++++++--- 3 files changed, 140 insertions(+), 19 deletions(-) diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index 5c0ea6b72f..e7bcdf5fb3 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -321,9 +321,20 @@ impl FileScanTaskReader { { let (spec, partition_data) = match (&task.partition_spec, &task.partition) { (Some(spec), Some(data)) => (spec.clone(), data.clone()), - // Unpartitioned table or missing spec: build_partition_constant - // handles the empty unified_type case with an early return. - _ => (Arc::new(PartitionSpec::unpartition_spec()), Struct::empty()), + // A missing spec/data is only acceptable when there are no partition + // fields to fill (unpartitioned table). If the unified type has fields + // but we lack a spec or data, the task is inconsistent and we cannot + // build the _partition column. + _ if unified_type.fields().is_empty() => { + (Arc::new(PartitionSpec::unpartition_spec()), Struct::empty()) + } + _ => { + return Err(Error::new( + ErrorKind::Unexpected, + "cannot build _partition column: unified partition type has fields \ + but the scan task is missing its partition spec or data", + )); + } }; let constant = build_partition_constant(unified_type, &spec, &partition_data)?; record_batch_transformer_builder = diff --git a/crates/iceberg/src/arrow/record_batch_transformer.rs b/crates/iceberg/src/arrow/record_batch_transformer.rs index c9334f38e7..bf7e89910a 100644 --- a/crates/iceberg/src/arrow/record_batch_transformer.rs +++ b/crates/iceberg/src/arrow/record_batch_transformer.rs @@ -628,6 +628,14 @@ impl RecordBatchTransformer { )); } Some(ColumnConstant::Scalar(datum)) => { + // Metadata/virtual fields (like _file, _spec_id) never exist in the + // data file, so they always use their pre-computed constant value. + // + // For identity-partitioned fields, the Iceberg spec's "Column + // Projection" rules only apply to "field ids which are not present in + // a data file". When the column IS present in the Parquet file, it + // must be read from the file; the partition metadata constant is only + // a fallback for when the column is absent (e.g. add_files). let is_metadata = get_metadata_field(*field_id).is_ok(); let present_in_file = field_id_to_source_schema_map.contains_key(field_id); diff --git a/crates/iceberg/src/partitioning.rs b/crates/iceberg/src/partitioning.rs index b0c6781057..ae8a58c6ae 100644 --- a/crates/iceberg/src/partitioning.rs +++ b/crates/iceberg/src/partitioning.rs @@ -21,7 +21,7 @@ use std::cmp::Reverse; use std::collections::{HashMap, HashSet}; use crate::spec::{ - NestedField, NestedFieldRef, PartitionSpec, Schema, StructType, Transform, Type, + NestedField, NestedFieldRef, PartitionField, PartitionSpec, Schema, StructType, Transform, Type, }; use crate::{Error, ErrorKind, Result}; @@ -36,6 +36,8 @@ use crate::{Error, ErrorKind, Result}; /// names take precedence when deduplicating by field_id. /// - Unknown transforms cause an error. /// - Fields whose source column was dropped from the schema are skipped. +/// - Two specs defining the same field_id must be compatible (same source, compatible +/// transforms); V1 tables do not guarantee field ids are unique across specs. /// - When a newer spec marks a field as Void (dropped) but an older spec has it with a /// real transform, the older spec's type is preserved while the newer spec's name is kept. /// - Fields are deduplicated by field_id; each unique field_id appears exactly once. @@ -53,7 +55,7 @@ pub fn compute_unified_partition_type<'a>( let active_field_ids = all_active_field_ids(specs.iter().copied(), schema); - let mut field_map: HashMap = HashMap::new(); + let mut field_map: HashMap = HashMap::new(); let mut type_map: HashMap = HashMap::new(); let mut name_map: HashMap = HashMap::new(); @@ -61,21 +63,25 @@ pub fn compute_unified_partition_type<'a>( for field in spec.fields() { let field_id = field.field_id; - if !active_field_ids.contains(&field_id) { - continue; - } - + // Reject unknown transforms up front: we cannot determine their result type, + // so we cannot build a partition column for them. This check must precede the + // active_field_ids filter below, otherwise an unknown transform could be + // silently skipped. if matches!(field.transform, Transform::Unknown) { return Err(Error::new( - ErrorKind::FeatureUnsupported, + ErrorKind::DataInvalid, format!( - "Partition field '{}' uses an unknown transform that is not \ - supported by this version of iceberg-rust", + "Partition field '{}' uses an unknown transform whose result type \ + cannot be determined", field.name ), )); } + if !active_field_ids.contains(&field_id) { + continue; + } + let source_field = match schema.field_by_id(field.source_id) { Some(f) => f, None => continue, @@ -89,6 +95,22 @@ pub fn compute_unified_partition_type<'a>( name_map.insert(field_id, field.name.clone()); } Some(existing) => { + // V1 tables do not guarantee field ids are unique across specs, so two + // specs may define the same field id. They must be compatible. + if !equivalent_ignoring_names(field, existing) { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "Conflicting partition fields for field id {field_id}: \ + '{}' and '{}'", + field.name, existing.name + ), + )); + } + + // Use the correct type for dropped partitions in v1 tables: if the + // newer spec voided the field but an older spec has a real transform, + // keep the older spec's type. if is_void_transform(existing) && !is_void_transform(field) { let res_type = field.transform.result_type(&source_field.field_type)?; field_map.insert(field_id, field); @@ -102,22 +124,47 @@ pub fn compute_unified_partition_type<'a>( let mut field_ids: Vec = field_map.keys().copied().collect(); field_ids.sort(); - let struct_fields: Vec = field_ids + let struct_fields = field_ids .into_iter() - .map(|fid| { - let name = &name_map[&fid]; - let ty = type_map.remove(&fid).unwrap(); - NestedField::optional(fid, name, ty).into() + .map(|fid| -> Result { + let name = name_map.get(&fid).ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + format!("Missing name for partition field {fid}"), + ) + })?; + let ty = type_map.remove(&fid).ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + format!("Missing type for partition field {fid}"), + ) + })?; + Ok(NestedField::optional(fid, name, ty).into()) }) - .collect(); + .collect::>>()?; Ok(StructType::new(struct_fields)) } -fn is_void_transform(field: &crate::spec::PartitionField) -> bool { +fn is_void_transform(field: &PartitionField) -> bool { matches!(field.transform, Transform::Void) } +/// Two partition fields with the same field id are compatible if they share the same +/// source id and have compatible transforms. Matches Java's +/// `Partitioning.equivalentIgnoringNames`. +fn equivalent_ignoring_names(field: &PartitionField, other: &PartitionField) -> bool { + field.field_id == other.field_id + && field.source_id == other.source_id + && compatible_transforms(&field.transform, &other.transform) +} + +/// Transforms are compatible if they are equal, or if either is Void (a dropped field). +/// Matches Java's `Partitioning.compatibleTransforms`. +fn compatible_transforms(t1: &Transform, t2: &Transform) -> bool { + t1 == t2 || matches!(t1, Transform::Void) || matches!(t2, Transform::Void) +} + fn all_active_field_ids<'a>( partition_specs: impl Iterator, schema: &Schema, @@ -344,4 +391,59 @@ mod tests { compute_unified_partition_type([&spec_v0, &spec_v1].into_iter(), &schema).unwrap(); assert_eq!(result.fields().len(), 2); } + + #[test] + fn test_unknown_transform_errors() { + let schema = test_schema(); + + // A spec using an unknown transform. Deserialize directly since the builder + // validates transforms. + let spec = serde_json::from_value::(serde_json::json!({ + "spec-id": 0, + "fields": [{ + "source-id": 4, + "field-id": 1000, + "name": "category", + "transform": "unknown" + }] + })) + .unwrap(); + + let err = compute_unified_partition_type([&spec].into_iter(), &schema).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + } + + #[test] + fn test_conflicting_partition_fields_error() { + let schema = test_schema(); + + // Spec 0: field id 1000 -> source 4 (category), identity + let spec_v0 = serde_json::from_value::(serde_json::json!({ + "spec-id": 0, + "fields": [{ + "source-id": 4, + "field-id": 1000, + "name": "category", + "transform": "identity" + }] + })) + .unwrap(); + + // Spec 1: field id 1000 reused for a different source (ts) and transform (year). + // This conflicts with spec 0 and must be rejected. + let spec_v1 = serde_json::from_value::(serde_json::json!({ + "spec-id": 1, + "fields": [{ + "source-id": 3, + "field-id": 1000, + "name": "ts_year", + "transform": "year" + }] + })) + .unwrap(); + + let err = + compute_unified_partition_type([&spec_v0, &spec_v1].into_iter(), &schema).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + } }