R/codec_parquet.R

In artoo: Lossless CDISC-Native Input and Output for Clinical Datasets

# codec_parquet.R — the Apache Parquet codec (nanoparquet engine + sidecar).
#
# Parquet stores the data natively via nanoparquet (a lightweight, zero-R-dep
# engine — arrow is banned), and artoo's full CDISC metadata rides alongside
# as the universal `metadata_json` sidecar (plan 4.0/4.1): the single
# Dataset-JSON-shaped string set_meta() stamps, embedded verbatim in the
# parquet file's key-value metadata under the key "metadata_json". This is
# exactly where artoo beats plain nanoparquet/arrow, which drop
# labels/formats/codelists — artoo round-trips the complete artoo_meta.
#
# Read precedence is META-FIRST (plan D1): types are reconstructed from the
# sidecar, with nanoparquet's native column types advisory. A parquet written
# by another tool (no sidecar) degrades gracefully to a bare frame, never an
# error (plan 9.B). nanoparquet round-trips a tens-of-KB metadata value
# verbatim, so the chunked-key fallback is unnecessary at clinical spec sizes.

.parquet_meta_key <- "metadata_json"

# Pull the metadata_json value out of a parquet file's key-value metadata, or
# NULL when the file carries none (a foreign / plain-nanoparquet file).
#' @noRd
.parquet_sidecar <- function(path) {
  md <- nanoparquet::read_parquet_metadata(path)
  kv <- md$file_meta_data$key_value_metadata
  tbl <- if (is.null(kv) || !length(kv)) NULL else kv[[1L]]
  if (is.null(tbl) || !(.parquet_meta_key %in% tbl$key)) {
    return(NULL)
  }
  # The key is present (checked above) and parquet KV values are always single
  # strings, so the first match is the sidecar.
  out <- tbl$value[tbl$key == .parquet_meta_key][[1L]]
  Encoding(out) <- "UTF-8"
  out
}

# encode contract: (x, meta, path, <codec args>, call) -> invisible(path).
#' @noRd
.encode_parquet <- function(
  x,
  meta,
  path,
  encoding = NULL,
  on_invalid = "error",
  compression = "snappy",
  call = rlang::caller_env()
) {
  rlang::check_installed("nanoparquet", reason = "to write Parquet files.")

  valid_comp <- c("snappy", "gzip", "zstd", "uncompressed")
  if (
    !is.character(compression) ||
      length(compression) != 1L ||
      !compression %in% valid_comp
  ) {
    .artoo_abort(
      c(
        "{.arg compression} must be one of {.val {valid_comp}}.",
        "x" = "You supplied {.val {compression}}."
      ),
      kind = "input",
      call = call
    )
  }

  # The sidecar must describe exactly the columns the parquet file carries;
  # reconcile a meta whose frame was mutated after apply_spec() (the read
  # path realizes columns from the sidecar, so a stale one mis-realizes).
  if (is_artoo_meta(meta)) {
    meta <- .meta_reconcile(meta, x)
  }

  # Parquet bytes stay UTF-8 (the format's STRING type is UTF-8 by spec); an
  # explicit `encoding` is recorded as the source-charset metadata so a later
  # write_xpt() can reproduce the original bytes. Validate the name loudly.
  if (!is.null(encoding) && is_artoo_meta(meta)) {
    .resolve_charset(encoding, call)
    meta <- .meta_set_encoding(meta, encoding)
  }

  # hms has no native parquet type; store the bare seconds (the read path
  # realizes it back from the sidecar dataType). as.numeric(units = "secs")
  # — never unclass(), which would keep a stray units attribute. Date and
  # POSIXct have native parquet types, so they pass through untouched.
  # Character columns are gated through the UTF-8 validity policy (parquet
  # STRING is UTF-8 by spec), then NFC-canonicalised (both no-ops on
  # ASCII / single-byte content).
  for (nm in names(x)) {
    if (inherits(x[[nm]], "difftime")) {
      x[[nm]] <- as.numeric(x[[nm]], units = "secs")
    } else if (is.character(x[[nm]])) {
      x[[nm]] <- .nfc(.to_target(x[[nm]], "UTF-8", on_invalid, call))
    }
  }
  # The frame-level metadata_json attribute (if any) is not a column; drop it
  # so it never leaks into the parquet schema. The sidecar below is the home.
  attr(x, "metadata_json") <- NULL

  kv <- if (is_artoo_meta(meta)) {
    stats::setNames(
      .meta_to_datasetjson(
        meta,
        extensions = TRUE,
        special = .collect_special_missings(x)
      ),
      .parquet_meta_key
    )
  } else {
    NULL
  }

  # Atomic write: build in a temp file, then rename over the target.
  .with_atomic_write(
    path,
    ".parquet.tmp",
    function(tmp) {
      if (is.null(kv)) {
        nanoparquet::write_parquet(x, tmp, compression = compression)
      } else {
        nanoparquet::write_parquet(
          x,
          tmp,
          compression = compression,
          metadata = kv
        )
      }
    },
    call
  )
}

