fix: added kernel specializations for `awkward_ListOffsetArray_reduce…

…_local_nextparents_64` (#2572)

awkward-cpp/pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "awkward_cpp"
-version = "20"
+version = "21"
 dependencies = [
     "numpy>=1.18.0"
 ]

awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_reduce_local_nextparents_64.cpp

@@ -4,17 +4,45 @@
 
 #include "awkward/kernels.h"
 
+template <typename C>
 ERROR awkward_ListOffsetArray_reduce_local_nextparents_64(
   int64_t* nextparents,
-  const int64_t* offsets,
+  const C* offsets,
   int64_t length) {
-  int64_t initialoffset = offsets[0];
+  int64_t initialoffset = (int64_t)(offsets[0]);
   for (int64_t i = 0;  i < length;  i++) {
-    for (int64_t j = offsets[i] - initialoffset;
+    for (int64_t j = (int64_t)(offsets[i]) - initialoffset;
          j < offsets[i + 1] - initialoffset;
          j++) {
       nextparents[j] = i;
     }
   }
   return success();
 }
+ERROR awkward_ListOffsetArray32_reduce_local_nextparents_64(
+  int64_t* nextparents,
+  const int32_t* offsets,
+  int64_t length) {
+  return awkward_ListOffsetArray_reduce_local_nextparents_64(
+    nextparents,
+    offsets,
+    length);
+}
+ERROR awkward_ListOffsetArrayU32_reduce_local_nextparents_64(
+  int64_t* nextparents,
+  const uint32_t* offsets,
+  int64_t length) {
+  return awkward_ListOffsetArray_reduce_local_nextparents_64(
+    nextparents,
+    offsets,
+    length);
+}
+ERROR awkward_ListOffsetArray64_reduce_local_nextparents_64(
+  int64_t* nextparents,
+  const int64_t* offsets,
+  int64_t length) {
+  return awkward_ListOffsetArray_reduce_local_nextparents_64(
+    nextparents,
+    offsets,
+    length);
+}

kernel-specification.yml

@@ -2217,11 +2217,21 @@ kernels:
 
   - name: awkward_ListOffsetArray_reduce_local_nextparents_64
     specializations:
-      - name: awkward_ListOffsetArray_reduce_local_nextparents_64
+      - name: awkward_ListOffsetArray32_reduce_local_nextparents_64
+        args:
+          - {name: nextparents, type: "List[int64_t]", dir: out}
+          - {name: offsets, type: "Const[List[int32_t]]", dir: in, role: reducer-offsets}
+          - {name: length, type: "int64_t", dir: in, role: reducer-length}
+      - name: awkward_ListOffsetArray64_reduce_local_nextparents_64
         args:
           - {name: nextparents, type: "List[int64_t]", dir: out}
           - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets}
           - {name: length, type: "int64_t", dir: in, role: reducer-length}
+      - name: awkward_ListOffsetArrayU32_reduce_local_nextparents_64
+        args:
+          - {name: nextparents, type: "List[int64_t]", dir: out}
+          - {name: offsets, type: "Const[List[uint32_t]]", dir: in, role: reducer-offsets}
+          - {name: length, type: "int64_t", dir: in, role: reducer-length}
     description: null
     definition: |
       def awkward_ListOffsetArray_reduce_local_nextparents_64(nextparents, offsets, length):

pyproject.toml

@@ -40,7 +40,7 @@ classifiers = [
     "Topic :: Utilities",
 ]
 dependencies = [
-    "awkward_cpp==20",
+    "awkward_cpp==21",
     "importlib_resources;python_version < \"3.9\"",
     "numpy>=1.18.0",
     "packaging",

src/awkward/contents/listoffsetarray.py

@@ -886,9 +886,10 @@ def _is_unique(self, negaxis, starts, parents, outlength):
         if not branch and (negaxis == depth):
             return self._content._is_unique(negaxis - 1, starts, parents, outlength)
         else:
-            nextparents = Index64.empty(
-                self._offsets[-1] - self._offsets[0], self._backend.index_nplike
+            nextlen = self._backend.index_nplike.index_as_shape_item(
+                self._offsets[-1] - self._offsets[0]
             )
+            nextparents = Index64.empty(nextlen, self._backend.index_nplike)
 
             assert (
                 nextparents.nplike is self._backend.index_nplike
@@ -981,9 +982,10 @@ def _unique(self, negaxis, starts, parents, outlength):
             )
 
         else:
-            nextparents = Index64.empty(
-                self._offsets[-1] - self._offsets[0], self._backend.index_nplike
+            nextlen = self._backend.index_nplike.index_as_shape_item(
+                self._offsets[-1] - self._offsets[0]
             )
+            nextparents = Index64.empty(nextlen, self._backend.index_nplike)
 
             assert (
                 nextparents.nplike is self._backend.index_nplike
@@ -1165,12 +1167,10 @@ def _argsort_next(
                 out_offsets, out, parameters=self._parameters
             )
         else:
-            nextparents = Index64.empty(
-                self._backend.index_nplike.index_as_shape_item(
-                    self._offsets[-1] - self._offsets[0]
-                ),
-                self._backend.index_nplike,
+            nextlen = self._backend.index_nplike.index_as_shape_item(
+                self._offsets[-1] - self._offsets[0]
             )
+            nextparents = Index64.empty(nextlen, self._backend.index_nplike)
 
             assert (
                 nextparents.nplike is self._backend.index_nplike
@@ -1299,10 +1299,10 @@ def _sort_next(self, negaxis, starts, parents, outlength, ascending, stable):
                 parameters=self._parameters,
             )
         else:
-            nextparents = Index64.empty(
-                index_nplike.index_as_shape_item(self._offsets[-1] - self._offsets[0]),
-                index_nplike,
+            nextlen = index_nplike.index_as_shape_item(
+                self._offsets[-1] - self._offsets[0]
             )
+            nextparents = Index64.empty(nextlen, index_nplike)
             lenstarts = self._offsets.length - 1
 
             assert (

tests/test_2571_awkward_ListOffsetArray_reduce_local_nextparents_64.py

@@ -0,0 +1,26 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import numpy as np
+import pytest
+
+import awkward as ak
+
+
+@pytest.mark.parametrize(
+    "indextype", [ak.index.Index32, ak.index.IndexU32, ak.index.Index64]
+)
+def test(indextype):
+    array = ak.Array(
+        ak.contents.ListOffsetArray(
+            indextype(np.array([0, 3, 3, 5, 6, 9])),
+            ak.contents.NumpyArray(np.array([6, 9, 9, 4, 4, 2, 5, 2, 7], np.int64)),
+        )
+    )
+    if indextype is ak.index.Index32:
+        assert array.layout.offsets.data.dtype == np.dtype(np.int32)
+    elif indextype is ak.index.IndexU32:
+        assert array.layout.offsets.data.dtype == np.dtype(np.uint32)
+    elif indextype is ak.index.Index64:
+        assert array.layout.offsets.data.dtype == np.dtype(np.int64)
+
+    assert ak.argsort(array).tolist() == [[0, 1, 2], [], [0, 1], [0], [1, 0, 2]]