Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: extend limit on number of regular arrays to concatenate #3312

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 39 additions & 0 deletions awkward-cpp/src/cpu-kernels/awkward_UnionArray_regular_index.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,45 @@ ERROR awkward_UnionArray_regular_index(
}
return success();
}
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);
}
ERROR awkward_UnionArray8_32_regular_index(
int32_t* toindex,
int32_t* current,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,15 @@ ERROR awkward_UnionArray_regular_index_getsize(
*size = *size + 1;
return success();
}
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);
}
ERROR awkward_UnionArray8_regular_index_getsize(
int64_t* size,
const int8_t* fromtags,
Expand Down
26 changes: 26 additions & 0 deletions kernel-specification.yml
Original file line number Diff line number Diff line change
Expand Up @@ -3458,6 +3458,27 @@ kernels:

- name: awkward_UnionArray_regular_index
specializations:
- 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}
- name: awkward_UnionArray8_32_regular_index
args:
- {name: toindex, type: "List[int32_t]", dir: out}
Expand Down Expand Up @@ -3493,6 +3514,11 @@ kernels:

- name: awkward_UnionArray_regular_index_getsize
specializations:
- 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}
- name: awkward_UnionArray8_regular_index_getsize
args:
- {name: size, type: "List[int64_t]", dir: out}
Expand Down
11 changes: 11 additions & 0 deletions tests/test_3312_concatenate_regular_arrays.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
from __future__ import annotations

import numpy as np

import awkward as ak


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)
Loading