From 6a7f1b0ee50f70bfe42ef828dbdbee09e477665a Mon Sep 17 00:00:00 2001 From: Rok Date: Wed, 9 Feb 2022 14:01:21 +0100 Subject: [PATCH] Adding AddDateAndDuration --- .../compute/kernels/scalar_arithmetic.cc | 19 +++++++++++++ .../compute/kernels/scalar_temporal_test.cc | 28 +++++++++++++++++++ docs/source/cpp/compute.rst | 4 +-- 3 files changed, 49 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc index 3c3fd92d53b..9459e397332 100644 --- a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc +++ b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc @@ -2536,12 +2536,31 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { // ---------------------------------------------------------------------- auto add = MakeArithmeticFunction("add", &add_doc); AddDecimalBinaryKernels("add", add.get()); + + // Add add(timestamp, duration) -> timestamp + for (auto unit : TimeUnit::values()) { + InputType in_type(match::TimestampTypeUnit(unit)); + auto exec = ScalarBinary::Exec; + DCHECK_OK(add->AddKernel({in_type, duration(unit)}, OutputType(FirstType), + std::move(exec))); + } + DCHECK_OK(registry->AddFunction(std::move(add))); // ---------------------------------------------------------------------- auto add_checked = MakeArithmeticFunctionNotNull("add_checked", &add_checked_doc); AddDecimalBinaryKernels("add_checked", add_checked.get()); + + // Add add_checked(timestamp, duration) -> timestamp + for (auto unit : TimeUnit::values()) { + InputType in_type(match::TimestampTypeUnit(unit)); + auto exec = + ScalarBinary::Exec; + DCHECK_OK(add_checked->AddKernel({in_type, duration(unit)}, OutputType(FirstType), + std::move(exec))); + } + 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 f4fa2403c1f..fdd2876304f 100644 --- a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc @@ -976,6 +976,34 @@ TEST_F(ScalarTemporalTest, TestTemporalDifference) { } } +TEST_F(ScalarTemporalTest, TestTemporalAddDateAndDuration) { + for (auto op : {"add", "add_checked"}) { + std::string milliseconds_between_date_and_time = + "[59000, 84203000, 3560000, 12800000, 3905000, 7810000, 11715000, 15620000, " + "19525000, 23430000, 27335000, 31240000, 35145000, 0, 0, 3723000, null]"; + std::string microseconds_between_date_and_time = + "[59000000, 84203000000, 3560000000, 12800000000, 3905000000, 7810000000, " + "11715000000, 15620000000, 19525000000, 23430000000, 27335000000, 31240000000, " + "35145000000, 0, 0, 3723000000, null]"; + auto dates32 = ArrayFromJSON(date32(), date32s); + auto dates64 = ArrayFromJSON(date64(), date64s); + + auto durations_ms = + ArrayFromJSON(duration(TimeUnit::MILLI), milliseconds_between_date_and_time); + 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_date_and_time); + 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, TestTemporalSubtractDateAndDuration) { for (auto op : {"subtract", "subtract_checked"}) { std::string milliseconds_between_time_and_date = diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst index 13402718772..425fda9ac11 100644 --- a/docs/source/cpp/compute.rst +++ b/docs/source/cpp/compute.rst @@ -430,9 +430,9 @@ Mixed time resolution temporal inputs will be cast to finest input resolution. +------------------+--------+----------------------------+----------------------------+-------+ | abs_checked | Unary | Numeric | Numeric | | +------------------+--------+----------------------------+----------------------------+-------+ -| add | Binary | Numeric | Numeric | \(1) | +| add | Binary | Numeric/Temporal | Numeric/Temporal | \(1) | +------------------+--------+----------------------------+----------------------------+-------+ -| add_checked | Binary | Numeric | Numeric | \(1) | +| add_checked | Binary | Numeric/Temporal | Numeric/Temporal | \(1) | +------------------+--------+----------------------------+----------------------------+-------+ | divide | Binary | Numeric | Numeric | \(1) | +------------------+--------+----------------------------+----------------------------+-------+