diff --git a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc index c248ee4caf7..b87328a3e6e 100644 --- a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc +++ b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc @@ -239,6 +239,62 @@ struct SubtractCheckedDate32 { } }; +template +struct AddTimeDuration { + template + static enable_if_t Call(KernelContext*, Arg0 left, Arg1 right, + Status* st) { + T result = arrow::internal::SafeSignedAdd(left, static_cast(right)); + if (result < 0 || multiple <= result) { + *st = Status::Invalid(result, " is not within the acceptable range of ", "[0, ", + multiple, ") s"); + } + return result; + } + + template + static enable_if_t Call(KernelContext*, Arg0 left, Arg1 right, + Status* st) { + T result = arrow::internal::SafeSignedAdd(left, right); + if (result < 0 || multiple <= result) { + *st = Status::Invalid(result, " is not within the acceptable range of ", "[0, ", + multiple, ") s"); + } + return result; + } +}; + +template +struct AddTimeDurationChecked { + template + static enable_if_t Call(KernelContext*, Arg0 left, Arg1 right, + Status* st) { + T result = 0; + if (ARROW_PREDICT_FALSE(AddWithOverflow(left, static_cast(right), &result))) { + *st = Status::Invalid("overflow"); + } + if (result < 0 || multiple <= result) { + *st = Status::Invalid(result, " is not within the acceptable range of ", "[0, ", + multiple, ") s"); + } + return result; + } + + template + static enable_if_t Call(KernelContext*, Arg0 left, Arg1 right, + Status* st) { + T result = 0; + if (ARROW_PREDICT_FALSE(AddWithOverflow(left, static_cast(right), &result))) { + *st = Status::Invalid("overflow"); + } + if (result < 0 || multiple <= result) { + *st = Status::Invalid(result, " is not within the acceptable range of ", "[0, ", + multiple, ") s"); + } + return result; + } +}; + template struct SubtractTimeDuration { template @@ -2553,6 +2609,8 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { DCHECK_OK(add->AddKernel({in_type, in_type}, duration(unit), std::move(exec))); } + AddArithmeticFunctionTimeDurations(add); + DCHECK_OK(registry->AddFunction(std::move(add))); // ---------------------------------------------------------------------- @@ -2577,6 +2635,8 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { add_checked->AddKernel({in_type, in_type}, duration(unit), std::move(exec))); } + AddArithmeticFunctionTimeDurations(add_checked); + DCHECK_OK(registry->AddFunction(std::move(add_checked))); // ---------------------------------------------------------------------- diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc index 522a942b04e..4cafa7953e8 100644 --- a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc @@ -1238,6 +1238,65 @@ TEST_F(ScalarTemporalTest, TestTemporalSubtractTime) { } } +TEST_F(ScalarTemporalTest, TestTemporalAddTimeAndDuration) { + for (auto op : {"add", "add_checked"}) { + auto arr_s = ArrayFromJSON(time32(TimeUnit::SECOND), times_s); + auto arr_s2 = ArrayFromJSON(time32(TimeUnit::SECOND), times_s2); + auto arr_ms = ArrayFromJSON(time32(TimeUnit::MILLI), times_ms); + auto arr_ms2 = ArrayFromJSON(time32(TimeUnit::MILLI), times_ms2); + auto arr_us = ArrayFromJSON(time64(TimeUnit::MICRO), times_us); + auto arr_us2 = ArrayFromJSON(time64(TimeUnit::MICRO), times_us2); + auto arr_ns = ArrayFromJSON(time64(TimeUnit::NANO), times_ns); + auto arr_ns2 = ArrayFromJSON(time64(TimeUnit::NANO), times_ns2); + + CheckScalarBinary(op, arr_s, + ArrayFromJSON(duration(TimeUnit::SECOND), seconds_between_time), + arr_s2); + CheckScalarBinary(op, arr_ms, + ArrayFromJSON(duration(TimeUnit::MILLI), milliseconds_between_time), + arr_ms2); + CheckScalarBinary(op, arr_us, + ArrayFromJSON(duration(TimeUnit::MICRO), microseconds_between_time), + arr_us2); + CheckScalarBinary(op, arr_ns, + ArrayFromJSON(duration(TimeUnit::NANO), nanoseconds_between_time), + arr_ns2); + + auto seconds_1 = ArrayFromJSON(time32(TimeUnit::SECOND), R"([1, null])"); + auto milliseconds_2k = ArrayFromJSON(duration(TimeUnit::MILLI), R"([2000, null])"); + auto milliseconds_3k = ArrayFromJSON(time32(TimeUnit::MILLI), R"([3000, null])"); + auto nanoseconds_1G = ArrayFromJSON(time64(TimeUnit::NANO), R"([1000000000, null])"); + auto microseconds_2M = ArrayFromJSON(duration(TimeUnit::MICRO), R"([2000000, null])"); + auto nanoseconds_3M = ArrayFromJSON(time64(TimeUnit::NANO), R"([3000000000, null])"); + auto microseconds_3M = ArrayFromJSON(time64(TimeUnit::MICRO), R"([3000000, null])"); + CheckScalarBinary(op, seconds_1, milliseconds_2k, milliseconds_3k); + CheckScalarBinary(op, nanoseconds_1G, microseconds_2M, nanoseconds_3M); + CheckScalarBinary(op, seconds_1, microseconds_2M, microseconds_3M); + + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, + ::testing::HasSubstr("-1 is not within the acceptable range of [0, 86400)"), + CallFunction(op, {ArrayFromJSON(time32(TimeUnit::SECOND), R"([0, null])"), + ArrayFromJSON(duration(TimeUnit::SECOND), R"([-1, null])")})); + + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, + ::testing::HasSubstr( + "86400000000001 is not within the acceptable range of [0, 86400000000000)"), + CallFunction(op, + {ArrayFromJSON(time64(TimeUnit::MICRO), R"([86400000000, null])"), + ArrayFromJSON(duration(TimeUnit::NANO), R"([1, null])")})); + + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, + ::testing::HasSubstr( + "86400000001 is not within the acceptable range of [0, 86400000000)"), + CallFunction(op, + {ArrayFromJSON(time64(TimeUnit::MICRO), R"([86400000000, null])"), + ArrayFromJSON(duration(TimeUnit::MICRO), R"([1, null])")})); + } +} + TEST_F(ScalarTemporalTest, TestTemporalSubtractTimeAndDuration) { for (auto op : {"subtract", "subtract_checked"}) { auto arr_s = ArrayFromJSON(time32(TimeUnit::SECOND), times_s);