Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 60 additions & 3 deletions datafusion_iceberg/src/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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::<u8>()
.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
Expand Down Expand Up @@ -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)))?;
Expand All @@ -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,
Expand Down
169 changes: 152 additions & 17 deletions datafusion_iceberg/src/pruning_statistics.rs
Original file line number Diff line number Diff line change
Expand Up @@ -391,6 +391,14 @@ fn any_iter_to_array(
opt.and_then(|value| Some(*value.downcast::<i64>().ok()?)),
)
})),
DataType::Timestamp(TimeUnit::Nanosecond, tz) => {
ScalarValue::iter_to_array(iter.map(move |opt| {
ScalarValue::TimestampNanosecond(
opt.and_then(|value| Some(*value.downcast::<i64>().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| {
Expand Down Expand Up @@ -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;
Expand All @@ -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,
Expand Down Expand Up @@ -580,7 +595,7 @@ pub(crate) fn transform_predicate(
Some(partition_comparison(column, op, literal))
}

fn fold_cast_literal(expr: Expr) -> Option<Expr> {
fn fold_cast_literal(expr: Expr, source_type: &Type) -> Option<Expr> {
match expr {
literal @ Expr::Literal(..) => Some(literal),
Expr::Cast(cast) => {
Expand All @@ -589,10 +604,18 @@ fn fold_cast_literal(expr: Expr) -> Option<Expr> {
};
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))
Expand All @@ -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<Expr> {
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 {
Expand Down Expand Up @@ -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"
)
}
Expand Down Expand Up @@ -710,6 +739,21 @@ fn transform_literal(expr: Expr, transform: &Transform) -> Option<Expr> {
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],
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -854,6 +906,14 @@ fn scalarvalue_to_value(scalar: &ScalarValue) -> Result<Value, Error> {
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"
))),
Expand All @@ -872,6 +932,11 @@ fn value_to_scalarvalue(value: Value) -> Result<ScalarValue, Error> {
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"
))),
Expand All @@ -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;
Expand Down Expand Up @@ -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),
_
Expand All @@ -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")),
Expand All @@ -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,
Expand Down Expand Up @@ -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::<TimestampNanosecondArray>()
.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)] {
Expand Down
Loading
Loading