Skip to content

Commit b8f9bc2

Browse files
IgosukiIgosuki
committed
fix decimal add because arrow2 doesn't include decimal add in arithmetics::add
1 parent 5ad5f7c commit b8f9bc2

File tree

5 files changed

+100
-44
lines changed

5 files changed

+100
-44
lines changed

‎datafusion/benches/data_utils/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ fn create_record_batch(
123123
schema,
124124
vec![
125125
Arc::new(Utf8Array::<i32>::from_slice(keys)),
126-
Arc::new(Float32Array::from_slice(&[i as f32; batch_size])),
126+
Arc::new(Float32Array::from_slice(vec![i as f32; batch_size])),
127127
Arc::new(Float64Array::from(values)),
128128
Arc::new(UInt64Array::from(integer_values_wide)),
129129
Arc::new(UInt64Array::from_slice(integer_values_narrow)),

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

Lines changed: 44 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ use std::{any::Any, convert::TryInto, sync::Arc};
2020
use crate::record_batch::RecordBatch;
2121
use arrow::array::*;
2222
use arrow::compute;
23+
use arrow::datatypes::DataType::Decimal;
2324
use arrow::datatypes::{DataType, Schema};
2425

2526
use crate::error::{DataFusionError, Result};
@@ -247,9 +248,24 @@ fn evaluate_regex_case_insensitive<O: Offset>(
247248

248249
fn evaluate(lhs: &dyn Array, op: &Operator, rhs: &dyn Array) -> Result<Arc<dyn Array>> {
249250
use Operator::*;
250-
if matches!(op, Plus | Minus | Divide | Multiply | Modulo | BitwiseAnd) {
251+
if matches!(op, Plus) {
252+
let arr: ArrayRef = match (lhs.data_type(), rhs.data_type()) {
253+
(Decimal(p1, s1), Decimal(p2, s2)) => {
254+
let left_array =
255+
lhs.as_any().downcast_ref::<PrimitiveArray<i128>>().unwrap();
256+
let right_array =
257+
rhs.as_any().downcast_ref::<PrimitiveArray<i128>>().unwrap();
258+
Arc::new(if *p1 == *p2 && *s1 == *s2 {
259+
compute::arithmetics::decimal::add(left_array, right_array)
260+
} else {
261+
compute::arithmetics::decimal::adaptive_add(left_array, right_array)?
262+
})
263+
}
264+
_ => compute::arithmetics::add(lhs, rhs).into(),
265+
};
266+
Ok(arr)
267+
} else if matches!(op, Minus | Divide | Multiply | Modulo) {
251268
let arr = match op {
252-
Operator::Plus => compute::arithmetics::add(lhs, rhs),
253269
Operator::Minus => compute::arithmetics::sub(lhs, rhs),
254270
Operator::Divide => compute::arithmetics::div(lhs, rhs),
255271
Operator::Multiply => compute::arithmetics::mul(lhs, rhs),
@@ -828,6 +844,7 @@ mod tests {
828844
use crate::error::Result;
829845
use crate::field_util::SchemaExt;
830846
use crate::physical_plan::expressions::{col, lit};
847+
use crate::test_util::create_decimal_array;
831848
use arrow::datatypes::{Field, SchemaRef};
832849
use arrow::error::ArrowError;
833850

@@ -1015,7 +1032,11 @@ mod tests {
10151032
}
10161033

10171034
fn add_decimal(left: &Int128Array, right: &Int128Array) -> Result<Int128Array> {
1018-
let mut decimal_builder = Int128Vec::with_capacity(left.len());
1035+
let mut decimal_builder = Int128Vec::from_data(
1036+
left.data_type().clone(),
1037+
Vec::<i128>::with_capacity(left.len()),
1038+
None,
1039+
);
10191040
for i in 0..left.len() {
10201041
if left.is_null(i) || right.is_null(i) {
10211042
decimal_builder.push(None);
@@ -1027,7 +1048,11 @@ mod tests {
10271048
}
10281049

10291050
fn subtract_decimal(left: &Int128Array, right: &Int128Array) -> Result<Int128Array> {
1030-
let mut decimal_builder = Int128Vec::with_capacity(left.len());
1051+
let mut decimal_builder = Int128Vec::from_data(
1052+
left.data_type().clone(),
1053+
Vec::<i128>::with_capacity(left.len()),
1054+
None,
1055+
);
10311056
for i in 0..left.len() {
10321057
if left.is_null(i) || right.is_null(i) {
10331058
decimal_builder.push(None);
@@ -1043,7 +1068,11 @@ mod tests {
10431068
right: &Int128Array,
10441069
scale: u32,
10451070
) -> Result<Int128Array> {
1046-
let mut decimal_builder = Int128Vec::with_capacity(left.len());
1071+
let mut decimal_builder = Int128Vec::from_data(
1072+
left.data_type().clone(),
1073+
Vec::<i128>::with_capacity(left.len()),
1074+
None,
1075+
);
10471076
let divide = 10_i128.pow(scale);
10481077
for i in 0..left.len() {
10491078
if left.is_null(i) || right.is_null(i) {
@@ -1061,7 +1090,11 @@ mod tests {
10611090
right: &Int128Array,
10621091
scale: i32,
10631092
) -> Result<Int128Array> {
1064-
let mut decimal_builder = Int128Vec::with_capacity(left.len());
1093+
let mut decimal_builder = Int128Vec::from_data(
1094+
left.data_type().clone(),
1095+
Vec::<i128>::with_capacity(left.len()),
1096+
None,
1097+
);
10651098
let mul = 10_f64.powi(scale);
10661099
for i in 0..left.len() {
10671100
if left.is_null(i) || right.is_null(i) {
@@ -1081,7 +1114,11 @@ mod tests {
10811114
}
10821115

10831116
fn modulus_decimal(left: &Int128Array, right: &Int128Array) -> Result<Int128Array> {
1084-
let mut decimal_builder = Int128Vec::with_capacity(left.len());
1117+
let mut decimal_builder = Int128Vec::from_data(
1118+
left.data_type().clone(),
1119+
Vec::<i128>::with_capacity(left.len()),
1120+
None,
1121+
);
10851122
for i in 0..left.len() {
10861123
if left.is_null(i) || right.is_null(i) {
10871124
decimal_builder.push(None);
@@ -2135,25 +2172,6 @@ mod tests {
21352172
assert_eq!(result.as_ref(), &expected as &dyn Array);
21362173
}
21372174

2138-
fn create_decimal_array(
2139-
array: &[Option<i128>],
2140-
_precision: usize,
2141-
_scale: usize,
2142-
) -> Result<Int128Array> {
2143-
let mut decimal_builder = Int128Vec::with_capacity(array.len());
2144-
for value in array {
2145-
match value {
2146-
None => {
2147-
decimal_builder.push(None);
2148-
}
2149-
Some(v) => {
2150-
decimal_builder.try_push(Some(*v))?;
2151-
}
2152-
}
2153-
}
2154-
Ok(decimal_builder.into())
2155-
}
2156-
21572175
#[test]
21582176
fn comparison_decimal_op_test() -> Result<()> {
21592177
let value_i128: i128 = 123;

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

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,6 @@ pub fn cast_with_error(
105105
) -> Result<Box<dyn Array>> {
106106
let result = cast::cast(array, cast_type, options)?;
107107
if result.null_count() != array.null_count() {
108-
println!("{result:?} : {array:?}");
109108
let casted_valids = result.validity().unwrap();
110109
let failed_casts = match array.validity() {
111110
Some(valids) => valids ^ casted_valids,
@@ -192,6 +191,7 @@ mod tests {
192191
use crate::error::Result;
193192
use crate::field_util::SchemaExt;
194193
use crate::physical_plan::expressions::col;
194+
use crate::test_util::create_decimal_array_from_slice;
195195
use arrow::{array::*, datatypes::*};
196196

197197
type StringArray = Utf8Array<i32>;
@@ -298,7 +298,7 @@ mod tests {
298298
#[test]
299299
fn test_cast_decimal_to_decimal() -> Result<()> {
300300
let array: Vec<i128> = vec![1234, 2222, 3, 4000, 5000];
301-
let decimal_array = Int128Array::from_slice(&array);
301+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
302302
generic_decimal_to_other_test_cast!(
303303
decimal_array,
304304
DataType::Decimal(10, 3),
@@ -315,7 +315,7 @@ mod tests {
315315
DEFAULT_DATAFUSION_CAST_OPTIONS
316316
);
317317

318-
let decimal_array = Int128Array::from_slice(&array);
318+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
319319
generic_decimal_to_other_test_cast!(
320320
decimal_array,
321321
DataType::Decimal(10, 3),
@@ -339,7 +339,7 @@ mod tests {
339339
fn test_cast_decimal_to_numeric() -> Result<()> {
340340
let array: Vec<i128> = vec![1, 2, 3, 4, 5];
341341
// decimal to i8
342-
let decimal_array = Int128Array::from_slice(&array);
342+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
343343
generic_decimal_to_other_test_cast!(
344344
decimal_array,
345345
DataType::Decimal(10, 0),
@@ -356,7 +356,7 @@ mod tests {
356356
DEFAULT_DATAFUSION_CAST_OPTIONS
357357
);
358358
// decimal to i16
359-
let decimal_array = Int128Array::from_slice(&array);
359+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
360360
generic_decimal_to_other_test_cast!(
361361
decimal_array,
362362
DataType::Decimal(10, 0),
@@ -373,7 +373,7 @@ mod tests {
373373
DEFAULT_DATAFUSION_CAST_OPTIONS
374374
);
375375
// decimal to i32
376-
let decimal_array = Int128Array::from_slice(&array);
376+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
377377
generic_decimal_to_other_test_cast!(
378378
decimal_array,
379379
DataType::Decimal(10, 0),
@@ -390,7 +390,7 @@ mod tests {
390390
DEFAULT_DATAFUSION_CAST_OPTIONS
391391
);
392392
// decimal to i64
393-
let decimal_array = Int128Array::from_slice(&array);
393+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
394394
generic_decimal_to_other_test_cast!(
395395
decimal_array,
396396
DataType::Decimal(10, 0),
@@ -408,7 +408,7 @@ mod tests {
408408
);
409409
// decimal to float32
410410
let array: Vec<i128> = vec![1234, 2222, 3, 4000, 5000];
411-
let decimal_array = Int128Array::from_slice(&array);
411+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
412412
generic_decimal_to_other_test_cast!(
413413
decimal_array,
414414
DataType::Decimal(10, 3),
@@ -425,7 +425,7 @@ mod tests {
425425
DEFAULT_DATAFUSION_CAST_OPTIONS
426426
);
427427
// decimal to float64
428-
let decimal_array = Int128Array::from_slice(&array);
428+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
429429
generic_decimal_to_other_test_cast!(
430430
decimal_array,
431431
DataType::Decimal(20, 6),

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

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,7 @@ mod tests {
129129
use crate::error::Result;
130130
use crate::field_util::SchemaExt;
131131
use crate::physical_plan::expressions::col;
132+
use crate::test_util::create_decimal_array_from_slice;
132133
use arrow::{array::*, datatypes::*};
133134

134135
type StringArray = Utf8Array<i32>;
@@ -234,7 +235,7 @@ mod tests {
234235
fn test_try_cast_decimal_to_decimal() -> Result<()> {
235236
// try cast one decimal data type to another decimal data type
236237
let array: Vec<i128> = vec![1234, 2222, 3, 4000, 5000];
237-
let decimal_array = Int128Array::from_slice(&array);
238+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
238239
generic_decimal_to_other_test_cast!(
239240
decimal_array,
240241
DataType::Decimal(10, 3),
@@ -250,7 +251,7 @@ mod tests {
250251
]
251252
);
252253

253-
let decimal_array = Int128Array::from_slice(&array);
254+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
254255
generic_decimal_to_other_test_cast!(
255256
decimal_array,
256257
DataType::Decimal(10, 3),
@@ -274,7 +275,7 @@ mod tests {
274275
// TODO we should add function to create Int128Array with value and metadata
275276
// https://github.com/apache/arrow-rs/issues/1009
276277
let array: Vec<i128> = vec![1, 2, 3, 4, 5];
277-
let decimal_array = Int128Array::from_slice(&array);
278+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
278279
// decimal to i8
279280
generic_decimal_to_other_test_cast!(
280281
decimal_array,
@@ -292,7 +293,7 @@ mod tests {
292293
);
293294

294295
// decimal to i16
295-
let decimal_array = Int128Array::from_slice(&array);
296+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
296297
generic_decimal_to_other_test_cast!(
297298
decimal_array,
298299
DataType::Decimal(10, 0),
@@ -309,7 +310,7 @@ mod tests {
309310
);
310311

311312
// decimal to i32
312-
let decimal_array = Int128Array::from_slice(&array);
313+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
313314
generic_decimal_to_other_test_cast!(
314315
decimal_array,
315316
DataType::Decimal(10, 0),
@@ -326,7 +327,7 @@ mod tests {
326327
);
327328

328329
// decimal to i64
329-
let decimal_array = Int128Array::from_slice(&array);
330+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
330331
generic_decimal_to_other_test_cast!(
331332
decimal_array,
332333
DataType::Decimal(10, 0),
@@ -344,7 +345,7 @@ mod tests {
344345

345346
// decimal to float32
346347
let array: Vec<i128> = vec![1234, 2222, 3, 4000, 5000];
347-
let decimal_array = Int128Array::from_slice(&array);
348+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
348349
generic_decimal_to_other_test_cast!(
349350
decimal_array,
350351
DataType::Decimal(10, 3),
@@ -360,7 +361,7 @@ mod tests {
360361
]
361362
);
362363
// decimal to float64
363-
let decimal_array = Int128Array::from_slice(&array);
364+
let decimal_array = create_decimal_array_from_slice(&array, 10, 3)?;
364365
generic_decimal_to_other_test_cast!(
365366
decimal_array,
366367
DataType::Decimal(20, 6),

‎datafusion/src/test_util.rs‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -334,3 +334,40 @@ mod tests {
334334
assert!(PathBuf::from(res).is_dir());
335335
}
336336
}
337+
338+
#[cfg(test)]
339+
pub fn create_decimal_array(
340+
array: &[Option<i128>],
341+
precision: usize,
342+
scale: usize,
343+
) -> crate::error::Result<arrow::array::Int128Array> {
344+
use arrow::array::{Int128Vec, TryPush};
345+
let mut decimal_builder = Int128Vec::from_data(
346+
DataType::Decimal(precision, scale),
347+
Vec::<i128>::with_capacity(array.len()),
348+
None,
349+
);
350+
351+
for value in array {
352+
match value {
353+
None => {
354+
decimal_builder.push(None);
355+
}
356+
Some(v) => {
357+
decimal_builder.try_push(Some(*v))?;
358+
}
359+
}
360+
}
361+
Ok(decimal_builder.into())
362+
}
363+
364+
#[cfg(test)]
365+
pub fn create_decimal_array_from_slice(
366+
array: &[i128],
367+
precision: usize,
368+
scale: usize,
369+
) -> crate::error::Result<arrow::array::Int128Array> {
370+
let decimal_array_values: Vec<Option<i128>> =
371+
array.into_iter().map(|v| Some(*v)).collect();
372+
create_decimal_array(&decimal_array_values, precision, scale)
373+
}

0 commit comments

Comments
 (0)