ARROW-14821: [R] Implement bindings for lubridate's floor_date, ceiling_date, and round_date

r/R/dplyr-datetime-helpers.R

@@ -417,3 +417,161 @@ build_strptime_exprs <- function(x, formats) {
     )
   )
 }
+
+# This function parses the "unit" argument to round_date, floor_date, and
+# ceiling_date. The input x is a single string like "second", "3 seconds",
+# "10 microseconds" or "2 secs" used to specify the size of the unit to
+# which the temporal data should be rounded. The matching rules implemented
+# are designed to mirror lubridate exactly: it extracts the numeric multiple
+# from the start of the string (presumed to be 1 if no number is present)
+# and selects the unit by looking at the first 3 characters only. This choice
+# ensures that "secs", "second", "microsecs" etc are all valid, but it is
+# very permissive and would interpret "mickeys" as microseconds. This
+# permissive implementation mirrors the corresponding implementation in
+# lubridate. The return value is a list with integer-valued components
+# "multiple" and  "unit"
+parse_period_unit <- function(x) {
+  # the regexp matches against fractional units, but per lubridate
+  # supports integer multiples of a known unit only
+  match_info <- regexpr(
+    pattern = " *(?<multiple>[0-9.,]+)? *(?<unit>[^ \t\n]+)",
+    text = x[[1]],
+    perl = TRUE
+  )
+
+  capture_start <- attr(match_info, "capture.start")
+  capture_length <- attr(match_info, "capture.length")
+  capture_end <- capture_start + capture_length - 1L
+
+  str_unit <- substr(x, capture_start[[2]], capture_end[[2]])
+  str_multiple <- substr(x, capture_start[[1]], capture_end[[1]])
+
+  known_units <- c("nanosecond", "microsecond", "millisecond", "second",
+                   "minute", "hour", "day", "week", "month", "quarter", "year")
+
+  # match the period unit
+  str_unit_start <- substr(str_unit, 1, 3)
+  unit <- as.integer(pmatch(str_unit_start, known_units)) - 1L
+
+  if (any(is.na(unit))) {
+    abort(
+      sprintf(
+        "Invalid period name: '%s'",
+        str_unit,
+        ". Known units are",
+        oxford_paste(known_units, "and")
+      )
+    )
+  }
+
+  # empty string in multiple interpreted as 1
+  if (capture_length[[1]] == 0) {
+    multiple <- 1L
+
+  # otherwise parse the multiple
+  } else {
+    multiple <- as.numeric(str_multiple)
+
+    # special cases: interpret fractions of 1 second as integer
+    # multiples of nanoseconds, microseconds, or milliseconds
+    # to mirror lubridate syntax
+    if (unit == 3L) {
+      if (multiple < 10^-6) {
+        unit <- 0L
+        multiple <- 10^9 * multiple
+      }
+      if (multiple < 10^-3) {
+        unit <- 1L
+        multiple <- 10^6 * multiple
+      }
+      if (multiple < 1) {
+        unit <- 2L
+        multiple <- 10^3 * multiple
+      }
+    }
+
+    multiple <- as.integer(multiple)
+  }
+
+  # more special cases: lubridate imposes sensible maximum
+  # values on the number of seconds, minutes and hours
+  if (unit == 3L && multiple > 60) {
+    abort("Rounding with second > 60 is not supported")
+  }
+  if (unit == 4L && multiple > 60) {
+    abort("Rounding with minute > 60 is not supported")
+  }
+  if (unit == 5L && multiple > 24) {
+    abort("Rounding with hour > 24 is not supported")
+  }
+
+  list(unit = unit, multiple = multiple)
+}
+
+# This function handles round/ceil/floor when unit is week. The fn argument
+# specifies which of the temporal rounding functions (round_date, etc) is to
+# be applied, x is the data argument to the rounding function, week_start is
+# an integer indicating which day of the week is the start date. The C++
+# library natively handles Sunday and Monday so in those cases we pass the
+# week_starts_monday option through. Other week_start values are handled here
+shift_temporal_to_week <- function(fn, x, week_start, options) {
+  if (week_start == 7) { # Sunday
+    options$week_starts_monday <- FALSE
+    return(Expression$create(fn, x, options = options))
+  }
+
+  if (week_start == 1) { # Monday
+    options$week_starts_monday <- TRUE
+    return(Expression$create(fn, x, options = options))
+  }
+
+  # other cases use offset-from-Monday: to ensure type-stable output there
+  # are two separate helpers, one to handle date32 input and the other to
+  # handle timestamps
+  options$week_starts_monday <- TRUE
+  offset <- as.integer(week_start) - 1
+
+  is_date32 <- inherits(x, "Date") ||
+    (inherits(x, "Expression") && x$type_id() == Type$DATE32)
+
+  if (is_date32) {
+    shifted_date <- shift_date32_to_week(fn, x, offset, options = options)
+  } else {
+    shifted_date <- shift_timestamp_to_week(fn, x, offset, options = options)
+  }
+
+  shifted_date
+}
+
+# timestamp input should remain timestamp
+shift_timestamp_to_week <- function(fn, x, offset, options) {
+  offset_seconds <- build_expr(
+    "cast",
+    Scalar$create(offset * 86400L, int64()),
+    options = cast_options(to_type = duration(unit = "s"))
+  )
+  shift_offset <- build_expr(fn, x - offset_seconds, options = options)
+
+  shift_offset + offset_seconds
+}
+
+# to avoid date32 types being cast to timestamp during the temporal
+# arithmetic, the offset logic needs to use the count in days and
+# use integer arithmetic: this feels inelegant, but it ensures that
+# temporal rounding functions remain type stable
+shift_date32_to_week <- function(fn, x, offset, options) {
+  # offset the date
+  offset <- Expression$scalar(Scalar$create(offset, int32()))
+  x_int <- build_expr("cast", x, options = cast_options(to_type = int32()))
+  x_int_offset <- x_int - offset
+  x_offset <- build_expr("cast", x_int_offset, options = cast_options(to_type = date32()))
+
+  # apply round/floor/ceil
+  shift_offset <- build_expr(fn, x_offset, options = options)
+
+  # undo offset and return
+  shift_int_offset <- build_expr("cast", shift_offset, options = cast_options(to_type = int32()))
+  shift_int <- shift_int_offset + offset
+
+  build_expr("cast", shift_int, options = cast_options(to_type = date32()))
+}

