From 44b0c3e22312291315e6306d2914e51ea48948ed Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 5 Jul 2023 13:40:12 -0500 Subject: [PATCH] fix: added kernel specializations for `awkward_ListOffsetArray_reduce_local_nextparents_64` (#2572) --- awkward-cpp/pyproject.toml | 2 +- ...ffsetArray_reduce_local_nextparents_64.cpp | 34 +++++++++++++++++-- kernel-specification.yml | 12 ++++++- pyproject.toml | 2 +- src/awkward/contents/listoffsetarray.py | 24 ++++++------- ...OffsetArray_reduce_local_nextparents_64.py | 26 ++++++++++++++ 6 files changed, 82 insertions(+), 18 deletions(-) create mode 100644 tests/test_2571_awkward_ListOffsetArray_reduce_local_nextparents_64.py diff --git a/awkward-cpp/pyproject.toml b/awkward-cpp/pyproject.toml index f818caf4df..a24bc0ceb0 100644 --- a/awkward-cpp/pyproject.toml +++ b/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" ] diff --git a/awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_reduce_local_nextparents_64.cpp b/awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_reduce_local_nextparents_64.cpp index 9d54246493..fe5f90e8f0 100644 --- a/awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_reduce_local_nextparents_64.cpp +++ b/awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_reduce_local_nextparents_64.cpp @@ -4,13 +4,14 @@ #include "awkward/kernels.h" +template 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; @@ -18,3 +19,30 @@ ERROR awkward_ListOffsetArray_reduce_local_nextparents_64( } 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); +} diff --git a/kernel-specification.yml b/kernel-specification.yml index d6fe83775d..437a22da63 100644 --- a/kernel-specification.yml +++ b/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): diff --git a/pyproject.toml b/pyproject.toml index 084ced4e90..fb58b46978 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/awkward/contents/listoffsetarray.py b/src/awkward/contents/listoffsetarray.py index 3a49bccf46..60349acf64 100644 --- a/src/awkward/contents/listoffsetarray.py +++ b/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 ( diff --git a/tests/test_2571_awkward_ListOffsetArray_reduce_local_nextparents_64.py b/tests/test_2571_awkward_ListOffsetArray_reduce_local_nextparents_64.py new file mode 100644 index 0000000000..86bd284ec4 --- /dev/null +++ b/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]]