diff --git a/cpp/src/arrow/compute/api_scalar.cc b/cpp/src/arrow/compute/api_scalar.cc index 425274043ed..dd399c69c9e 100644 --- a/cpp/src/arrow/compute/api_scalar.cc +++ b/cpp/src/arrow/compute/api_scalar.cc @@ -699,6 +699,8 @@ SCALAR_ARITHMETIC_BINARY(Add, "add", "add_checked") SCALAR_ARITHMETIC_BINARY(Divide, "divide", "divide_checked") SCALAR_ARITHMETIC_BINARY(Logb, "logb", "logb_checked") SCALAR_ARITHMETIC_BINARY(Multiply, "multiply", "multiply_checked") +SCALAR_ARITHMETIC_BINARY(FloorDivide, "floor_divide", "floor_divide_checked") +SCALAR_ARITHMETIC_BINARY(Modulo, "modulo", "modulo_checked") SCALAR_ARITHMETIC_BINARY(Power, "power", "power_checked") SCALAR_ARITHMETIC_BINARY(ShiftLeft, "shift_left", "shift_left_checked") SCALAR_ARITHMETIC_BINARY(ShiftRight, "shift_right", "shift_right_checked") diff --git a/cpp/src/arrow/compute/api_scalar.h b/cpp/src/arrow/compute/api_scalar.h index 1c27757fcfc..adfaedc2e39 100644 --- a/cpp/src/arrow/compute/api_scalar.h +++ b/cpp/src/arrow/compute/api_scalar.h @@ -605,6 +605,38 @@ Result Divide(const Datum& left, const Datum& right, ArithmeticOptions options = ArithmeticOptions(), ExecContext* ctx = NULLPTR); +/// \brief Divide two values and return the largest integer smaller or equal to the +/// quotient. +/// +/// 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] dividend the dividend +/// \param[in] divisor the divisor +/// \param[in] options arithmetic options (enable/disable overflow checking), optional +/// \param[in] ctx the function execution context, optional +/// \return the elementwise quotient without the fractional part +ARROW_EXPORT +Result FloorDivide(const Datum& left, const Datum& right, + ArithmeticOptions options = ArithmeticOptions(), + ExecContext* ctx = NULLPTR); + +/// \brief Calculate the remainder of dividing two values. +/// +/// Array values must be the same length. +/// If either argument is null the result will be null. +/// If the divisor is zero, 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 Negate values. /// /// If argument 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 1cccdca1481..c445baccc3b 100644 --- a/cpp/src/arrow/compute/kernels/base_arithmetic_internal.h +++ b/cpp/src/arrow/compute/kernels/base_arithmetic_internal.h @@ -425,6 +425,94 @@ struct DivideChecked { } }; +struct FloorDivide { + template + static enable_if_floating_point Call(KernelContext*, Arg0 left, Arg1 right, + Status*) { + return std::floor(left / right); + } + + template + static enable_if_integer_value Call(KernelContext*, Arg0 left, Arg1 right, + Status* st) { + T result; + if (ARROW_PREDICT_FALSE(DivideWithOverflow(left, right, &result))) { + if (ARROW_PREDICT_FALSE(right == 0)) { + *st = Status::Invalid("divide by zero"); + result = 0; + } else { + result = 0; + } + } + return result; + } +}; + +struct FloorDivideChecked { + 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(DivideWithOverflow(left, right, &result))) { + if (ARROW_PREDICT_FALSE(right == 0)) { + *st = Status::Invalid("divide by zero"); + result = 0; + } else { + *st = Status::Invalid("overflow"); + } + } + return result; + } + + template + static enable_if_floating_point Call(KernelContext*, Arg0 left, Arg1 right, + Status* st) { + static_assert(std::is_same::value && std::is_same::value, ""); + if (ARROW_PREDICT_FALSE(right == 0)) { + *st = Status::Invalid("divide by zero"); + return 0; + } + return std::floor(left / right); + } +}; + +struct Remainder { + template + static enable_if_floating_point Call(KernelContext*, Arg0 left, Arg1 right, + Status*) { + return std::fmod(left, right); + } + + template + static enable_if_integer_value Call(KernelContext*, Arg0 left, Arg1 right, + Status* st) { + if (ARROW_PREDICT_FALSE(right == 0)) { + *st = Status::Invalid("divide by zero"); + return 0; + } + return left % right; + } +}; + +struct RemainderChecked { + template + static enable_if_floating_point Call(KernelContext*, Arg0 left, Arg1 right, + Status*) { + return std::fmod(left, right); + } + + template + static enable_if_integer_value Call(KernelContext*, Arg0 left, Arg1 right, + Status* st) { + if (ARROW_PREDICT_FALSE(right == 0)) { + *st = Status::Invalid("divide by zero"); + return 0; + } + return left % right; + } +}; + 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 79e7d0fee1a..dade23cba13 100644 --- a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc +++ b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc @@ -1699,6 +1699,34 @@ const FunctionDoc div_checked_doc{ "integer overflow is encountered."), {"dividend", "divisor"}}; +const FunctionDoc floor_div_doc{ + "Calculate the quotient without the fractional part", + ("Integer division by zero returns an error. However, integer overflow\n" + "wraps around, and floating-point division by zero returns an infinite.\n" + "Use function \"floor_divide_checked\" if you want to get an error\n" + "in all the aforementioned cases."), + {"dividend", "divisor"}}; + +const FunctionDoc floor_div_checked_doc{ + "Calculate the quotient without the fractional part", + ("An error is returned when trying to divide by zero, or when\n" + "integer overflow is encountered."), + {"dividend", "divisor"}}; + +const FunctionDoc remainder_doc{ + "Calculate the remainder of dividing two values", + ("Integer division by zero returns an error. However, integer overflow\n" + "wraps around, and floating-point division by zero returns an infinite.\n" + "Use function \"remainder_checked\" if you want to get an error\n" + "in all the aforementioned cases."), + {"dividend", "divisor"}}; + +const FunctionDoc remainder_checked_doc{ + "Calculate the remainder of dividing two values", + ("An error is returned when trying to divide by zero, or when\n" + "integer overflow is encountered."), + {"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" @@ -2206,6 +2234,25 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { DCHECK_OK(registry->AddFunction(std::move(divide_checked))); + // ---------------------------------------------------------------------- + auto floor_divide = + MakeArithmeticFunctionNotNull("floor_divide", floor_div_doc); + DCHECK_OK(registry->AddFunction(std::move(floor_divide))); + + // ---------------------------------------------------------------------- + auto floor_divide_checked = MakeArithmeticFunctionNotNull( + "floor_divide_checked", floor_div_checked_doc); + DCHECK_OK(registry->AddFunction(std::move(floor_divide_checked))); + + // ---------------------------------------------------------------------- + auto remainder = MakeArithmeticFunctionNotNull("remainder", remainder_doc); + DCHECK_OK(registry->AddFunction(std::move(remainder))); + + // ---------------------------------------------------------------------- + auto remainder_checked = MakeArithmeticFunctionNotNull( + "remainder_checked", remainder_checked_doc); + DCHECK_OK(registry->AddFunction(std::move(remainder_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 fc0d5eb2e15..8c8fa16a249 100644 --- a/cpp/src/arrow/compute/kernels/scalar_arithmetic_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_arithmetic_test.cc @@ -937,6 +937,22 @@ TYPED_TEST(TestBinaryArithmeticSigned, DivideOverflowRaises) { this->AssertBinop(Divide, MakeArray(min), MakeArray(-1), "[0]"); } +TYPED_TEST(TestBinaryArithmeticIntegral, Remainder) { + // Empty arrays + this->AssertBinop(Remainder, "[]", "[]", "[]"); + // Ordinary arrays + this->AssertBinop(Remainder, "[3, 2, 6]", "[1, 1, 2]", "[0, 0, 0]"); + // Array with nulls + this->AssertBinop(Remainder, "[null, 10, 30, null, 20]", "[1, 4, 2, 5, 10]", + "[null, 2, 0, null, 0]"); + // Scalar divides by array + this->AssertBinop(Remainder, 33, "[null, 1, 3, null, 2]", "[null, 0, 0, null, 1]"); + // Array divides by scalar + this->AssertBinop(Remainder, "[null, 10, 30, null, 2]", 3, "[null, 1, 0, null, 2]"); + // Scalar divides by scalar + this->AssertBinop(Remainder, 16, 7, 2); +} + TYPED_TEST(TestBinaryArithmeticFloating, Power) { using CType = typename TestFixture::CType; auto max = std::numeric_limits::max();