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 1 commit
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 @@ -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);
}
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
26 changes: 26 additions & 0 deletions kernel-specification.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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):
Expand Down
10 changes: 10 additions & 0 deletions tests/test_3312_concatenate_regular_arrays.py
Original file line number Diff line number Diff line change
@@ -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)
Loading