Skip to content

Commit 493d3ee

Browse files
friendlymatthewfriendlymatthew
authored
Support casting Date to a time zone-specific timestamp (#7141)
* Date can cast to tz-aware timestamp * Cover other Date x TimeUnit cases * Preserve local time when casting date to timestamp
1 parent a0c3186 commit 493d3ee

File tree

1 file changed

+242
-32
lines changed

1 file changed

+242
-32
lines changed

‎arrow-cast/src/cast/mod.rs‎

Lines changed: 242 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -265,8 +265,8 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
265265
}
266266
(Timestamp(_, _), _) if to_type.is_numeric() => true,
267267
(_, Timestamp(_, _)) if from_type.is_numeric() => true,
268-
(Date64, Timestamp(_, None)) => true,
269-
(Date32, Timestamp(_, None)) => true,
268+
(Date64, Timestamp(_, _)) => true,
269+
(Date32, Timestamp(_, _)) => true,
270270
(
271271
Timestamp(_, _),
272272
Timestamp(_, _)
@@ -1806,44 +1806,63 @@ pub fn cast_with_options(
18061806
})?,
18071807
))
18081808
}
1809-
(Date64, Timestamp(TimeUnit::Second, None)) => Ok(Arc::new(
1810-
array
1809+
(Date64, Timestamp(TimeUnit::Second, _)) => {
1810+
let array = array
18111811
.as_primitive::<Date64Type>()
1812-
.unary::<_, TimestampSecondType>(|x| x / MILLISECONDS),
1813-
)),
1814-
(Date64, Timestamp(TimeUnit::Millisecond, None)) => {
1815-
cast_reinterpret_arrays::<Date64Type, TimestampMillisecondType>(array)
1812+
.unary::<_, TimestampSecondType>(|x| x / MILLISECONDS);
1813+
1814+
cast_with_options(&array, to_type, cast_options)
18161815
}
1817-
(Date64, Timestamp(TimeUnit::Microsecond, None)) => Ok(Arc::new(
1818-
array
1816+
(Date64, Timestamp(TimeUnit::Millisecond, _)) => {
1817+
let array = array
18191818
.as_primitive::<Date64Type>()
1820-
.unary::<_, TimestampMicrosecondType>(|x| x * (MICROSECONDS / MILLISECONDS)),
1821-
)),
1822-
(Date64, Timestamp(TimeUnit::Nanosecond, None)) => Ok(Arc::new(
1823-
array
1819+
.reinterpret_cast::<TimestampMillisecondType>();
1820+
1821+
cast_with_options(&array, to_type, cast_options)
1822+
}
1823+
1824+
(Date64, Timestamp(TimeUnit::Microsecond, _)) => {
1825+
let array = array
18241826
.as_primitive::<Date64Type>()
1825-
.unary::<_, TimestampNanosecondType>(|x| x * (NANOSECONDS / MILLISECONDS)),
1826-
)),
1827-
(Date32, Timestamp(TimeUnit::Second, None)) => Ok(Arc::new(
1828-
array
1827+
.unary::<_, TimestampMicrosecondType>(|x| x * (MICROSECONDS / MILLISECONDS));
1828+
1829+
cast_with_options(&array, to_type, cast_options)
1830+
}
1831+
(Date64, Timestamp(TimeUnit::Nanosecond, _)) => {
1832+
let array = array
1833+
.as_primitive::<Date64Type>()
1834+
.unary::<_, TimestampNanosecondType>(|x| x * (NANOSECONDS / MILLISECONDS));
1835+
1836+
cast_with_options(&array, to_type, cast_options)
1837+
}
1838+
(Date32, Timestamp(TimeUnit::Second, _)) => {
1839+
let array = array
18291840
.as_primitive::<Date32Type>()
1830-
.unary::<_, TimestampSecondType>(|x| (x as i64) * SECONDS_IN_DAY),
1831-
)),
1832-
(Date32, Timestamp(TimeUnit::Millisecond, None)) => Ok(Arc::new(
1833-
array
1841+
.unary::<_, TimestampSecondType>(|x| (x as i64) * SECONDS_IN_DAY);
1842+
1843+
cast_with_options(&array, to_type, cast_options)
1844+
}
1845+
(Date32, Timestamp(TimeUnit::Millisecond, _)) => {
1846+
let array = array
18341847
.as_primitive::<Date32Type>()
1835-
.unary::<_, TimestampMillisecondType>(|x| (x as i64) * MILLISECONDS_IN_DAY),
1836-
)),
1837-
(Date32, Timestamp(TimeUnit::Microsecond, None)) => Ok(Arc::new(
1838-
array
1848+
.unary::<_, TimestampMillisecondType>(|x| (x as i64) * MILLISECONDS_IN_DAY);
1849+
1850+
cast_with_options(&array, to_type, cast_options)
1851+
}
1852+
(Date32, Timestamp(TimeUnit::Microsecond, _)) => {
1853+
let array = array
18391854
.as_primitive::<Date32Type>()
1840-
.unary::<_, TimestampMicrosecondType>(|x| (x as i64) * MICROSECONDS_IN_DAY),
1841-
)),
1842-
(Date32, Timestamp(TimeUnit::Nanosecond, None)) => Ok(Arc::new(
1843-
array
1855+
.unary::<_, TimestampMicrosecondType>(|x| (x as i64) * MICROSECONDS_IN_DAY);
1856+
1857+
cast_with_options(&array, to_type, cast_options)
1858+
}
1859+
(Date32, Timestamp(TimeUnit::Nanosecond, _)) => {
1860+
let array = array
18441861
.as_primitive::<Date32Type>()
1845-
.unary::<_, TimestampNanosecondType>(|x| (x as i64) * NANOSECONDS_IN_DAY),
1846-
)),
1862+
.unary::<_, TimestampNanosecondType>(|x| (x as i64) * NANOSECONDS_IN_DAY);
1863+
1864+
cast_with_options(&array, to_type, cast_options)
1865+
}
18471866

18481867
(_, Duration(unit)) if from_type.is_numeric() => {
18491868
let array = cast_with_options(array, &Int64, cast_options)?;
@@ -5217,6 +5236,197 @@ mod tests {
52175236
}};
52185237
}
52195238

