Skip to content

Commit 435b53d

Browse files
viiryaviirya
authored
Cast from numeric/timestamp to timestamp/numeric (#5123)
* Casting between floating and timestamp * Fix * For decimals * Fix
1 parent d5a6cf4 commit 435b53d

1 file changed

Lines changed: 89 additions & 15 deletions

File tree

‎arrow-cast/src/cast.rs‎

Lines changed: 89 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -161,17 +161,16 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
161161
(Decimal128(_, _) | Decimal256(_, _), Utf8 | LargeUtf8) => true,
162162
// Utf8 to decimal
163163
(Utf8 | LargeUtf8, Decimal128(_, _) | Decimal256(_, _)) => true,
164-
(Decimal128(_, _) | Decimal256(_, _), _) => false,
165-
(_, Decimal128(_, _) | Decimal256(_, _)) => false,
166164
(Struct(_), _) => false,
167165
(_, Struct(_)) => false,
168166
(_, Boolean) => {
169-
DataType::is_numeric(from_type)
167+
DataType::is_integer(from_type) ||
168+
DataType::is_floating(from_type)
170169
|| from_type == &Utf8
171170
|| from_type == &LargeUtf8
172171
}
173172
(Boolean, _) => {
174-
DataType::is_numeric(to_type) || to_type == &Utf8 || to_type == &LargeUtf8
173+
DataType::is_integer(to_type) || DataType::is_floating(to_type) || to_type == &Utf8 || to_type == &LargeUtf8
175174
}
176175

177176
(Binary, LargeBinary | Utf8 | LargeUtf8 | FixedSizeBinary(_)) => true,
@@ -222,8 +221,8 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
222221
(Time64(_), Time32(to_unit)) => {
223222
matches!(to_unit, Second | Millisecond)
224223
}
225-
(Timestamp(_, _), _) if to_type.is_integer() => true,
226-
(_, Timestamp(_, _)) if from_type.is_integer() => true,
224+
(Timestamp(_, _), _) if to_type.is_numeric() && to_type != &Float16 => true,
225+
(_, Timestamp(_, _)) if from_type.is_numeric() && from_type != &Float16 => true,
227226
(Date64, Timestamp(_, None)) => true,
228227
(Date32, Timestamp(_, None)) => true,
229228
(
@@ -849,7 +848,7 @@ pub fn cast_with_options(
849848
cast_options,
850849
)
851850
}
852-
(Decimal128(_, scale), _) => {
851+
(Decimal128(_, scale), _) if !to_type.is_temporal() => {
853852
// cast decimal to other type
854853
match to_type {
855854
UInt8 => cast_decimal_to_integer::<Decimal128Type, UInt8Type>(
@@ -914,7 +913,7 @@ pub fn cast_with_options(
914913
))),
915914
}
916915
}
917-
(Decimal256(_, scale), _) => {
916+
(Decimal256(_, scale), _) if !to_type.is_temporal() => {
918917
// cast decimal to other type
919918
match to_type {
920919
UInt8 => cast_decimal_to_integer::<Decimal256Type, UInt8Type>(
@@ -979,7 +978,7 @@ pub fn cast_with_options(
979978
))),
980979
}
981980
}
982-
(_, Decimal128(precision, scale)) => {
981+
(_, Decimal128(precision, scale)) if !from_type.is_temporal() => {
983982
// cast data to decimal
984983
match from_type {
985984
UInt8 => cast_integer_to_decimal::<_, Decimal128Type, _>(
@@ -1068,7 +1067,7 @@ pub fn cast_with_options(
10681067
))),
10691068
}
10701069
}
1071-
(_, Decimal256(precision, scale)) => {
1070+
(_, Decimal256(precision, scale)) if !from_type.is_temporal() => {
10721071
// cast data to decimal
10731072
match from_type {
10741073
UInt8 => cast_integer_to_decimal::<_, Decimal256Type, _>(
@@ -1607,24 +1606,25 @@ pub fn cast_with_options(
16071606
.unary::<_, Time64MicrosecondType>(|x| x / (NANOSECONDS / MICROSECONDS)),
16081607
)),
16091608

1610-
(Timestamp(TimeUnit::Second, _), _) if to_type.is_integer() => {
1609+
// Timestamp to integer/floating/decimals
1610+
(Timestamp(TimeUnit::Second, _), _) if to_type.is_numeric() => {
16111611
let array = cast_reinterpret_arrays::<TimestampSecondType, Int64Type>(array)?;
16121612
cast_with_options(&array, to_type, cast_options)
16131613
}
1614-
(Timestamp(TimeUnit::Millisecond, _), _) if to_type.is_integer() => {
1614+
(Timestamp(TimeUnit::Millisecond, _), _) if to_type.is_numeric() => {
16151615
let array = cast_reinterpret_arrays::<TimestampMillisecondType, Int64Type>(array)?;
16161616
cast_with_options(&array, to_type, cast_options)
16171617
}
1618-
(Timestamp(TimeUnit::Microsecond, _), _) if to_type.is_integer() => {
1618+
(Timestamp(TimeUnit::Microsecond, _), _) if to_type.is_numeric() => {
16191619
let array = cast_reinterpret_arrays::<TimestampMicrosecondType, Int64Type>(array)?;
16201620
cast_with_options(&array, to_type, cast_options)
16211621
}
1622-
(Timestamp(TimeUnit::Nanosecond, _), _) if to_type.is_integer() => {
1622+
(Timestamp(TimeUnit::Nanosecond, _), _) if to_type.is_numeric() => {
16231623
let array = cast_reinterpret_arrays::<TimestampNanosecondType, Int64Type>(array)?;
16241624
cast_with_options(&array, to_type, cast_options)
16251625
}
16261626

1627-
(_, Timestamp(unit, tz)) if from_type.is_integer() => {
1627+
(_, Timestamp(unit, tz)) if from_type.is_numeric() => {
16281628
let array = cast_with_options(array, &Int64, cast_options)?;
16291629
Ok(make_timestamp_array(
16301630
array.as_primitive(),
@@ -4652,6 +4652,80 @@ mod tests {
46524652
assert_eq!(&actual, &expected);
46534653
}
46544654

4655+
#[test]
4656+
fn test_cast_floating_to_timestamp() {
4657+
let array = Int64Array::from(vec![Some(2), Some(10), None]);
4658+
let expected = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
4659+
4660+
let array = Float32Array::from(vec![Some(2.0), Some(10.6), None]);
4661+
let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
4662+
4663+
assert_eq!(&actual, &expected);
4664+
4665+
let array = Float64Array::from(vec![Some(2.1), Some(10.2), None]);
4666+
let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
4667+
4668+
assert_eq!(&actual, &expected);
4669+
}
4670+
4671+
#[test]
4672+
fn test_cast_timestamp_to_floating() {
4673+
let array = TimestampMillisecondArray::from(vec![Some(5), Some(1), None])
4674+
.with_timezone("UTC".to_string());
4675+
let expected = cast(&array, &DataType::Int64).unwrap();
4676+
4677+
let actual = cast(&cast(&array, &DataType::Float32).unwrap(), &DataType::Int64).unwrap();
4678+
assert_eq!(&actual, &expected);
4679+
4680+
let actual = cast(&cast(&array, &DataType::Float64).unwrap(), &DataType::Int64).unwrap();
4681+
assert_eq!(&actual, &expected);
4682+
}
4683+
4684+
#[test]
4685+
fn test_cast_decimal_to_timestamp() {
4686+
let array = Int64Array::from(vec![Some(2), Some(10), None]);
4687+
let expected = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
4688+
4689+
let array = Decimal128Array::from(vec![Some(200), Some(1000), None])
4690+
.with_precision_and_scale(4, 2)
4691+
.unwrap();
4692+
let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
4693+
4694+
assert_eq!(&actual, &expected);
4695+
4696+
let array = Decimal256Array::from(vec![
4697+
Some(i256::from_i128(2000)),
4698+
Some(i256::from_i128(10000)),
4699+
None,
4700+
])
4701+
.with_precision_and_scale(5, 3)
4702+
.unwrap();
4703+
let actual = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
4704+
4705+
assert_eq!(&actual, &expected);
4706+
}
4707+
4708+
#[test]
4709+
fn test_cast_timestamp_to_decimal() {
4710+
let array = TimestampMillisecondArray::from(vec![Some(5), Some(1), None])
4711+
.with_timezone("UTC".to_string());
4712+
let expected = cast(&array, &DataType::Int64).unwrap();
4713+
4714+
let actual = cast(
4715+
&cast(&array, &DataType::Decimal128(5, 2)).unwrap(),
4716+
&DataType::Int64,
4717+
)
4718+
.unwrap();
4719+
assert_eq!(&actual, &expected);
4720+
4721+
let actual = cast(
4722+
&cast(&array, &DataType::Decimal256(10, 5)).unwrap(),
4723+
&DataType::Int64,
4724+
)
4725+
.unwrap();
4726+
assert_eq!(&actual, &expected);
4727+
}
4728+
46554729
#[test]
46564730
fn test_cast_list_i32_to_list_u16() {
46574731
let value_data = Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2, 100000000]).into_data();

0 commit comments

Comments
 (0)