Add chr_transform argument to vec_order_radix() (#1334)

* Implement C side of `string_key` support

* Implement R side of `string_key` support, with documentation

* Add `string_key` tests

* Reorder loop to make reprotection easier

* Use a protection index to avoid conditional protection

* Tweak string-key error messages

* Rename `string_key` to `chr_transform`

* Use `proxy_chr_transform()` to apply character transforms up front

* Update documentation regarding the name change

* Mention common transformation functions in `chr_transform` param section

Author: DavisVaughan

R/order-radix.R

@@ -19,9 +19,16 @@
 #' unless `base::order(method = "radix")` is explicitly set, which also uses
 #' the C-locale. While sorting with the C-locale can be useful for
 #' algorithmic efficiency, in many real world uses it can be the cause of
-#' data analysis mistakes.
+#' data analysis mistakes. To balance these trade-offs, you can supply a
+#' `chr_transform` to transform character vectors into an alternative
+#' representation that orders in the C-locale in a less surprising way. For
+#' example, providing [base::tolower()] as a transform will order the original
+#' vector in a case-insensitive manner. Locale-aware ordering can be achieved
+#' by providing `stringi::stri_sort_key()` as a transform, setting the
+#' collation options as appropriate for your locale.
 #'
-#' Character vectors are always translated to UTF-8 before ordering.
+#' Character vectors are always translated to UTF-8 before ordering, and before
+#' any transform is applied by `chr_transform`.
 #'
 #' @param x A vector
 #' @param direction Direction to sort in.
@@ -36,6 +43,19 @@
 #'   - For data frames, a length `1` or `ncol(x)` character vector containing
 #'     only `"largest"` or `"smallest"`, specifying how `NA`s should be treated
 #'     in each column.
+#' @param chr_transform Transformation of character vectors for sorting in
+#'   alternate locales.
+#'   - If `NULL`, no transformation is done.
+#'   - Otherwise, this must be a function of one argument. The function will be
+#'     invoked with `x`, if it is a character vector, after it has been
+#'     translated to UTF-8, and should return a character vector with the same
+#'     length as `x`, also encoded as UTF-8.
+#'   - For data frames, `chr_transform` will be applied to all character
+#'     columns.
+#'
+#'   Common transformation functions include: `tolower()` for case-insensitive
+#'   ordering and `stringi::str_sort_key()` for locale-aware ordering. See the
+#'   Details section for more information.
 #' @return
 #' * `vec_order()` an integer vector the same size as `x`.
 #' * `vec_sort()` a vector with the same size and type as `x`.
@@ -67,9 +87,21 @@
 #'   direction = c("desc", "asc"),
 #'   na_value = c("largest", "smallest")
 #' )
+#'
+#' # Character vectors are ordered in the C locale, which orders capital letters
+#' # below lowercase ones
+#' y <- c("B", "A", "a")
+#' vec_sort(y)
+#'
+#' # To order in a case-insensitive manner, provide a `chr_transform` that
+#' # transforms the strings to all lowercase
+#' vec_sort(y, chr_transform = tolower)
 #' @noRd
-vec_order_radix <- function(x, direction = "asc", na_value = "largest") {
-  .Call(vctrs_order, x, direction, na_value)
+vec_order_radix <- function(x,
+                            direction = "asc",
+                            na_value = "largest",
+                            chr_transform = NULL) {
+  .Call(vctrs_order, x, direction, na_value, chr_transform)
 }
 
 #' Identify ordered groups
@@ -104,6 +136,9 @@ vec_order_radix <- function(x, direction = "asc", na_value = "largest") {
 #'
 #' vec_group_loc(df)
 #' @noRd
-vec_order_locs <- function(x, direction = "asc", na_value = "largest") {
-  .Call(vctrs_order_locs, x, direction, na_value)
+vec_order_locs <- function(x,
+                           direction = "asc",
+                           na_value = "largest",
+                           chr_transform = NULL) {
+  .Call(vctrs_order_locs, x, direction, na_value, chr_transform)
 }

src/init.c

@@ -129,8 +129,8 @@ extern SEXP vctrs_slice_complete(SEXP);
 extern SEXP vctrs_locate_complete(SEXP);
 extern SEXP vctrs_detect_complete(SEXP);
 extern SEXP vctrs_normalize_encoding(SEXP);
-extern SEXP vctrs_order(SEXP, SEXP, SEXP);
-extern SEXP vctrs_order_locs(SEXP, SEXP, SEXP);
+extern SEXP vctrs_order(SEXP, SEXP, SEXP, SEXP);
+extern SEXP vctrs_order_locs(SEXP, SEXP, SEXP, SEXP);
 extern SEXP vctrs_unrep(SEXP);
 extern SEXP vctrs_fill_missing(SEXP, SEXP, SEXP);
 extern SEXP vctrs_chr_paste_prefix(SEXP, SEXP, SEXP);
@@ -280,8 +280,8 @@ static const R_CallMethodDef CallEntries[] = {
   {"vctrs_locate_complete",            (DL_FUNC) &vctrs_locate_complete, 1},
   {"vctrs_detect_complete",            (DL_FUNC) &vctrs_detect_complete, 1},
   {"vctrs_normalize_encoding",         (DL_FUNC) &vctrs_normalize_encoding, 1},
-  {"vctrs_order",                      (DL_FUNC) &vctrs_order, 3},
-  {"vctrs_order_locs",                 (DL_FUNC) &vctrs_order_locs, 3},
+  {"vctrs_order",                      (DL_FUNC) &vctrs_order, 4},
+  {"vctrs_order_locs",                 (DL_FUNC) &vctrs_order_locs, 4},
   {"vctrs_unrep",                      (DL_FUNC) &vctrs_unrep, 1},
   {"vctrs_fill_missing",               (DL_FUNC) &vctrs_fill_missing, 3},
   {"vctrs_chr_paste_prefix",           (DL_FUNC) &vctrs_chr_paste_prefix, 3},

src/order-radix.c

@@ -19,6 +19,7 @@
 #include "order-groups.h"
 #include "order-truelength.h"
 #include "order-sortedness.h"
+#include "order-transform.h"
 
 // -----------------------------------------------------------------------------
 
@@ -183,46 +184,50 @@
 
 // -----------------------------------------------------------------------------
 
-static SEXP vec_order(SEXP x, SEXP decreasing, SEXP na_last);
+static SEXP vec_order(SEXP x, SEXP decreasing, SEXP na_last, SEXP chr_transform);
 
 // [[ register() ]]
-SEXP vctrs_order(SEXP x, SEXP direction, SEXP na_value) {
+SEXP vctrs_order(SEXP x, SEXP direction, SEXP na_value, SEXP chr_transform) {
   SEXP decreasing = PROTECT(parse_direction(direction));
   SEXP na_last = PROTECT(parse_na_value(na_value));
 
-  SEXP out = vec_order(x, decreasing, na_last);
+  SEXP out = vec_order(x, decreasing, na_last, chr_transform);
 
   UNPROTECT(2);
   return out;
 }
 
 
-static SEXP vec_order_impl(SEXP x, SEXP decreasing, SEXP na_last, bool locations);
+static SEXP vec_order_impl(SEXP x,
+                           SEXP decreasing,
+                           SEXP na_last,
+                           SEXP chr_transform,
+                           bool locations);
 
 static
-SEXP vec_order(SEXP x, SEXP decreasing, SEXP na_last) {
-  return vec_order_impl(x, decreasing, na_last, false);
+SEXP vec_order(SEXP x, SEXP decreasing, SEXP na_last, SEXP chr_transform) {
+  return vec_order_impl(x, decreasing, na_last, chr_transform, false);
 }
 
 // -----------------------------------------------------------------------------
 
-static SEXP vec_order_locs(SEXP x, SEXP decreasing, SEXP na_last);
+static SEXP vec_order_locs(SEXP x, SEXP decreasing, SEXP na_last, SEXP chr_transform);
 
 // [[ register() ]]
-SEXP vctrs_order_locs(SEXP x, SEXP direction, SEXP na_value) {
+SEXP vctrs_order_locs(SEXP x, SEXP direction, SEXP na_value, SEXP chr_transform) {
   SEXP decreasing = PROTECT(parse_direction(direction));
   SEXP na_last = PROTECT(parse_na_value(na_value));
 
-  SEXP out = vec_order_locs(x, decreasing, na_last);
+  SEXP out = vec_order_locs(x, decreasing, na_last, chr_transform);
 
   UNPROTECT(2);
   return out;
 }
 
 
 static
-SEXP vec_order_locs(SEXP x, SEXP decreasing, SEXP na_last) {
-  return vec_order_impl(x, decreasing, na_last, true);
+SEXP vec_order_locs(SEXP x, SEXP decreasing, SEXP na_last, SEXP chr_transform) {
+  return vec_order_impl(x, decreasing, na_last, chr_transform, true);
 }
 
 // -----------------------------------------------------------------------------
@@ -258,7 +263,7 @@ static void vec_order_switch(SEXP x,
  * the locations in `x` corresponding to each key.
  */
 static
-SEXP vec_order_impl(SEXP x, SEXP decreasing, SEXP na_last, bool locations) {
+SEXP vec_order_impl(SEXP x, SEXP decreasing, SEXP na_last, SEXP chr_transform, bool locations) {
   int n_prot = 0;
   int* p_n_prot = &n_prot;
 
@@ -269,6 +274,7 @@ SEXP vec_order_impl(SEXP x, SEXP decreasing, SEXP na_last, bool locations) {
 
   SEXP proxy = PROTECT_N(vec_proxy_order(x), p_n_prot);
   proxy = PROTECT_N(vec_normalize_encoding(proxy), p_n_prot);
+  proxy = PROTECT_N(proxy_chr_transform(proxy, chr_transform), p_n_prot);
 
   r_ssize size = vec_size(proxy);
   const enum vctrs_type type = vec_proxy_typeof(proxy);
@@ -3526,7 +3532,14 @@ void df_order_internal(SEXP x,
 
   // Iterate over remaining columns by group chunk
   for (r_ssize i = 1; i < n_cols; ++i) {
-    col = VECTOR_ELT(x, i);
+    // Get the number of group chunks from previous column group info
+    struct group_info* p_group_info_pre = groups_current(p_group_infos);
+    r_ssize n_groups = p_group_info_pre->n_groups;
+
+    // If there were no ties, we are completely done
+    if (n_groups == size) {
+      break;
+    }
 
     if (!recycle_decreasing) {
       col_decreasing = p_decreasing[i];
@@ -3541,15 +3554,7 @@ void df_order_internal(SEXP x,
     // processed at least one column.
     int* p_o_col = p_order->p_data;
 
-    // Get the number of group chunks from previous column group info
-    struct group_info* p_group_info_pre = groups_current(p_group_infos);
-    r_ssize n_groups = p_group_info_pre->n_groups;
-
-    // If there were no ties, we are completely done
-    if (n_groups == size) {
-      break;
-    }
-
+    col = VECTOR_ELT(x, i);
     type = vec_proxy_typeof(col);
 
     // If we are on the rerun pass, flip this back off so the
@@ -3558,7 +3563,7 @@ void df_order_internal(SEXP x,
       rerun_complex = rerun_complex ? false : true;
     }
 
-    // Pre sort unique characters once for the whole column
+    // Pre-sort unique characters once for the whole column
     if (type == vctrs_type_character) {
       const SEXP* p_col = STRING_PTR_RO(col);
 

src/order-transform.c

@@ -0,0 +1,113 @@
+/*
+ * The implementation of vec_order() is based on data.table’s forder() and their
+ * earlier contribution to R’s order(). See LICENSE.note for more information.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this file,
+ * You can obtain one at https://mozilla.org/MPL/2.0/.
+ *
+ * Copyright (c) 2020, RStudio
+ * Copyright (c) 2020, Data table team
+ */
+
+#include "order-transform.h"
+#include "utils.h"
+
+// -----------------------------------------------------------------------------
+
+static SEXP chr_apply_transform(SEXP x, SEXP chr_transform);
+static SEXP df_apply_transform(SEXP x, SEXP chr_transform);
+
+// [[ include("order-transform.h") ]]
+SEXP proxy_chr_transform(SEXP proxy, SEXP chr_transform) {
+  if (chr_transform == r_null) {
+    return proxy;
+  }
+
+  chr_transform = PROTECT(r_as_function(chr_transform, "chr_transform"));
+
+  SEXP out;
+
+  switch (vec_proxy_typeof(proxy)) {
+  case vctrs_type_character: out = chr_apply_transform(proxy, chr_transform); break;
+  case vctrs_type_dataframe: out = df_apply_transform(proxy, chr_transform); break;
+  default: out = proxy;
+  }
+
+  UNPROTECT(1);
+  return out;
+}
+
+// -----------------------------------------------------------------------------
+
+static
+SEXP chr_apply_transform(SEXP x, SEXP chr_transform) {
+  // Don't use vctrs dispatch utils because we match argument positionally
+  SEXP call = PROTECT(Rf_lang2(syms_chr_transform, syms_x));
+
+  SEXP mask = PROTECT(r_new_environment(R_GlobalEnv));
+  Rf_defineVar(syms_chr_transform, chr_transform, mask);
+  Rf_defineVar(syms_x, x, mask);
+
+  SEXP out = PROTECT(Rf_eval(call, mask));
+
+  if (vec_typeof(out) != vctrs_type_character) {
+    Rf_errorcall(
+      R_NilValue,
+      "`chr_transform` must return a character vector."
+    );
+  }
+
+  R_len_t x_size = vec_size(x);
+  R_len_t out_size = vec_size(out);
+
+  if (x_size != out_size) {
+    Rf_errorcall(
+      R_NilValue,
+      "`chr_transform` must return a vector of the same length (%i, not %i).",
+      x_size,
+      out_size
+    );
+  }
+
+  UNPROTECT(3);
+  return out;
+}
+
+// -----------------------------------------------------------------------------
+
+static
+SEXP df_apply_transform(SEXP x, SEXP chr_transform) {
+  const r_ssize n_cols = r_length(x);
+  const SEXP* v_x = VECTOR_PTR_RO(x);
+
+  r_ssize i = 0;
+
+  for (; i < n_cols; ++i) {
+    SEXP col = v_x[i];
+    if (vec_proxy_typeof(col) == vctrs_type_character) {
+      break;
+    }
+  }
+
+  if (i == n_cols) {
+    // No character columns
+    return x;
+  }
+
+  SEXP out = PROTECT(r_clone_referenced(x));
+
+  for (; i < n_cols; ++i) {
+    SEXP col = v_x[i];
+
+    if (vec_proxy_typeof(col) != vctrs_type_character) {
+      continue;
+    }
+
+    col = chr_apply_transform(col, chr_transform);
+    SET_VECTOR_ELT(out, i, col);
+  }
+
+  UNPROTECT(1);
+  return out;
+}

src/order-transform.h

@@ -0,0 +1,33 @@
+/*
+ * The implementation of vec_order() is based on data.table’s forder() and their
+ * earlier contribution to R’s order(). See LICENSE.note for more information.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this file,
+ * You can obtain one at https://mozilla.org/MPL/2.0/.
+ *
+ * Copyright (c) 2020, RStudio
+ * Copyright (c) 2020, Data table team
+ */
+
+#ifndef VCTRS_ORDER_TRANSFORM_H
+#define VCTRS_ORDER_TRANSFORM_H
+
+#include "vctrs.h"
+
+// -----------------------------------------------------------------------------
+
+/*
+ * `proxy_chr_transform()` iterates over `proxy`, applying `chr_transform`
+ * on any character vectors that it detects.
+ *
+ * It expects that:
+ * - If `proxy` is a data frame, it has been flattened by its corresponding
+ *   `vec_proxy_*()` function.
+ * - All character vectors in `proxy` have already been normalized to UTF-8
+ *   by `vec_normalize_encoding()`.
+ */
+SEXP proxy_chr_transform(SEXP proxy, SEXP chr_transform);
+
+// -----------------------------------------------------------------------------
+#endif

src/utils.c

@@ -1835,6 +1835,7 @@ SEXP syms_vctrs_common_class_fallback = NULL;
 SEXP syms_fallback_class = NULL;
 SEXP syms_abort = NULL;
 SEXP syms_message = NULL;
+SEXP syms_chr_transform = NULL;
 
 SEXP fns_bracket = NULL;
 SEXP fns_quote = NULL;
@@ -2110,6 +2111,7 @@ void vctrs_init_utils(SEXP ns) {
   syms_fallback_class = Rf_install("fallback_class");
   syms_abort = Rf_install("abort");
   syms_message = Rf_install("message");
+  syms_chr_transform = Rf_install("chr_transform");
 
   fns_bracket = Rf_findVar(syms_bracket, R_BaseEnv);
   fns_quote = Rf_findVar(Rf_install("quote"), R_BaseEnv);

src/utils.h

@@ -631,6 +631,7 @@ extern SEXP syms_vctrs_common_class_fallback;
 extern SEXP syms_fallback_class;
 extern SEXP syms_abort;
 extern SEXP syms_message;
+extern SEXP syms_chr_transform;
 
 static const char * const c_strs_vctrs_common_class_fallback = "vctrs:::common_class_fallback";