def test_emptyarray():
assert ak.from_buffers(*ak.to_buffers([])).tolist() == []
assert ak.from_buffers(*ak.to_buffers([[], [], []])).tolist() == [[], [],
[]]
assert pickle.loads(pickle.dumps(ak.Array([]), -1)).tolist() == []
assert pickle.loads(pickle.dumps(ak.Array([[], [], []]), -1)).tolist() == [
[],
[],
[],
]
def test():
array = ak.Array([[{
"a": 1,
"b": [1, 2, 3]
}], [{
"a": 1,
"b": [4, 5]
}, {
"a": 4,
"b": [2]
}]])
array_new = ak.Array(
ak.layout.ListOffsetArray64(
array.layout.offsets,
ak.layout.RecordArray({
"a":
array.layout.content["a"],
"b":
ak.layout.ListArray64(
array.layout.content["b"].offsets[:-1],
array.layout.content["b"].offsets[1:],
array.layout.content["b"].content,
),
}),
))
form, length, container = ak.to_buffers(array_new)
reconstituted = ak.from_buffers(form, length, container, lazy=True)
assert reconstituted.tolist() == array_new.tolist()
def test_longer_than_expected():
array = ak.Array(
ak.layout.ListOffsetArray64(
ak.layout.Index64([0, 2, 4]),
ak.layout.RecordArray({
"item1":
ak.layout.NumpyArray(np.arange(4)),
"longitem":
ak.layout.NumpyArray(np.arange(6)),
}),
))
out = ak.from_buffers(*ak.to_buffers(array), lazy=True)
assert ak.to_list(out) == [
[{
"item1": 0,
"longitem": 0
}, {
"item1": 1,
"longitem": 1
}],
[{
"item1": 2,
"longitem": 2
}, {
"item1": 3,
"longitem": 3
}],
]
def test_listoffsetarray():
assert ak.from_buffers(
*ak.to_buffers([[1, 2, 3], [], [4, 5]])).tolist() == [
[1, 2, 3],
[],
[4, 5],
]
assert ak.from_buffers(
*ak.to_buffers(["one", "two", "three", "four", "five"])).tolist() == [
"one", "two", "three", "four", "five"
]
assert ak.from_buffers(*ak.to_buffers([["one", "two", "three"], [],
["four", "five"]])).tolist() == [[
"one", "two", "three"
], [], ["four", "five"]]
assert pickle.loads(pickle.dumps(ak.Array([[1, 2, 3], [], [4, 5]]),
-1)).tolist() == [[1, 2, 3], [], [4, 5]]
def __getitem__(self, where):
out = self.layout[where]
if isinstance(out, ak.layout.Content):
form, length, indexes = ak.to_buffers(out,
form_key="getitem_node{id}",
virtual="pass")
aux_data = AuxData(form, length, indexes, self.aux_data.datakeys)
return DifferentiableArray(aux_data, self.tracers)
else:
return out
def test_lazy():
array = ak.Array([1, 2, 3, 4, 5])
form, length, container = ak.to_buffers(array)
assert ak.from_buffers(form, length, container, lazy=True).tolist() == [
1,
2,
3,
4,
5,
]
def test_indexedoptionarray():
assert ak.from_buffers(
*ak.to_buffers([1, 2, 3, None, None, 5])).tolist() == [
1,
2,
3,
None,
None,
5,
]
assert pickle.loads(pickle.dumps(ak.Array([1, 2, 3, None, None, 5]),
-1)).tolist() == [1, 2, 3, None, None, 5]
def test_unmaskedarray():
content = ak.Array([1, 2, 3, 4, 5]).layout
unmaskedarray = ak.layout.UnmaskedArray(content)
assert ak.from_buffers(*ak.to_buffers(unmaskedarray)).tolist() == [
1, 2, 3, 4, 5
]
assert pickle.loads(pickle.dumps(ak.Array(unmaskedarray),
-1)).tolist() == [
1,
2,
3,
4,
5,
]
def test_unionarray():
assert ak.from_buffers(*ak.to_buffers([[1, 2, 3], [], 4, 5])).tolist() == [
[1, 2, 3],
[],
4,
5,
]
assert pickle.loads(pickle.dumps(ak.Array([[1, 2, 3], [], 4, 5]),
-1)).tolist() == [
[1, 2, 3],
[],
4,
5,
]
def test_listarray():
listoffsetarray = ak.Array([[1, 2, 3], [], [4, 5]]).layout
listarray = ak.layout.ListArray64(listoffsetarray.starts,
listoffsetarray.stops,
listoffsetarray.content)
assert ak.from_buffers(*ak.to_buffers(listarray)).tolist() == [
[1, 2, 3],
[],
[4, 5],
]
assert pickle.loads(pickle.dumps(ak.Array(listarray), -1)).tolist() == [
[1, 2, 3],
[],
[4, 5],
]
def test_regulararray():
content = ak.Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]).layout
regulararray = ak.layout.RegularArray(content, 3, zeros_length=0)
assert ak.from_buffers(*ak.to_buffers(regulararray)).tolist() == [
[1, 2, 3],
[4, 5, 6],
[7, 8, 9],
[10, 11, 12],
]
assert pickle.loads(pickle.dumps(ak.Array(regulararray), -1)).tolist() == [
[1, 2, 3],
[4, 5, 6],
[7, 8, 9],
[10, 11, 12],
]
def test_numpyarray():
assert ak.from_buffers(*ak.to_buffers([1, 2, 3, 4, 5])).tolist() == [
1,
2,
3,
4,
5,
]
assert pickle.loads(pickle.dumps(ak.Array([1, 2, 3, 4, 5]),
-1)).tolist() == [
1,
2,
3,
4,
5,
]
def test_lazy_partitioned():
array = ak.repartition(ak.Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]), 3)
form, length, container = ak.to_buffers(array)
assert length == [3, 3, 3, 1]
assert ak.from_buffers(form, length, container, lazy=True).tolist() == [
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
]
def test_indexedarray():
content = ak.Array([0.0, 1.1, 2.2, 3.3, 4.4]).layout
index = ak.layout.Index64(np.array([3, 1, 1, 4, 2], dtype=np.int64))
indexedarray = ak.layout.IndexedArray64(index, content)
assert ak.from_buffers(*ak.to_buffers(indexedarray)).tolist() == [
3.3,
1.1,
1.1,
4.4,
2.2,
]
assert pickle.loads(pickle.dumps(ak.Array(indexedarray), -1)).tolist() == [
3.3,
1.1,
1.1,
4.4,
2.2,
]
def test_bytemaskedarray():
content = ak.Array([0.0, 1.1, 2.2, 3.3, 4.4]).layout
mask = ak.layout.Index8(
np.array([False, True, True, False, False], dtype=np.int8))
bytemaskedarray = ak.layout.ByteMaskedArray(mask, content, True)
assert ak.from_buffers(*ak.to_buffers(bytemaskedarray)).tolist() == [
None,
1.1,
2.2,
None,
None,
]
assert pickle.loads(pickle.dumps(ak.Array(bytemaskedarray),
-1)).tolist() == [
None,
1.1,
2.2,
None,
None,
]
def special_flatten(array):
if isinstance(array, DifferentiableArray):
aux_data, children = array.aux_data, array.tracers
else:
form, length, buffers = ak.to_buffers(array)
formjson = json.loads(form.tojson())
indexes = {k: v for k, v in buffers.items() if not k.endswith("-data")}
datakeys = []
for key in buffers:
partition, form_key, role = key.split("-")
if role == "data":
nodejson = find_datanode(formjson, form_key)
assert nodejson is not None
node = ak.forms.Form.fromjson(json.dumps(nodejson))
datakeys.append(key)
nextform = ak.forms.Form.fromjson(json.dumps(formjson))
aux_data = AuxData(nextform, length, indexes, datakeys)
children = [
jax.numpy.asarray(buffers[x], buffers[x].dtype) for x in datakeys
]
return children, aux_data
def test_from_buffers():
array = ak.Array([
[{
"x": 0.0,
"y": []
}, {
"x": 1.1,
"y": [1]
}, {
"x": 2.2,
"y": None
}],
[],
[{
"x": 3.3,
"y": [1, 2, 3]
}, None, {
"x": 4.4,
"y": [1, 2, 3, 4]
}],
])
assert str(
array.type) == '3 * var * ?{"x": float64, "y": option[var * int64]}'
pa_table = ak.to_arrow_table(array)
awkward_array = ak.from_arrow(pa_table)
form, length, container = ak.to_buffers(awkward_array)
reconstituted = ak.from_buffers(form, length, container, lazy=True)
assert reconstituted[2].tolist() == [
{
"x": 3.3,
"y": [1, 2, 3]
},
None,
{
"x": 4.4,
"y": [1, 2, 3, 4]
},
]
def test():
array = ak.from_numpy(np.zeros((3, 0), dtype=np.int32))
buffs = ak.to_buffers(array)
new_array = ak.from_buffers(*buffs)
assert ak.to_list(new_array) == [[], [], []]
def test_recordarray():
assert ak.from_buffers(
*ak.to_buffers([(1.1, [1]), (2.2, [1, 2]), (3.3,
[1, 2, 3])])).tolist() == [
(1.1, [1]),
(2.2, [1, 2]),
(3.3, [1, 2, 3])
]
assert ak.from_buffers(*ak.to_buffers([{
"x": 1.1,
"y": [1]
}, {
"x": 2.2,
"y": [1, 2]
}, {
"x": 3.3,
"y": [1, 2, 3]
}])).tolist() == [
{
"x": 1.1,
"y": [1]
},
{
"x": 2.2,
"y": [1, 2]
},
{
"x": 3.3,
"y": [1, 2, 3]
},
]
assert pickle.loads(
pickle.dumps(ak.Array([(1.1, [1]), (2.2, [1, 2]), (3.3, [1, 2, 3])]),
-1)).tolist() == [(1.1, [1]), (2.2, [1, 2]),
(3.3, [1, 2, 3])]
assert pickle.loads(
pickle.dumps(
ak.Array([
{
"x": 1.1,
"y": [1]
},
{
"x": 2.2,
"y": [1, 2]
},
{
"x": 3.3,
"y": [1, 2, 3]
},
]),
-1,
)).tolist() == [
{
"x": 1.1,
"y": [1]
},
{
"x": 2.2,
"y": [1, 2]
},
{
"x": 3.3,
"y": [1, 2, 3]
},
]
def test_lazy_buffers():
array = ak.from_json("""
[
{
"listcollection": [
{"item1": 1, "item2": 2},
{"item1": 2, "item2": 4},
{"item1": 3, "item2": 6}
],
"collection": {"item1": 3, "item2": 4},
"singleton": 5,
"listsingleton": [1, 2, 3],
"unioncollection": {"item1": 3},
"masked": null
},
{
"listcollection": [
{"item1": 1, "item2": 2},
{"item1": 2, "item2": 4},
{"item1": 3, "item2": 6}
],
"collection": {"item1": 3, "item2": 4},
"singleton": 5,
"listsingleton": [1, 2, 3],
"unioncollection": [{"item1": 2}],
"masked": 4
},
{
"listcollection": [
{"item1": 1, "item2": 2},
{"item1": 2, "item2": 4},
{"item1": 3, "item2": 6}
],
"collection": {"item1": 3, "item2": 4},
"singleton": 5,
"listsingleton": [1, 2, 3],
"unioncollection": {"item1": 4},
"masked": 4
}
]""")
canary = Canary()
key_format = "kitty-{form_key}-{attribute}"
form, length, container = ak.to_buffers(array,
container=canary,
key_format=key_format)
assert not any(op[0] == "get" for op in canary.ops)
canary.ops = []
cache = {}
out = ak.from_buffers(
form,
length,
container,
key_format=key_format,
lazy=True,
lazy_cache=cache,
lazy_cache_key="hello",
)
assert len(canary.ops) == 0
assert len(cache) == 0
assert len(out) == 3
assert len(canary.ops) == 0
assert len(cache) == 0
assert ak.to_list(ak.num(out.listcollection)) == [3, 3, 3]
assert set(canary.ops) == {("get", "kitty-node1-offsets")}
assert "hello" in cache
assert "hello(kitty-node1-virtual)" in cache
canary.ops = []
cache.clear()
assert ak.to_list(out.unioncollection) == [
{
"item1": 3
},
[{
"item1": 2
}],
{
"item1": 4
},
]
assert set(canary.ops) == {
("get", "kitty-node11-tags"),
("get", "kitty-node11-index"),
("get", "kitty-node14-offsets"),
("get", "kitty-node13-data"),
("get", "kitty-node16-data"),
}
assert "hello" in cache
assert "hello(kitty-node11-virtual)" in cache
assert "hello(kitty-node13-virtual)" in cache
assert "hello(kitty-node16-virtual)" in cache
canary.ops = []
cache.clear()
assert ak.to_list(out.masked) == [None, 4, 4]
assert set(canary.ops) == {
("get", "kitty-node17-index"),
("get", "kitty-node18-data"),
}
assert "hello" in cache
assert "hello(kitty-node17-virtual)" in cache
canary.ops = []
cache.clear()
def test_partitioned():
array = ak.repartition(ak.Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]), 3)
form, length, container = ak.to_buffers(array)
assert ak.from_buffers(form, length, container).tolist() == [
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
]
form, length, container = ak.to_buffers(array)
assert ak.from_buffers(form, length, container).tolist() == [
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
]
one = ak.Array([1, 2, 3, 4, 5])
two = ak.Array([6, 7, 8, 9, 10])
container = {}
lengths = []
form1, length, _ = ak.to_buffers(one, container, 0)
lengths.append(length)
form2, length, _ = ak.to_buffers(two, container, 1)
lengths.append(length)
assert form1 == form2
assert ak.from_buffers(form1, lengths, container).tolist() == [
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
]
assert pickle.loads(pickle.dumps(array, -1)).tolist() == [
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
]