From 0bddfb747ec7cf3cf3db39f6bb3e9dc31212b4fe Mon Sep 17 00:00:00 2001 From: Kun Liu Date: Sat, 15 Jan 2022 01:09:40 +0800 Subject: [PATCH] support cast/try_cast for decimal: signed numeric to decimal (#1442) * support cast/try_cast for decimal: signed numeric to decimal * add test case * change test case * change dependency * remove useless code * restore dependency --- .../src/physical_plan/expressions/cast.rs | 128 +++++++++++++++++- .../src/physical_plan/expressions/try_cast.rs | 123 ++++++++++++++++- 2 files changed, 248 insertions(+), 3 deletions(-) diff --git a/datafusion/src/physical_plan/expressions/cast.rs b/datafusion/src/physical_plan/expressions/cast.rs index bba125ebdcc9..df687b4dd98b 100644 --- a/datafusion/src/physical_plan/expressions/cast.rs +++ b/datafusion/src/physical_plan/expressions/cast.rs @@ -162,7 +162,10 @@ mod tests { use crate::physical_plan::expressions::col; use arrow::array::{StringArray, Time64NanosecondArray}; use arrow::{ - array::{Array, Int32Array, Int64Array, TimestampNanosecondArray, UInt32Array}, + array::{ + Array, DecimalArray, Float32Array, Float64Array, Int16Array, Int32Array, + Int64Array, Int8Array, TimestampNanosecondArray, UInt32Array, + }, datatypes::*, }; @@ -217,6 +220,129 @@ mod tests { }}; } + #[test] + fn test_cast_numeric_to_decimal() -> Result<()> { + // int8 + generic_test_cast!( + Int8Array, + DataType::Int8, + vec![1, 2, 3, 4, 5], + DecimalArray, + DataType::Decimal(3, 0), + vec![ + Some(1_i128), + Some(2_i128), + Some(3_i128), + Some(4_i128), + Some(5_i128), + ], + DEFAULT_DATAFUSION_CAST_OPTIONS + ); + + // int16 + generic_test_cast!( + Int16Array, + DataType::Int16, + vec![1, 2, 3, 4, 5], + DecimalArray, + DataType::Decimal(5, 0), + vec![ + Some(1_i128), + Some(2_i128), + Some(3_i128), + Some(4_i128), + Some(5_i128), + ], + DEFAULT_DATAFUSION_CAST_OPTIONS + ); + + // int32 + generic_test_cast!( + Int32Array, + DataType::Int32, + vec![1, 2, 3, 4, 5], + DecimalArray, + DataType::Decimal(10, 0), + vec![ + Some(1_i128), + Some(2_i128), + Some(3_i128), + Some(4_i128), + Some(5_i128), + ], + DEFAULT_DATAFUSION_CAST_OPTIONS + ); + + // int64 + generic_test_cast!( + Int64Array, + DataType::Int64, + vec![1, 2, 3, 4, 5], + DecimalArray, + DataType::Decimal(20, 0), + vec![ + Some(1_i128), + Some(2_i128), + Some(3_i128), + Some(4_i128), + Some(5_i128), + ], + DEFAULT_DATAFUSION_CAST_OPTIONS + ); + + // int64 to different scale + generic_test_cast!( + Int64Array, + DataType::Int64, + vec![1, 2, 3, 4, 5], + DecimalArray, + DataType::Decimal(20, 2), + vec![ + Some(100_i128), + Some(200_i128), + Some(300_i128), + Some(400_i128), + Some(500_i128), + ], + DEFAULT_DATAFUSION_CAST_OPTIONS + ); + + // float32 + generic_test_cast!( + Float32Array, + DataType::Float32, + vec![1.5, 2.5, 3.0, 1.123_456_8, 5.50], + DecimalArray, + DataType::Decimal(10, 2), + vec![ + Some(150_i128), + Some(250_i128), + Some(300_i128), + Some(112_i128), + Some(550_i128), + ], + DEFAULT_DATAFUSION_CAST_OPTIONS + ); + + // float64 + generic_test_cast!( + Float64Array, + DataType::Float64, + vec![1.5, 2.5, 3.0, 1.123_456_8, 5.50], + DecimalArray, + DataType::Decimal(20, 4), + vec![ + Some(15000_i128), + Some(25000_i128), + Some(30000_i128), + Some(11234_i128), + Some(55000_i128), + ], + DEFAULT_DATAFUSION_CAST_OPTIONS + ); + Ok(()) + } + #[test] fn test_cast_i32_u32() -> Result<()> { generic_test_cast!( diff --git a/datafusion/src/physical_plan/expressions/try_cast.rs b/datafusion/src/physical_plan/expressions/try_cast.rs index 1ba4a50260d4..d231ce276a77 100644 --- a/datafusion/src/physical_plan/expressions/try_cast.rs +++ b/datafusion/src/physical_plan/expressions/try_cast.rs @@ -119,9 +119,12 @@ mod tests { use super::*; use crate::error::Result; use crate::physical_plan::expressions::col; - use arrow::array::{StringArray, Time64NanosecondArray}; + use arrow::array::{Float32Array, Float64Array, StringArray, Time64NanosecondArray}; use arrow::{ - array::{Array, Int32Array, Int64Array, TimestampNanosecondArray, UInt32Array}, + array::{ + Array, DecimalArray, Int16Array, Int32Array, Int64Array, Int8Array, + TimestampNanosecondArray, UInt32Array, + }, datatypes::*, }; @@ -175,6 +178,122 @@ mod tests { }}; } + #[test] + fn test_try_cast_numeric_to_decimal() -> Result<()> { + // int8 + generic_test_cast!( + Int8Array, + DataType::Int8, + vec![1, 2, 3, 4, 5], + DecimalArray, + DataType::Decimal(3, 0), + vec![ + Some(1_i128), + Some(2_i128), + Some(3_i128), + Some(4_i128), + Some(5_i128), + ] + ); + + // int16 + generic_test_cast!( + Int16Array, + DataType::Int16, + vec![1, 2, 3, 4, 5], + DecimalArray, + DataType::Decimal(5, 0), + vec![ + Some(1_i128), + Some(2_i128), + Some(3_i128), + Some(4_i128), + Some(5_i128), + ] + ); + + // int32 + generic_test_cast!( + Int32Array, + DataType::Int32, + vec![1, 2, 3, 4, 5], + DecimalArray, + DataType::Decimal(10, 0), + vec![ + Some(1_i128), + Some(2_i128), + Some(3_i128), + Some(4_i128), + Some(5_i128), + ] + ); + + // int64 + generic_test_cast!( + Int64Array, + DataType::Int64, + vec![1, 2, 3, 4, 5], + DecimalArray, + DataType::Decimal(20, 0), + vec![ + Some(1_i128), + Some(2_i128), + Some(3_i128), + Some(4_i128), + Some(5_i128), + ] + ); + + // int64 to different scale + generic_test_cast!( + Int64Array, + DataType::Int64, + vec![1, 2, 3, 4, 5], + DecimalArray, + DataType::Decimal(20, 2), + vec![ + Some(100_i128), + Some(200_i128), + Some(300_i128), + Some(400_i128), + Some(500_i128), + ] + ); + + // float32 + generic_test_cast!( + Float32Array, + DataType::Float32, + vec![1.5, 2.5, 3.0, 1.123_456_8, 5.50], + DecimalArray, + DataType::Decimal(10, 2), + vec![ + Some(150_i128), + Some(250_i128), + Some(300_i128), + Some(112_i128), + Some(550_i128), + ] + ); + + // float64 + generic_test_cast!( + Float64Array, + DataType::Float64, + vec![1.5, 2.5, 3.0, 1.123_456_8, 5.50], + DecimalArray, + DataType::Decimal(20, 4), + vec![ + Some(15000_i128), + Some(25000_i128), + Some(30000_i128), + Some(11234_i128), + Some(55000_i128), + ] + ); + Ok(()) + } + #[test] fn test_cast_i32_u32() -> Result<()> { generic_test_cast!(