diff --git a/cpp/src/arrow/compute/api_scalar.cc b/cpp/src/arrow/compute/api_scalar.cc index 117426af44c7..b808b065d48c 100644 --- a/cpp/src/arrow/compute/api_scalar.cc +++ b/cpp/src/arrow/compute/api_scalar.cc @@ -787,6 +787,7 @@ Result RoundToMultiple(const Datum& arg, RoundToMultipleOptions options, SCALAR_ARITHMETIC_BINARY(Add, "add", "add_checked") SCALAR_ARITHMETIC_BINARY(Divide, "divide", "divide_checked") SCALAR_ARITHMETIC_BINARY(Logb, "logb", "logb_checked") +SCALAR_ARITHMETIC_BINARY(Modulo, "modulo", "modulo_checked") SCALAR_ARITHMETIC_BINARY(Multiply, "multiply", "multiply_checked") SCALAR_ARITHMETIC_BINARY(Power, "power", "power_checked") SCALAR_ARITHMETIC_BINARY(ShiftLeft, "shift_left", "shift_left_checked") diff --git a/cpp/src/arrow/compute/api_scalar.h b/cpp/src/arrow/compute/api_scalar.h index 492ea05f6d57..8df85b7464fb 100644 --- a/cpp/src/arrow/compute/api_scalar.h +++ b/cpp/src/arrow/compute/api_scalar.h @@ -632,6 +632,21 @@ Result Subtract(const Datum& left, const Datum& right, ArithmeticOptions options = ArithmeticOptions(), ExecContext* ctx = NULLPTR); +/// \brief Get the modulo of dividing two values. +/// Array values must be the same length. +/// If either argument is null the result will be null. +/// For integer types, if there is a zero divisor, an error will be raised. +/// +/// \param[in] left the dividend +/// \param[in] right the divisor +/// \param[in] options arithmetic options (enable/disable overflow checking), optional +/// \param[in] ctx the function execution context, optional +/// \return the elementwise remainder +ARROW_EXPORT +Result Modulo(const Datum& left, const Datum& right, + ArithmeticOptions options = ArithmeticOptions(), + ExecContext* ctx = NULLPTR); + /// \brief Multiply two values. Array values must be the same length. If either /// factor is null the result will be null. /// diff --git a/cpp/src/arrow/compute/kernels/base_arithmetic_internal.h b/cpp/src/arrow/compute/kernels/base_arithmetic_internal.h index 26c44a8ff077..f64083799307 100644 --- a/cpp/src/arrow/compute/kernels/base_arithmetic_internal.h +++ b/cpp/src/arrow/compute/kernels/base_arithmetic_internal.h @@ -31,6 +31,7 @@ namespace arrow { using internal::AddWithOverflow; using internal::DivideWithOverflow; +using internal::ModuloWithOverflow; using internal::MultiplyWithOverflow; using internal::NegateWithOverflow; using internal::SubtractWithOverflow; @@ -464,6 +465,62 @@ struct FloatingDivideChecked { // TODO: Add decimal }; +struct Modulo { + template + static enable_if_floating_value Call(KernelContext*, Arg0 left, Arg1 right, + Status* st) { + *st = Status::Invalid("Not implemented"); + return 0; + } + + template + static enable_if_integer_value Call(KernelContext*, Arg0 left, Arg1 right, + Status* st) { + if (ARROW_PREDICT_FALSE(right == 0)) { + *st = Status::Invalid("Modulo by zero"); + return 0; + } + + return left % right; + } + + template + static enable_if_decimal_value Call(KernelContext* ctx, Arg0 left, Arg1 right, + Status* st) { + return Divide::Call(ctx, left, right, st); + } +}; + +struct ModuloChecked { + template + static enable_if_floating_value Call(KernelContext*, Arg0 left, Arg1 right, + Status* st) { + *st = Status::Invalid("Not implemented"); + return 0; + } + + template + static enable_if_integer_value Call(KernelContext*, Arg0 left, Arg1 right, + Status* st) { + static_assert(std::is_same::value && std::is_same::value, ""); + T result; + if (ARROW_PREDICT_FALSE(ModuloWithOverflow(left, right, &result))) { + if (right == 0) { + *st = Status::Invalid("Modulo by zero"); + } else { + *st = Status::Invalid("Overflow"); + } + } + return result; + } + + template + static enable_if_decimal_value Call(KernelContext* ctx, Arg0 left, Arg1 right, + Status* st) { + return Divide::Call(ctx, left, right, st); + } +}; + struct Negate { template static constexpr enable_if_floating_value Call(KernelContext*, Arg arg, Status*) { diff --git a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc index c20cfc5688e9..ce9233f1ba73 100644 --- a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc +++ b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc @@ -673,7 +673,7 @@ void AddDecimalBinaryKernels(const std::string& name, ScalarFunction* func) { out_type = OutputType(ResolveDecimalAdditionOrSubtractionOutput); } else if (op == "multiply") { out_type = OutputType(ResolveDecimalMultiplicationOutput); - } else if (op == "divide") { + } else if (op == "divide" || op == "modulo") { out_type = OutputType(ResolveDecimalDivisionOutput); } else { DCHECK(false); @@ -764,7 +764,7 @@ struct ArithmeticFunction : ScalarFunction { return CastBinaryDecimalArgs(DecimalPromotion::kAdd, types); } else if (op == "multiply") { return CastBinaryDecimalArgs(DecimalPromotion::kMultiply, types); - } else if (op == "divide") { + } else if (op == "divide" || op == "modulo") { return CastBinaryDecimalArgs(DecimalPromotion::kDivide, types); } else { return Status::Invalid("Invalid decimal function: ", func_name); @@ -1149,6 +1149,18 @@ const FunctionDoc div_checked_doc{ "integer overflow is encountered."), {"dividend", "divisor"}}; +const FunctionDoc modulo_doc{ + "Get the modulo of dividing two values", + ("Integer division by zero returns an error.\n" + "Use function \"modulo_checked\" if you want to get an error\n" + "in all the aforementioned cases."), + {"dividend", "divisor"}}; + +const FunctionDoc modulo_checked_doc{ + "Get the modulo of dividing two values", + ("An error is returned when trying to divide by zero."), + {"dividend", "divisor"}}; + const FunctionDoc negate_doc{"Negate the argument element-wise", ("Results will wrap around on integer overflow.\n" "Use function \"negate_checked\" if you want overflow\n" @@ -1692,6 +1704,16 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { DCHECK_OK(registry->AddFunction(std::move(divide_checked))); + // ---------------------------------------------------------------------- + auto modulo = MakeArithmeticFunctionNotNull("modulo", modulo_doc); + AddDecimalBinaryKernels("modulo", modulo.get()); + DCHECK_OK(registry->AddFunction(std::move(modulo))); + + auto modulo_checked = + MakeArithmeticFunctionNotNull("modulo_checked", modulo_checked_doc); + AddDecimalBinaryKernels("modulo_checked", modulo_checked.get()); + DCHECK_OK(registry->AddFunction(std::move(modulo_checked))); + // ---------------------------------------------------------------------- auto negate = MakeUnaryArithmeticFunction("negate", negate_doc); AddDecimalUnaryKernels(negate.get()); diff --git a/cpp/src/arrow/compute/kernels/scalar_arithmetic_test.cc b/cpp/src/arrow/compute/kernels/scalar_arithmetic_test.cc index 1162dad855da..b4084a72674c 100644 --- a/cpp/src/arrow/compute/kernels/scalar_arithmetic_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_arithmetic_test.cc @@ -900,6 +900,47 @@ TYPED_TEST(TestBinaryArithmeticSigned, DivideOverflowRaises) { this->AssertBinop(Divide, MakeArray(min), MakeArray(-1), "[0]"); } +TYPED_TEST(TestBinaryArithmeticIntegral, Modulo) { + for (auto check_overflow : {false, true}) { + this->SetOverflowCheck(check_overflow); + + // Empty arrays + this->AssertBinop(Modulo, "[]", "[]", "[]"); + // Ordinary arrays + this->AssertBinop(Modulo, "[3, 2, 6]", "[1, 1, 2]", "[0, 0, 0]"); + // Array with nulls + this->AssertBinop(Modulo, "[null, 10, 30, null, 20]", "[1, 4, 2, 5, 10]", + "[null, 2, 0, null, 0]"); + // Scalar % Array + this->AssertBinop(Modulo, 33, "[null, 1, 3, null, 2]", "[null, 0, 0, null, 1]"); + // Array % Scalar + this->AssertBinop(Modulo, "[null, 10, 30, null, 2]", 3, "[null, 1, 0, null, 2]"); + // Scalar % Scalar + this->AssertBinop(Modulo, 16, 7, 2); + } +} + +TYPED_TEST(TestBinaryArithmeticSigned, Modulo) { + // Ordinary arrays + this->AssertBinop(Modulo, "[-3, 2, -7]", "[1, 1, 2]", "[0, 0, -1]"); + // Array with nulls + this->AssertBinop(Modulo, "[null, 10, 30, null, -21]", "[1, 4, 2, 5, 10]", + "[null, 2, 0, null, -1]"); + // Scalar % Array + this->AssertBinop(Modulo, 33, "[null, -1, -3, null, 2]", "[null, 0, 0, null, 1]"); + // Array % Scalar + this->AssertBinop(Modulo, "[null, 10, 30, null, 2]", 3, "[null, 1, 0, null, 2]"); + // Scalar % Scalar + this->AssertBinop(Modulo, -17, -8, -1); +} + +TYPED_TEST(TestBinaryArithmeticIntegral, ModuloByZero) { + for (auto check_overflow : {false, true}) { + this->SetOverflowCheck(check_overflow); + this->AssertBinopRaises(Modulo, "[3, 2, 6]", "[1, 1, 0]", "Modulo by zero"); + } +} + TYPED_TEST(TestBinaryArithmeticFloating, Power) { using CType = typename TestFixture::CType; auto max = std::numeric_limits::max(); diff --git a/cpp/src/arrow/util/basic_decimal.cc b/cpp/src/arrow/util/basic_decimal.cc index fc69bcf6f8ec..8c85fadc108f 100644 --- a/cpp/src/arrow/util/basic_decimal.cc +++ b/cpp/src/arrow/util/basic_decimal.cc @@ -1398,6 +1398,14 @@ BasicDecimal256 operator/(const BasicDecimal256& left, const BasicDecimal256& ri return result; } +BasicDecimal256 operator%(const BasicDecimal256& left, const BasicDecimal256& right) { + BasicDecimal256 remainder; + BasicDecimal256 result; + auto s = left.Divide(right, &result, &remainder); + DCHECK_EQ(s, DecimalStatus::kSuccess); + return remainder; +} + // Explicitly instantiate template base class, for DLL linking on Windows template class GenericBasicDecimal; template class GenericBasicDecimal; diff --git a/cpp/src/arrow/util/basic_decimal.h b/cpp/src/arrow/util/basic_decimal.h index 638c4870f1de..f35696f4ee5f 100644 --- a/cpp/src/arrow/util/basic_decimal.h +++ b/cpp/src/arrow/util/basic_decimal.h @@ -883,5 +883,7 @@ ARROW_EXPORT BasicDecimal256 operator*(const BasicDecimal256& left, const BasicDecimal256& right); ARROW_EXPORT BasicDecimal256 operator/(const BasicDecimal256& left, const BasicDecimal256& right); +ARROW_EXPORT BasicDecimal256 operator%(const BasicDecimal256& left, + const BasicDecimal256& right); } // namespace arrow diff --git a/cpp/src/arrow/util/int_util_overflow.h b/cpp/src/arrow/util/int_util_overflow.h index ffe78be2470d..85d8f2ddc78b 100644 --- a/cpp/src/arrow/util/int_util_overflow.h +++ b/cpp/src/arrow/util/int_util_overflow.h @@ -59,6 +59,7 @@ OPS_WITH_OVERFLOW(AddWithOverflow, add) OPS_WITH_OVERFLOW(SubtractWithOverflow, sub) OPS_WITH_OVERFLOW(MultiplyWithOverflow, mul) OPS_WITH_OVERFLOW(DivideWithOverflow, div) +OPS_WITH_OVERFLOW(ModuloWithOverflow, mod) #undef OP_WITH_OVERFLOW #undef OPS_WITH_OVERFLOW