def test_RecordArray_NumpyArray():
v1a = ak.layout.RecordArray(
[
ak.layout.NumpyArray(np.array([0, 1, 2, 3, 4], dtype=np.int64)),
ak.layout.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5])),
],
["x", "y"],
)
v2a = ak._v2.contents.recordarray.RecordArray(
[
ak._v2.contents.numpyarray.NumpyArray(
np.array([0, 1, 2, 3, 4], dtype=np.int64)
),
ak._v2.contents.numpyarray.NumpyArray(
np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5])
),
],
["x", "y"],
)
assert v1v2_equal(v1a, v2a)
assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a))
assert ak.to_list(v1a) == ak.to_list(v2a)
assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
v1b = ak.layout.RecordArray(
[
ak.layout.NumpyArray(np.array([0, 1, 2, 3, 4], dtype=np.int64)),
ak.layout.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5])),
],
None,
)
v2b = ak._v2.contents.recordarray.RecordArray(
[
ak._v2.contents.numpyarray.NumpyArray(
np.array([0, 1, 2, 3, 4], dtype=np.int64)
),
ak._v2.contents.numpyarray.NumpyArray(
np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5])
),
],
None,
)
assert v1v2_equal(v1b, v2b)
assert v1v2_equal(v2_to_v1(v2b), v1_to_v2(v1b))
assert ak.to_list(v1b) == ak.to_list(v2b)
assert newform(json.loads(v1b.form.tojson())) == v2b.form.tolist()
v1c = ak.layout.RecordArray([], [], 10)
v2c = ak._v2.contents.recordarray.RecordArray([], [], 10)
assert v1v2_equal(v1c, v2c)
assert v1v2_equal(v2_to_v1(v2c), v1_to_v2(v1c))
assert ak.to_list(v1c) == ak.to_list(v2c)
assert newform(json.loads(v1c.form.tojson())) == v2c.form.tolist()
v1d = ak.layout.RecordArray([], None, 10)
v2d = ak._v2.contents.recordarray.RecordArray([], None, 10)
assert v1v2_equal(v1d, v2d)
assert v1v2_equal(v2_to_v1(v2d), v1_to_v2(v1d))
assert ak.to_list(v1d) == ak.to_list(v2d)
assert newform(json.loads(v1c.form.tojson())) == v2c.form.tolist()
def test_ByteMaskedArray_NumpyArray():
v1a = ak.layout.ByteMaskedArray(
ak.layout.Index8(np.array([1, 0, 1, 0, 1], dtype=np.int8)),
ak.layout.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])),
valid_when=True,
)
v2a = ak._v2.contents.bytemaskedarray.ByteMaskedArray(
ak._v2.index.Index(np.array([1, 0, 1, 0, 1], dtype=np.int8)),
ak._v2.contents.numpyarray.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])),
valid_when=True,
)
assert v1v2_equal(v1a, v2a)
assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a))
assert ak.to_list(v1a) == ak.to_list(v2a)
assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
v1b = ak.layout.ByteMaskedArray(
ak.layout.Index8(np.array([0, 1, 0, 1, 0], dtype=np.int8)),
ak.layout.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])),
valid_when=False,
)
v2b = ak._v2.contents.bytemaskedarray.ByteMaskedArray(
ak._v2.index.Index(np.array([0, 1, 0, 1, 0], dtype=np.int8)),
ak._v2.contents.numpyarray.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])),
valid_when=False,
)
assert v1v2_equal(v1b, v2b)
assert v1v2_equal(v2_to_v1(v2b), v1_to_v2(v1b))
assert ak.to_list(v1b) == ak.to_list(v2b)
assert newform(json.loads(v1b.form.tojson())) == v2b.form.tolist()
def test_UnmaskedArray():
old = ak.layout.UnmaskedArray(
ak.layout.NumpyArray(
np.array([[0.0, 1.1, 2.2, 3.3], [0.0, 1.1, 2.2, 3.3]])))
new = v1_to_v2(old)
assert ak.to_list(old[0, 1:]) == [1.1, 2.2, 3.3]
assert ak.to_list(new[0, 1:]) == [1.1, 2.2, 3.3]
assert v1v2_equal(old[0, 1:], new[0, 1:])
with pytest.raises(IndexError):
new[1, "hello"]
with pytest.raises(IndexError):
new[1, ["hello", "there"]]
assert ak.to_list(new[1, np.newaxis, -2]) == [2.2]
assert ak.to_list(new[1, np.newaxis, np.newaxis, -2]) == [[2.2]]
assert old.minmax_depth == (2, 2)
assert new.minmax_depth == (2, 2)
assert ak.to_list(old[1, ..., -2]) == 2.2
assert ak.to_list(new[1, ..., -2]) == 2.2
expectation = [[0.0, 1.1, 2.2, 3.3], [0.0, 1.1, 2.2, 3.3]]
assert (ak.to_list(old[[1, 0], ]) == expectation)
assert (ak.to_list(new[[1, 0], ]) == expectation)
assert ak.to_list(new[[1, 0]]) == expectation
assert v1v2_equal(old[1, ...], new[1, ...])
assert ak.to_list(old[1, [1, 0]]) == [1.1, 0.0]
assert ak.to_list(new[1, [1, 0]]) == [1.1, 0.0]
def test_ListArray_RecordArray_NumpyArray():
v1a = ak.layout.ListArray64(
ak.layout.Index64(np.array([4, 100, 1], dtype=np.int64)),
ak.layout.Index64(np.array([7, 100, 3, 200], dtype=np.int64)),
ak.layout.RecordArray(
[ak.layout.NumpyArray(np.array([6.6, 4.4, 5.5, 7.7, 1.1, 2.2, 3.3, 8.8]))],
["nest"],
),
)
v2a = ak._v2.contents.listarray.ListArray(
ak._v2.index.Index(np.array([4, 100, 1], dtype=np.int64)),
ak._v2.index.Index(np.array([7, 100, 3, 200], dtype=np.int64)),
ak._v2.contents.recordarray.RecordArray(
[
ak._v2.contents.numpyarray.NumpyArray(
np.array([6.6, 4.4, 5.5, 7.7, 1.1, 2.2, 3.3, 8.8])
)
],
["nest"],
),
)
assert v1v2_equal(v1a, v2a)
assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a))
assert ak.to_list(v1a) == ak.to_list(v2a)
assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
def test_EmptyArray():
v1a = ak.layout.EmptyArray()
v2a = ak._v2.contents.emptyarray.EmptyArray()
assert v1v2_equal(v1a, v2a)
assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a))
assert ak.to_list(v1a) == ak.to_list(v2a)
assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
def test_NumpyArray():
a = ak._v2.contents.RegularArray(
v1_to_v2(ak.from_numpy(np.arange(2 * 3 * 5).reshape(-1, 5)).layout), 3)
assert ak.to_list(a[1]) == [
[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24],
[25, 26, 27, 28, 29],
]
assert ak.to_list(a[1, -2]) == [20, 21, 22, 23, 24]
assert a[1, -2, 2] == 22
with pytest.raises(IndexError):
a[1, -2, 2, 0]
assert ak.to_list(a[1, -2, 2:]) == [22, 23, 24]
with pytest.raises(IndexError):
a[1, -2, 2:, 0]
with pytest.raises(IndexError):
a[1, -2, "hello"]
with pytest.raises(IndexError):
a[1, -2, ["hello", "there"]]
assert ak.to_list(a[1, -2, np.newaxis, 2]) == [22]
assert ak.to_list(a[1, -2, np.newaxis, np.newaxis, 2]) == [[22]]
assert ak.to_list(a[1, -2, ...]) == [20, 21, 22, 23, 24]
assert a[1, -2, ..., 2] == 22
with pytest.raises(IndexError):
a[1, -2, ..., 2, 2]
b = ak.layout.RegularArray(
ak.from_numpy(np.arange(2 * 3 * 5).reshape(-1, 5)).layout, 3)
assert ak.to_list(b[1, -2, [3, 1, 1, 2]]) == [23, 21, 21, 22]
assert ak.to_list(a[1, -2, [3, 1, 1, 2]]) == [23, 21, 21, 22]
with pytest.raises(IndexError):
a[1, -2, [3, 1, 1, 2], 2]
def test_NumpyArray():
v1a = ak.layout.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3]))
v2a = ak._v2.contents.numpyarray.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3]))
assert v1v2_equal(v1a, v2a)
assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a))
assert ak.to_list(v1a) == ak.to_list(v2a)
assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
v1b = ak.layout.NumpyArray(np.arange(2 * 3 * 5, dtype=np.int64).reshape(2, 3, 5))
v2b = ak._v2.contents.numpyarray.NumpyArray(
np.arange(2 * 3 * 5, dtype=np.int64).reshape(2, 3, 5)
)
assert v1v2_equal(v1b, v2b)
assert v1v2_equal(v2_to_v1(v2b), v1_to_v2(v1b))
assert ak.to_list(v1b) == ak.to_list(v2b)
assert newform(json.loads(v1b.form.tojson())) == v2b.form.tolist()
def prepare_tuple_item(item):
if ak._util.isint(item):
return int(item)
elif isinstance(item, slice):
return item
elif ak._util.isstr(item):
return item
elif item is np.newaxis:
return item
elif item is Ellipsis:
return item
elif isinstance(item, ak.highlevel.Array):
return prepare_tuple_item(item.layout)
elif isinstance(item, ak.layout.Content):
return prepare_tuple_item(v1_to_v2(item))
elif isinstance(item, ak._v2.highlevel.Array):
return prepare_tuple_item(item.layout)
elif isinstance(item, ak._v2.contents.EmptyArray):
return prepare_tuple_item(item.toNumpyArray(np.int64))
elif isinstance(item, ak._v2.contents.NumpyArray):
return item.data
elif isinstance(item, ak._v2.contents.Content):
out = prepare_tuple_bool_to_int(prepare_tuple_nested(item))
if isinstance(out, ak._v2.contents.NumpyArray):
return out.data
else:
return out
elif isinstance(item, Iterable) and all(ak._util.isstr(x) for x in item):
return list(item)
elif isinstance(item, Iterable):
layout = ak._v2.operations.convert.to_layout(item)
as_array = layout.maybe_to_array(layout.nplike)
if as_array is None:
return prepare_tuple_item(layout)
else:
return as_array
else:
raise ak._v2._util.error(
TypeError(
"only integers, slices (`:`), ellipsis (`...`), np.newaxis (`None`), "
"integer/boolean arrays (possibly with variable-length nested "
"lists or missing values), field name (str) or names (non-tuple "
"iterable of str) are valid indices for slicing, not\n\n "
+ repr(item).replace("\n", "\n ")
)
)
def test_UnmaskedArray_NumpyArray():
v1a = ak.layout.UnmaskedArray(ak.layout.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3])))
v2a = ak._v2.contents.unmaskedarray.UnmaskedArray(
ak._v2.contents.numpyarray.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3]))
)
assert v1v2_equal(v1a, v2a)
assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a))
assert ak.to_list(v1a) == ak.to_list(v2a)
assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
def test_RegularArray_RecordArray_NumpyArray():
v1a = ak.layout.RegularArray(
ak.layout.RecordArray(
[ak.layout.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6]))],
["nest"],
),
3,
)
v2a = ak._v2.contents.regulararray.RegularArray(
ak._v2.contents.recordarray.RecordArray(
[
ak._v2.contents.numpyarray.NumpyArray(
np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6])
)
],
["nest"],
),
3,
)
assert v1v2_equal(v1a, v2a)
assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a))
assert ak.to_list(v1a) == ak.to_list(v2a)
assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
v1b = ak.layout.RegularArray(
ak.layout.RecordArray([ak.layout.EmptyArray()], ["nest"]),
0,
zeros_length=10,
)
v2b = ak._v2.contents.regulararray.RegularArray(
ak._v2.contents.recordarray.RecordArray(
[ak._v2.contents.emptyarray.EmptyArray()], ["nest"]
),
0,
zeros_length=10,
)
assert v1v2_equal(v1b, v2b)
assert v1v2_equal(v2_to_v1(v2b), v1_to_v2(v1b))
assert ak.to_list(v1b) == ak.to_list(v2b)
assert newform(json.loads(v1b.form.tojson())) == v2b.form.tolist()
def test_IndexedArray():
old = ak.layout.IndexedArray64(
ak.layout.Index64(np.array([1, 0], np.int64)),
ak.layout.RegularArray(
ak.from_numpy(np.arange(2 * 3 * 5).reshape(-1, 5)).layout, 3),
)
new = v1_to_v2(old)
assert v1v2_equal(old[1, 1:], new[1, 1:])
assert ak.to_list(old[1, 1:]) == [[5, 6, 7, 8, 9], [10, 11, 12, 13, 14]]
assert ak.to_list(new[1, 1:]) == [[5, 6, 7, 8, 9], [10, 11, 12, 13, 14]]
with pytest.raises(IndexError):
new[1, "hello"]
with pytest.raises(IndexError):
new[1, ["hello", "there"]]
assert ak.to_list(new[0, np.newaxis]) == [[[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24],
[25, 26, 27, 28, 29]]]
assert ak.to_list(old[0, np.newaxis]) == [[[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24],
[25, 26, 27, 28, 29]]]
assert old.minmax_depth == (3, 3)
assert new.minmax_depth == (3, 3)
assert v1v2_equal(old[1, ...], new[1, ...])
assert ak.to_list(old[1, ...]) == [
[0, 1, 2, 3, 4],
[5, 6, 7, 8, 9],
[10, 11, 12, 13, 14],
]
assert ak.to_list(new[1, ...]) == [
[0, 1, 2, 3, 4],
[5, 6, 7, 8, 9],
[10, 11, 12, 13, 14],
]
expectation = [
[[0, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, 14]],
[[15, 16, 17, 18, 19], [20, 21, 22, 23, 24], [25, 26, 27, 28, 29]],
]
assert (ak.to_list(old[[1, 0], ]) == expectation)
assert (ak.to_list(new[[1, 0], ]) == expectation)
assert ak.to_list(old[[1, 0]]) == expectation
assert ak.to_list(new[[1, 0]]) == expectation
assert ak.to_list(old[1, [1, 0]]) == [[5, 6, 7, 8, 9], [0, 1, 2, 3, 4]]
assert ak.to_list(new[1, [1, 0]]) == [[5, 6, 7, 8, 9], [0, 1, 2, 3, 4]]
assert v1v2_equal(old[1, [1, 0]], new[1, [1, 0]])
def test_ListOffsetArray_NumpyArray():
old = ak.layout.ListOffsetArray64(
ak.layout.Index64(np.array([0, 1, 2, 3], np.int64)),
ak.layout.NumpyArray(
np.array([
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[11.1, 22.2, 33.3, 44.4, 55.5, 66.6],
[21.1, 22.2, 23.3, 24.4, 25.5, 26.6],
[31.1, 32.2, 33.3, 34.4, 35.5, 36.6],
[41.1, 42.2, 43.3, 44.4, 45.5, 46.6],
])),
)
new = v1_to_v2(old)
assert ak.to_list(old[0, 0:]) == [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6]]
assert ak.to_list(new[0, 0:]) == [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6]]
assert v1v2_equal(old[0, 1:], new[0, 1:])
with pytest.raises(IndexError):
new[1, "hello"]
with pytest.raises(IndexError):
new[1, ["hello", "there"]]
assert ak.to_list(
new[1, np.newaxis]) == [[[11.1, 22.2, 33.3, 44.4, 55.5, 66.6]]]
assert ak.to_list(
old[1, np.newaxis]) == [[[11.1, 22.2, 33.3, 44.4, 55.5, 66.6]]]
assert old.minmax_depth == (3, 3)
assert new.minmax_depth == (3, 3)
assert ak.to_list(old[1, ...]) == [[11.1, 22.2, 33.3, 44.4, 55.5, 66.6]]
assert ak.to_list(new[1, ...]) == [[11.1, 22.2, 33.3, 44.4, 55.5, 66.6]]
expectation = [
[[11.1, 22.2, 33.3, 44.4, 55.5, 66.6]],
[[1.1, 2.2, 3.3, 4.4, 5.5, 6.6]],
]
assert (ak.to_list(old[[1, 0], ]) == expectation)
assert (ak.to_list(new[[1, 0], ]) == expectation)
assert ak.to_list(old[[1, 0]]) == expectation
assert ak.to_list(new[[1, 0]]) == expectation
assert ak.to_list(old[0, [0, 0]]) == [
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
]
assert ak.to_list(new[0, [0, 0]]) == [
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
]
def test_IndexedOptionArray_NumpyArray():
v1a = ak.layout.IndexedOptionArray64(
ak.layout.Index64(np.array([2, 2, -1, 1, -1, 5, 4], dtype=np.int64)),
ak.layout.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])),
)
v2a = ak._v2.contents.indexedoptionarray.IndexedOptionArray(
ak._v2.index.Index(np.array([2, 2, -1, 1, -1, 5, 4], dtype=np.int64)),
ak._v2.contents.numpyarray.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])),
)
assert v1v2_equal(v1a, v2a)
assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a))
assert ak.to_list(v1a) == ak.to_list(v2a)
assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
def test_ListOffsetArray_NumpyArray():
v1a = ak.layout.ListOffsetArray64(
ak.layout.Index64(np.array([1, 4, 4, 6], dtype=np.int64)),
ak.layout.NumpyArray([6.6, 1.1, 2.2, 3.3, 4.4, 5.5, 7.7]),
)
v2a = ak._v2.contents.listoffsetarray.ListOffsetArray(
ak._v2.index.Index(np.array([1, 4, 4, 6], dtype=np.int64)),
ak._v2.contents.numpyarray.NumpyArray([6.6, 1.1, 2.2, 3.3, 4.4, 5.5, 7.7]),
)
assert v1v2_equal(v1a, v2a)
assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a))
assert ak.to_list(v1a) == ak.to_list(v2a)
assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
def test():
array = ak.Array([[{
"x": [[1, 2], []]
}, {
"x": []
}, {
"x": [[3]]
}]],
with_name="Record")
array_v2 = ak._v2.highlevel.Array(v1_to_v2(array.layout))
# Flattening inside record does not drop parameters
assert "__record__" not in ak.parameters(ak.flatten(array, axis=3))
assert "__record__" not in array_v2.layout.flatten(axis=3).parameters
# Flattening before record does not drop parameters
assert "__record__" in ak.parameters(ak.flatten(array, axis=1))
assert "__record__" in array_v2.layout.flatten(axis=1).parameters
def test_ByteMaskedArray():
old = ak.layout.ByteMaskedArray(
ak.layout.Index8(np.array([1, 1, 1], np.int8)),
ak.layout.NumpyArray(
np.array([
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
])),
valid_when=True,
)
new = v1_to_v2(old)
assert ak.to_list(old[:, 5:]) == [[6.6], [6.6], [6.6]]
assert ak.to_list(new[:, 5:]) == [[6.6], [6.6], [6.6]]
assert v1v2_equal(old[:, 5:], new[:, 5:])
with pytest.raises(IndexError):
new[1, "hello"]
with pytest.raises(IndexError):
new[1, ["hello", "there"]]
assert ak.to_list(new[1, np.newaxis]) == [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6]]
assert ak.to_list(old[1, np.newaxis]) == [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6]]
assert old.minmax_depth == (2, 2)
assert new.minmax_depth == (2, 2)
assert ak.to_list(old[0, ..., 0]) == 1.1
assert ak.to_list(new[0, ..., 0]) == 1.1
expectation = [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6]]
assert (ak.to_list(old[[1, 0], ]) == expectation)
assert (ak.to_list(new[[1, 0], ]) == expectation)
assert ak.to_list(new[[1, 0]]) == expectation
assert ak.to_list(old[1, [1, 0]]) == [2.2, 1.1]
assert ak.to_list(new[1, [1, 0]]) == [2.2, 1.1]
def test_UnionArray_NumpyArray():
v1a = ak.layout.UnionArray8_64(
ak.layout.Index8(np.array([1, 1, 0, 0, 1, 0, 1], dtype=np.int8)),
ak.layout.Index64(np.array([4, 3, 0, 1, 2, 2, 4, 100], dtype=np.int64)),
[
ak.layout.NumpyArray(np.array([1, 2, 3], dtype=np.int64)),
ak.layout.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5])),
],
)
v2a = ak._v2.contents.unionarray.UnionArray(
ak._v2.index.Index(np.array([1, 1, 0, 0, 1, 0, 1], dtype=np.int8)),
ak._v2.index.Index(np.array([4, 3, 0, 1, 2, 2, 4, 100], dtype=np.int64)),
[
ak._v2.contents.numpyarray.NumpyArray(np.array([1, 2, 3], dtype=np.int64)),
ak._v2.contents.numpyarray.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5])),
],
)
assert v1v2_equal(v1a, v2a)
assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a))
assert ak.to_list(v1a) == ak.to_list(v2a)
assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
def test_RecordArray():
old = ak.layout.RecordArray(
[
ak.layout.NumpyArray(
np.array([[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], np.int64)),
ak.layout.NumpyArray(
np.array([[0.0, 1.1, 2.2, 3.3, 4.4, 5.5],
[0.0, 1.1, 2.2, 3.3, 4.4, 5.5]])),
],
["x", "y"],
)
new = v1_to_v2(old)
assert ak.to_list(old[:, 3:]) == [
{
"x": [3, 4],
"y": [3.3, 4.4, 5.5]
},
{
"x": [3, 4],
"y": [3.3, 4.4, 5.5]
},
]
assert ak.to_list(new[:, 3:]) == [
{
"x": [3, 4],
"y": [3.3, 4.4, 5.5]
},
{
"x": [3, 4],
"y": [3.3, 4.4, 5.5]
},
]
with pytest.raises(IndexError):
new[1, "hello"]
with pytest.raises(IndexError):
new[1, ["hello", "there"]]
assert ak.to_list(new[1, np.newaxis]) == [{
"x": [0, 1, 2, 3, 4],
"y": [0.0, 1.1, 2.2, 3.3, 4.4, 5.5]
}]
assert ak.to_list(old[1, np.newaxis]) == [{
"x": [0, 1, 2, 3, 4],
"y": [0.0, 1.1, 2.2, 3.3, 4.4, 5.5]
}]
assert old.minmax_depth == (2, 2)
assert new.minmax_depth == (2, 2)
assert ak.to_list(old[0, ..., 0]) == {"x": 0, "y": 0.0}
assert ak.to_list(new[0, ..., 0]) == {"x": 0, "y": 0.0}
expectation = [
{
"x": [0, 1, 2, 3, 4],
"y": [0.0, 1.1, 2.2, 3.3, 4.4, 5.5]
},
{
"x": [0, 1, 2, 3, 4],
"y": [0.0, 1.1, 2.2, 3.3, 4.4, 5.5]
},
]
assert (ak.to_list(old[[1, 0], ]) == expectation)
assert (ak.to_list(new[[1, 0], ]) == expectation)
assert ak.to_list(new[[1, 0]]) == expectation
assert ak.to_list(old[1, [1, 0]]) == [{
"x": 1,
"y": 1.1
}, {
"x": 0,
"y": 0.0
}]
assert ak.to_list(new[1, [1, 0]]) == [{
"x": 1,
"y": 1.1
}, {
"x": 0,
"y": 0.0
}]
def test_ListArray():
old = ak.layout.ListArray64(
ak.layout.Index64(np.array([0, 100, 1], np.int64)),
ak.layout.Index64(np.array([3, 100, 3, 200], np.int64)),
ak.layout.NumpyArray(
np.array([
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
])),
)
new = v1_to_v2(old)
assert ak.to_list(old[0, :2]) == [
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
]
assert ak.to_list(new[0, :2]) == [
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
]
with pytest.raises(IndexError):
new[1, "hello"]
with pytest.raises(IndexError):
new[1, ["hello", "there"]]
assert ak.to_list(new[0, np.newaxis]) == [[
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
]]
assert ak.to_list(old[0, np.newaxis]) == [[
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
]]
assert old.minmax_depth == (3, 3)
assert new.minmax_depth == (3, 3)
assert ak.to_list(old[0, ...]) == [
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
]
assert ak.to_list(new[0, ...]) == [
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
]
expectation = [
[],
[
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
],
]
assert (ak.to_list(old[[1, 0], ]) == expectation)
assert (ak.to_list(new[[1, 0], ]) == expectation)
assert ak.to_list(old[[1, 0]]) == expectation
assert ak.to_list(new[[1, 0]]) == expectation
assert ak.to_list(old[0, [1, 0]]) == [
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
]
assert ak.to_list(new[0, [1, 0]]) == [
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
]
def test_BitMaskedArray():
old = ak.layout.BitMaskedArray(
ak.layout.IndexU8(
np.packbits(
np.array(
[
[
1,
1,
1,
1,
0,
0,
0,
0,
1,
0,
1,
0,
1,
],
[
1,
1,
1,
1,
0,
0,
0,
0,
1,
0,
1,
0,
1,
],
],
np.uint8,
))),
ak.layout.NumpyArray(
np.array([
[
0.0,
1.0,
2.0,
3.0,
4.0,
5.0,
6.0,
7.0,
1.1,
2.2,
3.3,
4.4,
5.5,
6.6,
],
[
0.0,
1.0,
2.0,
3.0,
4.0,
5.0,
6.0,
7.0,
1.1,
2.2,
3.3,
4.4,
5.5,
6.6,
],
])),
valid_when=True,
length=2,
lsb_order=False,
)
new = v1_to_v2(old)
assert v1v2_equal(old[:, :], new[:, :])
assert ak.to_list(old[:, 5:]) == [
[5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
]
assert ak.to_list(new[:, 5:]) == [
[5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
]
assert v1v2_equal(old[:, 5:], new[:, 5:])
with pytest.raises(IndexError):
new[1, "hello"]
with pytest.raises(IndexError):
new[1, ["hello", "there"]]
assert ak.to_list(new[1, np.newaxis]) == [[
0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6
]]
assert ak.to_list(old[1, np.newaxis]) == [[
0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6
]]
assert old.minmax_depth == (2, 2)
assert new.minmax_depth == (2, 2)
assert ak.to_list(old[0, ..., 0]) == 0.0
assert ak.to_list(new[0, ..., 0]) == 0.0
expectation = [
[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6],
]
assert (ak.to_list(old[[1, 0], ]) == expectation)
assert (ak.to_list(new[[1, 0], ]) == expectation)
assert ak.to_list(old[[1, 0]]) == expectation
assert ak.to_list(new[[1, 0]]) == expectation
assert ak.to_list(old[1, [1, 0]]) == [1.0, 0.0]
assert ak.to_list(new[1, [1, 0]]) == [1.0, 0.0]
assert v1v2_equal(old[1, [1, 0]], new[1, [1, 0]])
def _getitem(self, where):
if ak._util.isint(where):
return self._getitem_at(where)
elif isinstance(where, slice) and where.step is None:
return self._getitem_range(where)
elif isinstance(where, slice):
return self._getitem((where,))
elif ak._util.isstr(where):
return self._getitem_field(where)
elif where is np.newaxis:
return self._getitem((where,))
elif where is Ellipsis:
return self._getitem((where,))
elif isinstance(where, tuple):
if len(where) == 0:
return self
items = [ak._v2._slicing.prepare_tuple_item(x) for x in where]
nextwhere = ak._v2._slicing.getitem_broadcast(items)
next = ak._v2.contents.RegularArray(
self,
self.length if self._nplike.known_shape else 1,
1,
None,
None,
self._nplike,
)
out = next._getitem_next(nextwhere[0], nextwhere[1:], None)
if out.length == 0:
return out._getitem_nothing()
else:
return out._getitem_at(0)
elif isinstance(where, ak.highlevel.Array):
return self._getitem(where.layout)
elif isinstance(where, ak.layout.Content):
return self._getitem(v1_to_v2(where))
elif isinstance(where, ak._v2.highlevel.Array):
return self._getitem(where.layout)
elif (
isinstance(where, Content)
and where._parameters is not None
and (where._parameters.get("__array__") in ("string", "bytestring"))
):
return self._getitem_fields(ak._v2.operations.convert.to_list(where))
elif isinstance(where, ak._v2.contents.emptyarray.EmptyArray):
return where.toNumpyArray(np.int64)
elif isinstance(where, ak._v2.contents.numpyarray.NumpyArray):
if issubclass(where.dtype.type, np.int64):
carry = ak._v2.index.Index64(where.data.reshape(-1))
allow_lazy = True
elif issubclass(where.dtype.type, np.integer):
carry = ak._v2.index.Index64(where.data.astype(np.int64).reshape(-1))
allow_lazy = "copied" # True, but also can be modified in-place
elif issubclass(where.dtype.type, (np.bool_, bool)):
if len(where.data.shape) == 1:
where = self._nplike.nonzero(where.data)[0]
carry = ak._v2.index.Index64(where)
allow_lazy = "copied" # True, but also can be modified in-place
else:
wheres = self._nplike.nonzero(where.data)
return self._getitem(wheres)
else:
raise ak._v2._util.error(
TypeError(
"array slice must be an array of integers or booleans, not\n\n {}".format(
repr(where.data).replace("\n", "\n ")
)
)
)
out = ak._v2._slicing.getitem_next_array_wrap(
self._carry(carry, allow_lazy), where.shape
)
if out.length == 0:
return out._getitem_nothing()
else:
return out._getitem_at(0)
elif isinstance(where, Content):
return self._getitem((where,))
elif isinstance(where, Iterable) and all(ak._util.isstr(x) for x in where):
return self._getitem_fields(where)
elif isinstance(where, Iterable):
layout = ak._v2.operations.convert.to_layout(where)
as_array = layout.maybe_to_array(layout.nplike)
if as_array is None:
return self._getitem(layout)
else:
return self._getitem(
ak._v2.contents.NumpyArray(as_array, None, None, layout.nplike)
)
else:
raise ak._v2._util.error(
TypeError(
"only integers, slices (`:`), ellipsis (`...`), np.newaxis (`None`), "
"integer/boolean arrays (possibly with variable-length nested "
"lists or missing values), field name (str) or names (non-tuple "
"iterable of str) are valid indices for slicing, not\n\n "
+ repr(where).replace("\n", "\n ")
)
)
def test_BitMaskedArray_RecordArray_NumpyArray():
v1a = ak.layout.BitMaskedArray(
ak.layout.IndexU8(
np.packbits(
np.array(
[
True,
True,
True,
True,
False,
False,
False,
False,
True,
False,
True,
False,
True,
]
)
)
),
ak.layout.RecordArray(
[
ak.layout.NumpyArray(
np.array(
[
0.0,
1.0,
2.0,
3.0,
4.0,
5.0,
6.0,
7.0,
1.1,
2.2,
3.3,
4.4,
5.5,
6.6,
]
)
)
],
["nest"],
),
valid_when=True,
length=13,
lsb_order=False,
)
v2a = ak._v2.contents.bitmaskedarray.BitMaskedArray(
ak._v2.index.Index(
np.packbits(
np.array(
[
True,
True,
True,
True,
False,
False,
False,
False,
True,
False,
True,
False,
True,
]
)
)
),
ak._v2.contents.recordarray.RecordArray(
[
ak._v2.contents.numpyarray.NumpyArray(
np.array(
[
0.0,
1.0,
2.0,
3.0,
4.0,
5.0,
6.0,
7.0,
1.1,
2.2,
3.3,
4.4,
5.5,
6.6,
]
)
)
],
["nest"],
),
valid_when=True,
length=13,
lsb_order=False,
)
assert v1v2_equal(v1a, v2a)
assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a))
assert ak.to_list(v1a) == ak.to_list(v2a)
assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
v1b = ak.layout.BitMaskedArray(
ak.layout.IndexU8(
np.packbits(
np.array(
[
0,
0,
0,
0,
1,
1,
1,
1,
0,
1,
0,
1,
0,
],
dtype=np.uint8,
)
)
),
ak.layout.RecordArray(
[
ak.layout.NumpyArray(
np.array(
[
0.0,
1.0,
2.0,
3.0,
4.0,
5.0,
6.0,
7.0,
1.1,
2.2,
3.3,
4.4,
5.5,
6.6,
]
)
)
],
["nest"],
),
valid_when=False,
length=13,
lsb_order=False,
)
v2b = ak._v2.contents.bitmaskedarray.BitMaskedArray(
ak._v2.index.Index(
np.packbits(
np.array(
[
0,
0,
0,
0,
1,
1,
1,
1,
0,
1,
0,
1,
0,
],
dtype=np.uint8,
)
)
),
ak._v2.contents.recordarray.RecordArray(
[
ak._v2.contents.numpyarray.NumpyArray(
np.array(
[
0.0,
1.0,
2.0,
3.0,
4.0,
5.0,
6.0,
7.0,
1.1,
2.2,
3.3,
4.4,
5.5,
6.6,
]
)
)
],
["nest"],
),
valid_when=False,
length=13,
lsb_order=False,
)
assert v1v2_equal(v1b, v2b)
assert v1v2_equal(v2_to_v1(v2b), v1_to_v2(v1b))
assert ak.to_list(v1b) == ak.to_list(v2b)
assert newform(json.loads(v1b.form.tojson())) == v2b.form.tolist()
v1c = ak.layout.BitMaskedArray(
ak.layout.IndexU8(
np.packbits(
np.array(
[
0,
0,
0,
0,
1,
1,
1,
1,
0,
0,
0,
1,
0,
1,
0,
1,
],
dtype=np.uint8,
)
)
),
ak.layout.RecordArray(
[
ak.layout.NumpyArray(
np.array(
[
0.0,
1.0,
2.0,
3.0,
4.0,
5.0,
6.0,
7.0,
1.1,
2.2,
3.3,
4.4,
5.5,
6.6,
]
)
)
],
["nest"],
),
valid_when=True,
length=13,
lsb_order=True,
)
v2c = ak._v2.contents.bitmaskedarray.BitMaskedArray(
ak._v2.index.Index(
np.packbits(
np.array(
[
0,
0,
0,
0,
1,
1,
1,
1,
0,
0,
0,
1,
0,
1,
0,
1,
],
dtype=np.uint8,
)
)
),
ak._v2.contents.recordarray.RecordArray(
[
ak._v2.contents.numpyarray.NumpyArray(
np.array(
[
0.0,
1.0,
2.0,
3.0,
4.0,
5.0,
6.0,
7.0,
1.1,
2.2,
3.3,
4.4,
5.5,
6.6,
]
)
)
],
["nest"],
),
valid_when=True,
length=13,
lsb_order=True,
)
assert v1v2_equal(v1c, v2c)
assert v1v2_equal(v2_to_v1(v2c), v1_to_v2(v1c))
assert ak.to_list(v1c) == ak.to_list(v2c)
assert newform(json.loads(v1c.form.tojson())) == v2c.form.tolist()
v1d = ak.layout.BitMaskedArray(
ak.layout.IndexU8(
np.packbits(
np.array(
[
1,
1,
1,
1,
0,
0,
0,
0,
0,
0,
0,
0,
1,
0,
1,
0,
],
dtype=np.uint8,
)
)
),
ak.layout.RecordArray(
[
ak.layout.NumpyArray(
np.array(
[
0.0,
1.0,
2.0,
3.0,
4.0,
5.0,
6.0,
7.0,
1.1,
2.2,
3.3,
4.4,
5.5,
6.6,
]
)
)
],
["nest"],
),
valid_when=False,
length=13,
lsb_order=True,
)
v2d = ak._v2.contents.bitmaskedarray.BitMaskedArray(
ak._v2.index.Index(
np.packbits(
np.array(
[
1,
1,
1,
1,
0,
0,
0,
0,
0,
0,
0,
0,
1,
0,
1,
0,
],
dtype=np.uint8,
)
)
),
ak._v2.contents.recordarray.RecordArray(
[
ak._v2.contents.numpyarray.NumpyArray(
np.array(
[
0.0,
1.0,
2.0,
3.0,
4.0,
5.0,
6.0,
7.0,
1.1,
2.2,
3.3,
4.4,
5.5,
6.6,
]
)
)
],
["nest"],
),
valid_when=False,
length=13,
lsb_order=True,
)
assert v1v2_equal(v1d, v2d)
assert v1v2_equal(v2_to_v1(v2d), v1_to_v2(v1d))
assert ak.to_list(v1d) == ak.to_list(v2d)
assert newform(json.loads(v1d.form.tojson())) == v2d.form.tolist()
def __getitem__(self, where):
try:
if ak._util.isint(where):
return self._getitem_at(where)
elif isinstance(where, slice) and where.step is None:
return self._getitem_range(where)
elif isinstance(where, slice):
return self.__getitem__((where,))
elif ak._util.isstr(where):
return self._getitem_field(where)
elif where is np.newaxis:
return self.__getitem__((where,))
elif where is Ellipsis:
return self.__getitem__((where,))
elif isinstance(where, tuple):
if len(where) == 0:
return self
nextwhere = self._getitem_broadcast(
[self._prepare_tuple_item(x) for x in where],
self.nplike,
)
next = ak._v2.contents.RegularArray(self, len(self), 1, None, None)
out = next._getitem_next(nextwhere[0], nextwhere[1:], None)
if len(out) == 0:
raise out._getitem_nothing()
else:
return out._getitem_at(0)
elif isinstance(where, ak.highlevel.Array):
return self.__getitem__(where.layout)
elif isinstance(where, ak.layout.Content):
return self.__getitem__(v1_to_v2(where))
elif isinstance(where, ak._v2.contents.numpyarray.NumpyArray):
if issubclass(where.dtype.type, np.int64):
carry = ak._v2.index.Index64(where.data.reshape(-1))
allow_lazy = True
elif issubclass(where.dtype.type, np.integer):
carry = ak._v2.index.Index64(
where.data.astype(np.int64).reshape(-1)
)
allow_lazy = "copied" # True, but also can be modified in-place
elif issubclass(where.dtype.type, (np.bool_, bool)):
carry = ak._v2.index.Index64(np.nonzero(where.data.reshape(-1))[0])
allow_lazy = "copied" # True, but also can be modified in-place
else:
raise TypeError(
"array slice must be an array of integers or booleans, not\n\n {0}".format(
repr(where.data).replace("\n", "\n ")
)
)
out = self._getitem_next_array_wrap(
self._carry(carry, allow_lazy, NestedIndexError), where.shape
)
if len(out) == 0:
return out._getitem_nothing()
else:
return out._getitem_at(0)
elif isinstance(where, Content):
return self.__getitem__((where,))
elif isinstance(where, Iterable) and all(ak._util.isstr(x) for x in where):
return self._getitem_fields(where)
elif isinstance(where, Iterable):
return self.__getitem__(
v1_to_v2(ak.operations.convert.to_layout(where))
)
else:
raise TypeError(
"only integers, slices (`:`), ellipsis (`...`), np.newaxis (`None`), "
"integer/boolean arrays (possibly with variable-length nested "
"lists or missing values), field name (str) or names (non-tuple "
"iterable of str) are valid indices for slicing, not\n\n "
+ repr(where).replace("\n", "\n ")
)
except NestedIndexError as err:
def format_slice(x):
if isinstance(x, slice):
if x.step is None:
return "{0}:{1}".format(
"" if x.start is None else x.start,
"" if x.stop is None else x.stop,
)
else:
return "{0}:{1}:{2}".format(
"" if x.start is None else x.start,
"" if x.stop is None else x.stop,
x.step,
)
elif isinstance(x, tuple):
return "(" + ", ".join(format_slice(y) for y in x) + ")"
elif isinstance(x, ak._v2.index.Index64):
return str(x.data)
elif isinstance(x, Content):
return str(ak.Array(v2_to_v1(x)))
else:
return repr(x)
raise IndexError(
"""cannot slice
{0}
with
{1}
at inner {2} of length {3}, using sub-slice {4}.{5}""".format(
repr(ak.Array(v2_to_v1(self))),
format_slice(where),
type(err.array).__name__,
len(err.array),
format_slice(err.slicer),
""
if err.details is None
else "\n\n{0} error: {1}.".format(
type(err.array).__name__, err.details
),
)
)
def _prepare_tuple_item(self, item):
if ak._util.isint(item):
return int(item)
elif isinstance(item, slice):
return item
elif ak._util.isstr(item):
return item
elif item is np.newaxis:
return item
elif item is Ellipsis:
return item
elif isinstance(item, ak.highlevel.Array):
return self._prepare_tuple_item(item.layout)
elif isinstance(item, ak.layout.Content):
return self._prepare_tuple_item(v1_to_v2(item))
elif isinstance(item, ak._v2.contents.NumpyArray):
return item.data
elif isinstance(
item,
(
ak._v2.contents.ListOffsetArray,
ak._v2.contents.ListArray,
ak._v2.contents.RegularArray,
),
):
return item.toListOffsetArray64(False)
elif isinstance(
item,
(
ak._v2.contents.IndexedOptionArray,
ak._v2.contents.ByteMaskedArray,
ak._v2.contents.BitMaskedArray,
ak._v2.contents.UnmaskedArray,
),
):
return item.toIndexedOptionArray64()
elif isinstance(item, Iterable) and all(ak._util.isstr(x) for x in item):
return list(item)
elif isinstance(item, Iterable):
return self._prepare_tuple_item(
v1_to_v2(ak.operations.convert.to_layout(item))
)
else:
raise TypeError(
"only integers, slices (`:`), ellipsis (`...`), np.newaxis (`None`), "
"integer/boolean arrays (possibly with variable-length nested "
"lists or missing values), field name (str) or names (non-tuple "
"iterable of str) are valid indices for slicing, not\n\n "
+ repr(item).replace("\n", "\n ")
)
