diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc index 60b2d791146..1b9d9254c58 100644 --- a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc @@ -3779,6 +3779,126 @@ TEST_F(ScalarTemporalTest, TestCeilFloorRoundTemporalDate) { CheckScalarUnary("ceil_temporal", arr_ns, arr_ns, &round_to_2_hours); } +TEST_F(ScalarTemporalTest, TestCeilFloorRoundTemporalDuration) { + auto check_unit = [&](TimeUnit::type time_unit, CalendarUnit calendar_unit) { + RoundTemporalOptions round_to_2_units(2, calendar_unit); + auto values = ArrayFromJSON(duration(time_unit), "[0, 1, 2, 3, -1, -2, -3, null]"); + + CheckScalarUnary("ceil_temporal", values, + ArrayFromJSON(duration(time_unit), "[0, 2, 2, 4, 0, -2, -2, null]"), + &round_to_2_units); + CheckScalarUnary("floor_temporal", values, + ArrayFromJSON(duration(time_unit), "[0, 0, 2, 2, -2, -2, -4, null]"), + &round_to_2_units); + CheckScalarUnary("round_temporal", values, + ArrayFromJSON(duration(time_unit), "[0, 2, 2, 4, 0, -2, -2, null]"), + &round_to_2_units); + }; + + check_unit(TimeUnit::SECOND, CalendarUnit::SECOND); + check_unit(TimeUnit::MILLI, CalendarUnit::MILLISECOND); + check_unit(TimeUnit::MICRO, CalendarUnit::MICROSECOND); + check_unit(TimeUnit::NANO, CalendarUnit::NANOSECOND); + + auto check_second_based_unit = [&](CalendarUnit calendar_unit, const char* values_json, + const char* ceil_round_json, + const char* floor_json) { + RoundTemporalOptions round_to_2_units(2, calendar_unit); + auto values = ArrayFromJSON(duration(TimeUnit::SECOND), values_json); + + CheckScalarUnary("ceil_temporal", values, + ArrayFromJSON(duration(TimeUnit::SECOND), ceil_round_json), + &round_to_2_units); + CheckScalarUnary("floor_temporal", values, + ArrayFromJSON(duration(TimeUnit::SECOND), floor_json), + &round_to_2_units); + CheckScalarUnary("round_temporal", values, + ArrayFromJSON(duration(TimeUnit::SECOND), ceil_round_json), + &round_to_2_units); + }; + + check_second_based_unit(CalendarUnit::MINUTE, + "[0, 60, 120, 180, -60, -120, -180, null]", + "[0, 120, 120, 240, 0, -120, -120, null]", + "[0, 0, 120, 120, -120, -120, -240, null]"); + + check_second_based_unit(CalendarUnit::HOUR, + "[0, 3600, 7200, 10800, -3600, -7200, -10800, null]", + "[0, 7200, 7200, 14400, 0, -7200, -7200, null]", + "[0, 0, 7200, 7200, -7200, -7200, -14400, null]"); + + // A day is treated as a physical 24-hour unit for duration values. + RoundTemporalOptions round_to_day(1, CalendarUnit::DAY); + auto day_values = ArrayFromJSON(duration(TimeUnit::SECOND), + "[0, 43200, 86399, 86400, -1, -43200, null]"); + + CheckScalarUnary( + "ceil_temporal", day_values, + ArrayFromJSON(duration(TimeUnit::SECOND), "[0, 86400, 86400, 86400, 0, 0, null]"), + &round_to_day); + CheckScalarUnary( + "floor_temporal", day_values, + ArrayFromJSON(duration(TimeUnit::SECOND), "[0, 0, 0, 86400, -86400, -86400, null]"), + &round_to_day); + CheckScalarUnary( + "round_temporal", day_values, + ArrayFromJSON(duration(TimeUnit::SECOND), "[0, 86400, 86400, 86400, 0, 0, null]"), + &round_to_day); + + // A week is treated as exactly seven physical days for duration values. + RoundTemporalOptions round_to_week(1, CalendarUnit::WEEK); + auto week_values = + ArrayFromJSON(duration(TimeUnit::SECOND), + "[0, 302400, 604799, 604800, -1, -302400, -604800, null]"); + + CheckScalarUnary("ceil_temporal", week_values, + ArrayFromJSON(duration(TimeUnit::SECOND), + "[0, 604800, 604800, 604800, 0, 0, -604800, null]"), + &round_to_week); + CheckScalarUnary("floor_temporal", week_values, + ArrayFromJSON(duration(TimeUnit::SECOND), + "[0, 0, 0, 604800, -604800, -604800, -604800, null]"), + &round_to_week); + CheckScalarUnary("round_temporal", week_values, + ArrayFromJSON(duration(TimeUnit::SECOND), + "[0, 604800, 604800, 604800, 0, 0, -604800, null]"), + &round_to_week); + + auto values = ArrayFromJSON(duration(TimeUnit::SECOND), "[0, 1, -1, null]"); + + RoundTemporalOptions round_to_month(1, CalendarUnit::MONTH); + for (const auto* function_name : + {"ceil_temporal", "floor_temporal", "round_temporal"}) { + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("Duration values can only be rounded using units"), + CallFunction(function_name, {values}, &round_to_month)); + } + + RoundTemporalOptions calendar_origin(2, CalendarUnit::SECOND); + calendar_origin.calendar_based_origin = true; + for (const auto* function_name : + {"ceil_temporal", "floor_temporal", "round_temporal"}) { + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, + ::testing::HasSubstr( + "calendar_based_origin is not supported for duration inputs"), + CallFunction(function_name, {values}, &calendar_origin)); + } +} + +TEST_F(ScalarTemporalTest, TestDurationTemporalWeekMultipleOverflow) { + auto values = ArrayFromJSON(duration(TimeUnit::SECOND), "[0]"); + RoundTemporalOptions options(std::numeric_limits::max(), CalendarUnit::WEEK); + + for (const auto* function_name : + {"ceil_temporal", "floor_temporal", "round_temporal"}) { + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, + ::testing::HasSubstr("Duration week multiple would not fit in 32-bit integer"), + CallFunction(function_name, {values}, &options)); + } +} + TEST_F(ScalarTemporalTest, DurationUnaryArithmetics) { auto arr = ArrayFromJSON(duration(TimeUnit::SECOND), "[2, -1, null, 3, 0]"); CheckScalarUnary("negate", arr, diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_unary.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_unary.cc index 1bad2d0a118..74a29081ae8 100644 --- a/cpp/src/arrow/compute/kernels/scalar_temporal_unary.cc +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_unary.cc @@ -18,6 +18,7 @@ #include #include #include +#include #include "arrow/builder.h" #include "arrow/compute/api_scalar.h" @@ -26,6 +27,7 @@ #include "arrow/compute/kernels/temporal_internal.h" #include "arrow/compute/registry_internal.h" #include "arrow/util/checked_cast.h" +#include "arrow/util/int_util_overflow.h" #include "arrow/util/logging_internal.h" #include "arrow/util/time.h" #include "arrow/util/value_parsing.h" @@ -177,6 +179,36 @@ struct TemporalComponentExtractRound static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { const RoundTemporalOptions& options = RoundTemporalState::Get(ctx); + + if constexpr (std::is_same_v) { + if (options.calendar_based_origin) { + return Status::Invalid( + "calendar_based_origin is not supported for duration inputs"); + } + + if (options.unit == CalendarUnit::WEEK) { + RoundTemporalOptions duration_options = options; + duration_options.unit = CalendarUnit::DAY; + if (::arrow::internal::MultiplyWithOverflow(options.multiple, 7, + &duration_options.multiple)) { + return Status::Invalid( + "Duration week multiple would not fit in 32-bit integer"); + } + return Base::ExecWithOptions(ctx, &duration_options, batch, out); + } + + switch (options.unit) { + case CalendarUnit::MONTH: + case CalendarUnit::QUARTER: + case CalendarUnit::YEAR: + return Status::Invalid( + "Duration values can only be rounded using units " + "from nanoseconds through weeks"); + default: + break; + } + } + return Base::ExecWithOptions(ctx, &options, batch, out); } }; @@ -2014,23 +2046,23 @@ void RegisterScalarTemporalUnary(FunctionRegistry* registry) { // output type. See TemporalComponentExtractRound for more. static const auto default_round_temporal_options = RoundTemporalOptions::Defaults(); - auto floor_temporal = UnaryTemporalFactory::Make( - "floor_temporal", OutputType(FirstType), floor_temporal_doc, - &default_round_temporal_options, RoundTemporalState::Init); + auto floor_temporal = + UnaryTemporalFactory:: + Make( + "floor_temporal", OutputType(FirstType), floor_temporal_doc, + &default_round_temporal_options, RoundTemporalState::Init); DCHECK_OK(registry->AddFunction(std::move(floor_temporal))); - auto ceil_temporal = UnaryTemporalFactory::Make( - "ceil_temporal", OutputType(FirstType), ceil_temporal_doc, - &default_round_temporal_options, RoundTemporalState::Init); + auto ceil_temporal = + UnaryTemporalFactory:: + Make( + "ceil_temporal", OutputType(FirstType), ceil_temporal_doc, + &default_round_temporal_options, RoundTemporalState::Init); DCHECK_OK(registry->AddFunction(std::move(ceil_temporal))); - auto round_temporal = UnaryTemporalFactory::Make( - "round_temporal", OutputType(FirstType), round_temporal_doc, - &default_round_temporal_options, RoundTemporalState::Init); + auto round_temporal = + UnaryTemporalFactory:: + Make( + "round_temporal", OutputType(FirstType), round_temporal_doc, + &default_round_temporal_options, RoundTemporalState::Init); DCHECK_OK(registry->AddFunction(std::move(round_temporal))); } diff --git a/cpp/src/arrow/compute/kernels/temporal_internal.h b/cpp/src/arrow/compute/kernels/temporal_internal.h index bc3b3888159..02f5965522c 100644 --- a/cpp/src/arrow/compute/kernels/temporal_internal.h +++ b/cpp/src/arrow/compute/kernels/temporal_internal.h @@ -200,6 +200,7 @@ struct TimestampFormatter { struct WithDates {}; struct WithTimes {}; struct WithTimestamps {}; +struct WithDurations {}; struct WithStringTypes {}; // This helper allows generating temporal kernels for selected type categories @@ -224,6 +225,18 @@ void AddTemporalKernels(Factory* fac, WithTimes, WithOthers... others) { AddTemporalKernels(fac, std::forward(others)...); } +template +void AddTemporalKernels(Factory* fac, WithDurations, WithOthers... others) { + fac->template AddKernel(duration(TimeUnit::SECOND)); + fac->template AddKernel( + duration(TimeUnit::MILLI)); + fac->template AddKernel( + duration(TimeUnit::MICRO)); + fac->template AddKernel( + duration(TimeUnit::NANO)); + AddTemporalKernels(fac, std::forward(others)...); +} + template void AddTemporalKernels(Factory* fac, WithTimestamps, WithOthers... others) { fac->template AddKernel( diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index 1e08e73668e..d350c811579 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -2974,6 +2974,49 @@ def test_round_temporal(unit): _check_temporal_rounding(ts_zoned, values, unit) +@pytest.mark.parametrize( + ("unit", "base_frequency", "round_frequency"), + ( + ("nanosecond", "1ns", "4ns"), + ("microsecond", "1us", "4us"), + ("millisecond", "1ms", "4ms"), + ("second", "1s", "4s"), + ("minute", "1min", "4min"), + ("hour", "1h", "4h"), + ("day", "1D", "4D"), + ("week", "7D", "28D"), + ), +) +@pytest.mark.pandas +def test_round_temporal_duration(unit, base_frequency, round_frequency): + base = pd.Timedelta(base_frequency) + values = pd.Series([ + -7 * base, + -4 * base, + -1 * base, + 0 * base, + 1 * base, + 4 * base, + 7 * base, + pd.NaT, + ]) + arrow_values = pa.array(values) + assert pa.types.is_duration(arrow_values.type) + + options = pc.RoundTemporalOptions(4, unit) + + for arrow_round, pandas_round in ( + (pc.ceil_temporal, values.dt.ceil), + (pc.floor_temporal, values.dt.floor), + (pc.round_temporal, values.dt.round), + ): + result = arrow_round( + arrow_values, options=options + ).to_pandas() + expected = pandas_round(round_frequency) + np.testing.assert_array_equal(result, expected) + + def test_count(): arr = pa.array([1, 2, 3, None, None]) assert pc.count(arr).as_py() == 3