From fafb16f3ffa71eb14429933ceccf4f1113191e66 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Thu, 15 Oct 2026 23:12:37 +0300 Subject: [PATCH] Support Iceberg v3 nanosecond timestamps end to end --- datafusion_iceberg/src/planner.rs | 63 +++- datafusion_iceberg/src/pruning_statistics.rs | 169 +++++++++-- datafusion_iceberg/src/statistics.rs | 5 + datafusion_iceberg/src/table/mod.rs | 5 + .../tests/integration_df_types.rs | 60 ++++ datafusion_iceberg/tests/roundtrip_types.rs | 174 ++++++++++- iceberg-rust-spec/src/arrow/schema.rs | 23 ++ iceberg-rust-spec/src/spec/manifest.rs | 30 +- iceberg-rust-spec/src/spec/values.rs | 275 +++++++++++++++--- iceberg-rust/src/arrow/partition.rs | 264 ++++++++++++++++- iceberg-rust/src/arrow/transform.rs | 95 +++++- iceberg-rust/src/file_format/parquet.rs | 16 + iceberg-rust/src/table/manifest.rs | 20 ++ iceberg-rust/src/table/transaction/append.rs | 34 ++- 14 files changed, 1154 insertions(+), 79 deletions(-) diff --git a/datafusion_iceberg/src/planner.rs b/datafusion_iceberg/src/planner.rs index be3ac696..c748d1a3 100644 --- a/datafusion_iceberg/src/planner.rs +++ b/datafusion_iceberg/src/planner.rs @@ -36,7 +36,8 @@ use iceberg_rust::{ namespace::Namespace, partition::{PartitionField, PartitionSpec, Transform}, schema::{Schema, DEFAULT_SCHEMA_ID}, - types::StructType, + table_metadata::FormatVersion, + types::{PrimitiveType, StructType, Type}, view_metadata::{Version, ViewRepresentation}, }, table::Table, @@ -170,8 +171,35 @@ async fn plan_create_table( let fields_with_ids = new_fields_with_ids(node.0.schema.as_arrow().fields(), &mut 0) .map_err(|err| DataFusionError::External(Box::new(err)))?; - let schema = StructType::try_from(&fields_with_ids) + let mut schema = StructType::try_from(&fields_with_ids) .map_err(|err| DataFusionError::External(Box::new(err)))?; + let mut properties = node.0.options.clone(); + if let Some(version) = properties.remove("format.format-version") { + properties.insert("format-version".to_owned(), version); + } + let format_version = properties + .get("format-version") + .map(|version| { + let number = version + .parse::() + .map_err(|error| DataFusionError::External(Box::new(error)))?; + FormatVersion::try_from(number) + .map_err(|error| DataFusionError::External(Box::new(error))) + }) + .transpose()? + .unwrap_or_default(); + if format_version != FormatVersion::V3 { + schema = StructType::new( + schema + .iter() + .cloned() + .map(|mut field| { + v2_compatible_timestamp_type(&mut field.field_type); + field + }) + .collect(), + ); + } let pacrtition_spec = node .0 @@ -214,7 +242,7 @@ async fn plan_create_table( .with_location(location) .with_schema(Schema::from_struct_type(schema, DEFAULT_SCHEMA_ID, None)) .with_partition_spec(partition_spec) - .with_properties(node.0.options.clone()) + .with_properties(properties) .build(&[namespace_name.as_ref().to_owned()], catalog) .await .map_err(|err| DataFusionError::External(Box::new(err)))?; @@ -224,6 +252,35 @@ async fn plan_create_table( ))))) } +fn v2_compatible_timestamp_type(data_type: &mut Type) { + match data_type { + Type::Primitive(PrimitiveType::TimestampNs) => { + *data_type = Type::Primitive(PrimitiveType::Timestamp); + } + Type::Primitive(PrimitiveType::TimestamptzNs) => { + *data_type = Type::Primitive(PrimitiveType::Timestamptz); + } + Type::Struct(fields) => { + *fields = StructType::new( + fields + .iter() + .cloned() + .map(|mut field| { + v2_compatible_timestamp_type(&mut field.field_type); + field + }) + .collect(), + ); + } + Type::List(list) => v2_compatible_timestamp_type(&mut list.element), + Type::Map(map) => { + v2_compatible_timestamp_type(&mut map.key); + v2_compatible_timestamp_type(&mut map.value); + } + Type::Primitive(_) => {} + } +} + async fn plan_create_view( node: &CreateIcebergView, session: &dyn Session, diff --git a/datafusion_iceberg/src/pruning_statistics.rs b/datafusion_iceberg/src/pruning_statistics.rs index 95c3ccf8..0df673fb 100644 --- a/datafusion_iceberg/src/pruning_statistics.rs +++ b/datafusion_iceberg/src/pruning_statistics.rs @@ -391,6 +391,14 @@ fn any_iter_to_array( opt.and_then(|value| Some(*value.downcast::().ok()?)), ) })), + DataType::Timestamp(TimeUnit::Nanosecond, tz) => { + ScalarValue::iter_to_array(iter.map(move |opt| { + ScalarValue::TimestampNanosecond( + opt.and_then(|value| Some(*value.downcast::().ok()?)), + tz.clone(), + ) + })) + } DataType::Timestamp(_, tz) => { // Make sure to preserve the column's timezone for the sake of comparisons. ScalarValue::iter_to_array(iter.map(move |opt| { @@ -498,13 +506,14 @@ pub(crate) fn transform_predicate( return None; }; let (column, literal, column_on_left) = match (*left, *right) { - (Expr::Column(column), literal) => (column, fold_cast_literal(literal)?, true), - (literal, Expr::Column(column)) => (column, fold_cast_literal(literal)?, false), + (Expr::Column(column), literal) => (column, literal, true), + (literal, Expr::Column(column)) => (column, literal, false), _ => return None, }; let field = partition_fields .iter() .find(|field| field.source_name() == column.name())?; + let literal = fold_cast_literal(literal, field.field_type())?; let legacy_negative_day = if field.transform() != &Transform::Identity { let Expr::Literal(value, _) = &literal else { return None; @@ -530,10 +539,16 @@ pub(crate) fn transform_predicate( } }; if field.transform() == &Transform::Hour { - let Expr::Literal(ScalarValue::TimestampMicrosecond(Some(micros), _), _) = &literal else { - return None; + let (timestamp, units_per_hour) = match &literal { + Expr::Literal(ScalarValue::TimestampMicrosecond(Some(value), _), _) => { + (*value, 3_600_000_000) + } + Expr::Literal(ScalarValue::TimestampNanosecond(Some(value), _), _) => { + (*value, 3_600_000_000_000) + } + _ => return None, }; - return project_hour_predicate(column, field.name(), op, *micros); + return project_hour_predicate(column, field.name(), op, timestamp, units_per_hour); } let op = match field.transform() { Transform::Identity => op, @@ -580,7 +595,7 @@ pub(crate) fn transform_predicate( Some(partition_comparison(column, op, literal)) } -fn fold_cast_literal(expr: Expr) -> Option { +fn fold_cast_literal(expr: Expr, source_type: &Type) -> Option { match expr { literal @ Expr::Literal(..) => Some(literal), Expr::Cast(cast) => { @@ -589,10 +604,18 @@ fn fold_cast_literal(expr: Expr) -> Option { }; let value = value.cast_to(cast.field.data_type()).ok()?; let value = match value { - ScalarValue::TimestampNanosecond(Some(nanos), None) if nanos % 1_000 == 0 => { + ScalarValue::TimestampNanosecond(Some(nanos), None) + if source_type == &Type::Primitive(PrimitiveType::Timestamp) + && nanos % 1_000 == 0 => + { ScalarValue::TimestampMicrosecond(Some(nanos / 1_000), None) } - ScalarValue::TimestampNanosecond(..) => return None, + ScalarValue::TimestampNanosecond(..) + if source_type != &Type::Primitive(PrimitiveType::TimestampNs) + && source_type != &Type::Primitive(PrimitiveType::TimestamptzNs) => + { + return None; + } value => value, }; Some(Expr::Literal(value, None)) @@ -609,19 +632,19 @@ fn project_hour_predicate( source_column: Column, partition_name: &str, op: Operator, - micros: i64, + timestamp: i64, + units_per_hour: i64, ) -> Option { - const MICROS_PER_HOUR: i64 = 3_600_000_000; const HOUR_ID_MODULUS: i64 = 1_i64 << 32; - let hour = micros.div_euclid(MICROS_PER_HOUR); + let hour = timestamp.div_euclid(units_per_hour); let (lower, upper) = match op { - Operator::Eq => (hour, hour.checked_add(i64::from(micros < 0))?), + Operator::Eq => (hour, hour.checked_add(i64::from(timestamp < 0))?), Operator::Lt | Operator::LtEq => ( - i64::MIN.div_euclid(MICROS_PER_HOUR), - hour.checked_add(i64::from(micros < 0))?, + i64::MIN.div_euclid(units_per_hour), + hour.checked_add(i64::from(timestamp < 0))?, ), - Operator::Gt | Operator::GtEq => (hour, i64::MAX.div_euclid(MICROS_PER_HOUR)), + Operator::Gt | Operator::GtEq => (hour, i64::MAX.div_euclid(units_per_hour)), _ => return None, }; if upper.checked_sub(lower)? >= HOUR_ID_MODULUS - 1 { @@ -677,12 +700,18 @@ fn literal_matches_source(value: &ScalarValue, field: &BoundPartitionField<'_>) ) | ( Type::Primitive(PrimitiveType::Timestamp), ScalarValue::TimestampMicrosecond(Some(_), None) + ) | ( + Type::Primitive(PrimitiveType::TimestampNs), + ScalarValue::TimestampNanosecond(Some(_), None) ) ) || matches!( (field.field_type(), value), ( Type::Primitive(PrimitiveType::Timestamptz), ScalarValue::TimestampMicrosecond(Some(_), Some(tz)) + ) | ( + Type::Primitive(PrimitiveType::TimestamptzNs), + ScalarValue::TimestampNanosecond(Some(_), Some(tz)) ) if tz.as_ref() == "UTC" ) } @@ -710,6 +739,21 @@ fn transform_literal(expr: Expr, transform: &Transform) -> Option { None, )) } + Transform::Day | Transform::Hour if matches!(value, ScalarValue::TimestampNanosecond(Some(nanos), _) if *nanos < 0) => + { + let ScalarValue::TimestampNanosecond(Some(nanos), _) = value else { + return None; + }; + let unit = if transform == &Transform::Day { + 86_400_000_000_000 + } else { + 3_600_000_000_000 + }; + Some(Expr::Literal( + ScalarValue::Int32(Some(i32::try_from(nanos.div_euclid(unit)).ok()?)), + None, + )) + } Transform::Year => Some(Expr::ScalarFunction(ScalarFunction::new_udf( Arc::new(ScalarUDF::new_from_impl(DateTransform::new())), vec![Expr::Literal(ScalarValue::new_utf8("year"), None), expr], @@ -769,6 +813,14 @@ impl DateTransform { DataType::Utf8, DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), ]), + TypeSignature::Exact(vec![ + DataType::Utf8, + DataType::Timestamp(TimeUnit::Nanosecond, None), + ]), + TypeSignature::Exact(vec![ + DataType::Utf8, + DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), + ]), ]), volatility: Volatility::Immutable, parameter_names: None, @@ -854,6 +906,14 @@ fn scalarvalue_to_value(scalar: &ScalarValue) -> Result { ScalarValue::TimestampMicrosecond(x, None) => Ok(Value::Timestamp(x.ok_or( Error::InvalidFormat("Value can't be null when converting to iceberg value".to_owned()), )?)), + ScalarValue::TimestampNanosecond(x, Some(tz)) if tz == &Arc::from("UTC") => { + Ok(Value::TimestampTZNano(x.ok_or(Error::InvalidFormat( + "Value can't be null when converting to iceberg value".to_owned(), + ))?)) + } + ScalarValue::TimestampNanosecond(x, None) => Ok(Value::TimestampNano(x.ok_or( + Error::InvalidFormat("Value can't be null when converting to iceberg value".to_owned()), + )?)), x => Err(Error::NotSupported(format!( "Transforming {x} to iceberg value" ))), @@ -872,6 +932,11 @@ fn value_to_scalarvalue(value: Value) -> Result { Some(x), Some("UTC".into()), )), + Value::TimestampNano(x) => Ok(ScalarValue::TimestampNanosecond(Some(x), None)), + Value::TimestampTZNano(x) => Ok(ScalarValue::TimestampNanosecond( + Some(x), + Some("UTC".into()), + )), x => Err(Error::NotSupported(format!( "Transforming {x} to iceberg value" ))), @@ -883,6 +948,7 @@ mod tests { use super::*; use datafusion::arrow::array::{ Array, Date32Array, Decimal128Array, Int64Array, TimestampMicrosecondArray, + TimestampNanosecondArray, }; use datafusion::arrow::datatypes::Field; use datafusion::common::config::ConfigOptions; @@ -1590,7 +1656,10 @@ mod tests { DataType::Timestamp(TimeUnit::Nanosecond, None), )); assert!(matches!( - fold_cast_literal(timestamp.clone()), + fold_cast_literal( + timestamp.clone(), + &Type::Primitive(PrimitiveType::Timestamp) + ), Some(Expr::Literal( ScalarValue::TimestampMicrosecond(Some(_), None), _ @@ -1613,7 +1682,11 @@ mod tests { )), DataType::Timestamp(TimeUnit::Nanosecond, None), )); - assert!(fold_cast_literal(sub_microsecond.clone()).is_none()); + assert!(fold_cast_literal( + sub_microsecond.clone(), + &Type::Primitive(PrimitiveType::Timestamp) + ) + .is_none()); assert!(transform_predicate( partition_comparison( Expr::Column(Column::from_name("ts")), @@ -1625,6 +1698,51 @@ mod tests { .is_none()); } + #[test] + fn nanosecond_partition_projection_keeps_sub_microsecond_precision() { + use datafusion_expr::expr::Cast; + + let source = StructField::new( + 1, + "ts", + false, + Type::Primitive(PrimitiveType::TimestampNs), + None, + ); + let literal = Expr::Cast(Cast::new( + Box::new(Expr::Literal( + ScalarValue::TimestampNanosecond(Some(-1), None), + None, + )), + DataType::Timestamp(TimeUnit::Nanosecond, None), + )); + assert_eq!( + fold_cast_literal(literal.clone(), &source.field_type), + Some(Expr::Literal( + ScalarValue::TimestampNanosecond(Some(-1), None), + None + )) + ); + for transform in [ + Transform::Year, + Transform::Month, + Transform::Day, + Transform::Hour, + ] { + let partition = PartitionField::new(1, 1000, "ts_part", transform); + let fields = [BoundPartitionField::new(&partition, &source)]; + assert!(transform_predicate( + partition_comparison( + Expr::Column(Column::from_name("ts")), + Operator::Eq, + literal.clone(), + ), + &fields, + ) + .is_some()); + } + } + fn project_id_predicate( transform: Transform, source_type: PrimitiveType, @@ -2104,6 +2222,23 @@ mod tests { } } + #[test] + fn any_iter_to_array_preserves_nanosecond_precision() { + for timezone in [None, Some("UTC".into())] { + let dt = DataType::Timestamp(TimeUnit::Nanosecond, timezone); + let iter = vec![Some(Value::TimestampNano(-1).into_any()), None].into_iter(); + let array = any_iter_to_array(iter, &dt).unwrap(); + + assert_eq!(array.data_type(), &dt); + let timestamps = array + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(timestamps.value(0), -1); + assert!(timestamps.is_null(1)); + } + } + #[test] fn scalar_value_roundtrip_preserves_timezone() { for val in [Value::Timestamp(TS_MICROS), Value::TimestampTZ(TS_MICROS)] { diff --git a/datafusion_iceberg/src/statistics.rs b/datafusion_iceberg/src/statistics.rs index 2325cc1a..12dd5839 100644 --- a/datafusion_iceberg/src/statistics.rs +++ b/datafusion_iceberg/src/statistics.rs @@ -133,6 +133,11 @@ fn convert_value_to_scalar_value(value: Value, field_type: &Type) -> Result Ok(ScalarValue::TimestampNanosecond(Some(ts), None)), + Value::TimestampTZNano(ts) => Ok(ScalarValue::TimestampNanosecond( + Some(ts), + Some("UTC".into()), + )), Value::String(s) => Ok(ScalarValue::Utf8(Some(s))), Value::UUID(u) => Ok(ScalarValue::FixedSizeBinary( 16, diff --git a/datafusion_iceberg/src/table/mod.rs b/datafusion_iceberg/src/table/mod.rs index 66c5cbc4..f635770e 100644 --- a/datafusion_iceberg/src/table/mod.rs +++ b/datafusion_iceberg/src/table/mod.rs @@ -2286,6 +2286,11 @@ fn value_to_scalarvalue(value: &Value) -> Result { Some(*ts), Some("UTC".into()), )), + Value::TimestampNano(ts) => Ok(ScalarValue::TimestampNanosecond(Some(*ts), None)), + Value::TimestampTZNano(ts) => Ok(ScalarValue::TimestampNanosecond( + Some(*ts), + Some("UTC".into()), + )), Value::String(s) => Ok(ScalarValue::Utf8(Some(s.clone()))), Value::UUID(u) => Ok(ScalarValue::FixedSizeBinary( 16, diff --git a/datafusion_iceberg/tests/integration_df_types.rs b/datafusion_iceberg/tests/integration_df_types.rs index 569d5c30..17759fc2 100644 --- a/datafusion_iceberg/tests/integration_df_types.rs +++ b/datafusion_iceberg/tests/integration_df_types.rs @@ -5,7 +5,67 @@ #[path = "df_common/mod.rs"] mod df_common; +use datafusion::arrow::datatypes::{DataType, TimeUnit}; +use datafusion_iceberg::catalog::catalog::IcebergCatalog; use df_common::{boot_df_stack, execute_scalar_f64, execute_scalar_i64, execute_sql}; +use iceberg_rust::catalog::{identifier::Identifier, tabular::Tabular}; +use iceberg_rust::spec::table_metadata::FormatVersion; + +#[tokio::test] +async fn integration_df_v3_create_table_infers_nanosecond_timestamp() { + let ctx = boot_df_stack().await; + execute_sql(&ctx, "CREATE SCHEMA warehouse.tp_ts_ns").await; + execute_sql( + &ctx, + "CREATE EXTERNAL TABLE warehouse.tp_ts_ns.v2 \ + (id BIGINT NOT NULL, ts TIMESTAMP NOT NULL) STORED AS ICEBERG \ + LOCATION '/warehouse/tp_ts_ns/v2'", + ) + .await; + let v2 = ctx.table("warehouse.tp_ts_ns.v2").await.unwrap(); + assert_eq!( + v2.schema() + .field_with_unqualified_name("ts") + .unwrap() + .data_type(), + &DataType::Timestamp(TimeUnit::Microsecond, None) + ); + let provider = ctx.catalog("warehouse").unwrap(); + let catalog = provider.downcast_ref::().unwrap().catalog(); + let v2_id = Identifier::new(&["tp_ts_ns".to_owned()], "v2"); + let Tabular::Table(v2_table) = catalog.clone().load_tabular(&v2_id).await.unwrap() else { + panic!("expected Iceberg table"); + }; + assert_eq!(v2_table.metadata().format_version, FormatVersion::V2); + for (name, version) in [("t", "3"), ("padded", "03")] { + execute_sql( + &ctx, + &format!( + "CREATE EXTERNAL TABLE warehouse.tp_ts_ns.{name} \ + (id BIGINT NOT NULL, ts TIMESTAMP NOT NULL) STORED AS ICEBERG \ + LOCATION '/warehouse/tp_ts_ns/{name}' OPTIONS ('format-version' '{version}')" + ), + ) + .await; + let table = ctx + .table(format!("warehouse.tp_ts_ns.{name}")) + .await + .unwrap(); + assert_eq!( + table + .schema() + .field_with_unqualified_name("ts") + .unwrap() + .data_type(), + &DataType::Timestamp(TimeUnit::Nanosecond, None) + ); + let v3_id = Identifier::new(&["tp_ts_ns".to_owned()], name); + let Tabular::Table(v3_table) = catalog.clone().load_tabular(&v3_id).await.unwrap() else { + panic!("expected Iceberg table"); + }; + assert_eq!(v3_table.metadata().format_version, FormatVersion::V3); + } +} #[tokio::test] async fn integration_df_date_equality_and_range() { diff --git a/datafusion_iceberg/tests/roundtrip_types.rs b/datafusion_iceberg/tests/roundtrip_types.rs index ff83a87f..96c7a3f5 100644 --- a/datafusion_iceberg/tests/roundtrip_types.rs +++ b/datafusion_iceberg/tests/roundtrip_types.rs @@ -10,10 +10,12 @@ use std::sync::Arc; use datafusion::arrow::array::{ ArrayRef, Date32Array, Decimal128Array, Float32Array, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray, Time64MicrosecondArray, TimestampMicrosecondArray, + TimestampNanosecondArray, }; use datafusion::arrow::datatypes::{DataType, Field, Schema as ArrowSchema, TimeUnit}; use datafusion::dataframe::DataFrameWriteOptions; -use datafusion::prelude::SessionContext; +use datafusion::prelude::{col, lit, SessionContext}; +use datafusion::scalar::ScalarValue; use datafusion_iceberg::catalog::catalog::IcebergCatalog; use iceberg_rust::catalog::identifier::Identifier; use iceberg_rust::catalog::Catalog; @@ -23,7 +25,9 @@ use iceberg_rust::object_store::ObjectStoreBuilder; use iceberg_rust::spec::decimal::decimal_from_i128_with_scale; use iceberg_rust::spec::manifest::DataFile; use iceberg_rust::spec::namespace::Namespace; -use iceberg_rust::spec::partition::{BoundPartitionField, PartitionField, Transform}; +use iceberg_rust::spec::partition::{ + BoundPartitionField, PartitionField, PartitionSpec, Transform, +}; use iceberg_rust::spec::schema::Schema; use iceberg_rust::spec::types::{PrimitiveType, StructField, Type}; use iceberg_rust::spec::values::Value; @@ -207,6 +211,16 @@ fn parquet_stats_and_partition_value_decode_correctly() { DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("UTC"))), false, ), + Field::new( + "ts_ns_col", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new( + "tstz_ns_col", + DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("UTC"))), + false, + ), Field::new("uuid_col", DataType::Utf8, false), ])); @@ -219,6 +233,8 @@ fn parquet_stats_and_partition_value_decode_correctly() { let time_val: i64 = 45_296_123_456; // 12:34:56.123456 let ts_val: i64 = 1_700_000_000_123_456; let tstz_val: i64 = 1_700_000_012_345_678; + let ts_ns_val: i64 = 1_700_000_000_123_456_789; + let tstz_ns_val: i64 = 1_700_000_012_345_678_901; let uuid_str = "550e8400-e29b-41d4-a716-446655440000"; let batch = RecordBatch::try_new( @@ -237,6 +253,8 @@ fn parquet_stats_and_partition_value_decode_correctly() { Arc::new(Time64MicrosecondArray::from(vec![time_val])), Arc::new(TimestampMicrosecondArray::from(vec![ts_val])), Arc::new(TimestampMicrosecondArray::from(vec![tstz_val]).with_timezone("UTC")), + Arc::new(TimestampNanosecondArray::from(vec![ts_ns_val])), + Arc::new(TimestampNanosecondArray::from(vec![tstz_ns_val]).with_timezone("UTC")), Arc::new(StringArray::from(vec![uuid_str])), ], ) @@ -259,6 +277,8 @@ fn parquet_stats_and_partition_value_decode_correctly() { ("time_col", Type::Primitive(PrimitiveType::Time)), ("ts_col", Type::Primitive(PrimitiveType::Timestamp)), ("tstz_col", Type::Primitive(PrimitiveType::Timestamptz)), + ("ts_ns_col", Type::Primitive(PrimitiveType::TimestampNs)), + ("tstz_ns_col", Type::Primitive(PrimitiveType::TimestamptzNs)), ("uuid_col", Type::Primitive(PrimitiveType::Uuid)), ] .into_iter() @@ -306,6 +326,154 @@ fn parquet_stats_and_partition_value_decode_correctly() { assert_eq!(bounds[&6], Value::Time(time_val)); assert_eq!(bounds[&7], Value::Timestamp(ts_val)); assert_eq!(bounds[&8], Value::TimestampTZ(tstz_val)); - assert_eq!(bounds[&9], uuid_val); + assert_eq!(bounds[&9], Value::TimestampNano(ts_ns_val)); + assert_eq!(bounds[&10], Value::TimestampTZNano(tstz_ns_val)); + assert_eq!(bounds[&11], uuid_val); + } +} + +#[tokio::test] +async fn v3_nanosecond_timestamps_round_trip_without_truncation() { + let object_store = ObjectStoreBuilder::memory(); + let catalog: Arc = Arc::new( + SqlCatalog::new("sqlite://", "test", object_store) + .await + .unwrap(), + ); + catalog + .create_namespace(&Namespace::try_new(&["public".to_string()]).unwrap(), None) + .await + .unwrap(); + let identifier = Identifier::new(&["public".to_string()], "ts_ns"); + let mut schema_builder = Schema::builder(); + for (id, (name, data_type)) in [ + ("id", PrimitiveType::Int), + ("ts_ns", PrimitiveType::TimestampNs), + ("tstz_ns", PrimitiveType::TimestamptzNs), + ] + .into_iter() + .enumerate() + { + schema_builder.with_struct_field(StructField::new( + id as i32 + 1, + name, + true, + Type::Primitive(data_type), + None, + )); + } + Table::builder() + .with_name("ts_ns") + .with_location("/ts_ns") + .with_schema(schema_builder.build().unwrap()) + .with_partition_spec( + PartitionSpec::builder() + .with_partition_field(PartitionField::new( + 2, + 1000, + "ts_ns_part", + Transform::Identity, + )) + .with_partition_field(PartitionField::new( + 3, + 1001, + "tstz_ns_part", + Transform::Identity, + )) + .build() + .unwrap(), + ) + .with_property(("format-version".to_owned(), "3".to_owned())) + .build(identifier.namespace(), catalog.clone()) + .await + .unwrap(); + + let ctx = SessionContext::new(); + ctx.register_catalog( + "warehouse", + Arc::new(IcebergCatalog::new(catalog, None).await.unwrap()), + ); + let values = [-1_i64, 0, 1_700_000_000_123_456_789]; + let batch = RecordBatch::try_from_iter(vec![ + ("id", Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef), + ( + "ts_ns", + Arc::new(TimestampNanosecondArray::from(values.to_vec())) as ArrayRef, + ), + ( + "tstz_ns", + Arc::new(TimestampNanosecondArray::from(values.to_vec()).with_timezone("UTC")) + as ArrayRef, + ), + ]) + .unwrap(); + ctx.read_batch(batch) + .unwrap() + .write_table("warehouse.public.ts_ns", DataFrameWriteOptions::default()) + .await + .unwrap(); + + let batches = ctx + .sql("SELECT ts_ns, tstz_ns FROM warehouse.public.ts_ns ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + let batch = &batches[0]; + assert_eq!(batch.num_rows(), values.len()); + for column in batch.columns() { + let timestamps = column + .as_any() + .downcast_ref::() + .expect("Iceberg v3 timestamp must keep nanosecond Arrow type"); + for (row, expected) in values.iter().enumerate() { + assert_eq!(timestamps.value(row), *expected); + } } + + let filtered = ctx + .table("warehouse.public.ts_ns") + .await + .unwrap() + .filter(col("ts_ns").eq(lit(ScalarValue::TimestampNanosecond(Some(-1), None)))) + .unwrap(); + let matching = filtered + .select_columns(&["id"]) + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(matching.iter().map(RecordBatch::num_rows).sum::(), 1); + let ids = matching[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ids.value(0), 0); + + let matching_tz = ctx + .table("warehouse.public.ts_ns") + .await + .unwrap() + .filter(col("tstz_ns").eq(lit(ScalarValue::TimestampNanosecond( + Some(-1), + Some("UTC".into()), + )))) + .unwrap() + .select_columns(&["id"]) + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + matching_tz.iter().map(RecordBatch::num_rows).sum::(), + 1 + ); + let ids = matching_tz[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ids.value(0), 0); } diff --git a/iceberg-rust-spec/src/arrow/schema.rs b/iceberg-rust-spec/src/arrow/schema.rs index 84046da8..76e07624 100644 --- a/iceberg-rust-spec/src/arrow/schema.rs +++ b/iceberg-rust-spec/src/arrow/schema.rs @@ -207,6 +207,14 @@ impl TryFrom<&DataType> for Type { })), DataType::Date32 => Ok(Type::Primitive(PrimitiveType::Date)), DataType::Time64(_) => Ok(Type::Primitive(PrimitiveType::Time)), + DataType::Timestamp(TimeUnit::Nanosecond, None) => { + Ok(Type::Primitive(PrimitiveType::TimestampNs)) + } + DataType::Timestamp(TimeUnit::Nanosecond, Some(tz)) + if tz.as_ref() == "UTC" || tz.as_ref() == "+00:00" => + { + Ok(Type::Primitive(PrimitiveType::TimestamptzNs)) + } DataType::Timestamp(_, None) => Ok(Type::Primitive(PrimitiveType::Timestamp)), DataType::Timestamp(_, Some(tz)) if tz.as_ref() == "UTC" || tz.as_ref() == "+00:00" => { Ok(Type::Primitive(PrimitiveType::Timestamptz)) @@ -1239,6 +1247,21 @@ mod tests { )]))]); let result: Result = (&arrow_schema).try_into(); assert!(matches!(result.unwrap_err(), Error::NotSupported(_))); + + for (arrow_type, iceberg_type) in [ + ( + DataType::Timestamp(TimeUnit::Nanosecond, None), + PrimitiveType::TimestampNs, + ), + ( + DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("UTC"))), + PrimitiveType::TimestamptzNs, + ), + ] { + let inferred = Type::try_from(&arrow_type).unwrap(); + assert_eq!(inferred, Type::Primitive(iceberg_type)); + assert_eq!(DataType::try_from(&inferred).unwrap(), arrow_type); + } } #[test] diff --git a/iceberg-rust-spec/src/spec/manifest.rs b/iceberg-rust-spec/src/spec/manifest.rs index 9318ade8..133e3755 100644 --- a/iceberg-rust-spec/src/spec/manifest.rs +++ b/iceberg-rust-spec/src/spec/manifest.rs @@ -605,6 +605,8 @@ fn avro_schema_datatype(data_type: &Type) -> Type { PrimitiveType::Time => Type::Primitive(PrimitiveType::Long), PrimitiveType::Timestamp => Type::Primitive(PrimitiveType::Long), PrimitiveType::Timestamptz => Type::Primitive(PrimitiveType::Long), + PrimitiveType::TimestampNs => Type::Primitive(PrimitiveType::Long), + PrimitiveType::TimestamptzNs => Type::Primitive(PrimitiveType::Long), p => Type::Primitive(p.clone()), }, t => t.clone(), @@ -693,8 +695,9 @@ impl AvroMap { /// /// Type promotion does not rewrite existing bounds, so a bound may be encoded at a narrower /// write-time type than the current schema reports. Per the Iceberg spec, the write-time type is -/// inferred from the byte length: a 4-byte bound for a current `long`/`double`/`timestamp` was -/// written as `int`/`float`/`date` and is promoted here. All other cases (including decimals, which +/// inferred from the byte length: a 4-byte bound for a current +/// `long`/`double`/`timestamp`/`timestamp_ns` was written as `int`/`float`/`date` and is +/// promoted here. All other cases (including decimals, which /// decode at any width from the scale in `data_type`) fall through to the plain single-value /// decoding. pub(crate) fn decode_bound(bytes: &[u8], data_type: &Type) -> Result { @@ -705,7 +708,7 @@ pub(crate) fn decode_bound(bytes: &[u8], data_type: &Type) -> Result Ok(Value::Double( ordered_float::OrderedFloat(f64::from(f32::from_le_bytes(bytes.try_into()?))), )), - (Type::Primitive(PrimitiveType::Timestamp), 4) => { + (Type::Primitive(PrimitiveType::Timestamp | PrimitiveType::TimestampNs), 4) => { Value::Date(i32::from_le_bytes(bytes.try_into()?)).cast(data_type) } _ => Value::try_from_bytes(bytes, data_type), @@ -721,7 +724,7 @@ pub(crate) fn decode_bound_for_copy(bytes: &[u8], data_type: &Type) -> Result { Value::try_from_bytes(bytes, &Type::Primitive(PrimitiveType::Float)) } - (Type::Primitive(PrimitiveType::Timestamp), 4) => { + (Type::Primitive(PrimitiveType::Timestamp | PrimitiveType::TimestampNs), 4) => { Value::try_from_bytes(bytes, &Type::Primitive(PrimitiveType::Date)) } _ => Value::try_from_bytes(bytes, data_type), @@ -2266,6 +2269,7 @@ mod tests { const LONG: Type = Type::Primitive(PrimitiveType::Long); const DOUBLE: Type = Type::Primitive(PrimitiveType::Double); const TIMESTAMP: Type = Type::Primitive(PrimitiveType::Timestamp); + const TIMESTAMP_NS: Type = Type::Primitive(PrimitiveType::TimestampNs); #[test] fn decode_bound_widens_int_written_before_long_promotion() { @@ -2299,6 +2303,24 @@ mod tests { assert_eq!(decoded, Value::Timestamp(259_200_000_000)); } + #[test] + fn decode_bound_promotes_date_written_before_timestamp_ns_promotion() { + let bytes = (-1_i32).to_le_bytes(); + assert_eq!( + decode_bound(&bytes, &TIMESTAMP_NS).unwrap(), + Value::TimestampNano(-86_400_000_000_000) + ); + assert_eq!( + decode_bound_for_copy(&bytes, &TIMESTAMP_NS).unwrap(), + Value::Date(-1) + ); + let nanos = 1_511_955_007_123_456_789_i64; + assert_eq!( + decode_bound(&nanos.to_le_bytes(), &TIMESTAMP_NS).unwrap(), + Value::TimestampNano(nanos) + ); + } + #[test] fn decode_bound_decodes_decimal_at_minimal_width() { let data_type = Type::Primitive(PrimitiveType::Decimal { diff --git a/iceberg-rust-spec/src/spec/values.rs b/iceberg-rust-spec/src/spec/values.rs index 912ef553..73ee5493 100644 --- a/iceberg-rust-spec/src/spec/values.rs +++ b/iceberg-rust-spec/src/spec/values.rs @@ -92,6 +92,10 @@ pub enum Value { Timestamp(i64), /// Stores microseconds from 1970-01-01 00:00:00.000000 in an 8-byte little-endian long TimestampTZ(i64), + /// Stores nanoseconds from 1970-01-01 00:00:00.000000000 in an 8-byte little-endian long + TimestampNano(i64), + /// Stores nanoseconds from 1970-01-01 00:00:00.000000000 UTC in an 8-byte little-endian long + TimestampTZNano(i64), /// UTF-8 bytes (without length) String(String), /// 16-byte big-endian value @@ -135,6 +139,8 @@ impl From for ByteBuf { Value::Time(val) => ByteBuf::from(val.to_le_bytes()), Value::Timestamp(val) => ByteBuf::from(val.to_le_bytes()), Value::TimestampTZ(val) => ByteBuf::from(val.to_le_bytes()), + Value::TimestampNano(val) => ByteBuf::from(val.to_le_bytes()), + Value::TimestampTZNano(val) => ByteBuf::from(val.to_le_bytes()), Value::String(val) => ByteBuf::from(val.as_bytes()), Value::UUID(val) => ByteBuf::from(val.as_u128().to_be_bytes()), Value::Fixed(_, val) => ByteBuf::from(val), @@ -157,6 +163,8 @@ impl fmt::Display for Value { Value::Time(t) => write!(f, "{t}"), Value::Timestamp(ts) => write!(f, "{ts}"), Value::TimestampTZ(ts) => write!(f, "{ts}"), + Value::TimestampNano(ts) => write!(f, "{ts}"), + Value::TimestampTZNano(ts) => write!(f, "{ts}"), Value::String(s) => write!(f, "{s}"), Value::UUID(u) => write!(f, "{u}"), Value::Fixed(size, data) => write!(f, "{data:?} ({size} bytes)"), @@ -276,17 +284,38 @@ impl Struct { // Int -> Date retain the existing value-cast behavior. let value = field .map(|value| { - let source_type = value.datatype(); - match (&source_type, datatype) { + match (value, datatype) { + // Avro stores timestamps as long, but untagged Value can decode + // small long values as Int. The partition schema disambiguates. + (Value::Int(value), Type::Primitive(PrimitiveType::Timestamp)) => { + Ok(Value::Timestamp(i64::from(value))) + } ( - Type::Primitive(PrimitiveType::Int), - Type::Primitive(PrimitiveType::Long), - ) - | ( - Type::Primitive(PrimitiveType::Float), - Type::Primitive(PrimitiveType::Double), - ) => value.promote_iceberg(&source_type, datatype), - _ => value.cast(datatype), + Value::Int(value), + Type::Primitive(PrimitiveType::Timestamptz), + ) => Ok(Value::TimestampTZ(i64::from(value))), + ( + Value::Int(value), + Type::Primitive(PrimitiveType::TimestampNs), + ) => Ok(Value::TimestampNano(i64::from(value))), + ( + Value::Int(value), + Type::Primitive(PrimitiveType::TimestamptzNs), + ) => Ok(Value::TimestampTZNano(i64::from(value))), + (value, datatype) => { + let source_type = value.datatype(); + match (&source_type, datatype) { + ( + Type::Primitive(PrimitiveType::Int), + Type::Primitive(PrimitiveType::Long), + ) + | ( + Type::Primitive(PrimitiveType::Float), + Type::Primitive(PrimitiveType::Double), + ) => value.promote_iceberg(&source_type, datatype), + _ => value.cast(datatype), + } + } } }) .transpose()?; @@ -404,6 +433,9 @@ impl Value { Value::Int(value) => ByteBuf::from((*value as i64).to_le_bytes()), // Apache Iceberg Java also hashes date as long, despite Appendix B. Value::Date(value) => ByteBuf::from((*value as i64).to_le_bytes()), + Value::TimestampNano(value) | Value::TimestampTZNano(value) => { + ByteBuf::from(value.div_euclid(1_000).to_le_bytes()) + } _ => >::into(self.clone()), }; let hash = murmur3::murmur3_32(&mut Cursor::new(bytes), 0)?; @@ -429,6 +461,9 @@ impl Value { Value::TimestampTZ(time) => Ok(Value::Int( micros_to_datetime(*time).year() - YEARS_BEFORE_UNIX_EPOCH, )), + Value::TimestampNano(time) | Value::TimestampTZNano(time) => Ok(Value::Int( + chrono::DateTime::from_timestamp_nanos(*time).year() - YEARS_BEFORE_UNIX_EPOCH, + )), _ => Err(Error::NotSupported( "Datatype for year partition transform.".to_string(), )), @@ -441,6 +476,9 @@ impl Value { Value::TimestampTZ(time) => { Ok(Value::Int(datetime_to_months(µs_to_datetime(*time)))) } + Value::TimestampNano(time) | Value::TimestampTZNano(time) => Ok(Value::Int( + datetime_to_months(&chrono::DateTime::from_timestamp_nanos(*time).naive_utc()), + )), _ => Err(Error::NotSupported( "Datatype for month partition transform.".to_string(), )), @@ -453,6 +491,9 @@ impl Value { Value::TimestampTZ(time) => Ok(Value::Int(datetime_to_days(µs_to_datetime( *time, )) as i32)), + Value::TimestampNano(time) | Value::TimestampTZNano(time) => { + Ok(Value::Int(time.div_euclid(86_400_000_000_000) as i32)) + } _ => Err(Error::NotSupported( "Datatype for day partition transform.".to_string(), )), @@ -464,6 +505,9 @@ impl Value { Value::TimestampTZ(time) => Ok(Value::Int(datetime_to_hours(µs_to_datetime( *time, )) as i32)), + Value::TimestampNano(time) | Value::TimestampTZNano(time) => { + Ok(Value::Int(time.div_euclid(3_600_000_000_000) as i32)) + } _ => Err(Error::NotSupported( "Datatype for hour partition transform.".to_string(), )), @@ -531,6 +575,12 @@ impl Value { PrimitiveType::Timestamptz => { Ok(Value::TimestampTZ(i64::from_le_bytes(bytes.try_into()?))) } + PrimitiveType::TimestampNs => { + Ok(Value::TimestampNano(i64::from_le_bytes(bytes.try_into()?))) + } + PrimitiveType::TimestamptzNs => Ok(Value::TimestampTZNano(i64::from_le_bytes( + bytes.try_into()?, + ))), PrimitiveType::String => Ok(Value::String(std::str::from_utf8(bytes)?.to_string())), // ByteArray-hinted Uuid bytes are its UTF-8 string form, not a big-endian integer. PrimitiveType::Uuid if physical_type_hint == Some(PhysicalTypeHint::ByteArray) => { @@ -556,9 +606,7 @@ impl Value { }; Ok(Value::Decimal(decimal_from_i128_with_scale(val, *scale)?)) } - PrimitiveType::TimestampNs - | PrimitiveType::TimestamptzNs - | PrimitiveType::Unknown + PrimitiveType::Unknown | PrimitiveType::Variant | PrimitiveType::Geometry(_) | PrimitiveType::Geography(_, _) => Err(Error::NotSupported(format!( @@ -628,6 +676,24 @@ impl Value { &NaiveDateTime::parse_from_str(&s, "%Y-%m-%dT%H:%M:%S%.f+00:00")?, )), ))), + (PrimitiveType::TimestampNs, JsonValue::String(s)) => { + let time = NaiveDateTime::parse_from_str(&s, "%Y-%m-%dT%H:%M:%S%.f")?; + Ok(Some(Value::TimestampNano( + time.and_utc().timestamp_nanos_opt().ok_or_else(|| { + Error::InvalidFormat("timestamp_ns exceeds i64 nanosecond range".into()) + })?, + ))) + } + (PrimitiveType::TimestamptzNs, JsonValue::String(s)) if s.ends_with("+00:00") => { + let time = chrono::DateTime::parse_from_rfc3339(&s)?; + Ok(Some(Value::TimestampTZNano( + time.timestamp_nanos_opt().ok_or_else(|| { + Error::InvalidFormat( + "timestamptz_ns exceeds i64 nanosecond range".into(), + ) + })?, + ))) + } (PrimitiveType::String, JsonValue::String(s)) => Ok(Some(Value::String(s))), (PrimitiveType::Uuid, JsonValue::String(s)) => { Ok(Some(Value::UUID(Uuid::parse_str(&s)?))) @@ -739,6 +805,8 @@ impl Value { Value::Time(_) => Type::Primitive(PrimitiveType::Time), Value::Timestamp(_) => Type::Primitive(PrimitiveType::Timestamp), Value::TimestampTZ(_) => Type::Primitive(PrimitiveType::Timestamptz), + Value::TimestampNano(_) => Type::Primitive(PrimitiveType::TimestampNs), + Value::TimestampTZNano(_) => Type::Primitive(PrimitiveType::TimestamptzNs), Value::Fixed(len, _) => Type::Primitive(PrimitiveType::Fixed(*len as u64)), Value::Binary(_) => Type::Primitive(PrimitiveType::Binary), Value::String(_) => Type::Primitive(PrimitiveType::String), @@ -770,6 +838,8 @@ impl Value { Value::Time(any) => Box::new(any), Value::Timestamp(any) => Box::new(any), Value::TimestampTZ(any) => Box::new(any), + Value::TimestampNano(any) => Box::new(any), + Value::TimestampTZNano(any) => Box::new(any), Value::Fixed(_, any) => Box::new(any), Value::Binary(any) => Box::new(any), Value::String(any) => Box::new(any), @@ -819,6 +889,36 @@ impl Value { .checked_mul(86_400_000_000) .map(Value::Timestamp) .ok_or_else(|| Error::InvalidFormat("date exceeds timestamp range".into())), + (Value::Date(input), Type::Primitive(PrimitiveType::TimestampNs)) => { + i64::from(input) + .checked_mul(86_400_000_000_000) + .map(Value::TimestampNano) + .ok_or_else(|| { + Error::InvalidFormat("date exceeds timestamp_ns range".into()) + }) + } + (Value::Timestamp(input), Type::Primitive(PrimitiveType::TimestampNs)) => input + .checked_mul(1_000) + .map(Value::TimestampNano) + .ok_or_else(|| { + Error::InvalidFormat("timestamp exceeds timestamp_ns range".into()) + }), + (Value::TimestampTZ(input), Type::Primitive(PrimitiveType::TimestamptzNs)) => input + .checked_mul(1_000) + .map(Value::TimestampTZNano) + .ok_or_else(|| { + Error::InvalidFormat("timestamptz exceeds timestamptz_ns range".into()) + }), + (Value::TimestampNano(input), Type::Primitive(PrimitiveType::Timestamp)) => { + Ok(Value::Timestamp(input.div_euclid(1_000))) + } + (Value::TimestampTZNano(input), Type::Primitive(PrimitiveType::Timestamptz)) => { + Ok(Value::TimestampTZ(input.div_euclid(1_000))) + } + ( + Value::TimestampNano(input) | Value::TimestampTZNano(input), + Type::Primitive(PrimitiveType::Date), + ) => Ok(Value::Date(input.div_euclid(86_400_000_000_000) as i32)), (Value::Int(input), Type::Primitive(PrimitiveType::Date)) => Ok(Value::Date(input)), (Value::LongInt(input), Type::Primitive(PrimitiveType::Time)) => { Ok(Value::Time(input)) @@ -829,6 +929,12 @@ impl Value { (Value::LongInt(input), Type::Primitive(PrimitiveType::Timestamptz)) => { Ok(Value::TimestampTZ(input)) } + (Value::LongInt(input), Type::Primitive(PrimitiveType::TimestampNs)) => { + Ok(Value::TimestampNano(input)) + } + (Value::LongInt(input), Type::Primitive(PrimitiveType::TimestamptzNs)) => { + Ok(Value::TimestampTZNano(input)) + } _ => Err(Error::NotSupported("cast".to_string())), } } @@ -856,6 +962,11 @@ impl Value { Type::Primitive(PrimitiveType::Date), Type::Primitive(PrimitiveType::Timestamp), ) => Value::Date(input).cast(target_type), + ( + Value::Date(input), + Type::Primitive(PrimitiveType::Date), + Type::Primitive(PrimitiveType::TimestampNs), + ) => Value::Date(input).cast(target_type), ( value @ Value::Decimal(_), Type::Primitive(PrimitiveType::Decimal { @@ -913,6 +1024,16 @@ impl From<&Value> for JsonValue { .format("%Y-%m-%dT%H:%M:%S%.f+00:00") .to_string(), ), + Value::TimestampNano(val) => JsonValue::String( + chrono::DateTime::from_timestamp_nanos(*val) + .format("%Y-%m-%dT%H:%M:%S%.9f") + .to_string(), + ), + Value::TimestampTZNano(val) => JsonValue::String( + chrono::DateTime::from_timestamp_nanos(*val) + .format("%Y-%m-%dT%H:%M:%S%.9f+00:00") + .to_string(), + ), Value::String(val) => JsonValue::String(val.clone()), Value::UUID(val) => JsonValue::String(val.to_string()), Value::Fixed(_, val) => { @@ -1097,6 +1218,12 @@ impl TrySub for Value { (Value::TimestampTZ(own), Value::TimestampTZ(other)) => { Ok(Value::TimestampTZ(own - other)) } + (Value::TimestampNano(own), Value::TimestampNano(other)) => Ok(Value::Decimal( + decimal_from_i128_with_scale(i128::from(*own) - i128::from(*other), 0)?, + )), + (Value::TimestampTZNano(own), Value::TimestampTZNano(other)) => Ok(Value::Decimal( + decimal_from_i128_with_scale(i128::from(*own) - i128::from(*other), 0)?, + )), (Value::String(own), Value::String(other)) => { Ok(Value::LongInt(sub_string(own, other) as i64)) } @@ -1686,6 +1813,30 @@ mod tests { assert_eq!(result, Value::Int(648)); } + #[test] + fn timestamp_ns_bucket_hashes_floor_microseconds() { + let micros = 1_510_871_468_000_001_i64; + let nanos = micros * 1_000 + 1; + let hash = murmur3::murmur3_32(&mut Cursor::new(micros.to_le_bytes()), 0).unwrap(); + assert_eq!(hash as i32, -1_207_196_810); + let expected = Value::Int(((hash & i32::MAX as u32) % 16) as i32); + for value in [ + Value::Timestamp(micros), + Value::TimestampNano(nanos), + Value::TimestampTZNano(nanos), + ] { + assert_eq!(value.transform(&Transform::Bucket(16)).unwrap(), expected); + } + for value in [Value::TimestampNano(-1), Value::TimestampTZNano(-1)] { + assert_eq!( + value.transform(&Transform::Bucket(16)).unwrap(), + Value::Timestamp(-1) + .transform(&Transform::Bucket(16)) + .unwrap() + ); + } + } + #[test] fn test_transform_month_unsupported() { let value = Value::Boolean(true); @@ -1818,10 +1969,16 @@ mod tests { // ----------------------------------------------------------------------- #[test] - #[ignore = "no Value::TimestampNano variant; Value::cast Timestamp->TimestampNano unimplemented"] fn test_timestamp_micros_promotes_to_timestamp_nanos_by_multiplying_by_one_thousand() { - // Value::Timestamp(micros).cast(TimestampNs) yields Value::TimestampNano(micros * 1000). - unimplemented!("Timestamp->TimestampNano cast"); + assert_eq!( + Value::Timestamp(-1) + .cast(&Type::Primitive(PrimitiveType::TimestampNs)) + .unwrap(), + Value::TimestampNano(-1_000) + ); + assert!(Value::Timestamp(i64::MAX) + .cast(&Type::Primitive(PrimitiveType::TimestampNs)) + .is_err()); } #[test] @@ -1840,25 +1997,36 @@ mod tests { } #[test] - #[ignore = "no Value::TimestampNano variant; Value::cast TimestampNano->Timestamp unimplemented"] fn test_timestamp_nanos_narrows_to_timestamp_micros_by_floor_division_by_one_thousand() { - // Value::TimestampNano(nanos).cast(Timestamp) yields Value::Timestamp(nanos.div_euclid(1000)). - unimplemented!("TimestampNano->Timestamp cast"); + assert_eq!( + Value::TimestampNano(-1) + .cast(&Type::Primitive(PrimitiveType::Timestamp)) + .unwrap(), + Value::Timestamp(-1) + ); } #[test] - #[ignore = "no Value::TimestampNano variant; Value::cast TimestampNano->Date unimplemented"] fn test_timestamp_nanos_casts_to_date_using_floor_division() { - // Value::TimestampNano(nanos).cast(Date) returns days-since-epoch with floor semantics. - unimplemented!("TimestampNano->Date cast"); + assert_eq!( + Value::TimestampNano(-1) + .cast(&Type::Primitive(PrimitiveType::Date)) + .unwrap(), + Value::Date(-1) + ); } #[test] - #[ignore = "no Value::TimestampNano variant; offset-bearing string rejected for without-zone"] fn test_timestamp_nanos_with_zone_strings_rejected_for_without_zone_targets() { - // `try_from_json` of `"2024-11-07T12:33:54.123456789+00:00"` is accepted for - // TimestampTzNs but rejected for TimestampNs (and similarly for TimestampTz/Timestamp). - unimplemented!("TimestampNano zone parsing"); + let value = JsonValue::String("2024-11-07T12:33:54.123456789+00:00".into()); + assert!( + Value::try_from_json(value.clone(), &Type::Primitive(PrimitiveType::TimestampNs)) + .is_err() + ); + assert!(matches!( + Value::try_from_json(value, &Type::Primitive(PrimitiveType::TimestamptzNs)), + Ok(Some(Value::TimestampTZNano(_))) + )); } #[test] @@ -1870,11 +2038,15 @@ mod tests { } #[test] - #[ignore = "Timestamptz parser hardcodes +00:00; offset-less string rejected for with-zone"] fn test_timestamp_nanos_without_zone_strings_rejected_for_with_zone_targets() { - // `try_from_json` of `"2024-11-07T12:33:54.123456789"` accepted for TimestampNs but - // rejected for TimestampTzNs. - unimplemented!("TimestampNano no-zone parsing"); + let value = JsonValue::String("2024-11-07T12:33:54.123456789".into()); + assert!(matches!( + Value::try_from_json(value.clone(), &Type::Primitive(PrimitiveType::TimestampNs)), + Ok(Some(Value::TimestampNano(_))) + )); + assert!( + Value::try_from_json(value, &Type::Primitive(PrimitiveType::TimestamptzNs)).is_err() + ); } #[test] @@ -2001,10 +2173,13 @@ mod tests { } #[test] - #[ignore = "no Value::TimestampNano variant"] fn test_timestamp_with_nanoseconds_casts_to_date_via_floor_division() { - // TimestampNano(nanos).cast(Date) returns days-since-epoch with floor semantics. - unimplemented!("TimestampNano->Date cast"); + assert_eq!( + Value::TimestampNano(-86_400_000_000_001) + .cast(&Type::Primitive(PrimitiveType::Date)) + .unwrap(), + Value::Date(-2) + ); } #[test] @@ -2291,12 +2466,36 @@ mod tests { } #[test] - #[ignore = "no Value::TimestampNano variant; try_from_json for timestamp_ns / timestamptz_ns absent"] fn test_literal_round_trips_via_json_for_timestamp_nano_variants() { - // Once Value::TimestampNano lands, two more cases extend the round-trip set: - // ("2017-11-29T11:30:07.123456789", TimestampNs, no-offset string) - // ("2017-11-29T11:30:07.123456789+00:00", TimestamptzNs, +00:00 offset) - unimplemented!("TimestampNano literal serde"); + for (value, data_type, literal) in [ + ( + Value::TimestampNano(1_511_955_007_123_456_789), + PrimitiveType::TimestampNs, + "2017-11-29T11:30:07.123456789", + ), + ( + Value::TimestampTZNano(1_511_955_007_123_456_789), + PrimitiveType::TimestamptzNs, + "2017-11-29T11:30:07.123456789+00:00", + ), + ] { + let json = JsonValue::from(&value); + assert_eq!(json, JsonValue::String(literal.into())); + assert_eq!( + Value::try_from_json(json, &Type::Primitive(data_type)).unwrap(), + Some(value.clone()) + ); + let bytes: ByteBuf = value.clone().into(); + assert_eq!( + Value::try_from_bytes(&bytes, &value.datatype()).unwrap(), + value + ); + } + assert!(Value::try_from_json( + JsonValue::String("2024-11-07T12:33:54.123456789Z".into()), + &Type::Primitive(PrimitiveType::TimestamptzNs) + ) + .is_err()); } // ----------------------------------------------------------------------- diff --git a/iceberg-rust/src/arrow/partition.rs b/iceberg-rust/src/arrow/partition.rs index 7e0f0e99..b0fca713 100644 --- a/iceberg-rust/src/arrow/partition.rs +++ b/iceberg-rust/src/arrow/partition.rs @@ -8,24 +8,30 @@ //! * Efficient handling of distinct partition values //! * Automatic management of partition streams and channels -use std::{collections::HashSet, hash::Hash}; +use std::{ + collections::{HashMap, HashSet}, + hash::Hash, +}; use arrow::{ array::{ as_primitive_array, as_string_array, Array, ArrayRef, BooleanArray, BooleanBufferBuilder, - PrimitiveArray, Scalar, StringArray, + PrimitiveArray, Scalar, StringArray, UInt32Array, }, compute::{ and, filter, filter_record_batch, is_not_null, is_null, kernels::cmp::{distinct, eq}, + take, }, datatypes::{ ArrowPrimitiveType, DataType, Int32Type, Int64Type, TimeUnit, TimestampMicrosecondType, + TimestampNanosecondType, }, error::ArrowError, record_batch::RecordBatch, + row::{RowConverter, SortField}, }; -use itertools::{iproduct, Itertools}; +use itertools::{iproduct, Either, Itertools}; use iceberg_rust_spec::{partition::BoundPartitionField, spec::values::Value}; @@ -34,6 +40,8 @@ use super::transform::transform_arrow; /// A record batch and the partition values shared by all its rows. pub type PartitionedRecordBatch = (Vec>, RecordBatch); +const MAX_BITMAP_COMPARISONS: usize = 1_000_000; + /// Partitions a record batch according to the given partition fields. /// /// This function takes a record batch and partition field specifications, then splits the batch into @@ -70,6 +78,20 @@ pub fn partition_record_batch<'a>( .iter() .map(|x| distinct_values(x.clone())) .collect::, ArrowError>>()?; + let candidate_groups = distinct_values.iter().zip(&partition_columns).fold( + 1_usize, + |product, (distinct, array)| { + product.saturating_mul(distinct.len() + usize::from(array.null_count() > 0)) + }, + ); + if candidate_groups > 16 + && candidate_groups.saturating_mul(record_batch.num_rows()) > MAX_BITMAP_COMPARISONS + { + return Ok(Either::Left(partition_record_batch_indexed( + record_batch, + partition_columns, + )?)); + } let mut true_buffer = BooleanBufferBuilder::new(record_batch.num_rows()); true_buffer.append_n(record_batch.num_rows(), true); let predicates = distinct_values @@ -111,6 +133,21 @@ pub fn partition_record_batch<'a>( Ok((Some(partition_value), eq(&scalar, value)?)) }) .collect::, ArrowError>>(), + DistinctValues::TimestampNano(set, timezone) => set + .into_iter() + .map(|x| { + let scalar = Scalar::new( + PrimitiveArray::::from(vec![x]) + .with_timezone_opt(timezone.clone()), + ); + let partition_value = if timezone.is_some() { + Value::TimestampTZNano(x) + } else { + Value::TimestampNano(x) + }; + Ok((Some(partition_value), eq(&scalar, value)?)) + }) + .collect::, ArrowError>>(), DistinctValues::String(set) => set .into_iter() .map(|x| { @@ -137,11 +174,87 @@ pub fn partition_record_batch<'a>( .collect::>, _)>, ArrowError>>() }, )?; - Ok(predicates.into_iter().map(move |(values, predicate)| { - Ok((values, filter_record_batch(record_batch, &predicate)?)) + Ok(Either::Right(predicates.into_iter().map( + move |(values, predicate)| Ok((values, filter_record_batch(record_batch, &predicate)?)), + ))) +} + +fn partition_record_batch_indexed<'a>( + record_batch: &'a RecordBatch, + partition_columns: Vec, +) -> Result> + 'a, ArrowError> { + let converter = RowConverter::new( + partition_columns + .iter() + .map(|array| SortField::new(array.data_type().clone())) + .collect(), + )?; + let rows = converter.convert_columns(&partition_columns)?; + let mut groups: HashMap<_, (usize, Vec)> = HashMap::new(); + for row in 0..record_batch.num_rows() { + let key = rows.row(row); + let index = u32::try_from(row).map_err(|_| { + ArrowError::ComputeError("Record batch exceeds u32 row-index range".into()) + })?; + groups + .entry(key) + .or_insert_with(|| (row, Vec::new())) + .1 + .push(index); + } + + let groups: Vec<_> = groups.into_values().collect(); + Ok(groups.into_iter().map(move |(first_row, indices)| { + let values = partition_columns + .iter() + .map(|array| partition_value_at(array, first_row)) + .collect::, _>>()?; + let indices = UInt32Array::from(indices); + let columns = record_batch + .columns() + .iter() + .map(|array| take(array.as_ref(), &indices, None)) + .collect::, _>>()?; + Ok(( + values, + RecordBatch::try_new(record_batch.schema(), columns)?, + )) })) } +fn partition_value_at(array: &ArrayRef, row: usize) -> Result, ArrowError> { + if array.is_null(row) { + return Ok(None); + } + let value = match array.data_type() { + DataType::Int32 => Value::Int(as_primitive_array::(array).value(row)), + DataType::Int64 => Value::LongInt(as_primitive_array::(array).value(row)), + DataType::Timestamp(TimeUnit::Microsecond, timezone) => { + let value = as_primitive_array::(array).value(row); + if timezone.is_some() { + Value::TimestampTZ(value) + } else { + Value::Timestamp(value) + } + } + DataType::Timestamp(TimeUnit::Nanosecond, timezone) => { + let value = as_primitive_array::(array).value(row); + if timezone.is_some() { + Value::TimestampTZNano(value) + } else { + Value::TimestampNano(value) + } + } + DataType::Utf8 => Value::String(as_string_array(array).value(row).to_owned()), + _ => { + return Err(ArrowError::ComputeError( + "Datatype not supported for partitioning".into(), + )) + } + }; + Ok(Some(value)) +} + /// Extracts distinct values from an Arrow array into a DistinctValues enum /// /// # Arguments @@ -174,6 +287,10 @@ fn distinct_values(array: ArrayRef) -> Result { distinct_values_primitive::(array.clone())?, timezone.clone(), )), + DataType::Timestamp(TimeUnit::Nanosecond, timezone) => Ok(DistinctValues::TimestampNano( + distinct_values_primitive::(array.clone())?, + timezone.clone(), + )), DataType::Utf8 => Ok(DistinctValues::String(distinct_values_string(array)?)), _ => Err(ArrowError::ComputeError( "Datatype not supported for transform.".to_string(), @@ -279,19 +396,32 @@ fn distinct_values_string(array: ArrayRef) -> Result, ArrowError /// * `Int` - Distinct 32-bit integer values /// * `Long` - Distinct 64-bit integer values /// * `Timestamp` - Distinct microsecond timestamps and their optional timezone +/// * `TimestampNano` - Distinct nanosecond timestamps and their optional timezone /// * `String` - Distinct string values enum DistinctValues { Int(HashSet), Long(HashSet), Timestamp(HashSet, Option>), + TimestampNano(HashSet, Option>), String(HashSet), } +impl DistinctValues { + fn len(&self) -> usize { + match self { + Self::Int(values) => values.len(), + Self::Long(values) => values.len(), + Self::Timestamp(values, _) | Self::TimestampNano(values, _) => values.len(), + Self::String(values) => values.len(), + } + } +} + #[cfg(test)] mod tests { use std::sync::Arc; - use arrow::array::{ArrayRef, TimestampMicrosecondArray}; + use arrow::array::{ArrayRef, TimestampMicrosecondArray, TimestampNanosecondArray}; use iceberg_rust_spec::spec::{ partition::{PartitionField, Transform}, types::{PrimitiveType, StructField, Type}, @@ -333,6 +463,128 @@ mod tests { Ok(()) } + #[test] + fn identity_partition_preserves_nanosecond_precision() -> Result<(), ArrowError> { + for timezone in [None, Some("UTC")] { + let array: ArrayRef = Arc::new( + TimestampNanosecondArray::from(vec![Some(10_i64), Some(11_i64), None]) + .with_timezone_opt(timezone), + ); + let batch = RecordBatch::try_from_iter(vec![("ts", array)])?; + let source_type = if timezone.is_some() { + PrimitiveType::TimestamptzNs + } else { + PrimitiveType::TimestampNs + }; + let source = StructField::new(1, "ts", false, Type::Primitive(source_type), None); + let partition = PartitionField::new(1, 1000, "ts", Transform::Identity); + let bound = BoundPartitionField::new(&partition, &source); + let mut groups = partition_record_batch(&batch, &[bound])? + .map(|group| group.map(|(values, rows)| (values[0].clone(), rows.num_rows()))) + .collect::, ArrowError>>()?; + groups.sort(); + let expected = if timezone.is_some() { + vec![ + (None, 1), + (Some(Value::TimestampTZNano(10)), 1), + (Some(Value::TimestampTZNano(11)), 1), + ] + } else { + vec![ + (None, 1), + (Some(Value::TimestampNano(10)), 1), + (Some(Value::TimestampNano(11)), 1), + ] + }; + assert_eq!(groups, expected); + } + Ok(()) + } + + #[test] + fn high_cardinality_nanosecond_partitions_use_indexed_groups() -> Result<(), ArrowError> { + let size = 2_048; + for timezone in [None, Some("UTC")] { + let timestamps: Vec<_> = (0..size) + .map(|row| if row == 0 { None } else { Some(i64::from(row)) }) + .collect(); + let array: ArrayRef = + Arc::new(TimestampNanosecondArray::from(timestamps).with_timezone_opt(timezone)); + let categories: ArrayRef = Arc::new(arrow::array::Int32Array::from_iter_values( + (0..size).map(|row| row % 2), + )); + let batch = RecordBatch::try_from_iter(vec![("ts", array), ("category", categories)])?; + let source_type = if timezone.is_some() { + PrimitiveType::TimestamptzNs + } else { + PrimitiveType::TimestampNs + }; + let ts = StructField::new(1, "ts", false, Type::Primitive(source_type), None); + let category = StructField::new( + 2, + "category", + true, + Type::Primitive(PrimitiveType::Int), + None, + ); + let ts_partition = PartitionField::new(1, 1000, "ts_part", Transform::Identity); + let category_partition = + PartitionField::new(2, 1001, "category_part", Transform::Identity); + let fields = [ + BoundPartitionField::new(&ts_partition, &ts), + BoundPartitionField::new(&category_partition, &category), + ]; + let groups = + partition_record_batch(&batch, &fields)?.collect::, ArrowError>>()?; + assert_eq!(groups.len(), size as usize); + for (values, batch) in groups { + assert_eq!(batch.num_rows(), 1); + let timestamp = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let category = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + let expected = if timestamp.is_null(0) { + None + } else if timezone.is_some() { + Some(Value::TimestampTZNano(timestamp.value(0))) + } else { + Some(Value::TimestampNano(timestamp.value(0))) + }; + assert_eq!(values, vec![expected, Some(Value::Int(category))]); + } + } + Ok(()) + } + + #[test] + fn indexed_grouping_accepts_more_than_ten_thousand_partitions() -> Result<(), ArrowError> { + let size = 10_001_i64; + let timestamps: ArrayRef = Arc::new(TimestampNanosecondArray::from_iter_values(0..size)); + let batch = RecordBatch::try_from_iter(vec![("ts", timestamps)])?; + let source = StructField::new( + 1, + "ts", + true, + Type::Primitive(PrimitiveType::TimestampNs), + None, + ); + let partition = PartitionField::new(1, 1000, "ts_part", Transform::Identity); + let rows = + partition_record_batch(&batch, &[BoundPartitionField::new(&partition, &source)])? + .try_fold(0, |total, group| { + group.map(|(_, batch)| total + batch.num_rows()) + })?; + assert_eq!(rows, size as usize); + Ok(()) + } + #[test] fn nullable_partition_values_keep_every_row() -> Result<(), ArrowError> { for (array, source_type) in [ diff --git a/iceberg-rust/src/arrow/transform.rs b/iceberg-rust/src/arrow/transform.rs index f2f47ec8..44d4f05b 100644 --- a/iceberg-rust/src/arrow/transform.rs +++ b/iceberg-rust/src/arrow/transform.rs @@ -19,7 +19,7 @@ use arrow::{ compute::{binary, cast, date_part, unary, DatePart}, datatypes::{ DataType, Date32Type, Int16Type, Int32Type, Int64Type, Time64MicrosecondType, TimeUnit, - TimestampMicrosecondType, + TimestampMicrosecondType, TimestampNanosecondType, }, error::ArrowError, }; @@ -28,6 +28,8 @@ use iceberg_rust_spec::{spec::partition::Transform, values::YEARS_BEFORE_UNIX_EP static MICROS_IN_HOUR: i64 = 3_600_000_000; static MICROS_IN_DAY: i64 = 86_400_000_000; +static NANOS_IN_HOUR: i64 = 3_600_000_000_000; +static NANOS_IN_DAY: i64 = 86_400_000_000_000; /// Applies an Iceberg partition transform to an Arrow array /// @@ -114,6 +116,65 @@ pub fn transform_arrow(array: ArrayRef, transform: &Transform) -> Result { + Ok(Arc::new(unary::<_, _, Int32Type>( + as_primitive_array::(&array), + |value| value.div_euclid(NANOS_IN_HOUR) as i32, + ))) + } + (DataType::Timestamp(TimeUnit::Nanosecond, _), Transform::Day) => { + Ok(Arc::new(unary::<_, _, Int32Type>( + as_primitive_array::(&array), + |value| value.div_euclid(NANOS_IN_DAY) as i32, + ))) + } + (DataType::Timestamp(TimeUnit::Nanosecond, _), Transform::Month) => { + let values = as_primitive_array::(&array); + // Iceberg timestamp partitions use UTC; Arrow needs chrono-tz for named timezones. + let utc_values = match array.data_type() { + DataType::Timestamp(_, Some(_)) => { + Some(values.clone().with_timezone_opt(None::>)) + } + _ => None, + }; + let values = utc_values.as_ref().unwrap_or(values); + let year = date_part(values, DatePart::Year)?; + let month = date_part(values, DatePart::Month)?; + Ok(Arc::new(binary::<_, _, _, Int32Type>( + as_primitive_array::(&year), + as_primitive_array::(&month), + datepart_to_months, + )?)) + } + (DataType::Timestamp(TimeUnit::Nanosecond, _), Transform::Year) => { + let values = as_primitive_array::(&array); + // Keep the epoch values unchanged while removing only the UTC display timezone. + let utc_values = match array.data_type() { + DataType::Timestamp(_, Some(_)) => { + Some(values.clone().with_timezone_opt(None::>)) + } + _ => None, + }; + Ok(Arc::new(unary::<_, _, Int32Type>( + as_primitive_array::(&date_part( + utc_values.as_ref().unwrap_or(values), + DatePart::Year, + )?), + datepart_to_years, + ))) + } + (DataType::Timestamp(TimeUnit::Nanosecond, _), Transform::Bucket(m)) => { + Ok(Arc::new(unary::<_, _, Int32Type>( + as_primitive_array::(&array), + |value| bucket(&value.div_euclid(1_000).to_le_bytes(), *m), + ))) + } + (DataType::Timestamp(TimeUnit::Microsecond, _), Transform::Bucket(m)) => { + Ok(Arc::new(unary::<_, _, Int32Type>( + as_primitive_array::(&array), + |value| bucket(&value.to_le_bytes(), *m), + ))) + } (DataType::Int16, Transform::Truncate(m)) => Ok(Arc::>::new( unary(as_primitive_array::(&array), |i| { i - i.rem_euclid(*m as i16) @@ -213,7 +274,9 @@ fn bucket(bytes: &[u8], width: u32) -> i32 { mod tests { use super::*; - use arrow::array::{ArrayRef, Date32Array, TimestampMicrosecondArray}; + use arrow::array::{ + ArrayRef, Date32Array, TimestampMicrosecondArray, TimestampNanosecondArray, + }; fn create_date32_array() -> ArrayRef { Arc::new(Date32Array::from(vec![ @@ -245,6 +308,34 @@ mod tests { ) as ArrayRef } + #[test] + fn test_timestamp_ns_transforms_preserve_sub_microsecond_values() { + for timezone in [None, Some("UTC")] { + let array: ArrayRef = Arc::new( + TimestampNanosecondArray::from(vec![Some(-1), Some(0), Some(NANOS_IN_DAY), None]) + .with_timezone_opt(timezone), + ); + for (transform, expected) in [ + (Transform::Year, vec![Some(-1), Some(0), Some(0), None]), + (Transform::Month, vec![Some(-1), Some(0), Some(0), None]), + (Transform::Day, vec![Some(-1), Some(0), Some(1), None]), + (Transform::Hour, vec![Some(-1), Some(0), Some(24), None]), + ] { + let result = transform_arrow(array.clone(), &transform).unwrap(); + let expected: ArrayRef = Arc::new(arrow::array::Int32Array::from(expected)); + assert_eq!(&result, &expected, "{transform:?}"); + } + let result = transform_arrow(array, &Transform::Bucket(17)).unwrap(); + let expected: ArrayRef = Arc::new(arrow::array::Int32Array::from(vec![ + Some(bucket(&(-1_i64).to_le_bytes(), 17)), + Some(bucket(&0_i64.to_le_bytes(), 17)), + Some(bucket(&(NANOS_IN_DAY / 1_000).to_le_bytes(), 17)), + None, + ])); + assert_eq!(&result, &expected); + } + } + #[test] fn test_identity_transform() { let array = create_date32_array(); diff --git a/iceberg-rust/src/file_format/parquet.rs b/iceberg-rust/src/file_format/parquet.rs index e5794015..4a27a16b 100644 --- a/iceberg-rust/src/file_format/parquet.rs +++ b/iceberg-rust/src/file_format/parquet.rs @@ -378,6 +378,14 @@ fn parquet_to_datafile_impl( { *entry = new } + ( + Value::TimestampNano(current), + Value::TimestampNano(new_val), + ) if *current > *new_val => *entry = new, + ( + Value::TimestampTZNano(current), + Value::TimestampTZNano(new_val), + ) if *current > *new_val => *entry = new, _ => (), } } @@ -442,6 +450,14 @@ fn parquet_to_datafile_impl( { *entry = new } + ( + Value::TimestampNano(current), + Value::TimestampNano(new_val), + ) if *current < *new_val => *entry = new, + ( + Value::TimestampTZNano(current), + Value::TimestampTZNano(new_val), + ) if *current < *new_val => *entry = new, _ => (), } } diff --git a/iceberg-rust/src/table/manifest.rs b/iceberg-rust/src/table/manifest.rs index 7d2ce334..2c824763 100644 --- a/iceberg-rust/src/table/manifest.rs +++ b/iceberg-rust/src/table/manifest.rs @@ -959,6 +959,16 @@ fn update_partitions( (Value::TimestampTZ(val), Value::TimestampTZ(current)) if *current > *val => { *current = *val } + (Value::TimestampNano(val), Value::TimestampNano(current)) + if *current > *val => + { + *current = *val + } + (Value::TimestampTZNano(val), Value::TimestampTZNano(current)) + if *current > *val => + { + *current = *val + } (Value::String(val), Value::String(current)) if *current > *val => { *current = val.clone() } @@ -1000,6 +1010,16 @@ fn update_partitions( (Value::TimestampTZ(val), Value::TimestampTZ(current)) if *current < *val => { *current = *val } + (Value::TimestampNano(val), Value::TimestampNano(current)) + if *current < *val => + { + *current = *val + } + (Value::TimestampTZNano(val), Value::TimestampTZNano(current)) + if *current < *val => + { + *current = *val + } (Value::String(val), Value::String(current)) if *current < *val => { *current = val.clone() } diff --git a/iceberg-rust/src/table/transaction/append.rs b/iceberg-rust/src/table/transaction/append.rs index 64194452..5c3bbd42 100644 --- a/iceberg-rust/src/table/transaction/append.rs +++ b/iceberg-rust/src/table/transaction/append.rs @@ -367,14 +367,15 @@ mod split_tests { /// Build a manifest entry sitting at partition `day = value`. fn entry_at(day: i32) -> ManifestEntry { + entry_at_partition(Value::Int(day), "day", FormatVersion::V2) + } + + fn entry_at_partition(value: Value, name: &str, version: FormatVersion) -> ManifestEntry { let data_file = DataFile::builder() .with_content(Content::Data) - .with_file_path(format!("s3://bucket/day={day}/data.parquet")) + .with_file_path(format!("s3://bucket/{name}={value}/data.parquet")) .with_file_format(FileFormat::Parquet) - .with_partition(Struct::from_iter(vec![( - "day".to_owned(), - Some(Value::Int(day)), - )])) + .with_partition(Struct::from_iter(vec![(name.to_owned(), Some(value))])) .with_record_count(1) .with_file_size_in_bytes(1024) .with_column_sizes(None) @@ -388,7 +389,7 @@ mod split_tests { .expect("data file builds"); ManifestEntry::builder() - .with_format_version(FormatVersion::V2) + .with_format_version(version) .with_status(Status::Added) .with_data_file(data_file) .with_sequence_number(1) @@ -443,6 +444,27 @@ mod split_tests { ); } + #[test] + fn nanosecond_timestamp_split_handles_distances_beyond_i64() { + let min = -5_000_000_000_000_000_000_i64; + let max = 5_000_000_000_000_000_000_i64; + let files = [min, 0, max].into_iter().map(|value| { + Ok(entry_at_partition( + Value::TimestampNano(value), + "ts_ns", + FormatVersion::V3, + )) + }); + let bounds = Rectangle::new( + smallvec![Value::TimestampNano(min)], + smallvec![Value::TimestampNano(max)], + ); + let [(lower, _), (upper, _)] = split_datafiles_once(files, bounds, &["ts_ns"]).unwrap(); + assert_eq!(lower.len() + upper.len(), 3); + assert!(!lower.is_empty()); + assert!(!upper.is_empty()); + } + #[test] fn empty_input_produces_no_groups() { let files: Vec> = Vec::new();