Skip to content

Commit 562b68f

Browse files
liukun4515liukun4515
committed
support cast/try_cast for decimal: signed numeric to decimal
1 parent 4b454d0 commit 562b68f

File tree

3 files changed

+227
-3
lines changed

3 files changed

+227
-3
lines changed

‎datafusion/Cargo.toml‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,10 @@ avro = ["avro-rs", "num-traits"]
5252
[dependencies]
5353
ahash = "0.7"
5454
hashbrown = { version = "0.11", features = ["raw"] }
55-
arrow = { version = "6.3.0", features = ["prettyprint"] }
56-
parquet = { version = "6.3.0", features = ["arrow"] }
55+
#arrow = { version = "6.3.0", features = ["prettyprint"] }
56+
#parquet = { version = "6.3.0", features = ["arrow"] }
57+
arrow = { path = "/Users/kliu3/Documents/github/arrow-rs/arrow", features = ["prettyprint"] }
58+
parquet = { path = "/Users/kliu3/Documents/github/arrow-rs/parquet", features = ["arrow"] }
5759
sqlparser = "0.13"
5860
paste = "^1.0"
5961
num_cpus = "1.13.0"

‎datafusion/src/physical_plan/expressions/cast.rs‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,9 @@ pub fn cast_with_options(
129129
if expr_type == cast_type {
130130
Ok(expr.clone())
131131
} else if can_cast_types(&expr_type, &cast_type) {
132+
// TODO
133+
// support numeric data type to decimal
134+
// support one type decimal to another type decimal
132135
Ok(Arc::new(CastExpr::new(expr, cast_type, cast_options)))
133136
} else {
134137
Err(DataFusionError::Internal(format!(
@@ -217,6 +220,95 @@ mod tests {
217220
}};
218221
}
219222

223+
#[test]
224+
fn test_cast_numeric_to_decimal() -> Result<()> {
225+
// int32
226+
generic_test_cast!(
227+
Int32Array,
228+
DataType::Int32,
229+
vec![1, 2, 3, 4, 5],
230+
// TODO
231+
UInt32Array,
232+
DataType::UInt32,
233+
vec![
234+
Some(1_u32),
235+
Some(2_u32),
236+
Some(3_u32),
237+
Some(4_u32),
238+
Some(5_u32)
239+
],
240+
DEFAULT_DATAFUSION_CAST_OPTIONS
241+
);
242+
// int64
243+
generic_test_cast!(
244+
Int32Array,
245+
DataType::Int64,
246+
vec![1, 2, 3, 4, 5],
247+
// TODO
248+
UInt32Array,
249+
DataType::UInt32,
250+
vec![
251+
Some(1_u32),
252+
Some(2_u32),
253+
Some(3_u32),
254+
Some(4_u32),
255+
Some(5_u32)
256+
],
257+
DEFAULT_DATAFUSION_CAST_OPTIONS
258+
);
259+
// float32
260+
generic_test_cast!(
261+
Int32Array,
262+
DataType::Float32,
263+
vec![1, 2, 3, 4, 5],
264+
// TODO
265+
UInt32Array,
266+
DataType::UInt32,
267+
vec![
268+
Some(1_u32),
269+
Some(2_u32),
270+
Some(3_u32),
271+
Some(4_u32),
272+
Some(5_u32)
273+
],
274+
DEFAULT_DATAFUSION_CAST_OPTIONS
275+
);
276+
// float64
277+
generic_test_cast!(
278+
Int32Array,
279+
DataType::Float64,
280+
vec![1, 2, 3, 4, 5],
281+
// TODO
282+
UInt32Array,
283+
DataType::UInt32,
284+
vec![
285+
Some(1_u32),
286+
Some(2_u32),
287+
Some(3_u32),
288+
Some(4_u32),
289+
Some(5_u32)
290+
],
291+
DEFAULT_DATAFUSION_CAST_OPTIONS
292+
);
293+
generic_test_cast!(
294+
Int32Array,
295+
DataType::Decimal(10, 4),
296+
vec![1, 2, 3, 4, 5],
297+
// TODO
298+
UInt32Array,
299+
DataType::UInt32,
300+
vec![
301+
Some(1_u32),
302+
Some(2_u32),
303+
Some(3_u32),
304+
Some(4_u32),
305+
Some(5_u32)
306+
],
307+
DEFAULT_DATAFUSION_CAST_OPTIONS
308+
);
309+
Ok(())
310+
}
311+
220312
#[test]
221313
fn test_cast_i32_u32() -> Result<()> {
222314
generic_test_cast!(

‎datafusion/src/physical_plan/expressions/try_cast.rs‎

Lines changed: 131 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,14 @@ impl PhysicalExpr for TryCastExpr {
7777

7878
fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
7979
let value = self.expr.evaluate(batch)?;
80+
// TODO cast signed numeric to decimal
81+
// INT32, INT64, FLOAT, DOUBLE
82+
// private[sql] val IntDecimal = DecimalType(3, 0)
83+
// private[sql] val IntDecimal = DecimalType(5, 0)
84+
// private[sql] val IntDecimal = DecimalType(10, 0)
85+
// private[sql] val LongDecimal = DecimalType(20, 0)
86+
// private[sql] val FloatDecimal = DecimalType(14, 7)
87+
// private[sql] val DoubleDecimal = DecimalType(30, 15)
8088
match value {
8189
ColumnarValue::Array(array) => Ok(ColumnarValue::Array(kernels::cast::cast(
8290
&array,
@@ -105,6 +113,11 @@ pub fn try_cast(
105113
if expr_type == cast_type {
106114
Ok(expr.clone())
107115
} else if can_cast_types(&expr_type, &cast_type) {
116+
// TODO
117+
// support numeric data type to decimal
118+
// support one type decimal to another type decimal
119+
// Now we just support signed numeric type to decimal,
120+
// and will implement decimal to signed numeric data type later.
108121
Ok(Arc::new(TryCastExpr::new(expr, cast_type)))
109122
} else {
110123
Err(DataFusionError::Internal(format!(
@@ -121,9 +134,10 @@ mod tests {
121134
use crate::physical_plan::expressions::col;
122135
use arrow::array::{StringArray, Time64NanosecondArray};
123136
use arrow::{
124-
array::{Array, Int32Array, Int64Array, TimestampNanosecondArray, UInt32Array},
137+
array::{Array, Int8Array, Int16Array, Int32Array, Int64Array, DecimalArray, TimestampNanosecondArray, UInt32Array},
125138
datatypes::*,
126139
};
140+
use crate::scalar::ScalarValue::Decimal128;
127141

128142
// runs an end-to-end test of physical type cast
129143
// 1. construct a record batch with a column "a" of type A
@@ -175,6 +189,122 @@ mod tests {
175189
}};
176190
}
177191

192+
// #[test]
193+
fn test_try_cast_numeric_to_decimal() -> Result<()> {
194+
// int8
195+
generic_test_cast!(
196+
Int8Array,
197+
DataType::Int8,
198+
vec![1, 2, 3, 4, 5],
199+
// TODO
200+
DecimalArray,
201+
DataType::Decimal(3,0),
202+
vec![
203+
Some(Decimal128(Some(1), 3, 0)),
204+
Some(Decimal128(Some(2), 3, 0)),
205+
Some(Decimal128(Some(3), 3, 0)),
206+
Some(Decimal128(Some(4), 3, 0)),
207+
Some(Decimal128(Some(5), 3, 0))
208+
]
209+
);
210+
// int16
211+
generic_test_cast!(
212+
Int16Array,
213+
DataType::Int16,
214+
vec![1, 2, 3, 4, 5],
215+
// TODO
216+
DecimalArray,
217+
DataType::Decimal(5,0),
218+
vec![
219+
Some(Decimal128(Some(1), 3, 0)),
220+
Some(Decimal128(Some(2), 3, 0)),
221+
Some(Decimal128(Some(3), 3, 0)),
222+
Some(Decimal128(Some(4), 3, 0)),
223+
Some(Decimal128(Some(5), 3, 0))
224+
]
225+
);
226+
// int32
227+
generic_test_cast!(
228+
Int32Array,
229+
DataType::Int32,
230+
vec![1, 2, 3, 4, 5],
231+
// TODO
232+
DecimalArray,
233+
DataType::Decimal(10,0),
234+
vec![
235+
Some(Decimal128(Some(1), 3, 0)),
236+
Some(Decimal128(Some(2), 3, 0)),
237+
Some(Decimal128(Some(3), 3, 0)),
238+
Some(Decimal128(Some(4), 3, 0)),
239+
Some(Decimal128(Some(5), 3, 0))
240+
]
241+
);
242+
// int64
243+
generic_test_cast!(
244+
Int32Array,
245+
DataType::Int64,
246+
vec![1, 2, 3, 4, 5],
247+
// TODO
248+
UInt32Array,
249+
DataType::UInt32,
250+
vec![
251+
Some(1_u32),
252+
Some(2_u32),
253+
Some(3_u32),
254+
Some(4_u32),
255+
Some(5_u32)
256+
]
257+
);
258+
// float32
259+
generic_test_cast!(
260+
Int32Array,
261+
DataType::Float32,
262+
vec![1, 2, 3, 4, 5],
263+
// TODO
264+
UInt32Array,
265+
DataType::UInt32,
266+
vec![
267+
Some(1_u32),
268+
Some(2_u32),
269+
Some(3_u32),
270+
Some(4_u32),
271+
Some(5_u32)
272+
]
273+
);
274+
// float64
275+
generic_test_cast!(
276+
Int32Array,
277+
DataType::Float64,
278+
vec![1, 2, 3, 4, 5],
279+
// TODO
280+
UInt32Array,
281+
DataType::UInt32,
282+
vec![
283+
Some(1_u32),
284+
Some(2_u32),
285+
Some(3_u32),
286+
Some(4_u32),
287+
Some(5_u32)
288+
]
289+
);
290+
generic_test_cast!(
291+
Int32Array,
292+
DataType::Decimal(10, 4),
293+
vec![1, 2, 3, 4, 5],
294+
// TODO
295+
UInt32Array,
296+
DataType::UInt32,
297+
vec![
298+
Some(1_u32),
299+
Some(2_u32),
300+
Some(3_u32),
301+
Some(4_u32),
302+
Some(5_u32)
303+
]
304+
);
305+
Ok(())
306+
}
307+
178308
#[test]
179309
fn test_cast_i32_u32() -> Result<()> {
180310
generic_test_cast!(

0 commit comments

Comments
 (0)