diff --git a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc index b27d494d864..b8dae54b67a 100644 --- a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc +++ b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc @@ -214,6 +214,31 @@ struct SubtractChecked { } }; +struct SubtractDate32 { + static constexpr int64_t kSecondsInDay = 86400; + + template + static constexpr T Call(KernelContext*, Arg0 left, Arg1 right, Status*) { + return arrow::internal::SafeSignedSubtract(left, right) * kSecondsInDay; + } +}; + +struct SubtractCheckedDate32 { + static constexpr int64_t kSecondsInDay = 86400; + + template + static T Call(KernelContext*, Arg0 left, Arg1 right, Status* st) { + T result = 0; + if (ARROW_PREDICT_FALSE(SubtractWithOverflow(left, right, &result))) { + *st = Status::Invalid("overflow"); + } + if (ARROW_PREDICT_FALSE(MultiplyWithOverflow(result, kSecondsInDay, &result))) { + *st = Status::Invalid("overflow"); + } + return result; + } +}; + struct Multiply { static_assert(std::is_same::value, ""); static_assert(std::is_same::value, ""); @@ -1611,6 +1636,23 @@ Result ResolveDecimalDivisionOutput(KernelContext*, }); } +Result ResolveTemporalOutput(KernelContext*, + const std::vector& args) { + DCHECK_EQ(args[0].type->id(), args[1].type->id()); + auto left_type = checked_cast(args[0].type.get()); + auto right_type = checked_cast(args[1].type.get()); + DCHECK_EQ(left_type->unit(), left_type->unit()); + + if ((left_type->timezone() == "" || right_type->timezone() == "") && + left_type->timezone() != right_type->timezone()) { + return Status::Invalid("Subtraction of zoned and non-zoned times is ambiguous. (", + left_type->timezone(), right_type->timezone(), ")."); + } + + auto type = duration(right_type->unit()); + return ValueDescr(std::move(type), GetBroadcastShape(args)); +} + template void AddDecimalUnaryKernels(ScalarFunction* func) { OutputType out_type(FirstType); @@ -2426,7 +2468,9 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { for (auto unit : TimeUnit::values()) { InputType in_type(match::TimestampTypeUnit(unit)); auto exec = ArithmeticExecFromOp(Type::TIMESTAMP); - DCHECK_OK(subtract->AddKernel({in_type, in_type}, duration(unit), std::move(exec))); + DCHECK_OK(subtract->AddKernel({in_type, in_type}, + OutputType::Resolver(ResolveTemporalOutput), + std::move(exec))); } // Add subtract(timestamp, duration) -> timestamp @@ -2437,6 +2481,18 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { std::move(exec))); } + // Add subtract(date32, date32) -> duration(TimeUnit::SECOND) + InputType in_type_date_32(date32()); + auto exec_date_32 = ScalarBinaryEqualTypes::Exec; + DCHECK_OK(subtract->AddKernel({in_type_date_32, in_type_date_32}, + duration(TimeUnit::SECOND), std::move(exec_date_32))); + + // Add subtract(date64, date64) -> duration(TimeUnit::MILLI) + InputType in_type_date_64(date64()); + auto exec_date_64 = ScalarBinaryEqualTypes::Exec; + DCHECK_OK(subtract->AddKernel({in_type_date_64, in_type_date_64}, + duration(TimeUnit::MILLI), std::move(exec_date_64))); + DCHECK_OK(registry->AddFunction(std::move(subtract))); // ---------------------------------------------------------------------- @@ -2444,7 +2500,17 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { "subtract_checked", &sub_checked_doc); AddDecimalBinaryKernels("subtract_checked", subtract_checked.get()); - // Add subtract(timestamp, duration) -> timestamp + // Add subtract_checked(timestamp, timestamp) -> duration + for (auto unit : TimeUnit::values()) { + InputType in_type(match::TimestampTypeUnit(unit)); + auto exec = + ArithmeticExecFromOp(Type::TIMESTAMP); + DCHECK_OK(subtract_checked->AddKernel({in_type, in_type}, + OutputType::Resolver(ResolveTemporalOutput), + std::move(exec))); + } + + // Add subtract_checked(timestamp, duration) -> timestamp for (auto unit : TimeUnit::values()) { InputType in_type(match::TimestampTypeUnit(unit)); auto exec = @@ -2453,6 +2519,20 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { OutputType(FirstType), std::move(exec))); } + // Add subtract_checked(date32, date32) -> duration(TimeUnit::SECOND) + auto exec_date_32_checked = + ScalarBinaryEqualTypes::Exec; + DCHECK_OK(subtract_checked->AddKernel({in_type_date_32, in_type_date_32}, + duration(TimeUnit::SECOND), + std::move(exec_date_32_checked))); + + // Add subtract_checked(date64, date64) -> duration(TimeUnit::MILLI) + auto exec_date_64_checked = + ScalarBinaryEqualTypes::Exec; + DCHECK_OK(subtract_checked->AddKernel({in_type_date_64, in_type_date_64}, + duration(TimeUnit::MILLI), + std::move(exec_date_64_checked))); + DCHECK_OK(registry->AddFunction(std::move(subtract_checked))); // ---------------------------------------------------------------------- diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc index 271254dbcc7..67a5a8b57f0 100644 --- a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc @@ -977,131 +977,144 @@ TEST_F(ScalarTemporalTest, TestTemporalDifference) { } TEST_F(ScalarTemporalTest, TestTemporalSubtractDateAndDuration) { - std::string op = "subtract"; - std::string milliseconds_between_time_and_date = - "[31535941000, -31706603000, 2674840000, -2604800000, 82495000," - "-180610000, -11715000, -15620000, -19525000, -23430000, -27335000," - "-31240000, -35145000, -86400000, -26352000000, 5180277000, null]"; - std::string microseconds_between_time_and_date = - "[31535941000000, -31706603000000, 2674840000000, -2604800000000, 82495000000," - "-180610000000, -11715000000, -15620000000, -19525000000, -23430000000, " - "-27335000000, -31240000000, -35145000000, -86400000000, -26352000000000, " - "5180277000000, null]"; - auto dates32 = ArrayFromJSON(date32(), date32s2); - auto dates64 = ArrayFromJSON(date64(), date64s2); - - auto durations_ms = - ArrayFromJSON(duration(TimeUnit::MILLI), milliseconds_between_time_and_date); - auto timestamps_ms = ArrayFromJSON(timestamp(TimeUnit::MILLI), times_seconds_precision); - CheckScalarBinary(op, dates32, durations_ms, timestamps_ms); - CheckScalarBinary(op, dates64, durations_ms, timestamps_ms); - - auto durations_us = - ArrayFromJSON(duration(TimeUnit::MICRO), microseconds_between_time_and_date); - auto timestamps_us = ArrayFromJSON(timestamp(TimeUnit::MICRO), times_seconds_precision); - CheckScalarBinary(op, dates32, durations_us, timestamps_us); - CheckScalarBinary(op, dates64, durations_us, timestamps_us); -} - -TEST_F(ScalarTemporalTest, TestTemporalSubtractDateAndDurationChecked) { - std::string op = "subtract_checked"; - std::string milliseconds_between_time_and_date = - "[31535941000, -31706603000, 2674840000, -2604800000, 82495000," - "-180610000, -11715000, -15620000, -19525000, -23430000, -27335000," - "-31240000, -35145000, -86400000, -26352000000, 5180277000, null]"; - std::string microseconds_between_time_and_date = - "[31535941000000, -31706603000000, 2674840000000, -2604800000000, 82495000000," - "-180610000000, -11715000000, -15620000000, -19525000000, -23430000000, " - "-27335000000, -31240000000, -35145000000, -86400000000, -26352000000000, " - "5180277000000, null]"; - auto dates32 = ArrayFromJSON(date32(), date32s2); - auto dates64 = ArrayFromJSON(date64(), date64s2); - - auto durations_ms = - ArrayFromJSON(duration(TimeUnit::MILLI), milliseconds_between_time_and_date); - auto timestamps_ms = ArrayFromJSON(timestamp(TimeUnit::MILLI), times_seconds_precision); - CheckScalarBinary(op, dates32, durations_ms, timestamps_ms); - CheckScalarBinary(op, dates64, durations_ms, timestamps_ms); - - auto durations_us = - ArrayFromJSON(duration(TimeUnit::MICRO), microseconds_between_time_and_date); - auto timestamps_us = ArrayFromJSON(timestamp(TimeUnit::MICRO), times_seconds_precision); - CheckScalarBinary(op, dates32, durations_us, timestamps_us); - CheckScalarBinary(op, dates64, durations_us, timestamps_us); + for (auto op : {"subtract", "subtract_checked"}) { + std::string milliseconds_between_time_and_date = + "[31535941000, -31706603000, 2674840000, -2604800000, 82495000," + "-180610000, -11715000, -15620000, -19525000, -23430000, -27335000," + "-31240000, -35145000, -86400000, -26352000000, 5180277000, null]"; + std::string microseconds_between_time_and_date = + "[31535941000000, -31706603000000, 2674840000000, -2604800000000, 82495000000," + "-180610000000, -11715000000, -15620000000, -19525000000, -23430000000, " + "-27335000000, -31240000000, -35145000000, -86400000000, -26352000000000, " + "5180277000000, null]"; + auto dates32 = ArrayFromJSON(date32(), date32s2); + auto dates64 = ArrayFromJSON(date64(), date64s2); + + auto durations_ms = + ArrayFromJSON(duration(TimeUnit::MILLI), milliseconds_between_time_and_date); + auto timestamps_ms = + ArrayFromJSON(timestamp(TimeUnit::MILLI), times_seconds_precision); + CheckScalarBinary(op, dates32, durations_ms, timestamps_ms); + CheckScalarBinary(op, dates64, durations_ms, timestamps_ms); + + auto durations_us = + ArrayFromJSON(duration(TimeUnit::MICRO), microseconds_between_time_and_date); + auto timestamps_us = + ArrayFromJSON(timestamp(TimeUnit::MICRO), times_seconds_precision); + CheckScalarBinary(op, dates32, durations_us, timestamps_us); + CheckScalarBinary(op, dates64, durations_us, timestamps_us); + } } TEST_F(ScalarTemporalTest, TestTemporalSubtractTimestampAndDuration) { - std::string op = "subtract"; - for (auto tz : {"", "UTC", "Pacific/Marquesas"}) { - auto timestamp_unit_s = timestamp(TimeUnit::SECOND, tz); - auto duration_unit_s = duration(TimeUnit::SECOND); - auto timestamp_unit_ms = timestamp(TimeUnit::MILLI, tz); - auto duration_unit_ms = duration(TimeUnit::MILLI); - auto timestamp_unit_us = timestamp(TimeUnit::MICRO, tz); - auto duration_unit_us = duration(TimeUnit::MICRO); - auto timestamp_unit_ns = timestamp(TimeUnit::NANO, tz); - auto duration_unit_ns = duration(TimeUnit::NANO); - - CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_s, times_seconds_precision2), - ArrayFromJSON(duration_unit_s, seconds_between), - ArrayFromJSON(timestamp_unit_s, times_seconds_precision)); - CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_ms, times_seconds_precision2), - ArrayFromJSON(duration_unit_ms, milliseconds_between), - ArrayFromJSON(timestamp_unit_ms, times_seconds_precision)); - CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_us, times_seconds_precision2), - ArrayFromJSON(duration_unit_us, microseconds_between), - ArrayFromJSON(timestamp_unit_us, times_seconds_precision)); - CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_ns, times_seconds_precision2), - ArrayFromJSON(duration_unit_ns, nanoseconds_between), - ArrayFromJSON(timestamp_unit_ns, times_seconds_precision)); - } + for (auto op : {"subtract", "subtract_checked"}) { + for (auto tz : {"", "UTC", "Pacific/Marquesas"}) { + auto timestamp_unit_s = timestamp(TimeUnit::SECOND, tz); + auto duration_unit_s = duration(TimeUnit::SECOND); + auto timestamp_unit_ms = timestamp(TimeUnit::MILLI, tz); + auto duration_unit_ms = duration(TimeUnit::MILLI); + auto timestamp_unit_us = timestamp(TimeUnit::MICRO, tz); + auto duration_unit_us = duration(TimeUnit::MICRO); + auto timestamp_unit_ns = timestamp(TimeUnit::NANO, tz); + auto duration_unit_ns = duration(TimeUnit::NANO); + + CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_s, times_seconds_precision2), + ArrayFromJSON(duration_unit_s, seconds_between), + ArrayFromJSON(timestamp_unit_s, times_seconds_precision)); + CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_ms, times_seconds_precision2), + ArrayFromJSON(duration_unit_ms, milliseconds_between), + ArrayFromJSON(timestamp_unit_ms, times_seconds_precision)); + CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_us, times_seconds_precision2), + ArrayFromJSON(duration_unit_us, microseconds_between), + ArrayFromJSON(timestamp_unit_us, times_seconds_precision)); + CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_ns, times_seconds_precision2), + ArrayFromJSON(duration_unit_ns, nanoseconds_between), + ArrayFromJSON(timestamp_unit_ns, times_seconds_precision)); + } - auto seconds_3 = ArrayFromJSON(timestamp(TimeUnit::SECOND), R"([3, null])"); - auto milliseconds_2k = ArrayFromJSON(duration(TimeUnit::MILLI), R"([2000, null])"); - auto milliseconds_1k = ArrayFromJSON(timestamp(TimeUnit::MILLI), R"([1000, null])"); - CheckScalarBinary(op, seconds_3, milliseconds_2k, milliseconds_1k); + auto seconds_3 = ArrayFromJSON(timestamp(TimeUnit::SECOND), R"([3, null])"); + auto milliseconds_2k = ArrayFromJSON(duration(TimeUnit::MILLI), R"([2000, null])"); + auto milliseconds_1k = ArrayFromJSON(timestamp(TimeUnit::MILLI), R"([1000, null])"); + CheckScalarBinary(op, seconds_3, milliseconds_2k, milliseconds_1k); - auto seconds_3_tz = ArrayFromJSON(timestamp(TimeUnit::SECOND, "UTC"), R"([3, null])"); - auto milliseconds_1k_tz = - ArrayFromJSON(timestamp(TimeUnit::MILLI, "UTC"), R"([1000, null])"); - CheckScalarBinary(op, seconds_3_tz, milliseconds_2k, milliseconds_1k_tz); + auto seconds_3_tz = ArrayFromJSON(timestamp(TimeUnit::SECOND, "UTC"), R"([3, null])"); + auto milliseconds_1k_tz = + ArrayFromJSON(timestamp(TimeUnit::MILLI, "UTC"), R"([1000, null])"); + CheckScalarBinary(op, seconds_3_tz, milliseconds_2k, milliseconds_1k_tz); + } } -TEST_F(ScalarTemporalTest, TestTemporalSubtractCheckedTimestampAndDuration) { - std::string op = "subtract_checked"; - for (auto tz : {"", "UTC", "Pacific/Marquesas"}) { - auto timestamp_unit_s = timestamp(TimeUnit::SECOND, tz); - auto duration_unit_s = duration(TimeUnit::SECOND); - auto timestamp_unit_ms = timestamp(TimeUnit::MILLI, tz); - auto duration_unit_ms = duration(TimeUnit::MILLI); - auto timestamp_unit_us = timestamp(TimeUnit::MICRO, tz); - auto duration_unit_us = duration(TimeUnit::MICRO); - auto timestamp_unit_ns = timestamp(TimeUnit::NANO, tz); - auto duration_unit_ns = duration(TimeUnit::NANO); - - CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_s, times_seconds_precision2), - ArrayFromJSON(duration_unit_s, seconds_between), - ArrayFromJSON(timestamp_unit_s, times_seconds_precision)); - CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_ms, times_seconds_precision2), - ArrayFromJSON(duration_unit_ms, milliseconds_between), - ArrayFromJSON(timestamp_unit_ms, times_seconds_precision)); - CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_us, times_seconds_precision2), - ArrayFromJSON(duration_unit_us, microseconds_between), - ArrayFromJSON(timestamp_unit_us, times_seconds_precision)); - CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_ns, times_seconds_precision2), - ArrayFromJSON(duration_unit_ns, nanoseconds_between), - ArrayFromJSON(timestamp_unit_ns, times_seconds_precision)); +TEST_F(ScalarTemporalTest, TestTemporalSubtractDate) { + for (auto op : {"subtract", "subtract_checked"}) { + auto arr_date32s = ArrayFromJSON(date32(), date32s); + auto arr_date32s2 = ArrayFromJSON(date32(), date32s2); + auto arr_date64s = ArrayFromJSON(date64(), date64s); + auto arr_date64s2 = ArrayFromJSON(date64(), date64s2); + + CheckScalarBinary(op, arr_date32s2, arr_date32s, + ArrayFromJSON(duration(TimeUnit::SECOND), seconds_between_date)); + CheckScalarBinary( + op, arr_date64s2, arr_date64s, + ArrayFromJSON(duration(TimeUnit::MILLI), milliseconds_between_date)); + CheckScalarBinary( + op, arr_date64s2, arr_date32s, + ArrayFromJSON(duration(TimeUnit::MILLI), milliseconds_between_date)); } +} + +TEST_F(ScalarTemporalTest, TestTemporalSubtractTimestamp) { + for (auto op : {"subtract", "subtract_checked"}) { + for (auto tz : {"", "UTC", "Pacific/Marquesas"}) { + auto timestamp_unit_s = timestamp(TimeUnit::SECOND, tz); + auto duration_unit_s = duration(TimeUnit::SECOND); + auto timestamp_unit_ms = timestamp(TimeUnit::MILLI, tz); + auto duration_unit_ms = duration(TimeUnit::MILLI); + auto timestamp_unit_us = timestamp(TimeUnit::MICRO, tz); + auto duration_unit_us = duration(TimeUnit::MICRO); + auto timestamp_unit_ns = timestamp(TimeUnit::NANO, tz); + auto duration_unit_ns = duration(TimeUnit::NANO); + + CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_s, times_seconds_precision2), + ArrayFromJSON(timestamp_unit_s, times_seconds_precision), + ArrayFromJSON(duration_unit_s, seconds_between)); + CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_ms, times_seconds_precision2), + ArrayFromJSON(timestamp_unit_ms, times_seconds_precision), + ArrayFromJSON(duration_unit_ms, milliseconds_between)); + CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_us, times_seconds_precision2), + ArrayFromJSON(timestamp_unit_us, times_seconds_precision), + ArrayFromJSON(duration_unit_us, microseconds_between)); + CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_ns, times_seconds_precision2), + ArrayFromJSON(timestamp_unit_ns, times_seconds_precision), + ArrayFromJSON(duration_unit_ns, nanoseconds_between)); + + CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_s, times_seconds_precision2), + ArrayFromJSON(timestamp_unit_ms, times_seconds_precision), + ArrayFromJSON(duration_unit_ms, milliseconds_between)); + + CheckScalarBinary(op, ArrayFromJSON(timestamp_unit_ms, times_seconds_precision2), + ArrayFromJSON(timestamp_unit_s, times_seconds_precision), + ArrayFromJSON(duration_unit_ms, milliseconds_between)); + } + + CheckScalarBinary( + op, ArrayFromJSON(timestamp(TimeUnit::NANO, "UTC"), times_seconds_precision2), + ArrayFromJSON(timestamp(TimeUnit::SECOND, "Pacific/Marquesas"), + times_seconds_precision), + ArrayFromJSON(duration(TimeUnit::NANO), nanoseconds_between)); - auto seconds_3 = ArrayFromJSON(timestamp(TimeUnit::SECOND), R"([3, null])"); - auto milliseconds_2k = ArrayFromJSON(duration(TimeUnit::MILLI), R"([2000, null])"); - auto milliseconds_1k = ArrayFromJSON(timestamp(TimeUnit::MILLI), R"([1000, null])"); - CheckScalarBinary(op, seconds_3, milliseconds_2k, milliseconds_1k); + CheckScalarBinary( + op, ArrayFromJSON(timestamp(TimeUnit::SECOND), times_seconds_precision2), + ArrayFromJSON(timestamp(TimeUnit::NANO), times_seconds_precision), + ArrayFromJSON(duration(TimeUnit::NANO), nanoseconds_between)); - auto seconds_3_tz = ArrayFromJSON(timestamp(TimeUnit::SECOND, "UTC"), R"([3, null])"); - auto milliseconds_1k_tz = - ArrayFromJSON(timestamp(TimeUnit::MILLI, "UTC"), R"([1000, null])"); - CheckScalarBinary(op, seconds_3_tz, milliseconds_2k, milliseconds_1k_tz); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, + ::testing::HasSubstr("Subtraction of zoned and non-zoned times is ambiguous."), + CallFunction( + op, + {ArrayFromJSON(timestamp(TimeUnit::SECOND, "UTC"), times_seconds_precision2), + ArrayFromJSON(timestamp(TimeUnit::SECOND), times_seconds_precision)})); + } } TEST_F(ScalarTemporalTest, TestTemporalDifferenceWeeks) {