fix: support all-None index in awkward_Index_nones_as_index.cpp

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 unique, contiguous integers (or -1), 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++;
+    }
   }
+  // Now set the null-value indices to by monotonically increasing and unique from the final index
   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();
 }

kernel-specification.yml

@@ -4500,14 +4500,14 @@ kernels:
     description: null
     definition: |
       def awkward_Index_nones_as_index(toindex, length):
-          last_index = 0
+          num_non_null = 0
           for i in range(length):
-              if toindex[i] > last_index:
-                  last_index = toindex[i]
+              if toindex[i] != -1:
+                  num_non_null += 1
           for i in range(length):
               if toindex[i] == -1:
-                  last_index = last_index + 1
-                  toindex[i] = last_index
+                  toindex[i] = num_non_null
+                  num_non_null += 1
     automatic-tests: false
     manual-tests: []
 

src/awkward/contents/indexedoptionarray.py

@@ -1265,10 +1265,15 @@ def _argsort_next(
                     nextoutindex.length,
                 )
             )
-
-        out = ak.contents.IndexedOptionArray.simplified(
-            nextoutindex, out, parameters=self._parameters
-        )
+            # Drop the option type: we have ensured that we don't have any
+            # -1 values in `nextoutindex` now!
+            out = out._carry(nextoutindex, False).copy(
+                parameters=parameters_union(out._parameters, self._parameters)
+            )
+        else:
+            out = ak.contents.IndexedOptionArray.simplified(
+                nextoutindex, out, parameters=self._parameters
+            )
 
         inject_nones = (
             True

tests/test_2769_argsort_all_none.py

@@ -0,0 +1,51 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import pytest  # noqa: F401
+
+import awkward as ak
+
+
+def test_all_none():
+    result = ak.argsort([None, None, None])
+    assert ak.is_valid(result)
+    assert result.to_list() == [0, 1, 2]
+    assert isinstance(result.layout, ak.contents.NumpyArray)
+    assert result.type == ak.types.ArrayType(ak.types.NumpyType("int64"), 3)
+
+
+def test_mixed_none_local():
+    result = ak.argsort([None, 1, None, 2, 0, None, -1])
+    assert ak.is_valid(result)
+    assert result.to_list() == [6, 4, 1, 3, 0, 2, 5]
+    assert result.type == ak.types.ArrayType(ak.types.NumpyType("int64"), 7)
+
+
+def test_mixed_none_2d_local():
+    result = ak.argsort(
+        [
+            [None, 1, None, 0, None, None, -1],
+            None,
+            [None, 2, None, 2, 0, None, -2],
+        ],
+        axis=1,
+    )
+    assert ak.is_valid(result)
+    assert result.to_list() == [[6, 3, 1, 0, 2, 4, 5], None, [6, 4, 1, 3, 0, 2, 5]]
+    assert result.type == ak.types.ArrayType(
+        ak.types.OptionType(ak.types.ListType(ak.types.NumpyType("int64"))), 3
+    )
+
+
+def test_mixed_none_2d_nonlocal():
+    result = ak.argsort(
+        [
+            [None, 1, None, 0, None, None, -1],
+            [None, 2, None, 2, 0, None, -2],
+        ],
+        axis=0,
+    )
+    assert ak.is_valid(result)
+    assert result.to_list() == [[0, 0, 0, 0, 1, 0, 1], [1, 1, 1, 1, 0, 1, 0]]
+    assert result.type == ak.types.ArrayType(
+        ak.types.ListType(ak.types.NumpyType("int64")), 2
+    )