From 7ecd4dee899bb51731e58b3292c016fe9b872249 Mon Sep 17 00:00:00 2001 From: Hung Date: Thu, 5 Nov 2026 23:57:09 +0800 Subject: [PATCH 1/3] fix(arrow): preserve INT96 timestamptz values in projections --- crates/iceberg/src/arrow/int96.rs | 327 +++++++--- crates/iceberg/src/arrow/reader/pipeline.rs | 647 ++++++++++++++++++-- 2 files changed, 819 insertions(+), 155 deletions(-) diff --git a/crates/iceberg/src/arrow/int96.rs b/crates/iceberg/src/arrow/int96.rs index 6add023098..338327a275 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,67 @@ use crate::{Error, ErrorKind}; pub(crate) fn coerce_int96_timestamps( arrow_schema: &ArrowSchemaRef, iceberg_schema: &Schema, + parquet_schema: &SchemaDescriptor, ) -> Option> { - let mut visitor = Int96CoercionVisitor::new(iceberg_schema); + if !parquet_schema + .columns() + .iter() + .any(|column| column.physical_type() == PhysicalType::INT96) + { + return None; + } + let mut visitor = Int96CoercionVisitor::new(iceberg_schema, parquet_schema); let coerced = visit_schema(arrow_schema, &mut visitor).ok()?; - if visitor.changed { + if visitor.changed && visitor.columns.len() == 0 { Some(Arc::new(coerced)) } else { None } } -/// 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 +167,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 +189,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 +197,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 +206,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 +221,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(), + // 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", ) - .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()), - ) - } + })?; + 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 +245,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 +269,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 +283,7 @@ 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(); assert_eq!( coerced.field(0).data_type(), &DataType::Timestamp(TimeUnit::Microsecond, None) @@ -292,6 +294,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 +304,20 @@ 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(); 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 +331,12 @@ 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).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 +349,18 @@ 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).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 +368,101 @@ 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).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), + None + ); + } + + #[test] + fn test_no_coercion_with_mismatched_leaf_count() { + let ts = Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true) + .with_metadata(field_id_meta(1)); + let id = Field::new("id", DataType::Int32, false).with_metadata(field_id_meta(2)); + let iceberg = iceberg_schema_with_timestamp(); + for (message, fields) in [ + ("message schema { optional int96 ts; }", vec![ + ts.clone(), + id, + ]), + ( + "message schema { optional int96 ts; required int32 id; }", + vec![ts], + ), + ] { + let parquet = parquet_schema(message); + let arrow_schema = Arc::new(ArrowSchema::new(fields)); + + assert_eq!( + coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet), + None + ); + } + } + + #[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), + 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 +470,7 @@ 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(); assert_eq!( coerced.field(0).data_type(), &DataType::Timestamp(TimeUnit::Microsecond, None) @@ -392,6 +482,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 +496,7 @@ 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(); assert_eq!( coerced.field(0).data_type(), &DataType::Timestamp(TimeUnit::Microsecond, None) @@ -414,6 +505,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 +515,14 @@ 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(); 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 +554,7 @@ 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(); let inner = match coerced.field(0).data_type() { DataType::Struct(fields) => fields, other => panic!("Expected Struct, got {other}"), @@ -473,6 +567,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 +603,7 @@ 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(); let element_dt = match coerced.field(0).data_type() { DataType::List(f) => f.data_type(), other => panic!("Expected List, got {other}"), @@ -517,6 +616,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 +656,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 +672,35 @@ 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(); + 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..888a687858 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,20 +982,60 @@ 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() } + fn assert_int96_projection_mask(task: &FileScanTask, expected_indices: &[usize]) { + use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; + use parquet::file::reader::{FileReader, SerializedFileReader}; + + use super::{ArrowReader, FileScanTaskReader}; + + let reader = SerializedFileReader::new(File::open(task.data_file_path()).unwrap()).unwrap(); + let metadata = ArrowReaderMetadata::try_new( + Arc::new(reader.metadata().clone()), + ArrowReaderOptions::default(), + ) + .unwrap(); + let missing_field_ids = metadata + .schema() + .fields() + .iter() + .next() + .is_some_and(|field| !field.metadata().contains_key(PARQUET_FIELD_ID_META_KEY)); + let use_position_fallback = missing_field_ids && task.name_mapping().is_none(); + let metadata = FileScanTaskReader::configure_arrow_reader_metadata( + metadata, + task, + missing_field_ids, + false, + ) + .unwrap(); + let mask = ArrowReader::get_arrow_projection_mask( + task.project_field_ids(), + task.schema(), + metadata.parquet_schema(), + metadata.schema(), + use_position_fallback, + ) + .unwrap(); + let actual: Vec = (0..metadata.parquet_schema().num_columns()) + .map(|index| mask.leaf_included(index)) + .collect(); + let expected: Vec = (0..metadata.parquet_schema().num_columns()) + .map(|index| expected_indices.contains(&index)) + .collect(); + assert_eq!( + actual, + expected, + "projection {:?}", + task.project_field_ids() + ); + } + // ArrowWriter cannot write INT96, so we use SerializedFileWriter directly. fn write_int96_parquet_file( table_location: &str, @@ -2871,6 +2943,437 @@ mod tests { assert_int96_read_matches(&file_path, schema, vec![1, 2], &expected_micros).await; } + #[tokio::test] + async fn test_read_int96_timestamptz_with_name_mapping() { + assert_int96_timestamp_projections( + PrimitiveType::Timestamptz, + Int96FieldIds::NameMapping, + false, + ) + .await; + } + + #[tokio::test] + async fn test_read_int96_timestamptz_ns_with_name_mapping() { + assert_int96_timestamp_projections( + PrimitiveType::TimestamptzNs, + Int96FieldIds::NameMapping, + false, + ) + .await; + } + + #[tokio::test] + async fn test_read_int96_timestamptz_with_field_ids_and_nulls() { + assert_int96_timestamp_projections( + PrimitiveType::Timestamptz, + Int96FieldIds::Embedded, + true, + ) + .await; + } + + #[tokio::test] + async fn test_read_int96_timestamptz_with_positional_fallback() { + assert_int96_timestamp_projections( + PrimitiveType::Timestamptz, + Int96FieldIds::Position, + false, + ) + .await; + } + + #[tokio::test] + async fn test_read_int96_timestamp_without_timezone_with_name_mapping() { + assert_int96_timestamp_projections( + PrimitiveType::Timestamp, + Int96FieldIds::NameMapping, + true, + ) + .await; + } + + #[tokio::test] + async fn test_read_int96_timestamp_ns_with_field_ids_and_nulls() { + assert_int96_timestamp_projections( + PrimitiveType::TimestampNs, + Int96FieldIds::Embedded, + true, + ) + .await; + } + + enum Int96FieldIds { + Embedded, + NameMapping, + Position, + } + + async fn assert_int96_timestamp_projections( + timestamp_type: PrimitiveType, + field_ids: Int96FieldIds, + with_nulls: bool, + ) { + use arrow_array::{TimestampMicrosecondArray, TimestampNanosecondArray, new_null_array}; + use parquet::basic::{LogicalType, Repetition, Type as PhysicalType}; + use parquet::data_type::{ByteArray, ByteArrayType, Int64Type, Int96, 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; + use crate::spec::Literal; + + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), + NestedField::optional(2, "ts", Type::Primitive(timestamp_type.clone())).into(), + NestedField::required(3, "s", Type::Primitive(PrimitiveType::String)).into(), + NestedField::optional(4, "missing_ts", Type::Primitive(timestamp_type.clone())) + .into(), + NestedField::optional(5, "default_s", Type::Primitive(PrimitiveType::String)) + .with_initial_default(Literal::string("default")) + .into(), + ]) + .build() + .unwrap(), + ); + let name_mapping = matches!(field_ids, Int96FieldIds::NameMapping).then(|| { + Arc::new(NameMapping::new(vec![ + MappedField::new(Some(1), vec!["id".to_string()], vec![]), + MappedField::new(Some(2), vec!["ts".to_string()], vec![]), + MappedField::new(Some(3), vec!["s".to_string()], vec![]), + ])) + }); + let with_field_ids = matches!(field_ids, Int96FieldIds::Embedded); + let parquet_schema = SchemaType::group_type_builder("schema") + .with_fields(vec![ + Arc::new( + SchemaType::primitive_type_builder("id", PhysicalType::INT64) + .with_repetition(Repetition::REQUIRED) + .with_id(with_field_ids.then_some(1)) + .build() + .unwrap(), + ), + Arc::new( + SchemaType::primitive_type_builder("ts", PhysicalType::INT96) + .with_repetition(Repetition::OPTIONAL) + .with_id(with_field_ids.then_some(2)) + .build() + .unwrap(), + ), + Arc::new( + SchemaType::primitive_type_builder("s", PhysicalType::BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_logical_type(Some(LogicalType::String)) + .with_id(with_field_ids.then_some(3)) + .build() + .unwrap(), + ), + ]) + .build() + .unwrap(); + let micros = vec![0_i64, 1_000_003, 2_000_006, 3_000_009, 4_000_012]; + let nanos: Vec = micros + .iter() + .map(|micros| { + micros * 1_000 + + i64::from(matches!( + timestamp_type, + PrimitiveType::TimestampNs | PrimitiveType::TimestamptzNs + )) + }) + .collect(); + let ids = vec![0_i64, 1, 2, 3, 4]; + let strings = vec!["0", "1", "2", "3", "4"]; + let definition_levels: Vec = (0..5) + .map(|index| i16::from(!with_nulls || index % 2 == 0)) + .collect(); + let int96_values: Vec = nanos + .iter() + .zip(&definition_levels) + .filter(|(_, level)| **level == 1) + .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 string_values: Vec = strings + .iter() + .map(|value| ByteArray::from(*value)) + .collect(); + + let tmp_dir = TempDir::new().unwrap(); + let file_path = tmp_dir.path().join("timestamptz_name_mapping.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(&ids, None, None) + .unwrap(); + column.close().unwrap(); + let mut column = row_group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch(&int96_values, Some(&definition_levels), None) + .unwrap(); + column.close().unwrap(); + let mut column = row_group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch(&string_values, None, None) + .unwrap(); + column.close().unwrap(); + row_group.close().unwrap(); + writer.close().unwrap(); + + let timezone = matches!( + timestamp_type, + PrimitiveType::Timestamptz | PrimitiveType::TimestamptzNs + ) + .then_some(UTC_TIME_ZONE); + let expected_timestamps: ArrayRef = match timestamp_type { + PrimitiveType::Timestamp | PrimitiveType::Timestamptz => { + let values: Vec> = micros + .into_iter() + .zip(&definition_levels) + .map(|(value, level)| (*level == 1).then_some(value)) + .collect(); + Arc::new(TimestampMicrosecondArray::from(values).with_timezone_opt(timezone)) + } + PrimitiveType::TimestampNs | PrimitiveType::TimestamptzNs => { + let values: Vec> = nanos + .into_iter() + .zip(&definition_levels) + .map(|(value, level)| (*level == 1).then_some(value)) + .collect(); + Arc::new(TimestampNanosecondArray::from(values).with_timezone_opt(timezone)) + } + other => panic!("Expected timestamp type, got {other}"), + }; + let missing_timestamps = new_null_array(expected_timestamps.data_type(), 5); + let expected_batch = + RecordBatch::try_new(Arc::new(schema_to_arrow_schema(&schema).unwrap()), vec![ + Arc::new(Int64Array::from(ids)), + expected_timestamps, + Arc::new(StringArray::from(strings)), + missing_timestamps, + Arc::new(StringArray::from(vec!["default"; 5])), + ]) + .unwrap(); + let mut actual = Vec::new(); + let mut expected = Vec::new(); + for field_ids in [ + vec![2], + vec![1, 2, 3], + vec![2, 1, 3], + vec![2, 4, 5, 1], + vec![4, 5], + ] { + let column_indices: Vec = + field_ids.iter().map(|id| (*id - 1) as usize).collect(); + expected.push(vec![expected_batch.project(&column_indices).unwrap()]); + if field_ids.iter().any(|id| *id <= 3) { + let task = int96_scan_task( + file_path.to_str().unwrap(), + schema.clone(), + field_ids.clone(), + name_mapping.clone(), + ); + let physical_indices: Vec = column_indices + .iter() + .copied() + .filter(|index| *index < 3) + .collect(); + assert_int96_projection_mask(&task, &physical_indices); + } + actual.push( + read_int96_batches_with_name_mapping( + file_path.to_str().unwrap(), + schema.clone(), + field_ids, + name_mapping.clone(), + ) + .await, + ); + } + assert_eq!(actual, expected); + } + + #[tokio::test] + 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_with_micros_type = |micros_type| { + Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), + NestedField::required( + 2, + "data", + Type::Struct(StructType::new(vec![ + NestedField::optional(3, "micros", Type::Primitive(micros_type)) + .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(), + ]) + .build() + .unwrap(), + ) + }; + let schema = schema_with_micros_type(PrimitiveType::Timestamp); + 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 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(); + + 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, physical_indices) in [ + (vec![1, 2, 5, 6], vec![0, 1, 2, 3], vec![0, 1, 2, 3, 4]), + (vec![2], vec![1], vec![1, 2]), + (vec![6, 2, 5], vec![3, 1, 2], vec![1, 2, 3, 4]), + ] { + let task = int96_scan_task( + file_path.to_str().unwrap(), + schema.clone(), + field_ids.clone(), + None, + ); + assert_int96_projection_mask(&task, &physical_indices); + 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); + } + + let incompatible_task = int96_scan_task( + file_path.to_str().unwrap(), + schema_with_micros_type(PrimitiveType::Timestamptz), + vec![1, 2, 6], + None, + ); + assert_int96_projection_mask(&incompatible_task, &[0, 2, 4]); + } + #[tokio::test] async fn test_read_int96_timestamps_with_fallback_ids_and_pos() { use arrow_array::TimestampMicrosecondArray; @@ -3057,12 +3560,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 +3645,7 @@ mod tests { element_field: NestedField::optional( 2, "element", - Type::Primitive(PrimitiveType::Timestamp), + Type::Primitive(timestamp_type.clone()), ) .into(), }), @@ -3140,36 +3656,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 +3788,7 @@ mod tests { value_field: NestedField::optional( 3, "value", - Type::Primitive(PrimitiveType::Timestamp), + Type::Primitive(timestamp_type.clone()), ) .into(), }), @@ -3273,26 +3799,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 From 020765a1d5fff8090176702ff36eef8ac9628460 Mon Sep 17 00:00:00 2001 From: Hung Date: Fri, 6 Nov 2026 00:10:18 +0800 Subject: [PATCH 2/3] test(arrow): remove unreachable leaf-count mismatch case --- crates/iceberg/src/arrow/int96.rs | 26 -------------------------- 1 file changed, 26 deletions(-) diff --git a/crates/iceberg/src/arrow/int96.rs b/crates/iceberg/src/arrow/int96.rs index 338327a275..ca0c57bc1b 100644 --- a/crates/iceberg/src/arrow/int96.rs +++ b/crates/iceberg/src/arrow/int96.rs @@ -385,32 +385,6 @@ mod tests { ); } - #[test] - fn test_no_coercion_with_mismatched_leaf_count() { - let ts = Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true) - .with_metadata(field_id_meta(1)); - let id = Field::new("id", DataType::Int32, false).with_metadata(field_id_meta(2)); - let iceberg = iceberg_schema_with_timestamp(); - for (message, fields) in [ - ("message schema { optional int96 ts; }", vec![ - ts.clone(), - id, - ]), - ( - "message schema { optional int96 ts; required int32 id; }", - vec![ts], - ), - ] { - let parquet = parquet_schema(message); - let arrow_schema = Arc::new(ArrowSchema::new(fields)); - - assert_eq!( - coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet), - None - ); - } - } - #[test] fn test_only_int96_timestamps_are_coerced() { let parquet = parquet_schema( From aecd9c9f51f19c58cb09166278475060ca11986c Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Tue, 10 Nov 2026 02:02:13 +0800 Subject: [PATCH 3/3] fix(arrow): propagate INT96 coercion errors and simplify tests --- crates/iceberg/src/arrow/int96.rs | 134 ++++- crates/iceberg/src/arrow/reader/pipeline.rs | 516 ++++++-------------- 2 files changed, 260 insertions(+), 390 deletions(-) diff --git a/crates/iceberg/src/arrow/int96.rs b/crates/iceberg/src/arrow/int96.rs index ca0c57bc1b..eea703a4fc 100644 --- a/crates/iceberg/src/arrow/int96.rs +++ b/crates/iceberg/src/arrow/int96.rs @@ -48,21 +48,23 @@ pub(crate) fn coerce_int96_timestamps( arrow_schema: &ArrowSchemaRef, iceberg_schema: &Schema, parquet_schema: &SchemaDescriptor, -) -> Option> { +) -> Result>> { if !parquet_schema .columns() .iter() .any(|column| column.physical_type() == PhysicalType::INT96) { - return None; + return Ok(None); } let mut visitor = Int96CoercionVisitor::new(iceberg_schema, parquet_schema); - let coerced = visit_schema(arrow_schema, &mut visitor).ok()?; - if visitor.changed && visitor.columns.len() == 0 { - Some(Arc::new(coerced)) - } else { - None - } + 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 INT96 Arrow fields to the Iceberg timestamp type. @@ -283,7 +285,9 @@ mod tests { ])); let iceberg = iceberg_schema_with_timestamp(); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).unwrap(); + let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .unwrap(); assert_eq!( coerced.field(0).data_type(), &DataType::Timestamp(TimeUnit::Microsecond, None) @@ -308,7 +312,9 @@ mod tests { .with_metadata(field_id_meta(1)), ])); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).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_TIME_ZONE.into())) @@ -331,7 +337,11 @@ mod tests { .with_metadata(field_id_meta(1)), ])); - assert!(coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).is_none()); + assert!( + coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .is_none() + ); } #[test] @@ -355,7 +365,11 @@ mod tests { .with_metadata(field_id_meta(1)), ])); - assert!(coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).is_none()); + assert!( + coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .is_none() + ); } #[test] @@ -368,7 +382,11 @@ mod tests { ])); let iceberg = iceberg_schema_with_timestamp(); - assert!(coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).is_none()); + assert!( + coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .is_none() + ); } #[test] @@ -380,11 +398,73 @@ mod tests { let iceberg = iceberg_schema_with_timestamp(); assert_eq!( - coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet), + 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( @@ -426,7 +506,7 @@ mod tests { )); assert_eq!( - coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet), + coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).unwrap(), Some(expected) ); } @@ -444,7 +524,9 @@ mod tests { )])); let iceberg = iceberg_schema_with_timestamp(); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).unwrap(); + let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .unwrap(); assert_eq!( coerced.field(0).data_type(), &DataType::Timestamp(TimeUnit::Microsecond, None) @@ -470,7 +552,9 @@ mod tests { .with_metadata(field_id_meta(1)), ])); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).unwrap(); + let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .unwrap(); assert_eq!( coerced.field(0).data_type(), &DataType::Timestamp(TimeUnit::Microsecond, None) @@ -489,7 +573,9 @@ mod tests { ])); let iceberg = iceberg_schema_with_timestamp(); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).unwrap(); + let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .unwrap(); assert_eq!(coerced.field(0).metadata(), &meta); } @@ -528,7 +614,9 @@ mod tests { .with_metadata(field_id_meta(1)), ])); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).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}"), @@ -577,7 +665,9 @@ mod tests { .with_metadata(field_id_meta(1)), ])); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).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}"), @@ -646,7 +736,9 @@ mod tests { .with_metadata(field_id_meta(1)), ])); - let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet).unwrap(); + let coerced = coerce_int96_timestamps(&arrow_schema, &iceberg, &parquet) + .unwrap() + .unwrap(); let expected_entries = Field::new( "key_value", DataType::Struct( diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index 888a687858..3918c486d5 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -699,13 +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(), arrow_metadata.parquet_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()) @@ -988,54 +988,6 @@ mod tests { .unwrap() } - fn assert_int96_projection_mask(task: &FileScanTask, expected_indices: &[usize]) { - use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; - use parquet::file::reader::{FileReader, SerializedFileReader}; - - use super::{ArrowReader, FileScanTaskReader}; - - let reader = SerializedFileReader::new(File::open(task.data_file_path()).unwrap()).unwrap(); - let metadata = ArrowReaderMetadata::try_new( - Arc::new(reader.metadata().clone()), - ArrowReaderOptions::default(), - ) - .unwrap(); - let missing_field_ids = metadata - .schema() - .fields() - .iter() - .next() - .is_some_and(|field| !field.metadata().contains_key(PARQUET_FIELD_ID_META_KEY)); - let use_position_fallback = missing_field_ids && task.name_mapping().is_none(); - let metadata = FileScanTaskReader::configure_arrow_reader_metadata( - metadata, - task, - missing_field_ids, - false, - ) - .unwrap(); - let mask = ArrowReader::get_arrow_projection_mask( - task.project_field_ids(), - task.schema(), - metadata.parquet_schema(), - metadata.schema(), - use_position_fallback, - ) - .unwrap(); - let actual: Vec = (0..metadata.parquet_schema().num_columns()) - .map(|index| mask.leaf_included(index)) - .collect(); - let expected: Vec = (0..metadata.parquet_schema().num_columns()) - .map(|index| expected_indices.contains(&index)) - .collect(); - assert_eq!( - actual, - expected, - "projection {:?}", - task.project_field_ids() - ); - } - // ArrowWriter cannot write INT96, so we use SerializedFileWriter directly. fn write_int96_parquet_file( table_location: &str, @@ -1135,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). @@ -2901,200 +2826,131 @@ mod tests { #[tokio::test] async fn test_read_int96_timestamps_with_field_ids() { - let schema = Arc::new( - Schema::builder() - .with_schema_id(1) - .with_fields(vec![ - NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamp)) - .into(), - NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(), - ]) - .build() - .unwrap(), - ); - - 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); - - assert_int96_read_matches(&file_path, schema, vec![1, 2], &expected_micros).await; + 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 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); - - assert_int96_read_matches(&file_path, schema, vec![1, 2], &expected_micros).await; - } - - #[tokio::test] - async fn test_read_int96_timestamptz_with_name_mapping() { - assert_int96_timestamp_projections( - PrimitiveType::Timestamptz, - Int96FieldIds::NameMapping, - false, - ) - .await; - } - - #[tokio::test] - async fn test_read_int96_timestamptz_ns_with_name_mapping() { - assert_int96_timestamp_projections( - PrimitiveType::TimestamptzNs, - Int96FieldIds::NameMapping, - false, - ) - .await; - } - - #[tokio::test] - async fn test_read_int96_timestamptz_with_field_ids_and_nulls() { - assert_int96_timestamp_projections( - PrimitiveType::Timestamptz, - Int96FieldIds::Embedded, - true, - ) - .await; - } - - #[tokio::test] - async fn test_read_int96_timestamptz_with_positional_fallback() { - assert_int96_timestamp_projections( - PrimitiveType::Timestamptz, - Int96FieldIds::Position, - false, - ) - .await; - } - - #[tokio::test] - async fn test_read_int96_timestamp_without_timezone_with_name_mapping() { - assert_int96_timestamp_projections( - PrimitiveType::Timestamp, - Int96FieldIds::NameMapping, - true, - ) - .await; + 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_timestamp_ns_with_field_ids_and_nulls() { - assert_int96_timestamp_projections( - PrimitiveType::TimestampNs, - Int96FieldIds::Embedded, - true, - ) - .await; - } - - enum Int96FieldIds { - Embedded, - NameMapping, - Position, + 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_timestamp_projections( + async fn assert_int96_nanosecond_projections( timestamp_type: PrimitiveType, - field_ids: Int96FieldIds, - with_nulls: bool, + with_field_ids: bool, ) { - use arrow_array::{TimestampMicrosecondArray, TimestampNanosecondArray, new_null_array}; - use parquet::basic::{LogicalType, Repetition, Type as PhysicalType}; - use parquet::data_type::{ByteArray, ByteArrayType, Int64Type, Int96, Int96Type}; + use arrow_array::TimestampNanosecondArray; + use parquet::data_type::{Int32Type, Int96, Int96Type}; use parquet::file::writer::SerializedFileWriter; - use parquet::schema::types::Type as SchemaType; + use parquet::schema::parser::parse_message_type; use crate::arrow::schema::UTC_TIME_ZONE; use crate::arrow::schema_to_arrow_schema; - use crate::spec::Literal; + let timezone = (timestamp_type == PrimitiveType::TimestamptzNs).then_some(UTC_TIME_ZONE); let schema = Arc::new( Schema::builder() .with_fields(vec![ - NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), - NestedField::optional(2, "ts", Type::Primitive(timestamp_type.clone())).into(), - NestedField::required(3, "s", Type::Primitive(PrimitiveType::String)).into(), - NestedField::optional(4, "missing_ts", Type::Primitive(timestamp_type.clone())) - .into(), - NestedField::optional(5, "default_s", Type::Primitive(PrimitiveType::String)) - .with_initial_default(Literal::string("default")) - .into(), + NestedField::optional(1, "ts", Type::Primitive(timestamp_type)).into(), + NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(), ]) .build() .unwrap(), ); - let name_mapping = matches!(field_ids, Int96FieldIds::NameMapping).then(|| { + 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!["id".to_string()], vec![]), - MappedField::new(Some(2), vec!["ts".to_string()], vec![]), - MappedField::new(Some(3), vec!["s".to_string()], vec![]), + MappedField::new(Some(1), vec!["ts".to_string()], vec![]), + MappedField::new(Some(2), vec!["id".to_string()], vec![]), ])) }); - let with_field_ids = matches!(field_ids, Int96FieldIds::Embedded); - let parquet_schema = SchemaType::group_type_builder("schema") - .with_fields(vec![ - Arc::new( - SchemaType::primitive_type_builder("id", PhysicalType::INT64) - .with_repetition(Repetition::REQUIRED) - .with_id(with_field_ids.then_some(1)) - .build() - .unwrap(), - ), - Arc::new( - SchemaType::primitive_type_builder("ts", PhysicalType::INT96) - .with_repetition(Repetition::OPTIONAL) - .with_id(with_field_ids.then_some(2)) - .build() - .unwrap(), - ), - Arc::new( - SchemaType::primitive_type_builder("s", PhysicalType::BYTE_ARRAY) - .with_repetition(Repetition::REQUIRED) - .with_logical_type(Some(LogicalType::String)) - .with_id(with_field_ids.then_some(3)) - .build() - .unwrap(), - ), - ]) - .build() - .unwrap(); - let micros = vec![0_i64, 1_000_003, 2_000_006, 3_000_009, 4_000_012]; - let nanos: Vec = micros - .iter() - .map(|micros| { - micros * 1_000 - + i64::from(matches!( - timestamp_type, - PrimitiveType::TimestampNs | PrimitiveType::TimestamptzNs - )) - }) - .collect(); - let ids = vec![0_i64, 1, 2, 3, 4]; - let strings = vec!["0", "1", "2", "3", "4"]; - let definition_levels: Vec = (0..5) - .map(|index| i16::from(!with_nulls || index % 2 == 0)) - .collect(); + let nanos = vec![Some(1), None, Some(2_000_006_001_i64)]; let int96_values: Vec = nanos .iter() - .zip(&definition_levels) - .filter(|(_, level)| **level == 1) - .map(|(nanos, _)| { + .flatten() + .map(|nanos| { let mut value = Int96::new(); value.set_data( *nanos as u32, @@ -3104,13 +2960,9 @@ mod tests { value }) .collect(); - let string_values: Vec = strings - .iter() - .map(|value| ByteArray::from(*value)) - .collect(); - + let ids = vec![0, 1, 2]; let tmp_dir = TempDir::new().unwrap(); - let file_path = tmp_dir.path().join("timestamptz_name_mapping.parquet"); + let file_path = tmp_dir.path().join("nanosecond_timestamps.parquet"); let mut writer = SerializedFileWriter::new( File::create(&file_path).unwrap(), Arc::new(parquet_schema), @@ -3119,97 +2971,41 @@ mod tests { .unwrap(); let mut row_group = writer.next_row_group().unwrap(); let mut column = row_group.next_column().unwrap().unwrap(); - column - .typed::() - .write_batch(&ids, None, None) - .unwrap(); - column.close().unwrap(); - let mut column = row_group.next_column().unwrap().unwrap(); column .typed::() - .write_batch(&int96_values, Some(&definition_levels), None) + .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(&string_values, None, None) + .typed::() + .write_batch(&ids, None, None) .unwrap(); column.close().unwrap(); row_group.close().unwrap(); writer.close().unwrap(); - let timezone = matches!( - timestamp_type, - PrimitiveType::Timestamptz | PrimitiveType::TimestamptzNs - ) - .then_some(UTC_TIME_ZONE); - let expected_timestamps: ArrayRef = match timestamp_type { - PrimitiveType::Timestamp | PrimitiveType::Timestamptz => { - let values: Vec> = micros - .into_iter() - .zip(&definition_levels) - .map(|(value, level)| (*level == 1).then_some(value)) - .collect(); - Arc::new(TimestampMicrosecondArray::from(values).with_timezone_opt(timezone)) - } - PrimitiveType::TimestampNs | PrimitiveType::TimestamptzNs => { - let values: Vec> = nanos - .into_iter() - .zip(&definition_levels) - .map(|(value, level)| (*level == 1).then_some(value)) - .collect(); - Arc::new(TimestampNanosecondArray::from(values).with_timezone_opt(timezone)) - } - other => panic!("Expected timestamp type, got {other}"), - }; - let missing_timestamps = new_null_array(expected_timestamps.data_type(), 5); let expected_batch = RecordBatch::try_new(Arc::new(schema_to_arrow_schema(&schema).unwrap()), vec![ - Arc::new(Int64Array::from(ids)), - expected_timestamps, - Arc::new(StringArray::from(strings)), - missing_timestamps, - Arc::new(StringArray::from(vec!["default"; 5])), + Arc::new(TimestampNanosecondArray::from(nanos).with_timezone_opt(timezone)), + Arc::new(Int32Array::from(ids)), ]) .unwrap(); - let mut actual = Vec::new(); - let mut expected = Vec::new(); - for field_ids in [ - vec![2], - vec![1, 2, 3], - vec![2, 1, 3], - vec![2, 4, 5, 1], - vec![4, 5], + for (field_ids, column_indices) in [ + (vec![1], vec![0]), + (vec![1, 2], vec![0, 1]), + (vec![2, 1], vec![1, 0]), ] { - let column_indices: Vec = - field_ids.iter().map(|id| (*id - 1) as usize).collect(); - expected.push(vec![expected_batch.project(&column_indices).unwrap()]); - if field_ids.iter().any(|id| *id <= 3) { - let task = int96_scan_task( - file_path.to_str().unwrap(), - schema.clone(), - field_ids.clone(), - name_mapping.clone(), - ); - let physical_indices: Vec = column_indices - .iter() - .copied() - .filter(|index| *index < 3) - .collect(); - assert_int96_projection_mask(&task, &physical_indices); - } - actual.push( - read_int96_batches_with_name_mapping( - file_path.to_str().unwrap(), - schema.clone(), - field_ids, - name_mapping.clone(), - ) - .await, - ); + 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); } - assert_eq!(actual, expected); } #[tokio::test] @@ -3235,44 +3031,41 @@ mod tests { }", ) .unwrap(); - let schema_with_micros_type = |micros_type| { - Arc::new( - Schema::builder() - .with_fields(vec![ - NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), - NestedField::required( - 2, - "data", - Type::Struct(StructType::new(vec![ - NestedField::optional(3, "micros", Type::Primitive(micros_type)) - .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), - ) + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + 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(), - ]) - .build() - .unwrap(), - ) - }; - let schema = schema_with_micros_type(PrimitiveType::Timestamp); + ]) + .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(); @@ -3347,31 +3140,16 @@ mod tests { ]) .unwrap(); - for (field_ids, column_indices, physical_indices) in [ - (vec![1, 2, 5, 6], vec![0, 1, 2, 3], vec![0, 1, 2, 3, 4]), - (vec![2], vec![1], vec![1, 2]), - (vec![6, 2, 5], vec![3, 1, 2], vec![1, 2, 3, 4]), + 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 task = int96_scan_task( - file_path.to_str().unwrap(), - schema.clone(), - field_ids.clone(), - None, - ); - assert_int96_projection_mask(&task, &physical_indices); 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); } - - let incompatible_task = int96_scan_task( - file_path.to_str().unwrap(), - schema_with_micros_type(PrimitiveType::Timestamptz), - vec![1, 2, 6], - None, - ); - assert_int96_projection_mask(&incompatible_task, &[0, 2, 4]); } #[tokio::test]