r/R/dplyr-funcs-datetime.R

@@ -25,6 +25,7 @@ register_bindings_datetime <- function() {
   register_bindings_duration_constructor()
   register_bindings_duration_helpers()
   register_bindings_datetime_parsers()
+  register_bindings_datetime_rounding()
 }
 
 register_bindings_datetime_utility <- function() {
@@ -622,4 +623,55 @@ register_bindings_datetime_parsers <- function() {
 
     build_expr("assume_timezone", coalesce_output, options = list(timezone = tz))
   })
+
+}
+
+register_bindings_datetime_rounding <- function() {
+  register_binding(
+    "round_date",
+    function(x,
+             unit = "second",
+             week_start = getOption("lubridate.week.start", 7)) {
+
+    opts <- parse_period_unit(unit)
+    if (opts$unit == 7L) { # weeks (unit = 7L) need to accommodate week_start
+      return(shift_temporal_to_week("round_temporal", x, week_start, options = opts))
+    }
+
+    Expression$create("round_temporal", x, options = opts)
+  })
+
+  register_binding(
+    "floor_date",
+    function(x,
+             unit = "second",
+             week_start = getOption("lubridate.week.start", 7)) {
+
+    opts <- parse_period_unit(unit)
+    if (opts$unit == 7L) { # weeks (unit = 7L) need to accommodate week_start
+      return(shift_temporal_to_week("floor_temporal", x, week_start, options = opts))
+    }
+
+    Expression$create("floor_temporal", x, options = opts)
+  })
+
+  register_binding(
+    "ceiling_date",
+    function(x,
+             unit = "second",
+             change_on_boundary = NULL,
+             week_start = getOption("lubridate.week.start", 7)) {
+    opts <- parse_period_unit(unit)
+    if (is.null(change_on_boundary)) {
+      change_on_boundary <- ifelse(call_binding("is.Date", x), TRUE, FALSE)
+    }
+    opts$ceil_is_strictly_greater <- change_on_boundary
+
+    if (opts$unit == 7L) { # weeks (unit = 7L) need to accommodate week_start
+      return(shift_temporal_to_week("ceil_temporal", x, week_start, options = opts))
+    }
+
+    Expression$create("ceil_temporal", x, options = opts)
+  })
+
 }

r/src/compute.cpp

@@ -519,6 +519,35 @@ std::shared_ptr<arrow::compute::FunctionOptions> make_compute_options(
     return out;
   }
 
+  if (func_name == "round_temporal" || func_name == "floor_temporal" ||
+      func_name == "ceil_temporal") {
+    using Options = arrow::compute::RoundTemporalOptions;
+
+    int64_t multiple = 1;
+    enum arrow::compute::CalendarUnit unit = arrow::compute::CalendarUnit::DAY;
+    bool week_starts_monday = true;
+    bool ceil_is_strictly_greater = true;
+    bool calendar_based_origin = true;
+
+    if (!Rf_isNull(options["multiple"])) {
+      multiple = cpp11::as_cpp<int64_t>(options["multiple"]);
+    }
+    if (!Rf_isNull(options["unit"])) {
+      unit = cpp11::as_cpp<enum arrow::compute::CalendarUnit>(options["unit"]);
+    }
+    if (!Rf_isNull(options["week_starts_monday"])) {
+      week_starts_monday = cpp11::as_cpp<bool>(options["week_starts_monday"]);
+    }
+    if (!Rf_isNull(options["ceil_is_strictly_greater"])) {
+      ceil_is_strictly_greater = cpp11::as_cpp<bool>(options["ceil_is_strictly_greater"]);
+    }
+    if (!Rf_isNull(options["calendar_based_origin"])) {
+      calendar_based_origin = cpp11::as_cpp<bool>(options["calendar_based_origin"]);
+    }
+    return std::make_shared<Options>(multiple, unit, week_starts_monday,
+                                     ceil_is_strictly_greater, calendar_based_origin);
+  }
+
   if (func_name == "round_to_multiple") {
     using Options = arrow::compute::RoundToMultipleOptions;
     auto out = std::make_shared<Options>(Options::Defaults());