Skip to content

Commit

Permalink
Cast from integer to timestamp
Browse files Browse the repository at this point in the history
  • Loading branch information
@viirya
viirya committed Nov 5, 2023
1 parent 7705aca commit 06084d0
Showing 1 changed file with 160 additions and 18 deletions.
178 changes: 160 additions & 18 deletions arrow-cast/src/cast.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -227,8 +227,8 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
(Time64(_), Time32(to_unit)) => {
matches!(to_unit, Second | Millisecond)
}
(Timestamp(_, _), Int64) => true,
(Int64, Timestamp(_, _)) => true,
(Timestamp(_, _), _) if to_type.is_integer() => true,
(_, Timestamp(_, _)) if from_type.is_integer() => true,
(Date64, Timestamp(_, None)) => true,
(Date32, Timestamp(_, None)) => true,
(
Expand Down Expand Up @@ -1621,24 +1621,104 @@ pub fn cast_with_options(
.unary::<_, Time64MicrosecondType>(|x| x / (NANOSECONDS / MICROSECONDS)),
)),

(Timestamp(TimeUnit::Second, _), Int64) => {
cast_reinterpret_arrays::<TimestampSecondType, Int64Type>(array)
(Timestamp(TimeUnit::Second, _), _) if to_type.is_integer() => {
let array = cast_reinterpret_arrays::<TimestampSecondType, Int64Type>(array)?;

match to_type {
Int8 => cast_numeric_arrays::<Int64Type, Int8Type>(&array, cast_options),
Int16 => cast_numeric_arrays::<Int64Type, Int16Type>(&array, cast_options),
Int32 => cast_numeric_arrays::<Int64Type, Int32Type>(&array, cast_options),
Int64 => Ok(array),
UInt8 => cast_numeric_arrays::<Int64Type, UInt8Type>(&array, cast_options),
UInt16 => cast_numeric_arrays::<Int64Type, UInt16Type>(&array, cast_options),
UInt32 => cast_numeric_arrays::<Int64Type, UInt32Type>(&array, cast_options),
UInt64 => cast_numeric_arrays::<Int64Type, UInt64Type>(&array, cast_options),
_ => Err(ArrowError::CastError(format!(
"Casting from {:?} to {:?} not supported",
from_type, to_type
))),
}
}
(Timestamp(TimeUnit::Millisecond, _), Int64) => {
cast_reinterpret_arrays::<TimestampMillisecondType, Int64Type>(array)
(Timestamp(TimeUnit::Millisecond, _), _) if to_type.is_integer() => {
let array = cast_reinterpret_arrays::<TimestampMillisecondType, Int64Type>(array)?;

match to_type {
Int8 => cast_numeric_arrays::<Int64Type, Int8Type>(&array, cast_options),
Int16 => cast_numeric_arrays::<Int64Type, Int16Type>(&array, cast_options),
Int32 => cast_numeric_arrays::<Int64Type, Int32Type>(&array, cast_options),
Int64 => Ok(array),
UInt8 => cast_numeric_arrays::<Int64Type, UInt8Type>(&array, cast_options),
UInt16 => cast_numeric_arrays::<Int64Type, UInt16Type>(&array, cast_options),
UInt32 => cast_numeric_arrays::<Int64Type, UInt32Type>(&array, cast_options),
UInt64 => cast_numeric_arrays::<Int64Type, UInt64Type>(&array, cast_options),
_ => Err(ArrowError::CastError(format!(
"Casting from {:?} to {:?} not supported",
from_type, to_type
))),
}
}
(Timestamp(TimeUnit::Microsecond, _), Int64) => {
cast_reinterpret_arrays::<TimestampMicrosecondType, Int64Type>(array)
(Timestamp(TimeUnit::Microsecond, _), _) if to_type.is_integer() => {
let array = cast_reinterpret_arrays::<TimestampMicrosecondType, Int64Type>(array)?;

match to_type {
Int8 => cast_numeric_arrays::<Int64Type, Int8Type>(&array, cast_options),
Int16 => cast_numeric_arrays::<Int64Type, Int16Type>(&array, cast_options),
Int32 => cast_numeric_arrays::<Int64Type, Int32Type>(&array, cast_options),
Int64 => Ok(array),
UInt8 => cast_numeric_arrays::<Int64Type, UInt8Type>(&array, cast_options),
UInt16 => cast_numeric_arrays::<Int64Type, UInt16Type>(&array, cast_options),
UInt32 => cast_numeric_arrays::<Int64Type, UInt32Type>(&array, cast_options),
UInt64 => cast_numeric_arrays::<Int64Type, UInt64Type>(&array, cast_options),
_ => Err(ArrowError::CastError(format!(
"Casting from {:?} to {:?} not supported",
from_type, to_type
))),
}
}
(Timestamp(TimeUnit::Nanosecond, _), Int64) => {
cast_reinterpret_arrays::<TimestampNanosecondType, Int64Type>(array)
(Timestamp(TimeUnit::Nanosecond, _), _) if to_type.is_integer() => {
let array = cast_reinterpret_arrays::<TimestampNanosecondType, Int64Type>(array)?;

match to_type {
Int8 => cast_numeric_arrays::<Int64Type, Int8Type>(&array, cast_options),
Int16 => cast_numeric_arrays::<Int64Type, Int16Type>(&array, cast_options),
Int32 => cast_numeric_arrays::<Int64Type, Int32Type>(&array, cast_options),
Int64 => Ok(array),
UInt8 => cast_numeric_arrays::<Int64Type, UInt8Type>(&array, cast_options),
UInt16 => cast_numeric_arrays::<Int64Type, UInt16Type>(&array, cast_options),
UInt32 => cast_numeric_arrays::<Int64Type, UInt32Type>(&array, cast_options),
UInt64 => cast_numeric_arrays::<Int64Type, UInt64Type>(&array, cast_options),
_ => Err(ArrowError::CastError(format!(
"Casting from {:?} to {:?} not supported",
from_type, to_type
))),
}
}

(Int64, Timestamp(unit, tz)) => Ok(make_timestamp_array(
array.as_primitive(),
unit.clone(),
tz.clone(),
)),
(_, Timestamp(unit, tz)) if from_type.is_integer() => {
let array = match from_type {
Int8 => cast_numeric_arrays::<Int8Type, Int64Type>(array, cast_options),
Int16 => cast_numeric_arrays::<Int16Type, Int64Type>(array, cast_options),
Int32 => cast_numeric_arrays::<Int32Type, Int64Type>(array, cast_options),
Int64 => {
return Ok(make_timestamp_array(
array.as_primitive(),
unit.clone(),
tz.clone(),
))
}
UInt8 => cast_numeric_arrays::<UInt8Type, Int64Type>(array, cast_options),
UInt16 => cast_numeric_arrays::<UInt16Type, Int64Type>(array, cast_options),
UInt32 => cast_numeric_arrays::<UInt32Type, Int64Type>(array, cast_options),
UInt64 => cast_numeric_arrays::<UInt64Type, Int64Type>(array, cast_options),
_ => unreachable!(),
}?;

Ok(make_timestamp_array(
array.as_primitive(),
unit.clone(),
tz.clone(),
))
}

(Timestamp(from_unit, from_tz), Timestamp(to_unit, to_tz)) => {
let array = cast_with_options(array, &Int64, cast_options)?;
Expand Down Expand Up @@ -4559,10 +4639,72 @@ mod tests {
}

#[test]
#[should_panic(expected = "Casting from Int32 to Timestamp(Microsecond, None) not supported")]
fn test_cast_int32_to_timestamp() {
fn test_cast_integer_to_timestamp() {
let array = Int64Array::from(vec![Some(2), Some(10), None]);
let expected = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();

let array = Int8Array::from(vec![Some(2), Some(10), None]);
let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();

assert_eq!(&actual, &expected);

let array = Int16Array::from(vec![Some(2), Some(10), None]);
let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();

assert_eq!(&actual, &expected);

let array = Int32Array::from(vec![Some(2), Some(10), None]);
cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();

assert_eq!(&actual, &expected);

let array = UInt8Array::from(vec![Some(2), Some(10), None]);
let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();

assert_eq!(&actual, &expected);

let array = UInt16Array::from(vec![Some(2), Some(10), None]);
let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();

assert_eq!(&actual, &expected);

let array = UInt32Array::from(vec![Some(2), Some(10), None]);
let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();

assert_eq!(&actual, &expected);

let array = UInt64Array::from(vec![Some(2), Some(10), None]);
let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();

assert_eq!(&actual, &expected);
}

#[test]
fn test_cast_timestamp_to_integer() {
let array = TimestampMillisecondArray::from(vec![Some(5), Some(1), None])
.with_timezone("UTC".to_string());
let expected = cast(&array, &DataType::Int64).unwrap();

let actual = cast(&cast(&array, &DataType::Int8).unwrap(), &DataType::Int64).unwrap();
assert_eq!(&actual, &expected);

let actual = cast(&cast(&array, &DataType::Int16).unwrap(), &DataType::Int64).unwrap();
assert_eq!(&actual, &expected);

let actual = cast(&cast(&array, &DataType::Int32).unwrap(), &DataType::Int64).unwrap();
assert_eq!(&actual, &expected);

let actual = cast(&cast(&array, &DataType::UInt8).unwrap(), &DataType::Int64).unwrap();
assert_eq!(&actual, &expected);

let actual = cast(&cast(&array, &DataType::UInt16).unwrap(), &DataType::Int64).unwrap();
assert_eq!(&actual, &expected);

let actual = cast(&cast(&array, &DataType::UInt32).unwrap(), &DataType::Int64).unwrap();
assert_eq!(&actual, &expected);

let actual = cast(&cast(&array, &DataType::UInt64).unwrap(), &DataType::Int64).unwrap();
assert_eq!(&actual, &expected);
}

#[test]
Expand Down

0 comments on commit 06084d0

Please sign in to comment.