Implement altrep mmap approach

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: grumpy
 Title: Read 'NumPy' '.npy' and '.npz' Files
-Version: 0.1.1.9001
+Version: 0.1.1.9002
 Authors@R: c(
     person("Hugo", "Gruson", , "hugo.gruson+R@normalesup.org", role = c("aut", "cre", "cph"),
            comment = c(ORCID = "0000-0002-4094-1476")),

NEWS.md

@@ -1,5 +1,12 @@
 # grumpy (development version)
 
+## Significant new features
+
+* `read_npy()` gains a new `lazy` argument. When `lazy = TRUE`, raw bytes are
+  read lazily using ALTREP and mmap against the `.npy` file on disk, resulting
+  in much better speed & memory performance. This is following a feature
+  request from @btraven00 in #11. 
+
 # grumpy 0.1.1
 
 * The package title and description have been revised based on CRAN feedback.

R/read_npy.R

@@ -1,6 +1,11 @@
 #' Read a .npy file
 #'
 #' @param file Path to the .npy file
+#' @param lazy If `TRUE`, and `file` is a path (not a connection), the data
+#'   payload is memory-mapped rather than read into memory upfront via
+#'   [readBin()]. This limits the number of copies in memory when the npy file
+#'   contains types native to R. Requires a POSIX platform; ignored (with a
+#'   warning) otherwise. Defaults to `FALSE`.
 #' @param ... Ignored. Reserved for future use.
 #'
 #' @returns An array containing the data from the .npy file
@@ -11,8 +16,7 @@
 #' read_npy(
 #'   system.file("extdata", "test.npy", package = "grumpy")
 #' )
-
-read_npy <- function(file, ...) {
+read_npy <- function(file, lazy = FALSE, ...) {
   chkDots(...)
 
   if (is.character(file)) {
@@ -22,6 +26,14 @@ read_npy <- function(file, ...) {
     con <- file(file, "rb")
     on.exit(close(con))
   } else if (inherits(file, "connection")) {
+    if (lazy) {
+      warning(
+        "`lazy = TRUE` is only supported when `file` is a path; ",
+        "ignoring it for connections.",
+        call. = FALSE
+      )
+      lazy <- FALSE
+    }
     con <- file
   } else {
     stop(
@@ -73,7 +85,19 @@ read_npy <- function(file, ...) {
 
   # Read the data
   num_elements <- prod(header$shape)
-  bytes <- readBin(con, "raw", n = sum(num_elements * header$nbytes))
+  n_bytes <- sum(num_elements * header$nbytes)
+
+  if (lazy) {
+    # Offset of the payload is wherever the connection currently sits:
+    # 6 (magic) + 2 (version) + header_len_field_size (2 or 4) + header_len
+    payload_offset <- seek(con, origin = "current")
+    # When lazy = TRUE, we know `file` is a path, not a connection, so we can
+    # mmap it directly.
+    bytes <- mmap_raw(file, offset = payload_offset, length = n_bytes)
+  } else {
+    bytes <- readBin(con, "raw", n = n_bytes)
+  }
+
   convert_bytes_to_array(
     bytes,
     what = header$base_type,
@@ -83,6 +107,31 @@ read_npy <- function(file, ...) {
   )
 }
 
+#' Create an ALTREP raw vector backed by a memory-mapped file slice
+#'
+#' Internal helper used by [read_npy()] when `lazy = TRUE`. The returned
+#' object behaves like an ordinary raw vector (as produced by [readBin()]),
+#' but its bytes are only paged in from disk on first access, and no R-level
+#' allocation happens for the full payload upfront. It can be passed to
+#' [convert_bytes_to_array()] exactly like a materialized raw vector.
+#'
+#' @param path Path to the file to map.
+#' @param offset Byte offset of the start of the desired slice.
+#' @param length Number of bytes to expose.
+#'
+#' @returns An ALTREP raw vector of length `length`.
+#'
+#' @noRd
+mmap_raw <- function(path, offset, length) {
+  .Call(
+    C_grumpy_make_mmap_raw,
+    path.expand(path),
+    as.double(offset),
+    as.double(length),
+    PACKAGE = "grumpy"
+  )
+}
+
 parse_npy_descr <- function(bytes) {
   # TODO: If I understand correctly, fortranarray in python are still displayed
   # the same way as regular arrays, but with a different order in memory.

src/altrep_mmap.c

@@ -0,0 +1,127 @@
+#include "altrep_mmap.h"
+
+#include <R_ext/Altrep.h>
+#include <string.h>
+
+// FIXME: we only support POSIX mmap() for now.
+// Windows has a different API (CreateFileMapping / MapViewOfFile) and would 
+// require a separate implementation.
+#ifndef _WIN32
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <unistd.h>
+#define GRUMPY_HAVE_MMAP 1
+#endif
+
+// ALTREP raw vector backed by a read-only mmap() of a slice of a file.
+//
+// Having a file-backed, lazy read of raw data is useful in the cases when
+// we don't need to do anything transformation (e.g., int32).
+//
+// It could also be helpful in the future if we ever implement an `index`
+// argument to `read_npy()` that allows reading a subset of the data. 
+
+typedef struct {
+  void *addr;      // mmap base address (page-aligned)
+  size_t map_len;  // bytes mapped (offset rounded down + length)
+  size_t skip;     // bytes to skip within the mapping to reach the logical start
+  R_xlen_t length; // logical length exposed to R
+} grumpy_mmap_info;
+
+static R_altrep_class_t grumpy_mmap_raw_class;
+
+static void grumpy_mmap_xptr_finalize(SEXP xp) {
+  grumpy_mmap_info *info = (grumpy_mmap_info *) R_ExternalPtrAddr(xp);
+  if (info == NULL) return;
+#ifdef GRUMPY_HAVE_MMAP
+  if (info->addr != NULL && info->addr != MAP_FAILED) {
+    munmap(info->addr, info->map_len);
+  }
+#endif
+  free(info);
+  R_ClearExternalPtr(xp);
+}
+
+static R_xlen_t grumpy_mmap_length(SEXP x) {
+  const grumpy_mmap_info *info = (grumpy_mmap_info *) R_ExternalPtrAddr(R_altrep_data1(x));
+  return info->length;
+}
+
+static void *grumpy_mmap_dataptr(SEXP x, Rboolean writeable) {
+  const grumpy_mmap_info *info = (grumpy_mmap_info *) R_ExternalPtrAddr(R_altrep_data1(x));
+  return (char *) info->addr + info->skip;
+}
+
+static const void *grumpy_mmap_dataptr_or_null(SEXP x) {
+  return grumpy_mmap_dataptr(x, FALSE);
+}
+
+static Rbyte grumpy_mmap_elt(SEXP x, R_xlen_t i) {
+  const Rbyte *p = (const Rbyte *) grumpy_mmap_dataptr_or_null(x);
+  return p[i];
+}
+
+static R_xlen_t grumpy_mmap_get_region(SEXP x, R_xlen_t start, R_xlen_t n, Rbyte *buf) {
+  const Rbyte *p = (const Rbyte *) grumpy_mmap_dataptr_or_null(x);
+  R_xlen_t len = grumpy_mmap_length(x);
+  R_xlen_t ncopy = (start + n > len) ? (len - start) : n;
+  memcpy(buf, p + start, ncopy);
+  return ncopy;
+}
+
+void grumpy_init_mmap_altrep(DllInfo *info) {
+  grumpy_mmap_raw_class = R_make_altraw_class("grumpy_mmap_raw", "grumpy", info);
+
+  R_set_altrep_Length_method(grumpy_mmap_raw_class, grumpy_mmap_length);
+  R_set_altvec_Dataptr_method(grumpy_mmap_raw_class, grumpy_mmap_dataptr);
+  R_set_altvec_Dataptr_or_null_method(grumpy_mmap_raw_class, grumpy_mmap_dataptr_or_null);
+  R_set_altraw_Elt_method(grumpy_mmap_raw_class, grumpy_mmap_elt);
+  R_set_altraw_Get_region_method(grumpy_mmap_raw_class, grumpy_mmap_get_region);
+}
+
+SEXP grumpy_make_mmap_raw(SEXP path_, SEXP offset_, SEXP length_) {
+#ifndef GRUMPY_HAVE_MMAP
+  error("mmap-backed reading is not supported on this platform");
+#else
+  const char *path = CHAR(STRING_ELT(path_, 0));
+  double offset = REAL(offset_)[0];
+  double length = REAL(length_)[0];
+
+  if (offset < 0 || length < 0)
+    error("offset and length must be non-negative");
+
+  int fd = open(path, O_RDONLY);
+  if (fd < 0)
+    error("Cannot open file for mmap: %s", path);
+
+  long page_size = sysconf(_SC_PAGE_SIZE);
+  size_t page_offset = ((size_t) offset) % page_size;
+  off_t map_start = (off_t) offset - (off_t) page_offset;
+  size_t map_len = (size_t) length + page_offset;
+
+  void *addr = mmap(NULL, map_len, PROT_READ, MAP_PRIVATE, fd, map_start);
+  // The fd is not needed once mapped.
+  close(fd);
+
+  if (addr == MAP_FAILED)
+    error("mmap() failed for file: %s", path);
+
+  grumpy_mmap_info *map_info = malloc(sizeof(grumpy_mmap_info));
+  if (map_info == NULL) {
+    munmap(addr, map_len);
+    error("Failed to allocate mmap bookkeeping struct");
+  }
+  map_info->addr = addr;
+  map_info->map_len = map_len;
+  map_info->skip = page_offset;
+  map_info->length = (R_xlen_t) length;
+
+  SEXP xp = PROTECT(R_MakeExternalPtr(map_info, R_NilValue, R_NilValue));
+  R_RegisterCFinalizerEx(xp, grumpy_mmap_xptr_finalize, TRUE);
+  SEXP ans = PROTECT(R_new_altrep(grumpy_mmap_raw_class, xp, R_NilValue));
+  UNPROTECT(2);
+
+  return ans;
+#endif /* GRUMPY_HAVE_MMAP */
+}

src/altrep_mmap.h

@@ -0,0 +1,13 @@
+#ifndef _GRUMPY_ALTREP_MMAP_H
+#define _GRUMPY_ALTREP_MMAP_H
+
+#include "grumpy.h"
+
+// Create an ALTREP raw vector backed by a read-only mmap() of `path`,
+// exposing `length` bytes starting at file offset `offset`.
+SEXP grumpy_make_mmap_raw(SEXP path, SEXP offset, SEXP length);
+
+// Register the ALTREP class and its methods. Called from R_init_grumpy().
+void grumpy_init_mmap_altrep(DllInfo *info);
+
+#endif /* _GRUMPY_ALTREP_MMAP_H */

src/grumpy.c

@@ -1,8 +1,10 @@
 #include "grumpy.h"
 #include "type_conversion.h"
+#include "altrep_mmap.h"
 
 static const R_CallMethodDef callMethods[] = {
   {"type_convert", (DL_FUNC) &type_convert, 5},
+  {"grumpy_make_mmap_raw", (DL_FUNC) &grumpy_make_mmap_raw, 3},
   {NULL, NULL, 0}
 };
 
@@ -11,4 +13,5 @@ void R_init_grumpy(DllInfo *info)
   R_registerRoutines(info, NULL, callMethods, NULL, NULL);
   R_useDynamicSymbols(info, FALSE);
   R_forceSymbols(info, TRUE);
+  grumpy_init_mmap_altrep(info);
 }