# Read the parquet frame, using nanoparquet's native column projection when it
# is available and col_select is given (a real columnar IO win). nanoparquet
# returns columns in the requested order, so the kept names are fed in FILE
# order (from the schema) to keep behavior identical to a full read. Any
# failure (old nanoparquet, odd schema) falls back to a full read; the generic
# filter in read_dataset() is the single source of selection correctness.
#' @noRd
.parquet_read_frame <- function(path, col_select) {
  full <- function() as.data.frame(nanoparquet::read_parquet(path))
  if (
    is.null(col_select) ||
      !"col_select" %in% names(formals(nanoparquet::read_parquet))
  ) {
    return(full())
  }
  tryCatch(
    {
      file_cols <- nanoparquet::read_parquet_schema(path)$name[-1L]
      keep <- file_cols[file_cols %in% col_select]
      if (!length(keep)) {
        return(full()) # all-unknown: let the generic filter raise the error
      }
      as.data.frame(nanoparquet::read_parquet(path, col_select = keep))
    },
    error = function(e) full()
  )
}

# decode contract: (path, <codec args>, call) -> list(data, meta).
#' @noRd
.decode_parquet <- function(
  path,
  col_select = NULL,
  encoding = NULL,
  call = rlang::caller_env()
) {
  rlang::check_installed("nanoparquet", reason = "to read Parquet files.")

  df <- .parquet_read_frame(path, col_select)
  # nanoparquet restores an R factor as a factor; artoo's canonical model holds
  # strings as character (the JSON codec returns character too), so a factor
  # column would make the same data non-identical across formats. Defactor
  # first, so the recode and meta-inference steps below see plain character.
  fct <- vapply(df, is.factor, logical(1))
  if (any(fct)) {
    df[fct] <- lapply(df[fct], as.character)
  }
  # Parquet STRING bytes are UTF-8; canonicalise to internal UTF-8 (NFC). An
  # explicit `encoding` instead reads a foreign file whose bytes are that
  # charset. Non-character columns pass through untouched.
  source_enc <- encoding %||% "UTF-8"
  for (nm in names(df)) {
    df[[nm]] <- .recode_col(df[[nm]], source_enc)
  }
  json <- .parquet_sidecar(path)
  # Parse the sidecar ONCE: the same parsed object feeds the meta rebuild and
  # the special-missing tag reattachment below. A foreign / plain-nanoparquet
  # file (no sidecar) gets metadata synthesized from the column types and
  # attributes — the same path a bare frame takes on write — so the result
  # still feeds get_meta() and write_xpt()/write_json(), never an abort
  # (plan 9.B).
  p <- if (is.null(json)) {
    NULL
  } else {
    jsonlite::fromJSON(json, simplifyVector = FALSE)
  }
  meta <- if (is.null(p)) .meta_from_frame(df) else .meta_from_parsed(p)
  if (is.null(meta)) {
    return(list(data = df, meta = NULL)) # 0-column foreign file
  }
  # Realize temporal columns from the meta dataType (plan D1: meta-first).
  # Date/POSIXct survive nanoparquet natively and realize idempotently (the
  # integer-backed DATE arrival is canonicalized to double); a time column
  # comes back a bare double and becomes hms here. A character ISO
  # 8601 column (the no-targetDataType --DTC form) stays text — realize is
  # for numeric storage, and text is already readable.
  for (nm in names(meta@columns)) {
    cm <- meta@columns[[nm]]
    dt <- cm$dataType %||% ""
    if (
      dt %in%
        c("date", "datetime", "time") &&
        nm %in% names(df) &&
        !is.character(df[[nm]])
    ) {
      df[[nm]] <- .realize_temporal(df[[nm]], dt, cm$displayFormat %||% NA)
    }
  }
  # Reattach special-missing tags AFTER realization (it rebuilds vectors).
  sm <- if (is.null(p)) NULL else .special_from_parsed(p)
  if (!is.null(sm)) {
    df <- .apply_special_missings(df, sm)
  }
  list(data = df, meta = meta)
}

