feat: add count_interval function with Rust and Arrow backend

.Rbuildignore

@@ -11,3 +11,8 @@
 ^\.github$
 ^client_secret.*\.json$
 ^.*auth-key\.json$
+^src/\.cargo$
+^src/rust/vendor$
+^src/rust/target$
+^src/Makevars$
+^src/Makevars\.win$

.gitignore

@@ -13,3 +13,6 @@ local/
 secrets/
 /Meta/
 docs
+src/rust/vendor
+src/Makevars
+src/Makevars.win

DESCRIPTION

@@ -10,7 +10,7 @@ Encoding: UTF-8
 Roxygen: list(markdown = TRUE)
 RoxygenNote: 7.3.3
 Imports:
-    arrow,
+    arrow (>= 24.0.0),
     cli,
     dbplyr,
     dplyr,
@@ -21,11 +21,14 @@ Imports:
     glue,
     googleCloudStorageR,
     janitor,
+    lubridate,
+    nanoarrow,
     readr,
     renv,
     rlang,
     stringr,
     targets,
+    tidyr,
     usethis,
     withr
 Suggests:
@@ -35,6 +38,12 @@ Suggests:
     tibble,
     knitr,
     rmarkdown
+Remotes:
+    apache/arrow/r
 Config/testthat/edition: 3
 VignetteBuilder: knitr
 URL: https://openjusticeok.github.io/ojoutils/
