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 special_unflatten(aux_data, children):
if any(isinstance(x, jax.core.Tracer) for x in children):
return DifferentiableArray(aux_data, children)
else:
buffers = dict(aux_data.indexes)
buffers.update(zip(aux_data.datakeys, children))
return ak.from_buffers(aux_data.form, aux_data.length, buffers)
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 layout(self):
buffers = dict(self.aux_data.indexes)
for key, tracer in zip(self.aux_data.datakeys, self.tracers):
if hasattr(tracer, "primal"):
buffers[key] = tracer.primal
return ak.from_buffers(self.aux_data.form,
self.aux_data.length,
buffers,
highlevel=False)
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 events(self):
"""Build events"""
events = self._events()
if events is None:
behavior = dict(self._schema.behavior)
behavior["__events_factory__"] = self
events = awkward.from_buffers(
self._schema.form,
len(self),
self._mapping,
key_format=partial(_key_formatter, self._partition_key),
lazy=True,
lazy_cache="new" if self._cache is None else self._cache,
behavior=behavior,
)
self._events = weakref.ref(events)
return events
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 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 events(self):
"""Build events"""
events = self._events()
if events is None:
prefix = self._partition_key
def key_formatter(partition, form_key, attribute):
return prefix + f"/{partition}/{form_key}/{attribute}"
behavior = dict(self._schema.behavior)
behavior["__events_factory__"] = self
events = awkward.from_buffers(
self._schema.form,
len(self),
self._mapping,
key_format=key_formatter,
lazy=True,
lazy_cache="new" if self._cache is None else self._cache,
behavior=behavior,
)
self._events = weakref.ref(events)
return events
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 snapshot(self):
return ak.from_buffers(self.form, self.vm.stack[0], self.vm.outputs)
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():
form = {
"class": "RecordArray",
"contents": {
"a": {
"class": "ListOffsetArray64",
"offsets": "i64",
"content": {
"class": "NumpyArray",
"itemsize": 8,
"format": "l",
"primitive": "int64",
"form_key": "node2",
},
"form_key": "node0",
},
"b": {
"class": "ListOffsetArray64",
"offsets": "i64",
"content": {
"class": "NumpyArray",
"itemsize": 8,
"format": "l",
"primitive": "int64",
"form_key": "node3",
},
"form_key": "node0",
},
},
"form_key": "node1",
}
container = {
"part0-node0-offsets": np.array([0, 2, 3, 3, 6], dtype=np.int64),
"part0-node2-data": np.array([1, 2, 3, 4, 5, 6], dtype=np.int64),
"part0-node3-data": np.array([10, 20, 30, 40, 50, 60], dtype=np.int64),
}
assert ak.from_buffers(form, 4, container).tolist() == [
{
"a": [1, 2],
"b": [10, 20]
},
{
"a": [3],
"b": [30]
},
{
"a": [],
"b": []
},
{
"a": [4, 5, 6],
"b": [40, 50, 60]
},
]
assert ak.from_buffers(form, 4, container, lazy=True).tolist() == [
{
"a": [1, 2],
"b": [10, 20]
},
{
"a": [3],
"b": [30]
},
{
"a": [],
"b": []
},
{
"a": [4, 5, 6],
"b": [40, 50, 60]
},
]
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_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,
]
def test_lazy():
form = {
"class": "RecordArray",
"contents": {
"Electron": {
"class": "ListOffsetArray64",
"content": {
"class": "RecordArray",
"contents": {
"charge": {
"class": "NumpyArray",
"form_key": "load_electron_charge",
"format": "i",
"has_identities": False,
"inner_shape": [],
"itemsize": 4,
"parameters": {},
"primitive": "int64",
},
"momentum": {
"class": "NumpyArray",
"form_key": "load_electron_momentum",
"format": "i",
"has_identities": False,
"inner_shape": [],
"itemsize": 4,
"parameters": {},
"primitive": "float64",
},
},
"form_key": "invalid",
"parameters": {
"__record__": "Electron"
},
},
"form_key": "load_electron_offsets",
"offsets": "i64",
},
"Muon": {
"class": "ListOffsetArray64",
"content": {
"class": "RecordArray",
"contents": {
"charge": {
"class": "NumpyArray",
"form_key": "load_muon_charge",
"format": "i",
"has_identities": False,
"inner_shape": [],
"itemsize": 4,
"parameters": {},
"primitive": "int64",
},
"momentum": {
"class": "NumpyArray",
"form_key": "load_muon_momentum",
"format": "i",
"has_identities": False,
"inner_shape": [],
"itemsize": 4,
"parameters": {},
"primitive": "float64",
},
},
"form_key": "invalid",
"parameters": {
"__record__": "Muon"
},
},
"form_key": "load_muon_offsets",
"offsets": "i64",
},
},
"form_key": "",
"parameters": {},
}
container = Verbose({
"part0-load_electron_offsets-offsets":
np.array([0, 3, 3, 5], dtype=np.int64),
"part0-load_muon_offsets-offsets":
np.array([0, 3, 3, 5], dtype=np.int64),
"part0-load_electron_charge-data":
np.array([1, 2, 3, 4, 5], dtype=np.int64),
"part0-load_electron_momentum-data":
np.array([1.1, 2.2, 3.3, 4.4, 5.5], dtype=np.float64),
"part0-load_muon_charge-data":
np.array([1, 2, 3, 4, 5], dtype=np.int64),
"part0-load_muon_momentum-data":
np.array([1.1, 2.2, 3.3, 4.4, 5.5], dtype=np.float64),
"part1-load_electron_offsets-offsets":
np.array([0, 3, 3, 5], dtype=np.int64),
"part1-load_muon_offsets-offsets":
np.array([0, 3, 3, 5], dtype=np.int64),
"part1-load_electron_charge-data":
np.array([1, 2, 3, 4, 5], dtype=np.int64),
"part1-load_electron_momentum-data":
np.array([1.1, 2.2, 3.3, 4.4, 5.5], dtype=np.float64),
"part1-load_muon_charge-data":
np.array([1, 2, 3, 4, 5], dtype=np.int64),
"part1-load_muon_momentum-data":
np.array([1.1, 2.2, 3.3, 4.4, 5.5], dtype=np.float64),
})
container.touched = False
lazy = ak.from_buffers(form, [3, 3], container, lazy=True)
one = ak.concatenate([lazy.Electron, lazy.Muon], axis=0)
lazy = ak.from_buffers(form, [3, 3], container, lazy=True)
two = ak.concatenate([lazy.Muon, lazy.Muon], axis=0)
lazy = ak.from_buffers(form, [3, 3], container, lazy=True)
three = ak.concatenate([lazy.Electron, lazy.Electron], axis=0)
assert not container.touched
assert one.tolist() == 4 * [
[
{
"charge": 1,
"momentum": 1.1
},
{
"charge": 2,
"momentum": 2.2
},
{
"charge": 3,
"momentum": 3.3
},
],
[],
[{
"charge": 4,
"momentum": 4.4
}, {
"charge": 5,
"momentum": 5.5
}],
]
assert two.tolist() == 4 * [
[
{
"charge": 1,
"momentum": 1.1
},
{
"charge": 2,
"momentum": 2.2
},
{
"charge": 3,
"momentum": 3.3
},
],
[],
[{
"charge": 4,
"momentum": 4.4
}, {
"charge": 5,
"momentum": 5.5
}],
]
assert three.tolist() == 4 * [
[
{
"charge": 1,
"momentum": 1.1
},
{
"charge": 2,
"momentum": 2.2
},
{
"charge": 3,
"momentum": 3.3
},
],
[],
[{
"charge": 4,
"momentum": 4.4
}, {
"charge": 5,
"momentum": 5.5
}],
]
assert container.touched
container.data_touched = False
lazy = ak.from_buffers(form, [3, 3], container, lazy=True)
one = ak.concatenate([lazy.Electron, lazy.Muon], axis=1)
lazy = ak.from_buffers(form, [3, 3], container, lazy=True)
two = ak.concatenate([lazy.Muon, lazy.Muon], axis=1)
lazy = ak.from_buffers(form, [3, 3], container, lazy=True)
three = ak.concatenate([lazy.Electron, lazy.Electron], axis=1)
assert not container.data_touched
assert one.tolist() == 2 * [
[
{
"charge": 1,
"momentum": 1.1
},
{
"charge": 2,
"momentum": 2.2
},
{
"charge": 3,
"momentum": 3.3
},
{
"charge": 1,
"momentum": 1.1
},
{
"charge": 2,
"momentum": 2.2
},
{
"charge": 3,
"momentum": 3.3
},
],
[],
[
{
"charge": 4,
"momentum": 4.4
},
{
"charge": 5,
"momentum": 5.5
},
{
"charge": 4,
"momentum": 4.4
},
{
"charge": 5,
"momentum": 5.5
},
],
]
assert two.tolist() == 2 * [
[
{
"charge": 1,
"momentum": 1.1
},
{
"charge": 2,
"momentum": 2.2
},
{
"charge": 3,
"momentum": 3.3
},
{
"charge": 1,
"momentum": 1.1
},
{
"charge": 2,
"momentum": 2.2
},
{
"charge": 3,
"momentum": 3.3
},
],
[],
[
{
"charge": 4,
"momentum": 4.4
},
{
"charge": 5,
"momentum": 5.5
},
{
"charge": 4,
"momentum": 4.4
},
{
"charge": 5,
"momentum": 5.5
},
],
]
assert three.tolist() == 2 * [
[
{
"charge": 1,
"momentum": 1.1
},
{
"charge": 2,
"momentum": 2.2
},
{
"charge": 3,
"momentum": 3.3
},
{
"charge": 1,
"momentum": 1.1
},
{
"charge": 2,
"momentum": 2.2
},
{
"charge": 3,
"momentum": 3.3
},
],
[],
[
{
"charge": 4,
"momentum": 4.4
},
{
"charge": 5,
"momentum": 5.5
},
{
"charge": 4,
"momentum": 4.4
},
{
"charge": 5,
"momentum": 5.5
},
],
]
form["contents"]["Electron"]["content"]["parameters"] = {}
form["contents"]["Muon"]["content"]["parameters"] = {}
container.data_touched = False
lazy = ak.from_buffers(form, [3, 3], container, lazy=True)
one = ak.concatenate([lazy.Electron, lazy.Muon], axis=1, merge=False)
assert not container.data_touched
assert one.tolist() == 2 * [
[
{
"charge": 1,
"momentum": 1.1
},
{
"charge": 2,
"momentum": 2.2
},
{
"charge": 3,
"momentum": 3.3
},
{
"charge": 1,
"momentum": 1.1
},
{
"charge": 2,
"momentum": 2.2
},
{
"charge": 3,
"momentum": 3.3
},
],
[],
[
{
"charge": 4,
"momentum": 4.4
},
{
"charge": 5,
"momentum": 5.5
},
{
"charge": 4,
"momentum": 4.4
},
{
"charge": 5,
"momentum": 5.5
},
],
]