Skip to content

Commit 5ce0ba3

Browse files
colinmarccolinmarc
committed
fix(iceberg-datafusion): handle timestamp predicates from DF
DataFusion sometimes passes dates as string literals, but can also pass timestamp ScalarValues, which need to be converted to predicates correctly in order to enable partition pruning.
1 parent c04b7ac commit 5ce0ba3

File tree

4 files changed

+369
-5
lines changed

4 files changed

+369
-5
lines changed

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/iceberg/src/spec/values.rs‎

Lines changed: 227 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1199,11 +1199,57 @@ impl Datum {
11991199
(PrimitiveLiteral::Long(val), _, PrimitiveType::Int) => {
12001200
Ok(Datum::i64_to_i32(*val))
12011201
}
1202-
(PrimitiveLiteral::Long(val), _, PrimitiveType::Timestamp) => {
1203-
Ok(Datum::timestamp_micros(*val))
1204-
}
1205-
(PrimitiveLiteral::Long(val), _, PrimitiveType::Timestamptz) => {
1206-
Ok(Datum::timestamptz_micros(*val))
1202+
(PrimitiveLiteral::Long(val), source_type, target_type) => {
1203+
match (source_type, target_type) {
1204+
(_, PrimitiveType::Long) => Ok(Datum::long(*val)),
1205+
(
1206+
PrimitiveType::Long
1207+
| PrimitiveType::Timestamp
1208+
| PrimitiveType::Timestamptz,
1209+
PrimitiveType::Timestamp,
1210+
) => Ok(Datum::timestamp_micros(*val)),
1211+
(
1212+
PrimitiveType::Long
1213+
| PrimitiveType::Timestamp
1214+
| PrimitiveType::Timestamptz,
1215+
PrimitiveType::Timestamptz,
1216+
) => Ok(Datum::timestamptz_micros(*val)),
1217+
(
1218+
PrimitiveType::Long
1219+
| PrimitiveType::TimestampNs
1220+
| PrimitiveType::TimestamptzNs,
1221+
PrimitiveType::TimestampNs,
1222+
) => Ok(Datum::timestamp_nanos(*val)),
1223+
(
1224+
PrimitiveType::Long
1225+
| PrimitiveType::TimestampNs
1226+
| PrimitiveType::TimestamptzNs,
1227+
PrimitiveType::TimestamptzNs,
1228+
) => Ok(Datum::timestamptz_nanos(*val)),
1229+
(
1230+
PrimitiveType::TimestampNs | PrimitiveType::TimestamptzNs,
1231+
PrimitiveType::Timestamp,
1232+
) => Ok(Datum::timestamp_micros(val / 1000)),
1233+
(
1234+
PrimitiveType::TimestampNs | PrimitiveType::TimestamptzNs,
1235+
PrimitiveType::Timestamptz,
1236+
) => Ok(Datum::timestamptz_micros(val / 1000)),
1237+
(
1238+
PrimitiveType::Timestamp | PrimitiveType::Timestamptz,
1239+
PrimitiveType::TimestampNs,
1240+
) => Ok(Datum::timestamp_nanos(val * 1000)),
1241+
(
1242+
PrimitiveType::Timestamp | PrimitiveType::Timestamptz,
1243+
PrimitiveType::TimestamptzNs,
1244+
) => Ok(Datum::timestamptz_nanos(val * 1000)),
1245+
_ => Err(Error::new(
1246+
ErrorKind::DataInvalid,
1247+
format!(
1248+
"Can't convert datum from {} type to {} type.",
1249+
self.r#type, target_primitive_type
1250+
),
1251+
)),
1252+
}
12071253
}
12081254
// Let's wait with nano's until this clears up: https://github.com/apache/iceberg/pull/11775
12091255
(PrimitiveLiteral::Int128(val), _, PrimitiveType::Long) => {
@@ -3954,4 +4000,180 @@ mod tests {
39544000

39554001
assert_eq!(double_sorted, double_expected);
39564002
}
4003+
4004+
#[test]
4005+
fn test_datum_timestamp_nanos_convert_to_timestamp_micros() {
4006+
let datum = Datum::timestamp_nanos(12345000);
4007+
4008+
let result = datum.to(&Primitive(PrimitiveType::Timestamp)).unwrap();
4009+
4010+
let expected = Datum::timestamp_micros(12345);
4011+
4012+
assert_eq!(result, expected);
4013+
}
4014+
4015+
#[test]
4016+
fn test_datum_timestamp_nanos_convert_to_timestamptz_micros() {
4017+
let datum = Datum::timestamp_nanos(12345000);
4018+
4019+
let result = datum.to(&Primitive(PrimitiveType::Timestamptz)).unwrap();
4020+
4021+
let expected = Datum::timestamptz_micros(12345);
4022+
4023+
assert_eq!(result, expected);
4024+
}
4025+
4026+
#[test]
4027+
fn test_datum_timestamptz_nanos_convert_to_timestamp_micros() {
4028+
let datum = Datum::timestamptz_nanos(12345000);
4029+
4030+
let result = datum.to(&Primitive(PrimitiveType::Timestamp)).unwrap();
4031+
4032+
let expected = Datum::timestamp_micros(12345);
4033+
4034+
assert_eq!(result, expected);
4035+
}
4036+
4037+
#[test]
4038+
fn test_datum_timestamptz_nanos_convert_to_timestamptz_micros() {
4039+
let datum = Datum::timestamptz_nanos(12345000);
4040+
4041+
let result = datum.to(&Primitive(PrimitiveType::Timestamptz)).unwrap();
4042+
4043+
let expected = Datum::timestamptz_micros(12345);
4044+
4045+
assert_eq!(result, expected);
4046+
}
4047+
4048+
#[test]
4049+
fn test_datum_timestamp_micros_convert_to_timestamp_nanos() {
4050+
let datum = Datum::timestamp_micros(12345);
4051+
4052+
let result = datum.to(&Primitive(PrimitiveType::TimestampNs)).unwrap();
4053+
4054+
let expected = Datum::timestamp_nanos(12345000);
4055+
4056+
assert_eq!(result, expected);
4057+
}
4058+
4059+
#[test]
4060+
fn test_datum_timestamp_micros_convert_to_timestamptz_nanos() {
4061+
let datum = Datum::timestamp_micros(12345);
4062+
4063+
let result = datum.to(&Primitive(PrimitiveType::TimestamptzNs)).unwrap();
4064+
4065+
let expected = Datum::timestamptz_nanos(12345000);
4066+
4067+
assert_eq!(result, expected);
4068+
}
4069+
4070+
#[test]
4071+
fn test_datum_timestamptz_micros_convert_to_timestamp_nanos() {
4072+
let datum = Datum::timestamptz_micros(12345);
4073+
4074+
let result = datum.to(&Primitive(PrimitiveType::TimestampNs)).unwrap();
4075+
4076+
let expected = Datum::timestamp_nanos(12345000);
4077+
4078+
assert_eq!(result, expected);
4079+
}
4080+
4081+
#[test]
4082+
fn test_datum_timestamptz_micros_convert_to_timestamptz_nanos() {
4083+
let datum = Datum::timestamptz_micros(12345);
4084+
4085+
let result = datum.to(&Primitive(PrimitiveType::TimestamptzNs)).unwrap();
4086+
4087+
let expected = Datum::timestamptz_nanos(12345000);
4088+
4089+
assert_eq!(result, expected);
4090+
}
4091+
4092+
#[test]
4093+
fn test_datum_timestamp_nanos_convert_to_timestamp_nanos() {
4094+
let datum = Datum::timestamp_nanos(12345);
4095+
4096+
let result = datum.to(&Primitive(PrimitiveType::TimestampNs)).unwrap();
4097+
4098+
let expected = Datum::timestamp_nanos(12345);
4099+
4100+
assert_eq!(result, expected);
4101+
}
4102+
4103+
#[test]
4104+
fn test_datum_timestamp_nanos_convert_to_timestamptz_nanos() {
4105+
let datum = Datum::timestamp_nanos(12345);
4106+
4107+
let result = datum.to(&Primitive(PrimitiveType::TimestamptzNs)).unwrap();
4108+
4109+
let expected = Datum::timestamptz_nanos(12345);
4110+
4111+
assert_eq!(result, expected);
4112+
}
4113+
4114+
#[test]
4115+
fn test_datum_timestamptz_nanos_convert_to_timestamp_nanos() {
4116+
let datum = Datum::timestamptz_nanos(12345);
4117+
4118+
let result = datum.to(&Primitive(PrimitiveType::TimestampNs)).unwrap();
4119+
4120+
let expected = Datum::timestamp_nanos(12345);
4121+
4122+
assert_eq!(result, expected);
4123+
}
4124+
4125+
#[test]
4126+
fn test_datum_timestamptz_nanos_convert_to_timestamptz_nanos() {
4127+
let datum = Datum::timestamptz_nanos(12345);
4128+
4129+
let result = datum.to(&Primitive(PrimitiveType::TimestamptzNs)).unwrap();
4130+
4131+
let expected = Datum::timestamptz_nanos(12345);
4132+
4133+
assert_eq!(result, expected);
4134+
}
4135+
4136+
#[test]
4137+
fn test_datum_long_convert_to_timestamp_nanos() {
4138+
let datum = Datum::long(12345);
4139+
4140+
let result = datum.to(&Primitive(PrimitiveType::TimestampNs)).unwrap();
4141+
4142+
let expected = Datum::timestamp_nanos(12345);
4143+
4144+
assert_eq!(result, expected);
4145+
}
4146+
4147+
#[test]
4148+
fn test_datum_long_convert_to_timestamptz_nanos() {
4149+
let datum = Datum::long(12345);
4150+
4151+
let result = datum.to(&Primitive(PrimitiveType::TimestamptzNs)).unwrap();
4152+
4153+
let expected = Datum::timestamptz_nanos(12345);
4154+
4155+
assert_eq!(result, expected);
4156+
}
4157+
4158+
#[test]
4159+
fn test_datum_timestamp_nanos_to_micros() {
4160+
let datum = Datum::timestamp_nanos(12345678);
4161+
4162+
let result = datum.to(&Primitive(PrimitiveType::Timestamp)).unwrap();
4163+
4164+
let expected = Datum::timestamp_micros(12345);
4165+
4166+
assert_eq!(result, expected);
4167+
}
4168+
4169+
#[test]
4170+
fn test_datum_timestamp_micros_to_nanos() {
4171+
let datum = Datum::timestamp_micros(12345);
4172+
4173+
let result = datum.to(&Primitive(PrimitiveType::TimestampNs)).unwrap();
4174+
4175+
let expected = Datum::timestamp_nanos(12345000);
4176+
4177+
assert_eq!(result, expected);
4178+
}
39574179
}

‎crates/integrations/datafusion/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ repository = { workspace = true }
3131
[dependencies]
3232
anyhow = { workspace = true }
3333
async-trait = { workspace = true }
34+
chrono = { workspace = true }
3435
datafusion = { workspace = true }
3536
futures = { workspace = true }
3637
iceberg = { workspace = true }

0 commit comments

Comments
 (0)