diff --git a/crates/iceberg/src/arrow/int96.rs b/crates/iceberg/src/arrow/int96.rs index 6add023098..eea703a4fc 100644 --- a/crates/iceberg/src/arrow/int96.rs +++ b/crates/iceberg/src/arrow/int96.rs @@ -23,8 +23,10 @@ use arrow_schema::{ DataType, Field, FieldRef, Fields, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, TimeUnit, }; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; +use parquet::basic::Type as PhysicalType; +use parquet::schema::types::{ColumnDescPtr, SchemaDescriptor}; -use crate::arrow::schema::{ArrowSchemaVisitor, DEFAULT_MAP_FIELD_NAME, visit_schema}; +use crate::arrow::schema::{ArrowSchemaVisitor, UTC_TIME_ZONE, visit_schema}; use crate::error::Result; use crate::spec::{PrimitiveType, Schema, Type}; use crate::{Error, ErrorKind}; @@ -34,6 +36,7 @@ use crate::{Error, ErrorKind}; /// arrow-rs defaults INT96 to `Timestamp(Nanosecond)`, which overflows i64 for dates outside /// ~1677-2262. We use arrow-rs's schema hint mechanism to read INT96 at the resolution /// specified by the Iceberg schema (`timestamp` → microsecond, `timestamp_ns` → nanosecond). +/// The Arrow schema must contain all physical columns in file order, before projection. /// /// Iceberg Java handles this differently: it bypasses parquet-mr with a custom column reader /// (`GenericParquetReaders.TimestampInt96Reader`). We achieve the same result via schema hints. @@ -44,66 +47,69 @@ use crate::{Error, ErrorKind}; pub(crate) fn coerce_int96_timestamps( arrow_schema: &ArrowSchemaRef, iceberg_schema: &Schema, -) -> Option> { - let mut visitor = Int96CoercionVisitor::new(iceberg_schema); - let coerced = visit_schema(arrow_schema, &mut visitor).ok()?; - if visitor.changed { - Some(Arc::new(coerced)) - } else { - None + parquet_schema: &SchemaDescriptor, +) -> Result>> { + if !parquet_schema + .columns() + .iter() + .any(|column| column.physical_type() == PhysicalType::INT96) + { + return Ok(None); } + let mut visitor = Int96CoercionVisitor::new(iceberg_schema, parquet_schema); + let coerced = visit_schema(arrow_schema, &mut visitor)?; + if visitor.columns.len() != 0 { + return Err(Error::new( + ErrorKind::Unexpected, + "Arrow schema has fewer primitive leaves than the Parquet schema has columns", + )); + } + Ok(visitor.changed.then(|| Arc::new(coerced))) } -/// Visitor that coerces `Timestamp(Nanosecond)` Arrow fields to the resolution -/// indicated by the Iceberg schema. +/// Visitor that coerces INT96 Arrow fields to the Iceberg timestamp type. struct Int96CoercionVisitor<'a> { iceberg_schema: &'a Schema, + columns: std::slice::Iter<'a, ColumnDescPtr>, field_stack: Vec, changed: bool, } impl<'a> Int96CoercionVisitor<'a> { - fn new(iceberg_schema: &'a Schema) -> Self { + fn new(iceberg_schema: &'a Schema, parquet_schema: &'a SchemaDescriptor) -> Self { Self { iceberg_schema, + columns: parquet_schema.columns().iter(), field_stack: Vec::new(), changed: false, } } - /// Determine the target TimeUnit for a Timestamp(Nanosecond) field based on the - /// Iceberg schema. Falls back to microsecond when field IDs are unavailable, - /// matching Iceberg Java behavior. - fn target_unit(&self, field: &FieldRef) -> Option { - if !matches!( - field.data_type(), - DataType::Timestamp(TimeUnit::Nanosecond, _) - ) { - return None; - } - - let target = field + fn target_type(&self, field: &FieldRef, tz: &Option>) -> DataType { + field .metadata() .get(PARQUET_FIELD_ID_META_KEY) .and_then(|id_str| id_str.parse::().ok()) .and_then(|field_id| self.iceberg_schema.field_by_id(field_id)) .and_then(|f| match &*f.field_type { - Type::Primitive(PrimitiveType::Timestamp | PrimitiveType::Timestamptz) => { - Some(TimeUnit::Microsecond) + Type::Primitive(PrimitiveType::Timestamp) => { + Some(DataType::Timestamp(TimeUnit::Microsecond, None)) } - Type::Primitive(PrimitiveType::TimestampNs | PrimitiveType::TimestamptzNs) => { - Some(TimeUnit::Nanosecond) + Type::Primitive(PrimitiveType::Timestamptz) => Some(DataType::Timestamp( + TimeUnit::Microsecond, + Some(UTC_TIME_ZONE.into()), + )), + Type::Primitive(PrimitiveType::TimestampNs) => { + Some(DataType::Timestamp(TimeUnit::Nanosecond, None)) } + Type::Primitive(PrimitiveType::TimestamptzNs) => Some(DataType::Timestamp( + TimeUnit::Nanosecond, + Some(UTC_TIME_ZONE.into()), + )), _ => None, }) // Iceberg Java reads INT96 as microseconds by default - .unwrap_or(TimeUnit::Microsecond); - - if target == TimeUnit::Nanosecond { - None - } else { - Some(target) - } + .unwrap_or_else(|| DataType::Timestamp(TimeUnit::Microsecond, tz.clone())) } } @@ -163,12 +169,10 @@ impl ArrowSchemaVisitor for Int96CoercionVisitor<'_> { .field_stack .last() .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Field stack underflow in struct"))?; - Ok(Field::new( - field_info.name(), - DataType::Struct(Fields::from(results)), - field_info.is_nullable(), - ) - .with_metadata(field_info.metadata().clone())) + Ok(field_info + .as_ref() + .clone() + .with_data_type(DataType::Struct(Fields::from(results)))) } fn list(&mut self, list: &DataType, value: Field) -> Result { @@ -187,10 +191,7 @@ impl ArrowSchemaVisitor for Int96CoercionVisitor<'_> { )); } }; - Ok( - Field::new(field_info.name(), list_type, field_info.is_nullable()) - .with_metadata(field_info.metadata().clone()), - ) + Ok(field_info.as_ref().clone().with_data_type(list_type)) } fn map(&mut self, map: &DataType, key_value: Field, value: Field) -> Result { @@ -198,8 +199,8 @@ impl ArrowSchemaVisitor for Int96CoercionVisitor<'_> { .field_stack .last() .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Field stack underflow in map"))?; - let sorted = match map { - DataType::Map(_, sorted) => *sorted, + let (entries, sorted) = match map { + DataType::Map(entries, sorted) => (entries, *sorted), _ => { return Err(Error::new( ErrorKind::Unexpected, @@ -207,17 +208,14 @@ impl ArrowSchemaVisitor for Int96CoercionVisitor<'_> { )); } }; - let struct_field = Field::new( - DEFAULT_MAP_FIELD_NAME, - DataType::Struct(Fields::from(vec![key_value, value])), - false, - ); - Ok(Field::new( - field_info.name(), - DataType::Map(Arc::new(struct_field), sorted), - field_info.is_nullable(), - ) - .with_metadata(field_info.metadata().clone())) + let struct_field = entries + .as_ref() + .clone() + .with_data_type(DataType::Struct(Fields::from(vec![key_value, value]))); + Ok(field_info + .as_ref() + .clone() + .with_data_type(DataType::Map(Arc::new(struct_field), sorted))) } fn primitive(&mut self, p: &DataType) -> Result { @@ -225,24 +223,22 @@ impl ArrowSchemaVisitor for Int96CoercionVisitor<'_> { Error::new(ErrorKind::Unexpected, "Field stack underflow in primitive") })?; - if let Some(target_unit) = self.target_unit(field_info) { - let tz = match field_info.data_type() { - DataType::Timestamp(_, tz) => tz.clone(), - _ => None, - }; - self.changed = true; - Ok(Field::new( - field_info.name(), - DataType::Timestamp(target_unit, tz), - field_info.is_nullable(), - ) - .with_metadata(field_info.metadata().clone())) - } else { - Ok( - Field::new(field_info.name(), p.clone(), field_info.is_nullable()) - .with_metadata(field_info.metadata().clone()), + // Before projection, every primitive (including list elements and map + // keys/values) consumes one physical column in file order. + let column = self.columns.next().ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + "Missing Parquet column for Arrow primitive", ) - } + })?; + let target = match p { + DataType::Timestamp(_, tz) if column.physical_type() == PhysicalType::INT96 => { + self.target_type(field_info, tz) + } + _ => field_info.data_type().clone(), + }; + self.changed |= target != *field_info.data_type(); + Ok(field_info.as_ref().clone().with_data_type(target)) } } @@ -251,10 +247,13 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; - use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit}; + use arrow_schema::{DataType, Field, FieldRef, Schema as ArrowSchema, TimeUnit}; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + use parquet::schema::parser::parse_message_type; + use parquet::schema::types::SchemaDescriptor; use super::coerce_int96_timestamps; + use crate::arrow::schema::UTC_TIME_ZONE; use crate::spec::{ListType, MapType, NestedField, PrimitiveType, Schema, StructType, Type}; fn iceberg_schema_with_timestamp() -> Schema { @@ -272,8 +271,13 @@ mod tests { HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), id.to_string())]) } + fn parquet_schema(message: &str) -> SchemaDescriptor { + SchemaDescriptor::new(Arc::new(parse_message_type(message).unwrap())) + } + #[test] fn test_coerce_timestamp_ns_to_us() { + let parquet = parquet_schema("message schema { optional int96 ts; required int32 id; }"); let arrow_schema = Arc::new(ArrowSchema::new(vec![ Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true) .with_metadata(field_id_meta(1)), @@ -281,7 +285,9 @@ mod tests { ])); let iceberg = iceberg_schema_with_timestamp(); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg).unwrap(); + let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .unwrap(); assert_eq!( coerced.field(0).data_type(), &DataType::Timestamp(TimeUnit::Microsecond, None) @@ -292,6 +298,7 @@ mod tests { #[test] fn test_coerce_timestamptz_ns_to_us() { + let parquet = parquet_schema("message schema { optional int96 ts; }"); let iceberg = Schema::builder() .with_schema_id(1) .with_fields(vec![ @@ -301,23 +308,22 @@ mod tests { .unwrap(); let arrow_schema = Arc::new(ArrowSchema::new(vec![ - Field::new( - "ts", - DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), - true, - ) - .with_metadata(field_id_meta(1)), + Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true) + .with_metadata(field_id_meta(1)), ])); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg).unwrap(); + let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .unwrap(); assert_eq!( coerced.field(0).data_type(), - &DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())) + &DataType::Timestamp(TimeUnit::Microsecond, Some(UTC_TIME_ZONE.into())) ); } #[test] fn test_no_coercion_when_iceberg_is_timestamp_ns() { + let parquet = parquet_schema("message schema { optional int96 ts; }"); let iceberg = Schema::builder() .with_schema_id(1) .with_fields(vec![ @@ -331,11 +337,16 @@ mod tests { .with_metadata(field_id_meta(1)), ])); - assert!(coerce_int96_timestamps(&arrow_schema, &iceberg).is_none()); + assert!( + coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .is_none() + ); } #[test] fn test_no_coercion_when_iceberg_is_timestamptz_ns() { + let parquet = parquet_schema("message schema { optional int96 ts; }"); let iceberg = Schema::builder() .with_schema_id(1) .with_fields(vec![ @@ -348,17 +359,22 @@ mod tests { let arrow_schema = Arc::new(ArrowSchema::new(vec![ Field::new( "ts", - DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), + DataType::Timestamp(TimeUnit::Nanosecond, Some(UTC_TIME_ZONE.into())), true, ) .with_metadata(field_id_meta(1)), ])); - assert!(coerce_int96_timestamps(&arrow_schema, &iceberg).is_none()); + assert!( + coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .is_none() + ); } #[test] fn test_no_coercion_when_already_microsecond() { + let parquet = parquet_schema("message schema { optional int96 ts; required int32 id; }"); let arrow_schema = Arc::new(ArrowSchema::new(vec![ Field::new("ts", DataType::Timestamp(TimeUnit::Microsecond, None), true) .with_metadata(field_id_meta(1)), @@ -366,13 +382,141 @@ mod tests { ])); let iceberg = iceberg_schema_with_timestamp(); - assert!(coerce_int96_timestamps(&arrow_schema, &iceberg).is_none()); + assert!( + coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .is_none() + ); + } + + #[test] + fn test_no_coercion_without_int96() { + let parquet = + parquet_schema("message schema { optional int64 ts (TIMESTAMP(NANOS,false)) = 1; }"); + let arrow_schema = + Arc::new(parquet::arrow::parquet_to_arrow_schema(&parquet, None).unwrap()); + let iceberg = iceberg_schema_with_timestamp(); + + assert_eq!( + coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).unwrap(), + None + ); + } + + #[test] + fn test_coercion_rejects_list_view() { + let parquet = parquet_schema( + "message schema { + optional group items (LIST) { + repeated group list { optional int32 element; } + } + optional int96 ts = 1; + }", + ); + let iceberg = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamptz)).into(), + ]) + .build() + .unwrap(); + let list_type = DataType::ListView(Arc::new(Field::new("element", DataType::Int32, true))); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("items", list_type.clone(), true), + Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true) + .with_metadata(field_id_meta(1)), + ])); + + let error = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).unwrap_err(); + assert_eq!( + (error.kind(), error.message()), + ( + crate::ErrorKind::DataInvalid, + format!("Cannot visit Arrow data type: {list_type}").as_str(), + ) + ); + } + + #[test] + fn test_coercion_rejects_mismatched_leaf_counts() { + let iceberg = iceberg_schema_with_timestamp(); + for (message, expected_message) in [ + ( + "message schema { optional int96 ts; required int32 id; }", + "Arrow schema has fewer primitive leaves than the Parquet schema has columns", + ), + ( + "message schema { optional int96 ts; }", + "Missing Parquet column for Arrow primitive", + ), + ] { + let parquet = parquet_schema(message); + let fields = [ + Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true) + .with_metadata(field_id_meta(1)), + Field::new("id", DataType::Int32, false).with_metadata(field_id_meta(2)), + ]; + let leaf_count = if parquet.num_columns() == 2 { 1 } else { 2 }; + let arrow_schema = Arc::new(ArrowSchema::new(fields[..leaf_count].to_vec())); + let error = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).unwrap_err(); + assert_eq!( + (error.kind(), error.message()), + (crate::ErrorKind::Unexpected, expected_message) + ); + } + } + + #[test] + fn test_only_int96_timestamps_are_coerced() { + let parquet = parquet_schema( + "message schema { + optional int64 micros (TIMESTAMP(MICROS,false)) = 1; + optional int64 nanos (TIMESTAMP(NANOS,false)) = 2; + optional int96 legacy = 3; + optional int64 adjusted (TIMESTAMP(NANOS,true)) = 4; + }", + ); + let arrow_schema = + Arc::new(parquet::arrow::parquet_to_arrow_schema(&parquet, None).unwrap()); + for (timestamp_type, unit) in [ + (PrimitiveType::Timestamptz, TimeUnit::Microsecond), + (PrimitiveType::TimestamptzNs, TimeUnit::Nanosecond), + ] { + let iceberg = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "micros", Type::Primitive(PrimitiveType::Timestamptz)) + .into(), + NestedField::optional(2, "nanos", Type::Primitive(PrimitiveType::Timestamptz)) + .into(), + NestedField::optional(3, "legacy", Type::Primitive(timestamp_type)).into(), + NestedField::optional(4, "adjusted", Type::Primitive(PrimitiveType::Timestamp)) + .into(), + ]) + .build() + .unwrap(); + let mut fields: Vec = arrow_schema.fields().iter().cloned().collect(); + fields[2] = Arc::new( + fields[2] + .as_ref() + .clone() + .with_data_type(DataType::Timestamp(unit, Some(UTC_TIME_ZONE.into()))), + ); + let expected = Arc::new(ArrowSchema::new_with_metadata( + fields, + arrow_schema.metadata().clone(), + )); + + assert_eq!( + coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).unwrap(), + Some(expected) + ); + } } // Without field IDs, the visitor can't look up the Iceberg type and falls back // to microsecond to match Iceberg Java behavior. #[test] fn test_defaults_to_us_without_field_ids() { + let parquet = parquet_schema("message schema { optional int96 ts; }"); let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new( "ts", DataType::Timestamp(TimeUnit::Nanosecond, None), @@ -380,7 +524,9 @@ mod tests { )])); let iceberg = iceberg_schema_with_timestamp(); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg).unwrap(); + let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .unwrap(); assert_eq!( coerced.field(0).data_type(), &DataType::Timestamp(TimeUnit::Microsecond, None) @@ -392,6 +538,7 @@ mod tests { // microsecond. #[test] fn test_defaults_to_us_when_iceberg_type_is_not_timestamp() { + let parquet = parquet_schema("message schema { optional int96 ts; }"); let iceberg = Schema::builder() .with_schema_id(1) .with_fields(vec![ @@ -405,7 +552,9 @@ mod tests { .with_metadata(field_id_meta(1)), ])); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg).unwrap(); + let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .unwrap(); assert_eq!( coerced.field(0).data_type(), &DataType::Timestamp(TimeUnit::Microsecond, None) @@ -414,6 +563,7 @@ mod tests { #[test] fn test_coerce_preserves_field_metadata() { + let parquet = parquet_schema("message schema { optional int96 ts; }"); let mut meta = field_id_meta(1); meta.insert("custom_key".to_string(), "custom_value".to_string()); @@ -423,12 +573,16 @@ mod tests { ])); let iceberg = iceberg_schema_with_timestamp(); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg).unwrap(); + let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .unwrap(); assert_eq!(coerced.field(0).metadata(), &meta); } #[test] fn test_coerce_timestamp_in_struct() { + let parquet = + parquet_schema("message schema { required group data { optional int96 ts; } }"); let iceberg = Schema::builder() .with_schema_id(1) .with_fields(vec![ @@ -460,7 +614,9 @@ mod tests { .with_metadata(field_id_meta(1)), ])); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg).unwrap(); + let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .unwrap(); let inner = match coerced.field(0).data_type() { DataType::Struct(fields) => fields, other => panic!("Expected Struct, got {other}"), @@ -473,6 +629,11 @@ mod tests { #[test] fn test_coerce_timestamp_in_list() { + let parquet = parquet_schema( + "message schema { optional group timestamps (LIST) { + repeated group list { optional int96 element; } + } }", + ); let iceberg = Schema::builder() .with_schema_id(1) .with_fields(vec![ @@ -504,7 +665,9 @@ mod tests { .with_metadata(field_id_meta(1)), ])); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg).unwrap(); + let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .unwrap(); let element_dt = match coerced.field(0).data_type() { DataType::List(f) => f.data_type(), other => panic!("Expected List, got {other}"), @@ -517,6 +680,11 @@ mod tests { #[test] fn test_coerce_timestamp_in_map_value() { + let parquet = parquet_schema( + "message schema { optional group ts_map (MAP) { + repeated group key_value { required binary key (STRING); optional int96 value; } + } }", + ); let iceberg = Schema::builder() .with_schema_id(1) .with_fields(vec![ @@ -552,9 +720,13 @@ mod tests { .with_metadata(field_id_meta(3)); let entries_field = Field::new( "key_value", - DataType::Struct(vec![key_field, value_field].into()), + DataType::Struct(vec![key_field.clone(), value_field].into()), false, - ); + ) + .with_metadata(HashMap::from([( + "custom_key".to_string(), + "custom_value".to_string(), + )])); let arrow_schema = Arc::new(ArrowSchema::new(vec![ Field::new( "ts_map", @@ -564,14 +736,37 @@ mod tests { .with_metadata(field_id_meta(1)), ])); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg).unwrap(); - let value_dt = match coerced.field(0).data_type() { - DataType::Map(entries, _) => match entries.data_type() { - DataType::Struct(fields) => fields[1].data_type().clone(), - other => panic!("Expected Struct inside Map, got {other}"), - }, - other => panic!("Expected Map, got {other}"), - }; - assert_eq!(value_dt, DataType::Timestamp(TimeUnit::Microsecond, None)); + let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .unwrap(); + let expected_entries = Field::new( + "key_value", + DataType::Struct( + vec![ + key_field, + Field::new( + "value", + DataType::Timestamp(TimeUnit::Microsecond, None), + true, + ) + .with_metadata(field_id_meta(3)), + ] + .into(), + ), + false, + ) + .with_metadata(HashMap::from([( + "custom_key".to_string(), + "custom_value".to_string(), + )])); + let expected = Arc::new(ArrowSchema::new(vec![ + Field::new( + "ts_map", + DataType::Map(Arc::new(expected_entries), false), + true, + ) + .with_metadata(field_id_meta(1)), + ])); + assert_eq!(coerced, expected); } } diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index 1ddee493c4..3918c486d5 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -699,10 +699,13 @@ impl FileScanTaskReader { // Coerce INT96 timestamp columns before building the stream reader to avoid // i64 overflow in arrow-rs. Apply this after assigning any missing field IDs // so the final schema contains both changes. - coerce_int96_timestamps(&schema, task.schema()).unwrap_or(schema) - } else if let Some(coerced) = - coerce_int96_timestamps(arrow_metadata.schema(), task.schema()) - { + coerce_int96_timestamps(&schema, task.schema(), arrow_metadata.parquet_schema())? + .unwrap_or(schema) + } else if let Some(coerced) = coerce_int96_timestamps( + arrow_metadata.schema(), + task.schema(), + arrow_metadata.parquet_schema(), + )? { coerced } else if install_row_number { Arc::clone(arrow_metadata.schema()) @@ -911,7 +914,10 @@ mod tests { RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_ROW_ID, }; use crate::scan::{FileScanTask, FileScanTaskDeleteFile, FileScanTaskStream}; - use crate::spec::{DataFileFormat, NestedField, PrimitiveType, Schema, SchemaRef, Type}; + use crate::spec::{ + DataFileFormat, MappedField, NameMapping, NestedField, PrimitiveType, Schema, SchemaRef, + Type, + }; // INT96 encoding: [nanos_low_u32, nanos_high_u32, julian_day_u32] // Julian day 2_440_588 = Unix epoch (1970-01-01) @@ -937,12 +943,38 @@ mod tests { file_path: &str, schema: SchemaRef, project_field_ids: Vec, + ) -> Vec { + read_int96_batches_with_name_mapping(file_path, schema, project_field_ids, None).await + } + + async fn read_int96_batches_with_name_mapping( + file_path: &str, + schema: SchemaRef, + project_field_ids: Vec, + name_mapping: Option>, ) -> Vec { let file_io = FileIO::new_with_fs(); let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build(); + let task = int96_scan_task(file_path, schema, project_field_ids, name_mapping); + let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; + reader + .read(tasks) + .unwrap() + .stream() + .try_collect() + .await + .unwrap() + } + + fn int96_scan_task( + file_path: &str, + schema: SchemaRef, + project_field_ids: Vec, + name_mapping: Option>, + ) -> FileScanTask { let file_size = std::fs::metadata(file_path).unwrap().len(); - let task = FileScanTask::builder() + FileScanTask::builder() .with_file_size_in_bytes(file_size) .with_start(0) .with_length(file_size) @@ -950,17 +982,9 @@ mod tests { .with_data_file_format(DataFileFormat::Parquet) .with_schema(schema) .with_project_field_ids(project_field_ids) + .with_name_mapping(name_mapping) .with_case_sensitive(false) .build() - .unwrap(); - - let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; - reader - .read(tasks) - .unwrap() - .stream() - .try_collect() - .await .unwrap() } @@ -1063,33 +1087,6 @@ mod tests { (file_path, expected_micros) } - async fn assert_int96_read_matches( - file_path: &str, - schema: SchemaRef, - project_field_ids: Vec, - expected_micros: &[i64], - ) { - use arrow_array::TimestampMicrosecondArray; - - let batches = read_int96_batches(file_path, schema, project_field_ids).await; - - assert_eq!(batches.len(), 1); - let ts_array = batches[0] - .column(0) - .as_any() - .downcast_ref::() - .expect("Expected TimestampMicrosecondArray"); - - for (i, expected) in expected_micros.iter().enumerate() { - assert_eq!( - ts_array.value(i), - *expected, - "Row {i}: got {}, expected {expected}", - ts_array.value(i) - ); - } - } - /// Writes a single-column Parquet file encrypted with `encryption_key`, then reads it /// back through `ArrowReader` and asserts the round-tripped values. The key length /// selects the AES-GCM variant in arrow-rs (16 -> AES-128, 32 -> AES-256). @@ -2829,46 +2826,330 @@ mod tests { #[tokio::test] async fn test_read_int96_timestamps_with_field_ids() { + for timestamp_type in [PrimitiveType::Timestamp, PrimitiveType::Timestamptz] { + assert_int96_microsecond_projections(timestamp_type, true, None).await; + } + } + + #[tokio::test] + async fn test_read_int96_timestamps_without_field_ids() { + for timestamp_type in [PrimitiveType::Timestamp, PrimitiveType::Timestamptz] { + assert_int96_microsecond_projections(timestamp_type, false, None).await; + } + } + + #[tokio::test] + async fn test_read_int96_timestamps_with_name_mapping() { + let name_mapping = Arc::new(NameMapping::new(vec![ + MappedField::new(Some(1), vec!["ts".to_string()], vec![]), + MappedField::new(Some(2), vec!["id".to_string()], vec![]), + ])); + for timestamp_type in [PrimitiveType::Timestamp, PrimitiveType::Timestamptz] { + assert_int96_microsecond_projections(timestamp_type, false, Some(name_mapping.clone())) + .await; + } + } + + async fn assert_int96_microsecond_projections( + timestamp_type: PrimitiveType, + with_field_ids: bool, + name_mapping: Option>, + ) { + use arrow_array::TimestampMicrosecondArray; + + use crate::arrow::schema::UTC_TIME_ZONE; + use crate::arrow::schema_to_arrow_schema; + + let timezone = (timestamp_type == PrimitiveType::Timestamptz).then_some(UTC_TIME_ZONE); let schema = Arc::new( Schema::builder() .with_schema_id(1) .with_fields(vec![ - NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamp)) - .into(), + NestedField::optional(1, "ts", Type::Primitive(timestamp_type)).into(), NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(), ]) .build() .unwrap(), ); + let tmp_dir = TempDir::new().unwrap(); + let (file_path, expected_micros) = write_int96_parquet_file( + tmp_dir.path().to_str().unwrap(), + "timestamps.parquet", + with_field_ids, + ); + let ids: Vec = (0..expected_micros.len() as i32).collect(); + let expected_batch = + RecordBatch::try_new(Arc::new(schema_to_arrow_schema(&schema).unwrap()), vec![ + Arc::new( + TimestampMicrosecondArray::from(expected_micros).with_timezone_opt(timezone), + ), + Arc::new(Int32Array::from(ids)), + ]) + .unwrap(); + for (field_ids, column_indices) in [ + (vec![1], vec![0]), + (vec![1, 2], vec![0, 1]), + (vec![2, 1], vec![1, 0]), + ] { + let actual = read_int96_batches_with_name_mapping( + &file_path, + schema.clone(), + field_ids, + name_mapping.clone(), + ) + .await; + let expected = vec![expected_batch.project(&column_indices).unwrap()]; + assert_eq!(actual, expected); + } + } + #[tokio::test] + async fn test_read_int96_nanosecond_timestamps() { + for timestamp_type in [PrimitiveType::TimestampNs, PrimitiveType::TimestamptzNs] { + for with_field_ids in [true, false] { + assert_int96_nanosecond_projections(timestamp_type.clone(), with_field_ids).await; + } + } + } + + async fn assert_int96_nanosecond_projections( + timestamp_type: PrimitiveType, + with_field_ids: bool, + ) { + use arrow_array::TimestampNanosecondArray; + use parquet::data_type::{Int32Type, Int96, Int96Type}; + use parquet::file::writer::SerializedFileWriter; + use parquet::schema::parser::parse_message_type; + + use crate::arrow::schema::UTC_TIME_ZONE; + use crate::arrow::schema_to_arrow_schema; + + let timezone = (timestamp_type == PrimitiveType::TimestamptzNs).then_some(UTC_TIME_ZONE); + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "ts", Type::Primitive(timestamp_type)).into(), + NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build() + .unwrap(), + ); + let message = if with_field_ids { + "message schema { optional int96 ts = 1; required int32 id = 2; }" + } else { + "message schema { optional int96 ts; required int32 id; }" + }; + let parquet_schema = parse_message_type(message).unwrap(); + let name_mapping = (!with_field_ids).then(|| { + Arc::new(NameMapping::new(vec![ + MappedField::new(Some(1), vec!["ts".to_string()], vec![]), + MappedField::new(Some(2), vec!["id".to_string()], vec![]), + ])) + }); + let nanos = vec![Some(1), None, Some(2_000_006_001_i64)]; + let int96_values: Vec = nanos + .iter() + .flatten() + .map(|nanos| { + let mut value = Int96::new(); + value.set_data( + *nanos as u32, + (nanos >> 32) as u32, + UNIX_EPOCH_JULIAN as u32, + ); + value + }) + .collect(); + let ids = vec![0, 1, 2]; let tmp_dir = TempDir::new().unwrap(); - let table_location = tmp_dir.path().to_str().unwrap().to_string(); - let (file_path, expected_micros) = - write_int96_parquet_file(&table_location, "with_ids.parquet", true); + let file_path = tmp_dir.path().join("nanosecond_timestamps.parquet"); + let mut writer = SerializedFileWriter::new( + File::create(&file_path).unwrap(), + Arc::new(parquet_schema), + Default::default(), + ) + .unwrap(); + let mut row_group = writer.next_row_group().unwrap(); + let mut column = row_group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch(&int96_values, Some(&[1, 0, 1]), None) + .unwrap(); + column.close().unwrap(); + let mut column = row_group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch(&ids, None, None) + .unwrap(); + column.close().unwrap(); + row_group.close().unwrap(); + writer.close().unwrap(); - assert_int96_read_matches(&file_path, schema, vec![1, 2], &expected_micros).await; + let expected_batch = + RecordBatch::try_new(Arc::new(schema_to_arrow_schema(&schema).unwrap()), vec![ + Arc::new(TimestampNanosecondArray::from(nanos).with_timezone_opt(timezone)), + Arc::new(Int32Array::from(ids)), + ]) + .unwrap(); + for (field_ids, column_indices) in [ + (vec![1], vec![0]), + (vec![1, 2], vec![0, 1]), + (vec![2, 1], vec![1, 0]), + ] { + let actual = read_int96_batches_with_name_mapping( + file_path.to_str().unwrap(), + schema.clone(), + field_ids, + name_mapping.clone(), + ) + .await; + let expected = vec![expected_batch.project(&column_indices).unwrap()]; + assert_eq!(actual, expected); + } } #[tokio::test] - async fn test_read_int96_timestamps_without_field_ids() { + async fn test_read_int96_with_mixed_timestamp_types_in_struct() { + use arrow_array::{StructArray, TimestampMicrosecondArray, TimestampNanosecondArray}; + use parquet::data_type::{Int64Type, Int96Type}; + use parquet::file::writer::SerializedFileWriter; + use parquet::schema::parser::parse_message_type; + + use crate::arrow::schema::UTC_TIME_ZONE; + use crate::arrow::schema_to_arrow_schema; + use crate::spec::StructType; + + let parquet_schema = parse_message_type( + "message schema { + required int64 id = 1; + required group data = 2 { + optional int64 micros (TIMESTAMP(MICROS,false)) = 3; + optional int96 legacy = 4; + } + optional int64 nanos (TIMESTAMP(NANOS,true)) = 5; + optional int96 trailing = 6; + }", + ) + .unwrap(); let schema = Arc::new( Schema::builder() - .with_schema_id(1) .with_fields(vec![ - NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamp)) + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), + NestedField::required( + 2, + "data", + Type::Struct(StructType::new(vec![ + NestedField::optional( + 3, + "micros", + Type::Primitive(PrimitiveType::Timestamp), + ) + .into(), + NestedField::optional( + 4, + "legacy", + Type::Primitive(PrimitiveType::Timestamptz), + ) + .into(), + ])), + ) + .into(), + NestedField::optional( + 5, + "nanos", + Type::Primitive(PrimitiveType::TimestamptzNs), + ) + .into(), + NestedField::optional(6, "trailing", Type::Primitive(PrimitiveType::Timestamp)) .into(), - NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(), ]) .build() .unwrap(), ); - + let (int96_value, expected_micros) = make_int96_test_value(); + let nanos = vec![1_000_003_001_i64, -1_000_003_001, 0]; let tmp_dir = TempDir::new().unwrap(); - let table_location = tmp_dir.path().to_str().unwrap().to_string(); - let (file_path, expected_micros) = - write_int96_parquet_file(&table_location, "no_ids.parquet", false); + let file_path = tmp_dir.path().join("mixed_timestamps.parquet"); + let mut writer = SerializedFileWriter::new( + File::create(&file_path).unwrap(), + Arc::new(parquet_schema), + Default::default(), + ) + .unwrap(); + let mut row_group = writer.next_row_group().unwrap(); + let mut column = row_group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch(&[0, 1, 2], None, None) + .unwrap(); + column.close().unwrap(); + let mut column = row_group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch(&[17, -17], Some(&[1, 0, 1]), None) + .unwrap(); + column.close().unwrap(); + let mut column = row_group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch(&[int96_value, int96_value], Some(&[1, 0, 1]), None) + .unwrap(); + column.close().unwrap(); + let mut column = row_group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch(&nanos, Some(&[1; 3]), None) + .unwrap(); + column.close().unwrap(); + let mut column = row_group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch(&[int96_value; 3], Some(&[1; 3]), None) + .unwrap(); + column.close().unwrap(); + row_group.close().unwrap(); + writer.close().unwrap(); - assert_int96_read_matches(&file_path, schema, vec![1, 2], &expected_micros).await; + let arrow_schema = Arc::new(schema_to_arrow_schema(&schema).unwrap()); + let DataType::Struct(fields) = arrow_schema.field(1).data_type() else { + panic!("Expected struct field"); + }; + let expected_batch = RecordBatch::try_new(arrow_schema.clone(), vec![ + Arc::new(Int64Array::from(vec![0, 1, 2])), + Arc::new(StructArray::new( + fields.clone(), + vec![ + Arc::new(TimestampMicrosecondArray::from(vec![ + Some(17), + None, + Some(-17), + ])), + Arc::new( + TimestampMicrosecondArray::from(vec![ + Some(expected_micros), + None, + Some(expected_micros), + ]) + .with_timezone(UTC_TIME_ZONE), + ), + ], + None, + )), + Arc::new(TimestampNanosecondArray::from(nanos).with_timezone(UTC_TIME_ZONE)), + Arc::new(TimestampMicrosecondArray::from(vec![expected_micros; 3])), + ]) + .unwrap(); + + for (field_ids, column_indices) in [ + (vec![1, 2, 5, 6], vec![0, 1, 2, 3]), + (vec![2], vec![1]), + (vec![6, 2, 5], vec![3, 1, 2]), + ] { + let actual = + read_int96_batches(file_path.to_str().unwrap(), schema.clone(), field_ids).await; + let expected = vec![expected_batch.project(&column_indices).unwrap()]; + assert_eq!(actual, expected); + } } #[tokio::test] @@ -3057,12 +3338,25 @@ mod tests { #[tokio::test] async fn test_read_int96_timestamps_in_list() { + assert_read_int96_timestamps_in_list(PrimitiveType::Timestamp).await; + } + + #[tokio::test] + async fn test_read_int96_timestamptz_in_list() { + assert_read_int96_timestamps_in_list(PrimitiveType::Timestamptz).await; + } + + async fn assert_read_int96_timestamps_in_list(timestamp_type: PrimitiveType) { use arrow_array::{ListArray, TimestampMicrosecondArray}; + use arrow_buffer::OffsetBuffer; use parquet::basic::{Repetition, Type as PhysicalType}; use parquet::data_type::Int96Type; use parquet::file::writer::SerializedFileWriter; use parquet::schema::types::Type as SchemaType; + use crate::arrow::schema::UTC_TIME_ZONE; + use crate::arrow::schema_to_arrow_schema; + let tmp_dir = TempDir::new().unwrap(); let table_location = tmp_dir.path().to_str().unwrap().to_string(); let file_path = format!("{table_location}/list_int96.parquet"); @@ -3129,7 +3423,7 @@ mod tests { element_field: NestedField::optional( 2, "element", - Type::Primitive(PrimitiveType::Timestamp), + Type::Primitive(timestamp_type.clone()), ) .into(), }), @@ -3140,36 +3434,46 @@ mod tests { .unwrap(), ); - let batches = read_int96_batches(&file_path, iceberg_schema, vec![1]).await; - - assert_eq!(batches.len(), 1); - let list_array = batches[0] - .column(0) - .as_any() - .downcast_ref::() - .expect("Expected ListArray"); - let ts_array = list_array - .values() - .as_any() - .downcast_ref::() - .expect("Expected TimestampMicrosecondArray inside list"); - - assert_eq!( - ts_array.value(0), - expected_micros, - "INT96 in list: got {}, expected {expected_micros}", - ts_array.value(0) + let batches = read_int96_batches(&file_path, iceberg_schema.clone(), vec![1]).await; + let arrow_schema = Arc::new(schema_to_arrow_schema(&iceberg_schema).unwrap()); + let DataType::List(element) = arrow_schema.field(0).data_type() else { + panic!("Expected list field"); + }; + let timezone = (timestamp_type == PrimitiveType::Timestamptz).then_some(UTC_TIME_ZONE); + let values = Arc::new( + TimestampMicrosecondArray::from(vec![expected_micros]).with_timezone_opt(timezone), + ); + let list = ListArray::new( + element.clone(), + OffsetBuffer::from_lengths([1]), + values, + None, ); + let expected = vec![RecordBatch::try_new(arrow_schema, vec![Arc::new(list)]).unwrap()]; + assert_eq!(batches, expected); } #[tokio::test] async fn test_read_int96_timestamps_in_map() { - use arrow_array::{MapArray, TimestampMicrosecondArray}; + assert_read_int96_timestamps_in_map(PrimitiveType::Timestamp).await; + } + + #[tokio::test] + async fn test_read_int96_timestamptz_in_map() { + assert_read_int96_timestamps_in_map(PrimitiveType::Timestamptz).await; + } + + async fn assert_read_int96_timestamps_in_map(timestamp_type: PrimitiveType) { + use arrow_array::{MapArray, StructArray, TimestampMicrosecondArray}; + use arrow_buffer::OffsetBuffer; use parquet::basic::{Repetition, Type as PhysicalType}; use parquet::data_type::{ByteArrayType, Int96Type}; use parquet::file::writer::SerializedFileWriter; use parquet::schema::types::Type as SchemaType; + use crate::arrow::schema::UTC_TIME_ZONE; + use crate::arrow::schema_to_arrow_schema; + let tmp_dir = TempDir::new().unwrap(); let table_location = tmp_dir.path().to_str().unwrap().to_string(); let file_path = format!("{table_location}/map_int96.parquet"); @@ -3262,7 +3566,7 @@ mod tests { value_field: NestedField::optional( 3, "value", - Type::Primitive(PrimitiveType::Timestamp), + Type::Primitive(timestamp_type.clone()), ) .into(), }), @@ -3273,26 +3577,35 @@ mod tests { .unwrap(), ); - let batches = read_int96_batches(&file_path, iceberg_schema, vec![1]).await; - - assert_eq!(batches.len(), 1); - let map_array = batches[0] - .column(0) - .as_any() - .downcast_ref::() - .expect("Expected MapArray"); - let ts_array = map_array - .values() - .as_any() - .downcast_ref::() - .expect("Expected TimestampMicrosecondArray as map values"); - - assert_eq!( - ts_array.value(0), - expected_micros, - "INT96 in map: got {}, expected {expected_micros}", - ts_array.value(0) + let batches = read_int96_batches(&file_path, iceberg_schema.clone(), vec![1]).await; + let arrow_schema = Arc::new(schema_to_arrow_schema(&iceberg_schema).unwrap()); + let DataType::Map(entry_field, sorted) = arrow_schema.field(0).data_type() else { + panic!("Expected map field"); + }; + let DataType::Struct(fields) = entry_field.data_type() else { + panic!("Expected map entries struct"); + }; + let timezone = (timestamp_type == PrimitiveType::Timestamptz).then_some(UTC_TIME_ZONE); + let entries = StructArray::new( + fields.clone(), + vec![ + Arc::new(StringArray::from(vec!["event_time"])), + Arc::new( + TimestampMicrosecondArray::from(vec![expected_micros]) + .with_timezone_opt(timezone), + ), + ], + None, + ); + let map = MapArray::new( + entry_field.clone(), + OffsetBuffer::from_lengths([1]), + entries, + None, + *sorted, ); + let expected = vec![RecordBatch::try_new(arrow_schema, vec![Arc::new(map)]).unwrap()]; + assert_eq!(batches, expected); } /// Writes `id` (Int32) plus a wide string column (field id 2) whose bytes dominate