+Config/rextendr/version: 0.5.0
+SystemRequirements: Cargo (Rust's package manager), rustc >= 1.65.0, xz
+Depends: 
+    R (>= 4.2)

NAMESPACE

@@ -1,5 +1,6 @@
 # Generated by roxygen2: do not edit by hand
 
+export(count_interval)
 export(describe_change)
 export(gcs_auth_bucket)
 export(gcs_list_objects)
@@ -11,10 +12,21 @@ export(ojo_create_project)
 export(ojo_parse_county)
 export(ojo_use_template)
 export(tar_gcs_csv)
+importFrom(arrow,as_arrow_table)
 importFrom(arrow,read_csv_arrow)
 importFrom(arrow,write_csv_arrow)
+importFrom(dplyr,all_of)
+importFrom(dplyr,coalesce)
+importFrom(dplyr,collect)
+importFrom(dplyr,compute)
+importFrom(dplyr,everything)
+importFrom(dplyr,filter)
+importFrom(dplyr,group_by)
 importFrom(dplyr,if_else)
+importFrom(dplyr,mutate)
 importFrom(dplyr,pull)
+importFrom(dplyr,select)
+importFrom(dplyr,ungroup)
 importFrom(fs,path_abs)
 importFrom(fs,path_wd)
 importFrom(gargle,token_fetch)
@@ -28,16 +40,28 @@ importFrom(googleCloudStorageR,gcs_get_object)
 importFrom(googleCloudStorageR,gcs_global_bucket)
 importFrom(googleCloudStorageR,gcs_list_objects)
 importFrom(janitor,clean_names)
+importFrom(lubridate,as_datetime)
+importFrom(lubridate,floor_date)
+importFrom(lubridate,period)
+importFrom(nanoarrow,as_nanoarrow_array_stream)
 importFrom(readr,write_lines)
 importFrom(renv,init)
 importFrom(renv,install)
+importFrom(rlang,":=")
 importFrom(rlang,abort)
 importFrom(rlang,as_name)
+importFrom(rlang,check_required)
 importFrom(rlang,enexpr)
 importFrom(rlang,ensym)
 importFrom(rlang,expr)
+importFrom(rlang,has_name)
 importFrom(rlang,is_interactive)
+importFrom(rlang,sym)
+importFrom(rlang,syms)
 importFrom(stringr,str_remove_all)
 importFrom(stringr,str_to_lower)
 importFrom(targets,tar_target_raw)
+importFrom(tidyr,complete)
+importFrom(tidyr,fill)
 importFrom(usethis,create_from_github)
+useDynLib(ojoutils, .registration = TRUE)

R/count_interval.R

@@ -0,0 +1,244 @@
+#' Count intervals over time periods
+#'
+#' @description
+#' Counts the number of active intervals for each time period (day, hour, etc.)
+#' given start and end dates. Useful for occupancy or population counts over time.
+#'
+#' @param data A data frame or Arrow table containing the interval data.
+#' @param start Character string. Name of the column containing interval start dates.
+#' @param end Character string. Name of the column containing interval end dates.
+#' @param period Character string. Time period for counting (e.g., "day", "hour",
+#'   "week", "month", "quarter", "year"). Defaults to "day".
+#' @param date_name Character string. Name for the output date column.
+#'   Defaults to "date".
+#' @param count_name Character string. Name for the output count column.
+#'   Defaults to "n".
+#' @param .by Character vector. Column names to group by. Defaults to empty
+#'   (no grouping).
+#' @param .fill Named list with "start" and "end" elements. Values to use when
+#'   filling NA start/end dates. If NULL (default), uses min/max of data.
+#' @param .inclusive Logical vector of length 2. Whether start and end boundaries
+#'   are inclusive. Defaults to `c(TRUE, TRUE)`.
+#'
+#' @return A tibble with columns for date, count, and any grouping variables.
+#'
+#' @examples
+#' \dontrun{
+#' # Basic usage
+#' df <- data.frame(
+#'   start = as.Date(c("2024-01-01", "2024-01-05")),
+#'   end = as.Date(c("2024-01-03", "2024-01-06"))
+#' )
+#' count_interval(df, start = "start", end = "end", period = "day")
+#'
+#' # With grouping
+#' df <- data.frame(
+#'   start = as.Date(c("2024-01-01", "2024-01-02")),
+#'   end = as.Date(c("2024-01-03", "2024-01-04")),
+#'   ward = c("A", "B")
+#' )
+#' count_interval(df, start = "start", end = "end", period = "day", .by = "ward")
+#' }
+#'
+#' @importFrom rlang check_required abort sym syms has_name :=
+#' @importFrom dplyr collect select mutate everything filter group_by ungroup
+#' @importFrom dplyr all_of coalesce compute
+#' @importFrom lubridate as_datetime floor_date period
+#' @importFrom arrow as_arrow_table
+#' @importFrom nanoarrow as_nanoarrow_array_stream
+#' @importFrom tidyr complete fill
+#' @export
+count_interval <- function(
+    data,
+    start,
+    end,
+    period = "day",
+    date_name = "date",
+    count_name = "n",
+    .by = character(),
+    .fill = list(start = NULL, end = NULL),
+    .inclusive = c(TRUE, TRUE)
+) {
+  rlang::check_required(data)
+  rlang::check_required(start)
+  rlang::check_required(end)
+
+  if (is.null(data)) {
+    rlang::abort("`data` must not be NULL.")
+  }
+
+  # Extract names for validation regardless of input type
+  data_names <- if (inherits(data, "ArrowTabular")) {
+    data$schema$names
+  } else {
+    names(data)
+  }
+
+  if (!start %in% data_names) {
+    rlang::abort(paste0("Column '", start, "' not found in `data`."))
+  }
+  if (!end %in% data_names) {
+    rlang::abort(paste0("Column '", end, "' not found in `data`."))
+  }
+  if (length(.by) > 0 && any(!.by %in% data_names)) {
+    missing <- setdiff(.by, data_names)
+    rlang::abort(paste0(
+      "Columns not found in `data`: ", paste(missing, collapse = ", "), "."
+    ))
+  }
+  if (date_name %in% c(start, end, count_name, .by)) {
+    rlang::abort(
+      "`date_name` must not match `start`, `end`, `count_name`, or any `.by` column."
+    )
+  }
+  if (count_name %in% c(start, end, date_name, .by)) {
+    rlang::abort(
+      "`count_name` must not match `start`, `end`, `date_name`, or any `.by` column."
+    )
+  }
+  if (any(.by %in% c(start, end))) {
+    rlang::abort("`.by` columns must not include `start` or `end`.")
+  }
+
+  if (!is.character(period) || length(period) != 1L || is.na(period)) {
+    rlang::abort("`period` must be a single non-NA character string.")
+  }
+
+  period_seq <- sub("minute", "min", period, fixed = TRUE)
+  period_seq <- sub("second", "sec", period_seq, fixed = TRUE)
+
+  if (
+    !is.list(.fill) ||
+      length(.fill) != 2L ||
+      !all(rlang::has_name(.fill, c("start", "end")))
+  ) {
+    rlang::abort(
+      "`.fill` must be a named list with 'start' and 'end' elements."
+    )
+  }
+
+  if (length(.inclusive) != 2L || !is.logical(.inclusive)) {
+    rlang::abort("`.inclusive` must be a length-2 logical vector.")
+  }
+
+  # Convert to arrow for processing
+  data <- arrow::as_arrow_table(data)
+
+  if (nrow(data) == 0) {
+    empty_res <- data |>
+      dplyr::collect() |>
+      dplyr::select(dplyr::all_of(.by)) |>
+      dplyr::mutate(
+        !!rlang::sym(date_name) := lubridate::as_datetime(character()),
+        !!rlang::sym(count_name) := integer()
+      ) |>
+      dplyr::select(!!rlang::sym(date_name), !!rlang::sym(count_name), dplyr::everything())
+    return(empty_res)
+  }
+
+  schema <- data$schema
+  start_type <- schema$GetFieldByName(start)$type
+  end_type <- schema$GetFieldByName(end)$type
+
+  # Compute fill defaults from data min/max when NULL
+  if (is.null(.fill$start)) {
+    start_vals <- data[[start]]$as_vector()
+    if (all(is.na(start_vals))) {
+      rlang::abort(
+        "All values in 'start' are NA and no `.fill$start` was provided."
+      )
+    }
+    fill_start <- lubridate::as_datetime(min(start_vals, na.rm = TRUE))
+  } else {
+    fill_start <- lubridate::as_datetime(.fill$start)
+  }
+
+  if (is.null(.fill$end)) {
+    end_vals <- data[[end]]$as_vector()
+    if (all(is.na(end_vals))) {
+      rlang::abort(
+        "All values in 'end' are NA and no `.fill$end` was provided."
+      )
+    }
+    fill_end <- lubridate::as_datetime(max(end_vals, na.rm = TRUE))
+  } else {
+    fill_end <- lubridate::as_datetime(.fill$end)
+  }
+
+  query <- data |>
+    dplyr::mutate(
+      !!rlang::sym(start) := lubridate::floor_date(
+        dplyr::coalesce(
+          !!rlang::sym(start),
+          arrow::as_arrow_array(fill_start)$cast(start_type)
+        ),
+        unit = period
+      ),
+      !!rlang::sym(end) := lubridate::floor_date(
+        dplyr::coalesce(
+          !!rlang::sym(end),
+          arrow::as_arrow_array(fill_end)$cast(end_type)
+        ),
+        unit = period
+      )
+    )
+
+  tab <- dplyr::compute(query)
+
+  end_vec <- tab[[end]]$as_vector()
+
+  max_expected_date <- max(end_vec, na.rm = TRUE)
+
+  tab$.end_plus_one <- end_vec + if (period == "quarter") {
+    lubridate::period(3, units = "months")
+  } else {
+    lubridate::period(1, units = period)
+  }
+
+  stream <- nanoarrow::as_nanoarrow_array_stream(tab)
+
+  res <- count_interval_(
+    stream = stream,
+    start = start,
+    end = end,
+    end_plus_one = ".end_plus_one",
+    date_name = date_name,
+    count_name = count_name,
+    by = .by,
+    inclusive = .inclusive
+  ) |>
+    arrow::as_arrow_table() |>
+    dplyr::collect()
+
+  if (nrow(res) == 0) {
+    return(
+      dplyr::select(
+        res,
+        !!rlang::sym(date_name),
+        !!rlang::sym(count_name),
+        !!!rlang::syms(.by)
+      )
+    )
+  }
+
+  res |>
+    dplyr::group_by(!!!rlang::syms(.by)) |>
+    tidyr::complete(
+      !!rlang::sym(date_name) := seq(
+        min(!!rlang::sym(date_name), na.rm = TRUE),
+        max(!!rlang::sym(date_name), na.rm = TRUE),
+        by = period_seq
+      )
+    ) |>
+    tidyr::fill(!!rlang::sym(count_name), .direction = "down") |>
+    dplyr::mutate(
+      !!rlang::sym(count_name) := dplyr::coalesce(!!rlang::sym(count_name), 0L)
+    ) |>
+    dplyr::filter(!!rlang::sym(date_name) <= max_expected_date) |>
+    dplyr::ungroup() |>
+    dplyr::select(
+      !!rlang::sym(date_name),
+      !!rlang::sym(count_name),
+      !!!rlang::syms(.by)
+    )
+}

R/extendr-wrappers.R

@@ -0,0 +1,10 @@
+# Generated by extendr: Do not edit by hand
+# nolint start
+
+#' @usage NULL
+#' @useDynLib ojoutils, .registration = TRUE
+NULL
+
+count_interval_ <- function(stream, start, end, end_plus_one, date_name, count_name, by, inclusive) .Call(wrap__count_interval_, stream, start, end, end_plus_one, date_name, count_name, by, inclusive)
+
+# nolint end

cleanup

@@ -0,0 +1 @@
+rm -f src/Makevars

cleanup.win

@@ -0,0 +1 @@
+rm -f src/Makevars.win

configure

@@ -0,0 +1,3 @@
+#!/usr/bin/env sh
+: "${R_HOME=`R RHOME`}"
+"${R_HOME}/bin/Rscript" tools/config.R 

configure.win

@@ -0,0 +1,2 @@
+#!/usr/bin/env sh
+"${R_HOME}/bin${R_ARCH_BIN}/Rscript.exe" tools/config.R