fix limit on number of regular arrays

scikit-hep · Nov 21, 2024 · f03a113 · f03a113
1 parent 4db7973
commit f03a113
Show file tree

Hide file tree

Showing 4 changed files with 84 additions and 0 deletions.
diff --git a/awkward-cpp/src/cpu-kernels/awkward_UnionArray_regular_index.cpp b/awkward-cpp/src/cpu-kernels/awkward_UnionArray_regular_index.cpp
@@ -60,3 +60,42 @@ ERROR awkward_UnionArray8_64_regular_index(
     fromtags,
     length);
 }
+ERROR awkward_UnionArray64_32_regular_index(
+  int32_t* toindex,
+  int32_t* current,
+  int64_t size,
+  const int64_t* fromtags,
+  int64_t length) {
+  return awkward_UnionArray_regular_index<int64_t, int32_t>(
+    toindex,
+    current,
+    size,
+    fromtags,
+    length);
+}
+ERROR awkward_UnionArray64_U32_regular_index(
+  uint32_t* toindex,
+  uint32_t* current,
+  int64_t size,
+  const int64_t* fromtags,
+  int64_t length) {
+  return awkward_UnionArray_regular_index<int64_t, uint32_t>(
+    toindex,
+    current,
+    size,
+    fromtags,
+    length);
+}
+ERROR awkward_UnionArray64_64_regular_index(
+  int64_t* toindex,
+  int64_t* current,
+  int64_t size,
+  const int64_t* fromtags,
+  int64_t length) {
+  return awkward_UnionArray_regular_index<int64_t, int64_t>(
+    toindex,
+    current,
+    size,
+    fromtags,
+    length);
+}
diff --git a/awkward-cpp/src/cpu-kernels/awkward_UnionArray_regular_index_getsize.cpp b/awkward-cpp/src/cpu-kernels/awkward_UnionArray_regular_index_getsize.cpp
@@ -28,3 +28,12 @@ ERROR awkward_UnionArray8_regular_index_getsize(
     fromtags,
     length);
 }
+ERROR awkward_UnionArray64_regular_index_getsize(
+  int64_t* size,
+  const int64_t* fromtags,
+  int64_t length) {
+  return awkward_UnionArray_regular_index_getsize<int64_t>(
+    size,
+    fromtags,
+    length);
+}
diff --git a/kernel-specification.yml b/kernel-specification.yml
@@ -3479,6 +3479,27 @@ kernels:
           - {name: size, type: "int64_t", dir: in, role: default}
           - {name: fromtags, type: "Const[List[int8_t]]", dir: in, role: UnionArray-tags}
           - {name: length, type: "int64_t", dir: in, role: default}
+      - name: awkward_UnionArray64_32_regular_index
+        args:
+          - {name: toindex, type: "List[int32_t]", dir: out}
+          - {name: current, type: "List[int32_t]", dir: out}
+          - {name: size, type: "int64_t", dir: in, role: default}
+          - {name: fromtags, type: "Const[List[int64_t]]", dir: in, role: UnionArray-tags}
+          - {name: length, type: "int64_t", dir: in, role: default}
+      - name: awkward_UnionArray64_64_regular_index
+        args:
+          - {name: toindex, type: "List[int64_t]", dir: out}
+          - {name: current, type: "List[int64_t]", dir: out}
+          - {name: size, type: "int64_t", dir: in, role: default}
+          - {name: fromtags, type: "Const[List[int64_t]]", dir: in, role: UnionArray-tags}
+          - {name: length, type: "int64_t", dir: in, role: default}
+      - name: awkward_UnionArray64_U32_regular_index
+        args:
+          - {name: toindex, type: "List[uint32_t]", dir: out}
+          - {name: current, type: "List[uint32_t]", dir: out}
+          - {name: size, type: "int64_t", dir: in, role: default}
+          - {name: fromtags, type: "Const[List[int64_t]]", dir: in, role: UnionArray-tags}
+          - {name: length, type: "int64_t", dir: in, role: default}
     description: null
     definition: |
       def awkward_UnionArray_regular_index(toindex, current, size, fromtags, length):
@@ -3498,6 +3519,11 @@ kernels:
           - {name: size, type: "List[int64_t]", dir: out}
           - {name: fromtags, type: "Const[List[int8_t]]", dir: in, role: UnionArray-tags}
           - {name: length, type: "int64_t", dir: in, role: default}
+      - name: awkward_UnionArray64_regular_index_getsize
+        args:
+          - {name: size, type: "List[int64_t]", dir: out}
+          - {name: fromtags, type: "Const[List[int64_t]]", dir: in, role: UnionArray-tags}
+          - {name: length, type: "int64_t", dir: in, role: default}
     description: null
     definition: |
       def awkward_UnionArray_regular_index_getsize(size, fromtags, length):

diff --git a/tests/test_3312_concatenate_regular_arrays.py b/tests/test_3312_concatenate_regular_arrays.py
@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+import awkward as ak
+import numpy as np
+
+
+def test():
+    ak.concatenate([ak.Array([i])[:, np.newaxis] for i in range(127)], axis=1)
+
+    ak.concatenate([ak.Array([i])[:, np.newaxis] for i in range(128)], axis=1)