# ---- exported wrappers ------------------------------------------------------

#' Write a dataset to Apache Parquet
#'
#' Serialize a data frame to an Apache Parquet (`.parquet`) file, storing the
#' data natively while preserving the full `artoo_meta` as a CDISC-shaped
#' sidecar in the file's key-value metadata. The emit end of the artoo
#' workflow (spec -> apply_spec -> write_parquet); a thin wrapper over
#' [write_dataset()] with `format = "parquet"`. Requires the lightweight
#' `nanoparquet` package.
#'
#' @details
#' **Metadata where plain Parquet has none.** A bare nanoparquet/arrow file
#' drops labels, formats, and codelists; `write_parquet()` embeds the complete
#' `artoo_meta` as a single Dataset-JSON-shaped string under the
#' `metadata_json` key, so [read_parquet()] restores every CDISC attribute.
#' The same string is what a `.json` file or an rds carries, so conversion
#' between any two formats stays lossless. A reader without artoo still opens
#' the data and can see the `metadata_json` block.
#'
#' @param x *The dataset to write.* `<data.frame>: required`. Typically the
#'   output of [apply_spec()], carrying `artoo_meta`.
#' @param path *Destination `.parquet` path.* `<character(1)>: required`.
#' @param encoding *Source charset to record.* `<character(1)> | NULL`. The
#'   parquet bytes are always written as UTF-8 (the format's STRING type is
#'   UTF-8 by spec); `encoding` only records the data's original charset in the
#'   `artoo_meta`, so a later [write_xpt()] can reproduce the source bytes.
#'   `NULL` (default) leaves the recorded encoding untouched.
#'
#'   **Tip:** any SAS or IANA spelling listed by [artoo_encodings()] is
#'   accepted.
#' @param on_invalid *Policy for values that are not valid UTF-8.*
#'   `<character(1)>: default "error"`. One of `"error"` (abort with
#'   `artoo_error_codec`), `"replace"` (substitute `?` and warn with
#'   `artoo_warning_encoding`), `"ignore"` (drop the invalid bytes), or
#'   `"translit"` / `"fold"` (accepted for pipeline symmetry; behave as
#'   `"error"` here, since a byte-level invalidity has no character fold).
#'   See [write_json()] for when this fires; parquet STRING bytes are
#'   UTF-8 by spec, exactly like Dataset-JSON.
#' @param compression *Column compression codec.* `<character(1)>: default
#'   "snappy"`. One of:
#'
#'   - `"snappy"` (default) — fast, the parquet ecosystem default.
#'   - `"gzip"` — smaller files, slower.
#'   - `"zstd"` — the best size/speed trade-off where supported.
#'   - `"uncompressed"` — raw pages.
#'
#' @return *The input `x`*, invisibly, so a write can sit mid-pipeline.
#'
#' @examples
#' spec <- artoo_spec(cdisc_adam_datasets, cdisc_adam_variables, codelists = cdisc_codelists)
#'
#' # ---- Example 1: write a conformed dataset to Parquet ----
#' #
#' # apply_spec() attaches the metadata; write_parquet() stores the data
#' # natively and the metadata as a CDISC-shaped sidecar.
#' adsl <- apply_spec(cdisc_adsl, spec, "ADSL", conformance = "off")
#' path <- tempfile(fileext = ".parquet")
#' write_parquet(adsl, path)
#'
#' # ---- Example 2: round-trip and confirm the metadata survived ----
#' #
#' # Reading it back yields an identical artoo_meta.
#' back <- read_parquet(path)
#' identical(get_meta(back)@columns, get_meta(adsl)@columns)
#'
#' @seealso [read_parquet()] for the inverse; [write_dataset()] for the
#'   generic dispatcher.
#' @export
write_parquet <- function(
  x,
  path,
  encoding = NULL,
  on_invalid = c("error", "translit", "fold", "replace", "ignore"),
  compression = "snappy"
) {
  on_invalid <- match.arg(on_invalid)
  write_dataset(
    x,
    path,
    format = "parquet",
    encoding = encoding,
    on_invalid = on_invalid,
    compression = compression
  )
}

