diff --git a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc index 53d4751b8c4f..c7d27077176b 100644 --- a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc +++ b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc @@ -1126,6 +1126,40 @@ void AddArithmeticFunctionDurationTime(std::shared_ptr func) { std::move(exec_4))); } +// Apply Op in double precision and truncate toward zero, like numpy. +// NaN and out of range results are errors. +template +struct DurationByFloat { + template + static T Call(KernelContext* ctx, Arg0 left, Arg1 right, Status* st) { + double value = Op::template Call( + ctx, static_cast(left), static_cast(right), st); + double truncated = std::trunc(value); + // The int64 range is [-2^63, 2^63); the comparison is false for NaN + if (ARROW_PREDICT_FALSE(!(truncated >= -0x1p63 && truncated < 0x1p63))) { + *st = Status::Invalid(std::isnan(truncated) ? "result is NaN" : "overflow"); + return 0; + } + return static_cast(truncated); + } +}; + +// Add Op(duration, float64) -> duration, and Op(float64, duration) -> duration for +// Multiply. The same kernels are used for the checked and unchecked functions. +template +void AddDurationByFloatKernels(ScalarFunction* func) { + for (auto unit : TimeUnit::values()) { + auto exec = + ScalarBinaryNotNull>::Exec; + DCHECK_OK(func->AddKernel({duration(unit), float64()}, duration(unit), exec)); + if constexpr (std::is_same_v) { + exec = ScalarBinaryNotNull>::Exec; + DCHECK_OK(func->AddKernel({float64(), duration(unit)}, duration(unit), exec)); + } + } +} + const FunctionDoc absolute_value_doc{ "Calculate the absolute value of the argument element-wise", ("Results will wrap around on integer overflow.\n" @@ -1711,6 +1745,7 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { DCHECK_OK(multiply->AddKernel({duration(unit), int64()}, duration(unit), exec)); DCHECK_OK(multiply->AddKernel({int64(), duration(unit)}, duration(unit), exec)); } + AddDurationByFloatKernels(multiply.get()); DCHECK_OK(registry->AddFunction(std::move(multiply))); @@ -1727,6 +1762,7 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { DCHECK_OK( multiply_checked->AddKernel({int64(), duration(unit)}, duration(unit), exec)); } + AddDurationByFloatKernels(multiply_checked.get()); DCHECK_OK(registry->AddFunction(std::move(multiply_checked))); @@ -1740,6 +1776,7 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { DCHECK_OK( divide->AddKernel({duration(unit), int64()}, duration(unit), std::move(exec))); } + AddDurationByFloatKernels(divide.get()); // Add divide(duration, duration) -> float64 for (auto unit : TimeUnit::values()) { @@ -1761,6 +1798,7 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { DCHECK_OK(divide_checked->AddKernel({duration(unit), int64()}, duration(unit), std::move(exec))); } + AddDurationByFloatKernels(divide_checked.get()); // Add divide_checked(duration, duration) -> float64 for (auto unit : TimeUnit::values()) { diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc index 86a81ffdd384..f292722aaf38 100644 --- a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc @@ -1749,6 +1749,48 @@ TEST_F(ScalarTemporalTest, TestTemporalDivideDuration) { {duration(TimeUnit::NANO), duration(TimeUnit::NANO)}); } +TEST_F(ScalarTemporalTest, TestTemporalMultiplyDivideDurationByFloat) { + // Results are truncated, slots with a null are not computed, -2^63 is in range + for (auto u : TimeUnit::values()) { + auto durations = + ArrayFromJSON(duration(u), "[10, 5, 7, -5, 2, null, 1, -4611686018427387904]"); + for (auto op : {"multiply", "multiply_checked"}) { + CheckScalarBinaryCommutative( + op, durations, + ArrayFromJSON(float64(), "[2.5, 1.5, 1.5, 1.5, 0.4, NaN, null, 2]"), + ArrayFromJSON(duration(u), + "[25, 7, 10, -7, 0, null, null, -9223372036854775808]")); + } + for (auto op : {"divide", "divide_checked"}) { + CheckScalarBinary( + op, durations, ArrayFromJSON(float64(), "[0.5, 2, 2, 2, 3, 0, null, 0.5]"), + ArrayFromJSON(duration(u), + "[20, 2, 3, -2, 0, null, null, -9223372036854775808]")); + } + } + // NaN and out of range results, and division by zero, are errors in all variants + auto one = ArrayFromJSON(duration(TimeUnit::SECOND), "[1]"); + auto big = ArrayFromJSON(duration(TimeUnit::SECOND), "[4611686018427387904]"); + auto nan = ArrayFromJSON(float64(), "[NaN]"); + for (auto op : {"multiply", "multiply_checked"}) { + EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, ::testing::HasSubstr("NaN"), + CallFunction(op, {nan, one})); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("overflow"), + CallFunction(op, {big, ArrayFromJSON(float64(), "[2]")})); + } + for (auto op : {"divide", "divide_checked"}) { + EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, ::testing::HasSubstr("NaN"), + CallFunction(op, {one, nan})); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("overflow"), + CallFunction(op, {big, ArrayFromJSON(float64(), "[-0.25]")})); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("divide by zero"), + CallFunction(op, {one, ArrayFromJSON(float64(), "[0]")})); + } +} + TEST_F(ScalarTemporalTest, TestTemporalDifferenceWeeks) { auto raw_days = ArrayFromJSON(timestamp(TimeUnit::SECOND), R"([ "2021-08-09", "2021-08-10", "2021-08-11", "2021-08-12", "2021-08-13", "2021-08-14", "2021-08-15", diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst index a328d7d07655..3fa3cfa44f01 100644 --- a/docs/source/cpp/compute.rst +++ b/docs/source/cpp/compute.rst @@ -493,6 +493,9 @@ appropriately. Mixed decimal and floating-point arguments will cast all arguments to floating-point, while mixed decimal and integer arguments will cast all arguments to decimals. Mixed time resolution temporal inputs will be cast to finest input resolution. +A duration multiplied or divided by a float64 is computed in float64, then +truncated toward zero (like numpy); NaN, overflow and division by zero are errors +in all variants. +-------------------+--------+-------------------------+-------------------------------+-------+ | Function name | Arity | Input types | Output type | Notes | diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index 691d2ec98f33..5be3f095f5d3 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -1941,6 +1941,16 @@ def test_arithmetic_multiply(): assert result.equals(expected) +def test_arithmetic_multiply_divide_duration_by_float(): + # Example from GH-48003. Results are truncated toward zero, like numpy. + x = pa.array([0, 1, 2, 3, 4, 5], type=pa.duration("s")) + half = pa.array([0, 0, 1, 1, 2, 2], type=pa.duration("s")) + assert pc.multiply(x, 0.5).equals(half) + assert pc.divide(x, 2.0).equals(half) + assert pc.multiply(x, 1.5).equals( + pa.array([0, 1, 3, 4, 6, 7], type=pa.duration("s"))) + + def test_arithmetic_remainder_modulo(): left = pa.array([7, -7, 7, -7, None]) right = pa.array([3, 3, -3, -3, 3])