fix: support all-None index

scikit-hep · Oct 24, 2023 · 1d9dc4a · 1d9dc4a
1 parent 8927adb
commit 1d9dc4a
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/awkward-cpp/src/cpu-kernels/awkward_Index_nones_as_index.cpp b/awkward-cpp/src/cpu-kernels/awkward_Index_nones_as_index.cpp
@@ -8,12 +8,17 @@ template <typename T>
 ERROR awkward_Index_nones_as_index(
   T* toindex,
   int64_t length) {
-  int64_t last_index = 0;
+  int64_t n_non_null = 0;
+  // Assuming that `toindex` comprises of contiguous integers, and is zero-based
+  // Compute the number of non-null values to determine our starting index
   for (int64_t i = 0; i < length; i++) {
-    toindex[i] > last_index ? last_index = toindex[i] : last_index;
+    if (toindex[i] != -1) {
+        n_non_null++;
+    }
   }
+  //
   for (int64_t i = 0; i < length; i++) {
-    toindex[i] == -1 ? toindex[i] = ++last_index : toindex[i];
+    toindex[i] == -1 ? toindex[i] = n_non_null++ : toindex[i];
   }
   return success();
 }