#' Read a dataset from Apache Parquet
#'
#' Read an Apache Parquet (`.parquet`) file back to a data frame, restoring the
#' `artoo_meta` from its `metadata_json` sidecar and realizing SAS
#' date/datetime/time variables to R `Date` / `POSIXct` / `hms::hms`. A
#' parquet written by another tool (with no artoo sidecar) reads back as a
#' bare frame. A thin wrapper over [read_dataset()] with `format = "parquet"`.
#' Requires the lightweight `nanoparquet` package.
#'
#' @param path *Source `.parquet` path.* `<character(1)>: required`.
#' @param encoding *Source charset of the string columns.* `<character(1)> |
#'   NULL`. `NULL` (default) reads the UTF-8 bytes parquet stores. Pass a
#'   charset name only to read a foreign file whose string columns hold that
#'   charset's bytes; they are transcoded to UTF-8 on read.
#'
#'   **Tip:** any SAS or IANA spelling listed by [artoo_encodings()] is
#'   accepted.
#' @inheritParams read_dataset
#'
#' @return *A `<data.frame>`* carrying `artoo_meta` when the file recorded it
#'   (read it with [get_meta()]); otherwise a plain data frame.
#'
#' @examples
#' spec <- artoo_spec(cdisc_adam_datasets, cdisc_adam_variables, codelists = cdisc_codelists)
#'
#' # ---- Example 1: round-trip a conformed dataset through Parquet ----
#' #
#' # The variable labels, types, and keys survive the round-trip.
#' adsl <- apply_spec(cdisc_adsl, spec, "ADSL", conformance = "off")
#' path <- tempfile(fileext = ".parquet")
#' write_parquet(adsl, path)
#' back <- read_parquet(path)
#' get_meta(back)@columns$STUDYID$label
#'
#' # ---- Example 2: the metadata names the dataset and row count ----
#' #
#' # The restored artoo_meta exposes the dataset-level attributes.
#' get_meta(back)@dataset$records
#'
#' @seealso [write_parquet()] for the inverse; [read_dataset()] for the
#'   generic dispatcher.
#' @export
read_parquet <- function(
  path,
  col_select = NULL,
  n_max = Inf,
  encoding = NULL
) {
  read_dataset(
    path,
    format = "parquet",
    col_select = col_select,
    n_max = n_max,
    encoding = encoding
  )
}

.register_codec(
  "parquet",
  encode = ".encode_parquet",
  decode = ".decode_parquet",
  extensions = c("parquet", "pq"),
  mode = "rw",
  engine = "nanoparquet"
)