Skip to content

feat: support Decimal256 for the abs function #7904

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Oct 23, 2023
Merged

feat: support Decimal256 for the abs function #7904

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
cd56987
feat: support Decimal256 for the abs function
jonahgao Oct 22, 2023
f8f3fbb
Remove useless comment
jonahgao Oct 23, 2023
5671aee
use wrapping_abs
jonahgao Oct 23, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 17 additions & 11 deletions datafusion/physical-expr/src/math_expressions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,8 @@

use arrow::array::ArrayRef;
use arrow::array::{
BooleanArray, Decimal128Array, Float32Array, Float64Array, Int16Array, Int32Array,
Int64Array, Int8Array,
BooleanArray, Decimal128Array, Decimal256Array, Float32Array, Float64Array,
Int16Array, Int32Array, Int64Array, Int8Array,
};
use arrow::datatypes::DataType;
use arrow::error::ArrowError;
Expand Down Expand Up @@ -701,6 +701,18 @@ macro_rules! make_try_abs_function {
}};
}

macro_rules! make_decimal_abs_function {
($ARRAY_TYPE:ident) => {{
|args: &[ArrayRef]| {
let array = downcast_arg!(&args[0], "abs arg", $ARRAY_TYPE);
let res: $ARRAY_TYPE = array
.unary(|x| x.wrapping_abs())
.with_data_type(args[0].data_type().clone());
Ok(Arc::new(res) as ArrayRef)
}
}};
}

/// Abs SQL function
/// Return different implementations based on input datatype to reduce branches during execution
pub(super) fn create_abs_function(
Expand All @@ -723,15 +735,9 @@ pub(super) fn create_abs_function(
| DataType::UInt32
| DataType::UInt64 => Ok(|args: &[ArrayRef]| Ok(args[0].clone())),

// Decimal should keep the same precision and scale by using `with_data_type()`.
// https://github.com/apache/arrow-rs/issues/4644
DataType::Decimal128(_, _) => Ok(|args: &[ArrayRef]| {
let array = downcast_arg!(&args[0], "abs arg", Decimal128Array);
let res: Decimal128Array = array
.unary(i128::abs)
.with_data_type(args[0].data_type().clone());
Ok(Arc::new(res) as ArrayRef)
}),
// Decimal types
DataType::Decimal128(_, _) => Ok(make_decimal_abs_function!(Decimal128Array)),
DataType::Decimal256(_, _) => Ok(make_decimal_abs_function!(Decimal256Array)),

other => not_impl_err!("Unsupported data type {other:?} for function abs"),
}
Expand Down
63 changes: 41 additions & 22 deletions datafusion/sqllogictest/test_files/math.slt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -395,7 +395,7 @@ NaN NaN

# abs: return type
query TT rowsort
SELECT arrow_typeof(c1), arrow_typeof(c2) FROM test_nullable_float limit 1
SELECT arrow_typeof(abs(c1)), arrow_typeof(abs(c2)) FROM test_nullable_float limit 1
----
Float32 Float64

Expand Down Expand Up @@ -466,34 +466,48 @@ drop table test_non_nullable_float

statement ok
CREATE TABLE test_nullable_decimal(
c1 DECIMAL(10, 2),
c2 DECIMAL(38, 10)
) AS VALUES (0, 0), (NULL, NULL);

query RR
c1 DECIMAL(10, 2), /* Decimal128 */
c2 DECIMAL(38, 10), /* Decimal128 with max precision */
c3 DECIMAL(40, 2), /* Decimal256 */
c4 DECIMAL(76, 10) /* Decimal256 with max precision */
) AS VALUES
(0, 0, 0, 0),
(NULL, NULL, NULL, NULL);

query RRRR
INSERT into test_nullable_decimal values
(-99999999.99, '-9999999999999999999999999999.9999999999'),
(99999999.99, '9999999999999999999999999999.9999999999');
(
-99999999.99,
'-9999999999999999999999999999.9999999999',
'-99999999999999999999999999999999999999.99',
'-999999999999999999999999999999999999999999999999999999999999999999.9999999999'
),
(
99999999.99,
'9999999999999999999999999999.9999999999',
'99999999999999999999999999999999999999.99',
'999999999999999999999999999999999999999999999999999999999999999999.9999999999'
)
----
2


query R rowsort
query R
SELECT c1*0 FROM test_nullable_decimal WHERE c1 IS NULL;
----
NULL

query R rowsort
query R
SELECT c1/0 FROM test_nullable_decimal WHERE c1 IS NULL;
----
NULL

query R rowsort
query R
SELECT c1%0 FROM test_nullable_decimal WHERE c1 IS NULL;
----
NULL

query R rowsort
query R
SELECT c1*0 FROM test_nullable_decimal WHERE c1 IS NOT NULL;
----
0
Expand All @@ -507,19 +521,24 @@ query error DataFusion error: Arrow error: Divide by zero error
SELECT c1%0 FROM test_nullable_decimal WHERE c1 IS NOT NULL;

# abs: return type
query TT rowsort
SELECT arrow_typeof(c1), arrow_typeof(c2) FROM test_nullable_decimal limit 1
query TTTT
SELECT
arrow_typeof(abs(c1)),
arrow_typeof(abs(c2)),
arrow_typeof(abs(c3)),
arrow_typeof(abs(c4))
FROM test_nullable_decimal limit 1
----
Decimal128(10, 2) Decimal128(38, 10)
Decimal128(10, 2) Decimal128(38, 10) Decimal256(40, 2) Decimal256(76, 10)

# abs: Decimal128
query RR rowsort
SELECT abs(c1), abs(c2) FROM test_nullable_decimal
# abs: decimals
query RRRR rowsort
SELECT abs(c1), abs(c2), abs(c3), abs(c4) FROM test_nullable_decimal
----
0 0
99999999.99 9999999999999999999999999999.9999999999
99999999.99 9999999999999999999999999999.9999999999
NULL NULL
0 0 0 0
99999999.99 9999999999999999999999999999.9999999999 99999999999999999999999999999999999999.99 999999999999999999999999999999999999999999999999999999999999999999.9999999999
99999999.99 9999999999999999999999999999.9999999999 99999999999999999999999999999999999999.99 999999999999999999999999999999999999999999999999999999999999999999.9999999999
NULL NULL NULL NULL

statement ok
drop table test_nullable_decimal
Expand Down