You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2023/09/18 21:34:41 UTC
[arrow-datafusion] branch main updated: feat: natively support more data types for the `abs` function. (#7568)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new d512bee11b feat: natively support more data types for the `abs` function. (#7568)
d512bee11b is described below
commit d512bee11b3458d751c05b35baebecdb469dda2c
Author: Jonah Gao <jo...@gmail.com>
AuthorDate: Tue Sep 19 05:34:35 2023 +0800
feat: natively support more data types for the `abs` function. (#7568)
* refactor: the return type of `abs()` should be the same as the input type
* add sqllogic tests
---
datafusion/core/tests/sql/mod.rs | 2 +-
datafusion/expr/src/built_in_function.rs | 10 +-
datafusion/optimizer/src/analyzer/type_coercion.rs | 4 +-
datafusion/physical-expr/src/functions.rs | 6 +-
datafusion/physical-expr/src/math_expressions.rs | 74 ++++++++++-
datafusion/proto/src/physical_plan/mod.rs | 4 +-
datafusion/sqllogictest/test_files/decimal.slt | 2 +-
datafusion/sqllogictest/test_files/math.slt | 138 ++++++++++++++++++---
datafusion/sqllogictest/test_files/scalar.slt | 6 +-
9 files changed, 216 insertions(+), 30 deletions(-)
diff --git a/datafusion/core/tests/sql/mod.rs b/datafusion/core/tests/sql/mod.rs
index 48fb69e51c..522cdce1a4 100644
--- a/datafusion/core/tests/sql/mod.rs
+++ b/datafusion/core/tests/sql/mod.rs
@@ -635,7 +635,7 @@ pub fn make_partition(sz: i32) -> RecordBatch {
/// Specialised String representation
fn col_str(column: &ArrayRef, row_index: usize) -> String {
- if column.is_null(row_index) {
+ if column.data_type() == &DataType::Null || column.is_null(row_index) {
return "NULL".to_string();
}
diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs
index 0fb23be5af..f38c7f12a8 100644
--- a/datafusion/expr/src/built_in_function.rs
+++ b/datafusion/expr/src/built_in_function.rs
@@ -789,8 +789,9 @@ impl BuiltinScalarFunction {
BuiltinScalarFunction::ArrowTypeof => Ok(Utf8),
- BuiltinScalarFunction::Abs
- | BuiltinScalarFunction::Acos
+ BuiltinScalarFunction::Abs => Ok(input_expr_types[0].clone()),
+
+ BuiltinScalarFunction::Acos
| BuiltinScalarFunction::Asin
| BuiltinScalarFunction::Atan
| BuiltinScalarFunction::Acosh
@@ -1162,8 +1163,9 @@ impl BuiltinScalarFunction {
Signature::uniform(2, vec![Int64], self.volatility())
}
BuiltinScalarFunction::ArrowTypeof => Signature::any(1, self.volatility()),
- BuiltinScalarFunction::Abs
- | BuiltinScalarFunction::Acos
+ BuiltinScalarFunction::Abs => Signature::any(1, self.volatility()),
+
+ BuiltinScalarFunction::Acos
| BuiltinScalarFunction::Asin
| BuiltinScalarFunction::Atan
| BuiltinScalarFunction::Acosh
diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs
index 6f258e6df3..f976a3468d 100644
--- a/datafusion/optimizer/src/analyzer/type_coercion.rs
+++ b/datafusion/optimizer/src/analyzer/type_coercion.rs
@@ -883,14 +883,14 @@ mod test {
fn scalar_function() -> Result<()> {
let empty = empty();
let lit_expr = lit(10i64);
- let fun: BuiltinScalarFunction = BuiltinScalarFunction::Abs;
+ let fun: BuiltinScalarFunction = BuiltinScalarFunction::Acos;
let scalar_function_expr =
Expr::ScalarFunction(ScalarFunction::new(fun, vec![lit_expr]));
let plan = LogicalPlan::Projection(Projection::try_new(
vec![scalar_function_expr],
empty,
)?);
- let expected = "Projection: abs(CAST(Int64(10) AS Float64))\n EmptyRelation";
+ let expected = "Projection: acos(CAST(Int64(10) AS Float64))\n EmptyRelation";
assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), &plan, expected)
}
diff --git a/datafusion/physical-expr/src/functions.rs b/datafusion/physical-expr/src/functions.rs
index 233b903d11..503cfe412b 100644
--- a/datafusion/physical-expr/src/functions.rs
+++ b/datafusion/physical-expr/src/functions.rs
@@ -171,6 +171,11 @@ pub fn create_physical_expr(
)))))
})
}
+ BuiltinScalarFunction::Abs => {
+ let input_data_type = input_phy_exprs[0].data_type(input_schema)?;
+ let abs_fun = math_expressions::create_abs_function(&input_data_type)?;
+ Arc::new(move |args| make_scalar_function(abs_fun)(args))
+ }
// These don't need args and input schema
_ => create_physical_fun(fun, execution_props)?,
};
@@ -360,7 +365,6 @@ pub fn create_physical_fun(
) -> Result<ScalarFunctionImplementation> {
Ok(match fun {
// math functions
- BuiltinScalarFunction::Abs => Arc::new(math_expressions::abs),
BuiltinScalarFunction::Acos => Arc::new(math_expressions::acos),
BuiltinScalarFunction::Asin => Arc::new(math_expressions::asin),
BuiltinScalarFunction::Atan => Arc::new(math_expressions::atan),
diff --git a/datafusion/physical-expr/src/math_expressions.rs b/datafusion/physical-expr/src/math_expressions.rs
index d1bfe4970c..a9dc0bd58f 100644
--- a/datafusion/physical-expr/src/math_expressions.rs
+++ b/datafusion/physical-expr/src/math_expressions.rs
@@ -18,11 +18,15 @@
//! Math expressions
use arrow::array::ArrayRef;
-use arrow::array::{BooleanArray, Float32Array, Float64Array, Int64Array};
+use arrow::array::{
+ BooleanArray, Decimal128Array, Float32Array, Float64Array, Int16Array, Int32Array,
+ Int64Array, Int8Array,
+};
use arrow::datatypes::DataType;
-use datafusion_common::internal_err;
+use arrow::error::ArrowError;
use datafusion_common::ScalarValue;
use datafusion_common::ScalarValue::{Float32, Int64};
+use datafusion_common::{internal_err, not_impl_err};
use datafusion_common::{DataFusionError, Result};
use datafusion_expr::ColumnarValue;
use rand::{thread_rng, Rng};
@@ -31,6 +35,8 @@ use std::iter;
use std::mem::swap;
use std::sync::Arc;
+type MathArrayFunction = fn(&[ArrayRef]) -> Result<ArrayRef>;
+
macro_rules! downcast_compute_op {
($ARRAY:expr, $NAME:expr, $FUNC:ident, $TYPE:ident) => {{
let n = $ARRAY.as_any().downcast_ref::<$TYPE>();
@@ -667,6 +673,70 @@ fn compute_truncate64(x: f64, y: i64) -> f64 {
(x * factor).round() / factor
}
+macro_rules! make_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.abs());
+ Ok(Arc::new(res) as ArrayRef)
+ }
+ }};
+}
+
+macro_rules! make_try_abs_function {
+ ($ARRAY_TYPE:ident) => {{
+ |args: &[ArrayRef]| {
+ let array = downcast_arg!(&args[0], "abs arg", $ARRAY_TYPE);
+ let res: $ARRAY_TYPE = array.try_unary(|x| {
+ x.checked_abs().ok_or_else(|| {
+ ArrowError::ComputeError(format!(
+ "{} overflow on abs({})",
+ stringify!($ARRAY_TYPE),
+ x
+ ))
+ })
+ })?;
+ 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(
+ input_data_type: &DataType,
+) -> Result<MathArrayFunction> {
+ match input_data_type {
+ DataType::Float32 => Ok(make_abs_function!(Float32Array)),
+ DataType::Float64 => Ok(make_abs_function!(Float64Array)),
+
+ // Types that may overflow, such as abs(-128_i8).
+ DataType::Int8 => Ok(make_try_abs_function!(Int8Array)),
+ DataType::Int16 => Ok(make_try_abs_function!(Int16Array)),
+ DataType::Int32 => Ok(make_try_abs_function!(Int32Array)),
+ DataType::Int64 => Ok(make_try_abs_function!(Int64Array)),
+
+ // Types of results are the same as the input.
+ DataType::Null
+ | DataType::UInt8
+ | DataType::UInt16
+ | 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)
+ }),
+
+ other => not_impl_err!("Unsupported data type {other:?} for function abs"),
+ }
+}
+
#[cfg(test)]
mod tests {
diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs
index 88558d1be2..71d5a35569 100644
--- a/datafusion/proto/src/physical_plan/mod.rs
+++ b/datafusion/proto/src/physical_plan/mod.rs
@@ -1835,12 +1835,12 @@ mod roundtrip_tests {
let execution_props = ExecutionProps::new();
let fun_expr = functions::create_physical_fun(
- &BuiltinScalarFunction::Abs,
+ &BuiltinScalarFunction::Acos,
&execution_props,
)?;
let expr = ScalarFunctionExpr::new(
- "abs",
+ "acos",
fun_expr,
vec![col("a", &schema)?],
&DataType::Int64,
diff --git a/datafusion/sqllogictest/test_files/decimal.slt b/datafusion/sqllogictest/test_files/decimal.slt
index 3405401f7c..a326a0cc49 100644
--- a/datafusion/sqllogictest/test_files/decimal.slt
+++ b/datafusion/sqllogictest/test_files/decimal.slt
@@ -469,7 +469,7 @@ select c1%c5 from decimal_simple;
query T
select arrow_typeof(abs(c1)) from decimal_simple limit 1;
----
-Float64
+Decimal128(10, 6)
query R rowsort
diff --git a/datafusion/sqllogictest/test_files/math.slt b/datafusion/sqllogictest/test_files/math.slt
index a56d4d2ecf..a3ee307f49 100644
--- a/datafusion/sqllogictest/test_files/math.slt
+++ b/datafusion/sqllogictest/test_files/math.slt
@@ -112,6 +112,18 @@ SELECT iszero(1.0), iszero(0.0), iszero(-0.0), iszero(NULL)
----
false true true NULL
+# abs: empty argumnet
+statement error DataFusion error: Error during planning: No function matches the given name and argument types 'abs\(\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tabs\(Any\)
+SELECT abs();
+
+# abs: wrong number of arguments
+statement error DataFusion error: Error during planning: No function matches the given name and argument types 'abs\(Int64, Int64\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tabs\(Any\)
+SELECT abs(1, 2);
+
+# abs: unsupported argument type
+statement error DataFusion error: Optimizer rule 'simplify_expressions' failed\ncaused by\nThis feature is not implemented: Unsupported data type Utf8 for function abs
+SELECT abs('foo');
+
statement ok
CREATE TABLE test_nullable_integer(
@@ -123,35 +135,45 @@ CREATE TABLE test_nullable_integer(
c6 SMALLINT UNSIGNED,
c7 INT UNSIGNED,
c8 BIGINT UNSIGNED,
+ dataset TEXT
)
AS VALUES
- (NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL);
+ (NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, 'nulls'),
+ (0, 0, 0, 0, 0, 0, 0, 0, 'zeros'),
+ (1, 1, 1, 1, 1, 1, 1, 1, 'ones');
-query IIIIIIII
-SELECT c1*0, c2*0, c3*0, c4*0, c5*0, c6*0, c7*0, c8*0 FROM test_nullable_integer
+query IIIIIIIIT
+INSERT into test_nullable_integer values(-128, -32768, -2147483648, -9223372036854775808, 0, 0, 0, 0, 'mins');
----
-NULL NULL NULL NULL NULL NULL NULL NULL
+1
+
+query IIIIIIIIT
+INSERT into test_nullable_integer values(127, 32767, 2147483647, 9223372036854775807, 255, 65535, 4294967295, 18446744073709551615, 'maxs');
+----
+1
query IIIIIIII
-SELECT c1/0, c2/0, c3/0, c4/0, c5/0, c6/0, c7/0, c8/0 FROM test_nullable_integer
+SELECT c1*0, c2*0, c3*0, c4*0, c5*0, c6*0, c7*0, c8*0 FROM test_nullable_integer where dataset = 'nulls'
----
NULL NULL NULL NULL NULL NULL NULL NULL
query IIIIIIII
-SELECT c1%0, c2%0, c3%0, c4%0, c5%0, c6%0, c7%0, c8%0 FROM test_nullable_integer
+SELECT c1/0, c2/0, c3/0, c4/0, c5/0, c6/0, c7/0, c8/0 FROM test_nullable_integer where dataset = 'nulls'
----
NULL NULL NULL NULL NULL NULL NULL NULL
query IIIIIIII
-INSERT INTO test_nullable_integer VALUES(1, 1, 1, 1, 1, 1, 1, 1)
+SELECT c1%0, c2%0, c3%0, c4%0, c5%0, c6%0, c7%0, c8%0 FROM test_nullable_integer where dataset = 'nulls'
----
-1
+NULL NULL NULL NULL NULL NULL NULL NULL
query IIIIIIII rowsort
-select c1*0, c2*0, c3*0, c4*0, c5*0, c6*0, c7*0, c8*0 from test_nullable_integer
+select c1*0, c2*0, c3*0, c4*0, c5*0, c6*0, c7*0, c8*0 from test_nullable_integer where dataset != 'maxs'
----
0 0 0 0 0 0 0 0
-NULL NULL NULL NULL NULL NULL NULL NULL
+0 0 0 0 0 0 0 0
+0 0 0 0 0 0 0 0
+NULL NULL NULL NULL NULL NULL NULL NULL
query error DataFusion error: Arrow error: Divide by zero error
SELECT c1/0 FROM test_nullable_integer
@@ -201,6 +223,50 @@ SELECT c7%0 FROM test_nullable_integer
query error DataFusion error: Arrow error: Divide by zero error
SELECT c8%0 FROM test_nullable_integer
+# abs: return type
+query TTTTTTTT rowsort
+select
+ arrow_typeof(abs(c1)), arrow_typeof(abs(c2)), arrow_typeof(abs(c3)), arrow_typeof(abs(c4)),
+ arrow_typeof(abs(c5)), arrow_typeof(abs(c6)), arrow_typeof(abs(c7)), arrow_typeof(abs(c8))
+from test_nullable_integer limit 1
+----
+Int8 Int16 Int32 Int64 UInt8 UInt16 UInt32 UInt64
+
+# abs: unsigned integers
+query IIII rowsort
+select abs(c5), abs(c6), abs(c7), abs(c8) from test_nullable_integer
+----
+0 0 0 0
+0 0 0 0
+1 1 1 1
+255 65535 4294967295 18446744073709551615
+NULL NULL NULL NULL
+
+# abs: signed integers
+query IIII rowsort
+select abs(c1), abs(c2), abs(c3), abs(c4) from test_nullable_integer where dataset != 'mins'
+----
+0 0 0 0
+1 1 1 1
+127 32767 2147483647 9223372036854775807
+NULL NULL NULL NULL
+
+# abs: Int8 overlow
+statement error DataFusion error: Arrow error: Compute error: Int8Array overflow on abs\(-128\)
+select abs(c1) from test_nullable_integer where dataset = 'mins'
+
+# abs: Int16 overlow
+statement error DataFusion error: Arrow error: Compute error: Int16Array overflow on abs\(-32768\)
+select abs(c2) from test_nullable_integer where dataset = 'mins'
+
+# abs: Int32 overlow
+statement error DataFusion error: Arrow error: Compute error: Int32Array overflow on abs\(-2147483648\)
+select abs(c3) from test_nullable_integer where dataset = 'mins'
+
+# abs: Int64 overlow
+statement error DataFusion error: Arrow error: Compute error: Int64Array overflow on abs\(-9223372036854775808\)
+select abs(c4) from test_nullable_integer where dataset = 'mins'
+
statement ok
drop table test_nullable_integer
@@ -327,6 +393,22 @@ SELECT c1%1, c2%1 FROM test_nullable_float
NULL NULL
NaN NaN
+# abs: return type
+query TT rowsort
+SELECT arrow_typeof(c1), arrow_typeof(c2) FROM test_nullable_float limit 1
+----
+Float32 Float64
+
+# abs: floats
+query RR rowsort
+SELECT abs(c1), abs(c2) from test_nullable_float
+----
+0 0
+1 1
+1 1
+NULL NULL
+NaN NaN
+
statement ok
drop table test_nullable_float
@@ -383,7 +465,18 @@ drop table test_non_nullable_float
statement ok
-CREATE TABLE test_nullable_decimal(c1 DECIMAL(9, 2)) AS VALUES (1), (NULL);
+CREATE TABLE test_nullable_decimal(
+ c1 DECIMAL(10, 2),
+ c2 DECIMAL(38, 10)
+ ) AS VALUES (0, 0), (NULL, NULL);
+
+query RR
+INSERT into test_nullable_decimal values
+ (-99999999.99, '-9999999999999999999999999999.9999999999'),
+ (99999999.99, '9999999999999999999999999999.9999999999');
+----
+2
+
query R rowsort
SELECT c1*0 FROM test_nullable_decimal WHERE c1 IS NULL;
@@ -404,6 +497,8 @@ query R rowsort
SELECT c1*0 FROM test_nullable_decimal WHERE c1 IS NOT NULL;
----
0
+0
+0
query error DataFusion error: Arrow error: Divide by zero error
SELECT c1/0 FROM test_nullable_decimal WHERE c1 IS NOT NULL;
@@ -411,6 +506,21 @@ SELECT c1/0 FROM test_nullable_decimal WHERE c1 IS NOT NULL;
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
+----
+Decimal128(10, 2) Decimal128(38, 10)
+
+# abs: Decimal128
+query RR rowsort
+SELECT abs(c1), abs(c2) FROM test_nullable_decimal
+----
+0 0
+99999999.99 9999999999999999999999999999.9999999999
+99999999.99 9999999999999999999999999999.9999999999
+NULL NULL
+
statement ok
drop table test_nullable_decimal
@@ -428,11 +538,11 @@ SELECT c1*0 FROM test_non_nullable_decimal
----
0
-query error DataFusion error: Optimizer rule 'simplify_expressions' failed\ncaused by\nArrow error: Divide by zero error
+query error DataFusion error: Optimizer rule 'simplify_expressions' failed\ncaused by\nArrow error: Divide by zero error
SELECT c1/0 FROM test_non_nullable_decimal
-query error DataFusion error: Optimizer rule 'simplify_expressions' failed\ncaused by\nArrow error: Divide by zero error
+query error DataFusion error: Optimizer rule 'simplify_expressions' failed\ncaused by\nArrow error: Divide by zero error
SELECT c1%0 FROM test_non_nullable_decimal
statement ok
-drop table test_non_nullable_decimal
\ No newline at end of file
+drop table test_non_nullable_decimal
diff --git a/datafusion/sqllogictest/test_files/scalar.slt b/datafusion/sqllogictest/test_files/scalar.slt
index 7d7f269606..e5c1a82849 100644
--- a/datafusion/sqllogictest/test_files/scalar.slt
+++ b/datafusion/sqllogictest/test_files/scalar.slt
@@ -67,19 +67,19 @@ CREATE TABLE small_floats(
## abs
# abs scalar function
-query RRR rowsort
+query III rowsort
select abs(64), abs(0), abs(-64);
----
64 0 64
# abs scalar nulls
-query R rowsort
+query ? rowsort
select abs(null);
----
NULL
# abs with columns
-query RRR rowsort
+query III rowsort
select abs(a), abs(b), abs(c) from signed_integers;
----
1 100 567