Add support for categorical columns with NA

R/read_h5mu.R

@@ -25,11 +25,7 @@ read_dataframe <- function(group) {
             } else {
                 labels <- H5Rdereference(labels, h5loc=col)
                 labels_items <- H5Dread(labels)
-                n_labels <- length(unique(values))
-                if (length(labels_items) > n_labels) {
-                    labels_items <- labels_items[seq_len(n_labels)]
-                }
-                values <- factor(as.integer(values), labels=labels_items)
+                values <- convert_categoricals(values, labels_items)
                 H5Dclose(labels)
             }
             H5Aclose(attr)
@@ -44,6 +40,30 @@ read_dataframe <- function(group) {
     do.call(data.frame, args=col_list)
 }
 
+#' Helper function to convert values + labels into factors
+#'
+#' @description  A helper function to convert categories into factors.
+#'  Assumptions: 
+#'      - values correspond to the zero indexed categories
+#'          (i.e. value 0 is the first category) 
+#'      - NA are encoded with a value -1
+#'  Categories not uses will be dropped.
+#'
+#' @param values Vector of integer level numbers (zero indexed). -1 indicate NA
+#' @param categories Labels for level numbers (zero indexed).
+#'
+#' @returns factor with categorical values
+#'
+#' @keywords internal
+#' @noRd 
+convert_categoricals <- function(values, categories) {
+    # The levels are 0 indexed integers
+    levels <- seq_len(length(categories))-1
+    value_factor <- factor(as.integer(values), levels, labels=categories)
+    # Drop unused levels
+    droplevels(value_factor)
+}
+
 #' @importFrom rhdf5 H5Dread H5Aexists H5Aopen H5Aread H5Aclose
 read_dataframe_legacy <- function(dataset) {
     table <- H5Dread(dataset)

tests/testthat/test_readh5mu.R

@@ -74,3 +74,25 @@ test_that("a SE object with a sparse matrix written to H5AD can be read", {
     expect_true(inherits(assay(se_b), "DelayedArray"))
   }
 })
+
+test_that("Categoricals columns with NA are loaded correctly", {
+  values <- c(0, -1)
+  categories <- c('a')
+  res <- convert_categoricals(values, categories)
+  expect_equal(as.character(res[1]), 'a')
+  expect_true(is.na(res[2]))
+})
+
+test_that("Extra levels in categoricals are ignored", {
+  values <- c(0)
+  categories <- c('a', 'b')
+  res <- convert_categoricals(values, categories)
+  expect_equal(levels(res), c('a'))
+})
+
+test_that("Extra levels in categoricals are ignored when NA are present", {
+  values <- c(0, -1)
+  categories <- c('a', 'b')
+  res <- convert_categoricals(values, categories)
+  expect_equal(levels(res), c('a'))
+})