cop: add abs rpn function

src/coprocessor/dag/rpn_expr/impl_math.rs

@@ -0,0 +1,115 @@
+use cop_codegen::rpn_fn;
+
+use crate::coprocessor::codec::data_type::*;
+use crate::coprocessor::codec::{self, Error};
+use crate::coprocessor::Result;
+
+#[rpn_fn]
+#[inline]
+fn abs_int(arg: &Option<Int>) -> Result<Option<Int>> {
+    match arg {
+        None => Ok(None),
+        Some(arg) => match (*arg).checked_abs() {
+            None => Err(Error::overflow("BIGINT", &format!("abs({})", *arg)))?,
+            Some(arg_abs) => Ok(Some(arg_abs)),
+        },
+    }
+}
+
+#[rpn_fn]
+#[inline]
+fn abs_uint(arg: &Option<Int>) -> Result<Option<Int>> {
+    Ok(*arg)
+}
+
+#[rpn_fn]
+#[inline]
+fn abs_real(arg: &Option<Real>) -> Result<Option<Real>> {
+    match arg {
+        Some(arg) => Ok(Some(num_traits::Signed::abs(arg))),
+        None => Ok(None),
+    }
+}
+
+#[rpn_fn]
+#[inline]
+fn abs_decimal(arg: &Option<Decimal>) -> Result<Option<Decimal>> {
+    match arg {
+        Some(arg) => {
+            let res: codec::Result<Decimal> = arg.to_owned().abs().into();
+            Ok(Some(res?))
+        }
+        None => Ok(None),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use tipb::expression::ScalarFuncSig;
+
+    use crate::coprocessor::dag::rpn_expr::types::test_util::RpnFnScalarEvaluator;
+
+    #[test]
+    fn test_abs_int() {
+        let test_cases = vec![
+            (ScalarFuncSig::AbsInt, -3, Some(3), false),
+            (
+                ScalarFuncSig::AbsInt,
+                std::i64::MAX,
+                Some(std::i64::MAX),
+                false,
+            ),
+            (
+                ScalarFuncSig::AbsUInt,
+                std::u64::MAX as i64,
+                Some(std::u64::MAX as i64),
+                false,
+            ),
+            (ScalarFuncSig::AbsInt, std::i64::MIN, Some(0), true),
+        ];
+
+        for (sig, arg, expect_output, is_err) in test_cases {
+            let output = RpnFnScalarEvaluator::new().push_param(arg).evaluate(sig);
+
+            if is_err {
+                assert!(output.is_err());
+            } else {
+                let output = output.unwrap();
+                assert_eq!(output, expect_output, "{:?}", arg);
+            }
+        }
+    }
+
+    #[test]
+    fn test_abs_real() {
+        let test_cases: Vec<(Real, Option<Real>)> = vec![
+            (Real::new(3.5).unwrap(), Real::new(3.5).ok()),
+            (Real::new(-3.5).unwrap(), Real::new(3.5).ok()),
+        ];
+
+        for (arg, expect_output) in test_cases {
+            let output = RpnFnScalarEvaluator::new()
+                .push_param(arg)
+                .evaluate(ScalarFuncSig::AbsReal)
+                .unwrap();
+            assert_eq!(output, expect_output, "{:?}", arg);
+        }
+    }
+
+    #[test]
+    fn test_abs_decimal() {
+        let test_cases = vec![("1.1", "1.1"), ("-1.1", "1.1")];
+
+        for (arg, expect_output) in test_cases {
+            let arg = arg.parse::<Decimal>().ok();
+            let expect_output = expect_output.parse::<Decimal>().ok();
+            let output = RpnFnScalarEvaluator::new()
+                .push_param(arg.clone())
+                .evaluate(ScalarFuncSig::AbsDecimal)
+                .unwrap();
+            assert_eq!(output, expect_output, "{:?}", arg);
+        }
+    }
+}

src/coprocessor/dag/rpn_expr/mod.rs

@@ -7,6 +7,7 @@ pub mod impl_cast;
 pub mod impl_compare;
 pub mod impl_control;
 pub mod impl_like;
+pub mod impl_math;
 pub mod impl_op;
 
 pub use self::types::*;
@@ -21,6 +22,7 @@ use self::impl_arithmetic::*;
 use self::impl_compare::*;
 use self::impl_control::*;
 use self::impl_like::*;
+use self::impl_math::*;
 use self::impl_op::*;
 
 fn map_int_sig<F>(value: ScalarFuncSig, children: &[Expr], mapper: F) -> Result<RpnFnMeta>
@@ -168,6 +170,10 @@ fn map_pb_sig_to_rpn_func(value: ScalarFuncSig, children: &[Expr]) -> Result<Rpn
         ScalarFuncSig::IfNullTime => if_null_fn_meta::<DateTime>(),
         ScalarFuncSig::IfNullDuration => if_null_fn_meta::<Duration>(),
         ScalarFuncSig::IfNullJson => if_null_fn_meta::<Json>(),
+        ScalarFuncSig::AbsInt => abs_int_fn_meta(),
+        ScalarFuncSig::AbsUInt => abs_uint_fn_meta(),
+        ScalarFuncSig::AbsReal => abs_real_fn_meta(),
+        ScalarFuncSig::AbsDecimal => abs_decimal_fn_meta(),
         ScalarFuncSig::CoalesceInt => coalesce_fn_meta::<Int>(),
         ScalarFuncSig::CoalesceReal => coalesce_fn_meta::<Real>(),
         ScalarFuncSig::CoalesceString => coalesce_fn_meta::<Bytes>(),