rok commented on a change in pull request #11818:
URL: https://github.com/apache/arrow/pull/11818#discussion_r778170813
##########
File path: cpp/src/arrow/compute/kernels/scalar_temporal_unary.cc
##########
@@ -452,6 +474,296 @@ struct Nanosecond {
}
};
+// ----------------------------------------------------------------------
+// Round temporal values to given frequency
+
+template <typename Duration, typename Localizer>
+year_month_day GetFlooredYmd(int64_t arg, int multiple, Localizer localizer_) {
+ year_month_day ymd{floor<days>(localizer_.template
ConvertTimePoint<Duration>(arg))};
+
+ if (multiple == 1) {
+ return year_month_day(ymd.year() / ymd.month() / 1);
+ } else {
+ int32_t total_months_origin = 1970 * 12;
+ int32_t total_months = static_cast<int32_t>(ymd.year()) * 12 +
+
static_cast<int32_t>(static_cast<uint32_t>(ymd.month())) - 1 -
+ total_months_origin;
+
+ if (total_months >= 0) {
+ total_months = total_months / multiple * multiple;
+ } else {
+ total_months = (total_months - multiple + 1) / multiple * multiple;
+ }
+ return year_month_day(year{1970} / jan / 0) + months{total_months};
+ }
+}
+
+template <typename Duration, typename Unit, typename Localizer>
+const Duration FloorTimePoint(const int64_t arg, const int64_t multiple,
+ Localizer localizer_, Status* st) {
+ const auto t = localizer_.template ConvertTimePoint<Duration>(arg);
+ const Unit d = floor<Unit>(t.time_since_epoch());
+
+ if (multiple == 1) {
+ return localizer_.template
ConvertLocalToSys<Duration>(duration_cast<Duration>(d),
+ st);
+ } else {
+ const Unit unit = Unit{multiple};
+ auto m = (d.count() >= 0) ? d / unit * unit : (d - unit + Unit{1}) / unit
* unit;
+ return localizer_.template
ConvertLocalToSys<Duration>(duration_cast<Duration>(m),
+ st);
+ }
+}
+
+template <typename Duration, typename Unit, typename Localizer>
+Duration CeilTimePoint(const int64_t arg, const int64_t multiple, Localizer
localizer_,
+ Status* st) {
+ const Duration result =
+ FloorTimePoint<Duration, Unit, Localizer>(arg, multiple, localizer_, st);
+ const auto c =
+ localizer_.template
ConvertTimePoint<Duration>(result.count()).time_since_epoch() +
+ duration_cast<Duration>(Unit{multiple});
+ return localizer_.template
ConvertLocalToSys<Duration>(duration_cast<Duration>(c), st);
+}
+
+template <typename Duration, typename Unit, typename Localizer>
+Duration RoundTimePoint(const int64_t arg, const int64_t multiple, Localizer
localizer_,
+ Status* st) {
+ const Duration f =
+ FloorTimePoint<Duration, Unit, Localizer>(arg, multiple, localizer_, st);
+ const Duration c =
+ CeilTimePoint<Duration, Unit, Localizer>(arg, multiple, localizer_, st);
+ return (Duration{arg} - f > c - Duration{arg}) ? c : f;
+}
+
+template <typename Duration, typename Localizer>
+struct CeilTemporal {
+ explicit CeilTemporal(const RoundTemporalOptions* options, Localizer&&
localizer)
+ : localizer_(std::move(localizer)), options(*options) {}
+
+ template <typename T, typename Arg0>
+ T Call(KernelContext*, Arg0 arg, Status* st) const {
+ Duration t;
+ switch (options.unit) {
+ case compute::CalendarUnit::NANOSECOND:
+ t = CeilTimePoint<Duration, std::chrono::nanoseconds, Localizer>(
+ arg, options.multiple, localizer_, st);
+ break;
+ case compute::CalendarUnit::MICROSECOND:
+ t = CeilTimePoint<Duration, std::chrono::microseconds, Localizer>(
+ arg, options.multiple, localizer_, st);
+ break;
+ case compute::CalendarUnit::MILLISECOND:
+ t = CeilTimePoint<Duration, std::chrono::milliseconds, Localizer>(
+ arg, options.multiple, localizer_, st);
+ break;
+ case compute::CalendarUnit::SECOND:
+ t = CeilTimePoint<Duration, std::chrono::seconds, Localizer>(
+ arg, options.multiple, localizer_, st);
+ break;
+ case compute::CalendarUnit::MINUTE:
+ t = CeilTimePoint<Duration, minutes, Localizer>(arg, options.multiple,
localizer_,
+ st);
+ break;
+ case compute::CalendarUnit::HOUR:
+ t = CeilTimePoint<Duration, std::chrono::hours, Localizer>(arg,
options.multiple,
+ localizer_,
st);
+ break;
+ case compute::CalendarUnit::DAY:
+ t = CeilTimePoint<Duration, days, Localizer>(arg, options.multiple,
localizer_,
+ st);
+ break;
+ case compute::CalendarUnit::WEEK:
+ t = CeilTimePoint<Duration, weeks, Localizer>(arg, options.multiple,
localizer_,
+ st);
+ break;
+ case compute::CalendarUnit::MONTH: {
+ year_month_day ymd =
+ GetFlooredYmd<Duration, Localizer>(arg, options.multiple,
localizer_);
+ ymd += months{options.multiple};
+ t = localizer_.ConvertDays(ymd.year() / ymd.month() /
1).time_since_epoch();
+ break;
+ }
+ case compute::CalendarUnit::QUARTER: {
+ year_month_day ymd =
+ GetFlooredYmd<Duration, Localizer>(arg, 3 * options.multiple,
localizer_);
+ ymd += months{3 * options.multiple};
+ t = localizer_.ConvertDays(ymd.year() / ymd.month() /
1).time_since_epoch();
+ break;
+ }
+ case compute::CalendarUnit::YEAR: {
+ year_month_day ymd(
+ floor<days>(localizer_.template ConvertTimePoint<Duration>(arg)));
+ year y{(static_cast<int32_t>(ymd.year()) / options.multiple + 1) *
+ options.multiple};
+ t = localizer_.ConvertDays(y / jan / 1).time_since_epoch();
+ break;
+ }
+ default:
+ t = Duration{arg};
+ }
+ return static_cast<T>(t.count());
+ }
+
+ Localizer localizer_;
+ RoundTemporalOptions options;
+};
+
+template <typename Duration, typename Localizer>
+struct FloorTemporal {
+ explicit FloorTemporal(const RoundTemporalOptions* options, Localizer&&
localizer)
+ : localizer_(std::move(localizer)), options(*options) {}
+
+ template <typename T, typename Arg0>
+ T Call(KernelContext*, Arg0 arg, Status* st) const {
+ Duration t;
+ switch (options.unit) {
+ case compute::CalendarUnit::NANOSECOND:
+ t = FloorTimePoint<Duration, std::chrono::nanoseconds, Localizer>(
+ arg, options.multiple, localizer_, st);
+ break;
+ case compute::CalendarUnit::MICROSECOND:
+ t = FloorTimePoint<Duration, std::chrono::microseconds, Localizer>(
+ arg, options.multiple, localizer_, st);
+ break;
+ case compute::CalendarUnit::MILLISECOND:
+ t = FloorTimePoint<Duration, std::chrono::milliseconds, Localizer>(
+ arg, options.multiple, localizer_, st);
+ break;
+ case compute::CalendarUnit::SECOND:
+ t = FloorTimePoint<Duration, std::chrono::seconds, Localizer>(
+ arg, options.multiple, localizer_, st);
+ break;
+ case compute::CalendarUnit::MINUTE:
+ t = FloorTimePoint<Duration, minutes, Localizer>(arg, options.multiple,
+ localizer_, st);
+ break;
+ case compute::CalendarUnit::HOUR:
+ t = FloorTimePoint<Duration, std::chrono::hours, Localizer>(arg,
options.multiple,
+
localizer_, st);
+ break;
+ case compute::CalendarUnit::DAY:
+ t = FloorTimePoint<Duration, days, Localizer>(arg, options.multiple,
localizer_,
+ st);
+ break;
+ case compute::CalendarUnit::WEEK:
+ t = FloorTimePoint<Duration, weeks, Localizer>(arg, options.multiple,
localizer_,
+ st);
+ break;
+ case compute::CalendarUnit::MONTH: {
+ year_month_day ymd =
+ GetFlooredYmd<Duration, Localizer>(arg, options.multiple,
localizer_);
+ t = localizer_.ConvertDays(ymd.year() / ymd.month() /
1).time_since_epoch();
+ break;
+ }
+ case compute::CalendarUnit::QUARTER: {
+ year_month_day ymd =
+ GetFlooredYmd<Duration, Localizer>(arg, 3 * options.multiple,
localizer_);
+ t = localizer_.ConvertDays(ymd.year() / ymd.month() /
1).time_since_epoch();
+ break;
+ }
+ case compute::CalendarUnit::YEAR: {
+ year_month_day ymd(
+ floor<days>(localizer_.template ConvertTimePoint<Duration>(arg)));
+ year y{(static_cast<int32_t>(ymd.year()) / options.multiple) *
options.multiple};
+ t = localizer_.ConvertDays(y / jan / 1).time_since_epoch();
+ break;
+ }
+ default:
+ t = Duration{arg};
+ }
+ return static_cast<T>(t.count());
+ }
+
+ Localizer localizer_;
+ RoundTemporalOptions options;
+};
+
+template <typename Duration, typename Localizer>
+struct RoundTemporal {
+ explicit RoundTemporal(const RoundTemporalOptions* options, Localizer&&
localizer)
+ : localizer_(std::move(localizer)), options(*options) {}
+
+ template <typename T, typename Arg0>
+ T Call(KernelContext*, Arg0 arg, Status* st) const {
+ Duration t;
+ switch (options.unit) {
+ case compute::CalendarUnit::NANOSECOND:
+ t = RoundTimePoint<Duration, std::chrono::nanoseconds, Localizer>(
+ arg, options.multiple, localizer_, st);
+ break;
+ case compute::CalendarUnit::MICROSECOND:
+ t = RoundTimePoint<Duration, std::chrono::microseconds, Localizer>(
+ arg, options.multiple, localizer_, st);
+ break;
+ case compute::CalendarUnit::MILLISECOND:
+ t = RoundTimePoint<Duration, std::chrono::milliseconds, Localizer>(
+ arg, options.multiple, localizer_, st);
+ break;
+ case compute::CalendarUnit::SECOND:
+ t = RoundTimePoint<Duration, std::chrono::seconds, Localizer>(
+ arg, options.multiple, localizer_, st);
+ break;
+ case compute::CalendarUnit::MINUTE:
+ t = RoundTimePoint<Duration, minutes, Localizer>(arg, options.multiple,
+ localizer_, st);
+ break;
+ case compute::CalendarUnit::HOUR:
+ t = RoundTimePoint<Duration, std::chrono::hours, Localizer>(arg,
options.multiple,
+
localizer_, st);
+ break;
+ case compute::CalendarUnit::DAY:
+ t = RoundTimePoint<Duration, days, Localizer>(arg, options.multiple,
localizer_,
+ st);
+ break;
+ case compute::CalendarUnit::WEEK:
+ t = RoundTimePoint<Duration, weeks, Localizer>(arg, options.multiple,
localizer_,
+ st);
+ break;
+ case compute::CalendarUnit::MONTH: {
+ auto t0 = localizer_.template ConvertTimePoint<Duration>(arg);
+ year_month_day ymd =
+ GetFlooredYmd<Duration, Localizer>(arg, options.multiple,
localizer_);
+
+ auto f = localizer_.ConvertDays(ymd.year() / ymd.month() / 1);
+ ymd += months{options.multiple};
+ auto c = localizer_.ConvertDays(ymd.year() / ymd.month() / 1);
+
+ t = (t0 - f > c - t0) ? c.time_since_epoch() : f.time_since_epoch();
Review comment:
I'm not sure how to round half by default. Switching to up (`>=`).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