5239+
#[test]
5240+
fn test_cast_date32_to_timestamp_and_timestamp_with_timezone() {
5241+
let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu
5242+
let a = Date32Array::from(vec![Some(18628), None, None]); // 2021-1-1, 2022-1-1
5243+
let array = Arc::new(a) as ArrayRef;
5244+
5245+
let b = cast(
5246+
&array,
5247+
&DataType::Timestamp(TimeUnit::Second, Some(tz.into())),
5248+
)
5249+
.unwrap();
5250+
let c = b.as_primitive::<TimestampSecondType>();
5251+
let string_array = cast(&c, &DataType::Utf8).unwrap();
5252+
let result = string_array.as_string::<i32>();
5253+
assert_eq!("2021-01-01T00:00:00+05:45", result.value(0));
5254+
5255+
let b = cast(&array, &DataType::Timestamp(TimeUnit::Second, None)).unwrap();
5256+
let c = b.as_primitive::<TimestampSecondType>();
5257+
let string_array = cast(&c, &DataType::Utf8).unwrap();
5258+
let result = string_array.as_string::<i32>();
5259+
assert_eq!("2021-01-01T00:00:00", result.value(0));
5260+
}
5261+
5262+
#[test]
5263+
fn test_cast_date32_to_timestamp_with_timezone() {
5264+
let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu
5265+
let a = Date32Array::from(vec![Some(18628), Some(18993), None]); // 2021-1-1, 2022-1-1
5266+
let array = Arc::new(a) as ArrayRef;
5267+
let b = cast(
5268+
&array,
5269+
&DataType::Timestamp(TimeUnit::Second, Some(tz.into())),
5270+
)
5271+
.unwrap();
5272+
let c = b.as_primitive::<TimestampSecondType>();
5273+
assert_eq!(1609438500, c.value(0));
5274+
assert_eq!(1640974500, c.value(1));
5275+
assert!(c.is_null(2));
5276+
5277+
let string_array = cast(&c, &DataType::Utf8).unwrap();
5278+
let result = string_array.as_string::<i32>();
5279+
assert_eq!("2021-01-01T00:00:00+05:45", result.value(0));
5280+
assert_eq!("2022-01-01T00:00:00+05:45", result.value(1));
5281+
}
5282+
5283+
#[test]
5284+
fn test_cast_date32_to_timestamp_with_timezone_ms() {
5285+
let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu
5286+
let a = Date32Array::from(vec![Some(18628), Some(18993), None]); // 2021-1-1, 2022-1-1
5287+
let array = Arc::new(a) as ArrayRef;
5288+
let b = cast(
5289+
&array,
5290+
&DataType::Timestamp(TimeUnit::Millisecond, Some(tz.into())),
5291+
)
5292+
.unwrap();
5293+
let c = b.as_primitive::<TimestampMillisecondType>();
5294+
assert_eq!(1609438500000, c.value(0));
5295+
assert_eq!(1640974500000, c.value(1));
5296+
assert!(c.is_null(2));
5297+
5298+
let string_array = cast(&c, &DataType::Utf8).unwrap();
5299+
let result = string_array.as_string::<i32>();
5300+
assert_eq!("2021-01-01T00:00:00+05:45", result.value(0));
5301+
assert_eq!("2022-01-01T00:00:00+05:45", result.value(1));
5302+
}
5303+
5304+
#[test]
5305+
fn test_cast_date32_to_timestamp_with_timezone_us() {
5306+
let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu
5307+
let a = Date32Array::from(vec![Some(18628), Some(18993), None]); // 2021-1-1, 2022-1-1
5308+
let array = Arc::new(a) as ArrayRef;
5309+
let b = cast(
5310+
&array,
5311+
&DataType::Timestamp(TimeUnit::Microsecond, Some(tz.into())),
5312+
)
5313+
.unwrap();
5314+
let c = b.as_primitive::<TimestampMicrosecondType>();
5315+
assert_eq!(1609438500000000, c.value(0));
5316+
assert_eq!(1640974500000000, c.value(1));
5317+
assert!(c.is_null(2));
5318+
5319+
let string_array = cast(&c, &DataType::Utf8).unwrap();
5320+
let result = string_array.as_string::<i32>();
5321+
assert_eq!("2021-01-01T00:00:00+05:45", result.value(0));
5322+
assert_eq!("2022-01-01T00:00:00+05:45", result.value(1));
5323+
}
5324+
5325+
#[test]
5326+
fn test_cast_date32_to_timestamp_with_timezone_ns() {
5327+
let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu
5328+
let a = Date32Array::from(vec![Some(18628), Some(18993), None]); // 2021-1-1, 2022-1-1
5329+
let array = Arc::new(a) as ArrayRef;
5330+
let b = cast(
5331+
&array,
5332+
&DataType::Timestamp(TimeUnit::Nanosecond, Some(tz.into())),
5333+
)
5334+
.unwrap();
5335+
let c = b.as_primitive::<TimestampNanosecondType>();
5336+
assert_eq!(1609438500000000000, c.value(0));
5337+
assert_eq!(1640974500000000000, c.value(1));
5338+
assert!(c.is_null(2));
5339+
5340+
let string_array = cast(&c, &DataType::Utf8).unwrap();
5341+
let result = string_array.as_string::<i32>();
5342+
assert_eq!("2021-01-01T00:00:00+05:45", result.value(0));
5343+
assert_eq!("2022-01-01T00:00:00+05:45", result.value(1));
5344+
}
5345+
5346+
#[test]
5347+
fn test_cast_date64_to_timestamp_with_timezone() {
5348+
let array = Date64Array::from(vec![Some(864000000005), Some(1545696000001), None]);
5349+
let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu
5350+
let b = cast(
5351+
&array,
5352+
&DataType::Timestamp(TimeUnit::Second, Some(tz.into())),
5353+
)
5354+
.unwrap();
5355+
5356+
let c = b.as_primitive::<TimestampSecondType>();
5357+
assert_eq!(863979300, c.value(0));
5358+
assert_eq!(1545675300, c.value(1));
5359+
assert!(c.is_null(2));
5360+
5361+
let string_array = cast(&c, &DataType::Utf8).unwrap();
5362+
let result = string_array.as_string::<i32>();
5363+
assert_eq!("1997-05-19T00:00:00+05:45", result.value(0));
5364+
assert_eq!("2018-12-25T00:00:00+05:45", result.value(1));
5365+
}
5366+
5367+
#[test]
5368+
fn test_cast_date64_to_timestamp_with_timezone_ms() {
5369+
let array = Date64Array::from(vec![Some(864000000005), Some(1545696000001), None]);
5370+
let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu
5371+
let b = cast(
5372+
&array,
5373+
&DataType::Timestamp(TimeUnit::Millisecond, Some(tz.into())),
5374+
)
5375+
.unwrap();
5376+
5377+
let c = b.as_primitive::<TimestampMillisecondType>();
5378+
assert_eq!(863979300005, c.value(0));
5379+
assert_eq!(1545675300001, c.value(1));
5380+
assert!(c.is_null(2));
5381+
5382+
let string_array = cast(&c, &DataType::Utf8).unwrap();
5383+
let result = string_array.as_string::<i32>();
5384+
assert_eq!("1997-05-19T00:00:00.005+05:45", result.value(0));
5385+
assert_eq!("2018-12-25T00:00:00.001+05:45", result.value(1));
5386+
}
5387+
5388+
#[test]
5389+
fn test_cast_date64_to_timestamp_with_timezone_us() {
5390+
let array = Date64Array::from(vec![Some(864000000005), Some(1545696000001), None]);
5391+
let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu
5392+
let b = cast(
5393+
&array,
5394+
&DataType::Timestamp(TimeUnit::Microsecond, Some(tz.into())),
5395+
)
5396+
.unwrap();
5397+
5398+
let c = b.as_primitive::<TimestampMicrosecondType>();
5399+
assert_eq!(863979300005000, c.value(0));
5400+
assert_eq!(1545675300001000, c.value(1));
5401+
assert!(c.is_null(2));
5402+
5403+
let string_array = cast(&c, &DataType::Utf8).unwrap();
5404+
let result = string_array.as_string::<i32>();
5405+
assert_eq!("1997-05-19T00:00:00.005+05:45", result.value(0));
5406+
assert_eq!("2018-12-25T00:00:00.001+05:45", result.value(1));
5407+
}
5408+
5409+
#[test]
5410+
fn test_cast_date64_to_timestamp_with_timezone_ns() {
5411+
let array = Date64Array::from(vec![Some(864000000005), Some(1545696000001), None]);
5412+
let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu
5413+
let b = cast(
5414+
&array,
5415+
&DataType::Timestamp(TimeUnit::Nanosecond, Some(tz.into())),
5416+
)
5417+
.unwrap();
5418+
5419+
let c = b.as_primitive::<TimestampNanosecondType>();
5420+
assert_eq!(863979300005000000, c.value(0));
5421+
assert_eq!(1545675300001000000, c.value(1));
5422+
assert!(c.is_null(2));
5423+
5424+
let string_array = cast(&c, &DataType::Utf8).unwrap();
5425+
let result = string_array.as_string::<i32>();
5426+
assert_eq!("1997-05-19T00:00:00.005+05:45", result.value(0));
5427+
assert_eq!("2018-12-25T00:00:00.001+05:45", result.value(1));
5428+
}
5429+
52205430
#[test]
52215431
fn test_cast_timestamp_to_strings() {
52225432
// "2018-12-25T00:00:02.001", "1997-05-19T00:00:03.005", None

0 commit comments

Comments
 (0)