def test_with_name_field(): array = awkward1.Array([{ "x": 0.0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": [2, 2] }, { "x": 3.3, "y": [3, 3, 3] }]) array2 = awkward1.repartition(array, 2) z = awkward1.Array([100, 200, 300, 400]) z2 = awkward1.repartition(z, 3) one = awkward1.with_name(array, "Wilbur") assert isinstance(one.layout, awkward1.layout.Content) assert one.layout.parameters["__record__"] == "Wilbur" one = awkward1.with_name(array2, "Wilbur") assert isinstance(one.layout, awkward1.partition.PartitionedArray) assert one.layout.partition(0).parameters["__record__"] == "Wilbur" assert one.layout.partition(1).parameters["__record__"] == "Wilbur" two = awkward1.with_field(array, z, "z") assert isinstance(two.layout, awkward1.layout.Content) assert awkward1.to_list(two) == [{ "x": 0.0, "y": [], "z": 100 }, { "x": 1.1, "y": [1], "z": 200 }, { "x": 2.2, "y": [2, 2], "z": 300 }, { "x": 3.3, "y": [3, 3, 3], "z": 400 }] two = awkward1.with_field(array2, z, "z") assert isinstance(two.layout, awkward1.partition.PartitionedArray) assert awkward1.to_list(two) == [{ "x": 0.0, "y": [], "z": 100 }, { "x": 1.1, "y": [1], "z": 200 }, { "x": 2.2, "y": [2, 2], "z": 300 }, { "x": 3.3, "y": [3, 3, 3], "z": 400 }] two = awkward1.with_field(array2, z2, "z") assert isinstance(two.layout, awkward1.partition.PartitionedArray) assert awkward1.to_list(two) == [{ "x": 0.0, "y": [], "z": 100 }, { "x": 1.1, "y": [1], "z": 200 }, { "x": 2.2, "y": [2, 2], "z": 300 }, { "x": 3.3, "y": [3, 3, 3], "z": 400 }] two = awkward1.with_field(array2, z2, "z") assert isinstance(two.layout, awkward1.partition.PartitionedArray) assert awkward1.to_list(two) == [{ "x": 0.0, "y": [], "z": 100 }, { "x": 1.1, "y": [1], "z": 200 }, { "x": 2.2, "y": [2, 2], "z": 300 }, { "x": 3.3, "y": [3, 3, 3], "z": 400 }]
def test_RegularArray_getitem(): array = awkward1.Array(numpy.array([[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]]), check_valid=True) @numba.njit def f1(x, i): return x[i] assert f1(array, -2).tolist() == [1.1, 2.2, 3.3] assert f1(array, 0).tolist() == [1.1, 2.2, 3.3] assert f1(array, 1).tolist() == [4.4, 5.5, 6.6] assert f1(array, -1).tolist() == [4.4, 5.5, 6.6] @numba.njit def f2(x, i, j): return x[i][j] assert f2(array, 1, 0) == 4.4 assert f2(array, 1, 1) == 5.5 assert f2(array, 1, 2) == 6.6 assert f2(array, -1, 0) == 4.4 assert f2(array, -1, 1) == 5.5 assert f2(array, -1, 2) == 6.6 assert f2(array, 1, -3) == 4.4 assert f2(array, 1, -2) == 5.5 assert f2(array, 1, -1) == 6.6 array = awkward1.Array(numpy.array([[1.1, 2.2], [3.3, 4.4], [5.5, 6.6]]), check_valid=True) @numba.njit def f3(x, i1, i2): return x[i1:i2] assert awkward1.to_list(f3(array, -1, 3)) == [[5.5, 6.6]] assert awkward1.to_list(f3(array, -2, 3)) == [[3.3, 4.4], [5.5, 6.6]] assert awkward1.to_list(f3(array, -3, 3)) == [[1.1, 2.2], [3.3, 4.4], [5.5, 6.6]] assert awkward1.to_list(f3(array, 0, 3)) == [[1.1, 2.2], [3.3, 4.4], [5.5, 6.6]] assert awkward1.to_list(f3(array, 1, 3)) == [[3.3, 4.4], [5.5, 6.6]] assert awkward1.to_list(f3(array, 2, 3)) == [[5.5, 6.6]] assert awkward1.to_list(f3(array, 3, 3)) == [] assert awkward1.to_list(f3(array, 0, 0)) == [] assert awkward1.to_list(f3(array, 0, 1)) == [[1.1, 2.2]] assert awkward1.to_list(f3(array, 0, 2)) == [[1.1, 2.2], [3.3, 4.4]] assert awkward1.to_list(f3(array, 0, 3)) == [[1.1, 2.2], [3.3, 4.4], [5.5, 6.6]] assert awkward1.to_list(f3(array, 0, -1)) == [[1.1, 2.2], [3.3, 4.4]] assert awkward1.to_list(f3(array, 0, -2)) == [[1.1, 2.2]] assert awkward1.to_list(f3(array, 0, -3)) == []
def test_ListArray_getitem_field(): array = awkward1.Array([[{ "x": 0.0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": [2, 2] }], [], [{ "x": 3.3, "y": [3, 3, 3] }, { "x": 4.4, "y": [4, 4, 4, 4] }], [{ "x": 5.5, "y": [5, 5, 5, 5, 5] }], [{ "x": 6.6, "y": [6, 6, 6, 6, 6, 6] }, { "x": 7.7, "y": [7, 7, 7, 7, 7, 7, 7] }, { "x": 8.8, "y": [8, 8, 8, 8, 8, 8, 8, 8] }, { "x": 9.9, "y": [9, 9, 9, 9, 9, 9, 9, 9, 9] }]], check_valid=True) @numba.njit def f1(x): return x["x"] assert awkward1.to_list(f1(array)) == [[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]] @numba.njit def f2(x): return x.y assert awkward1.to_list(f2(array)) == [[[], [1], [2, 2]], [], [[3, 3, 3], [4, 4, 4, 4]], [[5, 5, 5, 5, 5]], [[6, 6, 6, 6, 6, 6], [7, 7, 7, 7, 7, 7, 7], [8, 8, 8, 8, 8, 8, 8, 8], [9, 9, 9, 9, 9, 9, 9, 9, 9]]] @numba.njit def f3(x): return x[1:4].x assert awkward1.to_list(f3(array)) == [[], [3.3, 4.4], [5.5]] @numba.njit def f4(x): return x[1:4]["y"] assert awkward1.to_list(f4(array)) == [[], [[3, 3, 3], [4, 4, 4, 4]], [[5, 5, 5, 5, 5]]] @numba.njit def f5(x): return x["x"][1:4] assert awkward1.to_list(f5(array)) == [[], [3.3, 4.4], [5.5]] @numba.njit def f6(x): return x.y[1:4] assert awkward1.to_list(f6(array)) == [[], [[3, 3, 3], [4, 4, 4, 4]], [[5, 5, 5, 5, 5]]] @numba.njit def f7(x): return x[4]["x"] assert awkward1.to_list(awkward1.to_list( f7(array))) == [6.6, 7.7, 8.8, 9.9] @numba.njit def f8(x): return x[4].y assert awkward1.to_list(awkward1.to_list( f8(array))) == [[6, 6, 6, 6, 6, 6], [7, 7, 7, 7, 7, 7, 7], [8, 8, 8, 8, 8, 8, 8, 8], [9, 9, 9, 9, 9, 9, 9, 9, 9]] @numba.njit def f9(x): return x.x[4] assert awkward1.to_list(awkward1.to_list( f9(array))) == [6.6, 7.7, 8.8, 9.9] @numba.njit def f10(x): return x["y"][4] assert awkward1.to_list(awkward1.to_list( f10(array))) == [[6, 6, 6, 6, 6, 6], [7, 7, 7, 7, 7, 7, 7], [8, 8, 8, 8, 8, 8, 8, 8], [9, 9, 9, 9, 9, 9, 9, 9, 9]] @numba.njit def f11(x): return x[4]["x"][1] assert f11(array) == 7.7 @numba.njit def f12(x): return x[4].y[1] assert awkward1.to_list(awkward1.to_list( f12(array))) == [7, 7, 7, 7, 7, 7, 7] @numba.njit def f12b(x): return x[4].y[1][6] assert awkward1.to_list(awkward1.to_list(f12b(array))) == 7 @numba.njit def f13(x): return x.x[4][1] assert f13(array) == 7.7 @numba.njit def f14(x): return x["y"][4][1] assert awkward1.to_list(awkward1.to_list( f14(array))) == [7, 7, 7, 7, 7, 7, 7] @numba.njit def f14b(x): return x["y"][4][1][6] assert awkward1.to_list(f14b(array)) == 7
def test_RecordArray_deep_field(): array = awkward1.Array([{ "x": { "y": { "z": 1.1 } } }, { "x": { "y": { "z": 2.2 } } }, { "x": { "y": { "z": 3.3 } } }], check_valid=True) @numba.njit def f1(x): return x[1]["x"].y["z"] assert f1(array) == 2.2 @numba.njit def f2(x): return x["x"][1].y["z"] assert f2(array) == 2.2 @numba.njit def f3(x): return x["x"].y[1]["z"] assert f3(array) == 2.2 @numba.njit def f4(x): return x["x"].y["z"][1] assert f4(array) == 2.2 @numba.njit def f5(x): return x["x"].y["z"] assert awkward1.to_list(f5(array)) == [1.1, 2.2, 3.3] @numba.njit def f6(x): return x.x["y"].z assert awkward1.to_list(f6(array)) == [1.1, 2.2, 3.3] @numba.njit def f7(x): return x.x["y"] assert awkward1.to_list(f7(array)) == [{"z": 1.1}, {"z": 2.2}, {"z": 3.3}] @numba.njit def f8(x): return x.x assert awkward1.to_list(f8(array)) == [{ "y": { "z": 1.1 } }, { "y": { "z": 2.2 } }, { "y": { "z": 3.3 } }]
def test_views(): assert awkward1.to_list( awkward1_numba_arrayview.ArrayView.fromarray( awkward1.Array( [1.1, 2.2, 3.3, 4.4, 5.5], check_valid=True)).toarray()) == [1.1, 2.2, 3.3, 4.4, 5.5] assert awkward1.to_list( awkward1_numba_arrayview.ArrayView.fromarray( awkward1.Array(numpy.array([[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]]), check_valid=True)).toarray()) == [[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]] assert awkward1.to_list( awkward1_numba_arrayview.ArrayView.fromarray( awkward1.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]], check_valid=True)).toarray()) == [[1.1, 2.2, 3.3], [], [4.4, 5.5]] assert awkward1.to_list( awkward1_numba_arrayview.ArrayView.fromarray( awkward1.Array([1.1, 2.2, None, 3.3, None, 4.4, 5.5], check_valid=True)).toarray()) == [ 1.1, 2.2, None, 3.3, None, 4.4, 5.5 ] assert awkward1.to_list( awkward1_numba_arrayview.ArrayView.fromarray( awkward1.Array([{ "x": 0.0, "y": [] }, { "x": 1.1, "y": [1, 1] }, { "x": 2.2, "y": [2, 2, 2] }], check_valid=True)).toarray()) == [{ "x": 0.0, "y": [] }, { "x": 1.1, "y": [1, 1] }, { "x": 2.2, "y": [2, 2, 2] }] assert awkward1.to_list( awkward1_numba_arrayview.ArrayView.fromarray( awkward1.Array([(0.0, []), (1.1, [1, 1]), (2.2, [2, 2, 2])], check_valid=True)).toarray()) == [(0.0, []), (1.1, [1, 1]), (2.2, [2, 2, 2])] assert awkward1.to_list( awkward1_numba_arrayview.ArrayView.fromarray( awkward1.Array([1.1, 2.2, 3.3, [], [1], [2, 2]], check_valid=True)).toarray()) == [ 1.1, 2.2, 3.3, [], [1], [2, 2] ]
def test_RecordArray(): array = awkward1.Array([{ "x": 0.0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": [2, 2] }, { "x": 3.3, "y": [3, 3, 3] }, { "x": 4.4, "y": [4, 4, 4, 4] }, { "x": 5.5, "y": [5, 5, 5] }, { "x": 6.6, "y": [6, 6] }, { "x": 7.7, "y": [7] }, { "x": 8.8, "y": [] }]) assert awkward1.to_list(array) == [{ 'x': 0.0, 'y': [] }, { 'x': 1.1, 'y': [1] }, { 'x': 2.2, 'y': [2, 2] }, { 'x': 3.3, 'y': [3, 3, 3] }, { 'x': 4.4, 'y': [4, 4, 4, 4] }, { 'x': 5.5, 'y': [5, 5, 5] }, { 'x': 6.6, 'y': [6, 6] }, { 'x': 7.7, 'y': [7] }, { 'x': 8.8, 'y': [] }] assert awkward1.to_list(array.layout.sort(-1, True, False)) == { 'x': [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8], 'y': [[], [1], [2, 2], [3, 3, 3], [4, 4, 4, 4], [5, 5, 5], [6, 6], [7], []] } assert awkward1.to_list(array.layout.sort(-1, False, False)) == { 'x': [8.8, 7.7, 6.6, 5.5, 4.4, 3.3, 2.2, 1.1, 0.0], 'y': [[], [1], [2, 2], [3, 3, 3], [4, 4, 4, 4], [5, 5, 5], [6, 6], [7], []] } assert awkward1.to_list(array.layout.argsort(-1, True, False)) == { 'x': [0, 1, 2, 3, 4, 5, 6, 7, 8], 'y': [[], [0], [0, 1], [0, 1, 2], [0, 1, 2, 3], [0, 1, 2], [0, 1], [0], []] } assert awkward1.to_list(array.x.layout.argsort( 0, True, False)) == [0, 1, 2, 3, 4, 5, 6, 7, 8] assert awkward1.to_list(array.x.layout.argsort( 0, False, False)) == [8, 7, 6, 5, 4, 3, 2, 1, 0] array_y = array.y assert awkward1.to_list(array_y) == [[], [1], [2, 2], [3, 3, 3], [4, 4, 4, 4], [5, 5, 5], [6, 6], [7], []] assert awkward1.to_list(array.y.layout.argsort(0, True, False)) == [[], [0], [1, 0], [2, 1, 0], [3, 2, 1, 0], [4, 3, 2], [5, 4], [6], []] assert awkward1.to_list(array.y.layout.argsort(1, True, True)) == [[], [0], [0, 1], [0, 1, 2], [0, 1, 2, 3], [0, 1, 2], [0, 1], [0], []]
def test_string_equal(): one = awkward1.Array(["one", "two", "three"], check_valid=True) two = awkward1.Array(["ONE", "two", "four"], check_valid=True) assert awkward1.to_list(one == two) == [False, True, False]
def test_IndexedArray(): array = awkward1.Array([[0.0, 1.1, 2.2, 3.3], [], [4.4, 5.5, 6.6], None, [7.7], None, [8.8, 9.9, 10.0, 11.1, 12.2]]) assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False)) == [[(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3), (2.2, 3.3)], [], [(4.4, 5.5), (4.4, 6.6), (5.5, 6.6)], None, [], None, [(8.8, 9.9), (8.8, 10.0), (8.8, 11.1), (8.8, 12.2), (9.9, 10.0), (9.9, 11.1), (9.9, 12.2), (10.0, 11.1), (10.0, 12.2), (11.1, 12.2)]]
def test_axis2(): array = awkward1.Array([[[0.0, 1.1, 2.2, 3.3], [], [4.4, 5.5, 6.6]], [], [[7.7], [8.8, 9.9, 10.0, 11.1, 12.2]]]) assert awkward1.to_list(awkward1.combinations(array, 2, axis=1, replacement=False)) == [[([0.0, 1.1, 2.2, 3.3], []), ([0.0, 1.1, 2.2, 3.3], [4.4, 5.5, 6.6]), ([], [4.4, 5.5, 6.6])], [], [([7.7], [8.8, 9.9, 10.0, 11.1, 12.2])]] assert awkward1.to_list(awkward1.combinations(array, 2, axis=2, replacement=False)) == [[[(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3), (2.2, 3.3)], [], [(4.4, 5.5), (4.4, 6.6), (5.5, 6.6)]], [], [[], [(8.8, 9.9), (8.8, 10.0), (8.8, 11.1), (8.8, 12.2), (9.9, 10.0), (9.9, 11.1), (9.9, 12.2), (10.0, 11.1), (10.0, 12.2), (11.1, 12.2)]]]
def test_mask(): array = awkward1.Array([1, 2, 3, 4, 5]) mask = awkward1.Array([False, True, True, True, False]) assert awkward1.to_list(awkward1.mask(array, mask)) == [None, 2, 3, 4, None]
def test_zip(): x = awkward1.Array([[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]]) y = awkward1.Array([1.1, 2.2, 3.3, 4.4, 5.5]) one = awkward1.zip({"x": x, "y": y}) two = awkward1.zip({"x": x, "y": y}, depth_limit=1) xx, yy = awkward1.unzip(two) assert isinstance(one.layout, awkward1.layout.Content) assert isinstance(two.layout, awkward1.layout.Content) assert isinstance(xx.layout, awkward1.layout.Content) assert isinstance(yy.layout, awkward1.layout.Content) assert awkward1.to_list(one) == [[{ "x": 1, "y": 1.1 }, { "x": 2, "y": 1.1 }, { "x": 3, "y": 1.1 }], [], [{ "x": 4, "y": 3.3 }, { "x": 5, "y": 3.3 }], [{ "x": 6, "y": 4.4 }], [{ "x": 7, "y": 5.5 }, { "x": 8, "y": 5.5 }, { "x": 9, "y": 5.5 }, { "x": 10, "y": 5.5 }]] assert awkward1.to_list(two) == [{ "x": [1, 2, 3], "y": 1.1 }, { "x": [], "y": 2.2 }, { "x": [4, 5], "y": 3.3 }, { "x": [6], "y": 4.4 }, { "x": [7, 8, 9, 10], "y": 5.5 }] if not awkward1._util.py27 and not awkward1._util.py35: assert awkward1.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] assert awkward1.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5] x = awkward1.repartition(x, 3) assert isinstance(x.layout, awkward1.partition.PartitionedArray) assert awkward1.to_list(x) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] one = awkward1.zip({"x": x, "y": y}) two = awkward1.zip({"x": x, "y": y}, depth_limit=1) xx, yy = awkward1.unzip(two) assert isinstance(one.layout, awkward1.partition.PartitionedArray) assert isinstance(two.layout, awkward1.partition.PartitionedArray) assert isinstance(xx.layout, awkward1.partition.PartitionedArray) assert isinstance(yy.layout, awkward1.partition.PartitionedArray) assert awkward1.to_list(one) == [[{ "x": 1, "y": 1.1 }, { "x": 2, "y": 1.1 }, { "x": 3, "y": 1.1 }], [], [{ "x": 4, "y": 3.3 }, { "x": 5, "y": 3.3 }], [{ "x": 6, "y": 4.4 }], [{ "x": 7, "y": 5.5 }, { "x": 8, "y": 5.5 }, { "x": 9, "y": 5.5 }, { "x": 10, "y": 5.5 }]] assert awkward1.to_list(two) == [{ "x": [1, 2, 3], "y": 1.1 }, { "x": [], "y": 2.2 }, { "x": [4, 5], "y": 3.3 }, { "x": [6], "y": 4.4 }, { "x": [7, 8, 9, 10], "y": 5.5 }] if not awkward1._util.py27 and not awkward1._util.py35: assert awkward1.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] assert awkward1.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5] y = awkward1.repartition(y, 2) assert isinstance(x.layout, awkward1.partition.PartitionedArray) assert awkward1.to_list(y) == [1.1, 2.2, 3.3, 4.4, 5.5] one = awkward1.zip({"x": x, "y": y}) two = awkward1.zip({"x": x, "y": y}, depth_limit=1) xx, yy = awkward1.unzip(two) assert isinstance(one.layout, awkward1.partition.PartitionedArray) assert isinstance(two.layout, awkward1.partition.PartitionedArray) assert isinstance(xx.layout, awkward1.partition.PartitionedArray) assert isinstance(yy.layout, awkward1.partition.PartitionedArray) assert awkward1.to_list(one) == [[{ "x": 1, "y": 1.1 }, { "x": 2, "y": 1.1 }, { "x": 3, "y": 1.1 }], [], [{ "x": 4, "y": 3.3 }, { "x": 5, "y": 3.3 }], [{ "x": 6, "y": 4.4 }], [{ "x": 7, "y": 5.5 }, { "x": 8, "y": 5.5 }, { "x": 9, "y": 5.5 }, { "x": 10, "y": 5.5 }]] assert awkward1.to_list(two) == [{ "x": [1, 2, 3], "y": 1.1 }, { "x": [], "y": 2.2 }, { "x": [4, 5], "y": 3.3 }, { "x": [6], "y": 4.4 }, { "x": [7, 8, 9, 10], "y": 5.5 }] if not awkward1._util.py27 and not awkward1._util.py35: assert awkward1.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] assert awkward1.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5] x = awkward1.repartition(x, None) assert isinstance(x.layout, awkward1.layout.Content) assert awkward1.to_list(x) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] one = awkward1.zip({"x": x, "y": y}) two = awkward1.zip({"x": x, "y": y}, depth_limit=1) xx, yy = awkward1.unzip(two) assert isinstance(one.layout, awkward1.partition.PartitionedArray) assert isinstance(two.layout, awkward1.partition.PartitionedArray) assert isinstance(xx.layout, awkward1.partition.PartitionedArray) assert isinstance(yy.layout, awkward1.partition.PartitionedArray) assert awkward1.to_list(one) == [[{ "x": 1, "y": 1.1 }, { "x": 2, "y": 1.1 }, { "x": 3, "y": 1.1 }], [], [{ "x": 4, "y": 3.3 }, { "x": 5, "y": 3.3 }], [{ "x": 6, "y": 4.4 }], [{ "x": 7, "y": 5.5 }, { "x": 8, "y": 5.5 }, { "x": 9, "y": 5.5 }, { "x": 10, "y": 5.5 }]] assert awkward1.to_list(two) == [{ "x": [1, 2, 3], "y": 1.1 }, { "x": [], "y": 2.2 }, { "x": [4, 5], "y": 3.3 }, { "x": [6], "y": 4.4 }, { "x": [7, 8, 9, 10], "y": 5.5 }] if not awkward1._util.py27 and not awkward1._util.py35: assert awkward1.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] assert awkward1.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5] y = awkward1.repartition(y, None) assert isinstance(y.layout, awkward1.layout.Content) assert awkward1.to_list(y) == [1.1, 2.2, 3.3, 4.4, 5.5] one = awkward1.zip({"x": x, "y": y}) two = awkward1.zip({"x": x, "y": y}, depth_limit=1) xx, yy = awkward1.unzip(two) assert isinstance(one.layout, awkward1.layout.Content) assert isinstance(two.layout, awkward1.layout.Content) assert isinstance(xx.layout, awkward1.layout.Content) assert isinstance(yy.layout, awkward1.layout.Content) assert awkward1.to_list(one) == [[{ "x": 1, "y": 1.1 }, { "x": 2, "y": 1.1 }, { "x": 3, "y": 1.1 }], [], [{ "x": 4, "y": 3.3 }, { "x": 5, "y": 3.3 }], [{ "x": 6, "y": 4.4 }], [{ "x": 7, "y": 5.5 }, { "x": 8, "y": 5.5 }, { "x": 9, "y": 5.5 }, { "x": 10, "y": 5.5 }]] assert awkward1.to_list(two) == [{ "x": [1, 2, 3], "y": 1.1 }, { "x": [], "y": 2.2 }, { "x": [4, 5], "y": 3.3 }, { "x": [6], "y": 4.4 }, { "x": [7, 8, 9, 10], "y": 5.5 }] if not awkward1._util.py27 and not awkward1._util.py35: assert awkward1.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] assert awkward1.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5]
def test_getitem_basic(): one = awkward1.from_iter([[1.1, 2.2, 3.3], [], [4.4, 5.5]], highlevel=False) two = awkward1.from_iter([[6.6], [], [], [], [7.7, 8.8, 9.9]], highlevel=False) array = awkward1.partition.IrregularlyPartitionedArray([one, two]) assert awkward1.to_list(array[0]) == [1.1, 2.2, 3.3] assert awkward1.to_list(array[1]) == [] assert awkward1.to_list(array[2]) == [4.4, 5.5] assert awkward1.to_list(array[3]) == [6.6] assert awkward1.to_list(array[4]) == [] assert awkward1.to_list(array[5]) == [] assert awkward1.to_list(array[6]) == [] assert awkward1.to_list(array[7]) == [7.7, 8.8, 9.9] assert awkward1.to_list(array[-1]) == [7.7, 8.8, 9.9] assert awkward1.to_list(array[-2]) == [] assert awkward1.to_list(array[-3]) == [] assert awkward1.to_list(array[-4]) == [] assert awkward1.to_list(array[-5]) == [6.6] assert awkward1.to_list(array[-6]) == [4.4, 5.5] assert awkward1.to_list(array[-7]) == [] assert awkward1.to_list(array[-8]) == [1.1, 2.2, 3.3] assert array[:].tojson( ) == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert array[2:6].tojson() == "[[4.4,5.5],[6.6],[],[]]" assert array[1:].tojson() == "[[],[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert array[2:].tojson() == "[[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert array[3:].tojson() == "[[6.6],[],[],[],[7.7,8.8,9.9]]" assert array[4:].tojson() == "[[],[],[],[7.7,8.8,9.9]]" assert array[5:].tojson() == "[[],[],[7.7,8.8,9.9]]" assert array[6:].tojson() == "[[],[7.7,8.8,9.9]]" assert array[7:].tojson() == "[[7.7,8.8,9.9]]" assert array[8:].tojson() == "[]" assert array[-1:].tojson() == "[[7.7,8.8,9.9]]" assert array[-2:].tojson() == "[[],[7.7,8.8,9.9]]" assert array[-3:].tojson() == "[[],[],[7.7,8.8,9.9]]" assert array[-4:].tojson() == "[[],[],[],[7.7,8.8,9.9]]" assert array[-5:].tojson() == "[[6.6],[],[],[],[7.7,8.8,9.9]]" assert array[-6:].tojson() == "[[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert array[-7:].tojson() == "[[],[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert array[-8:].tojson( ) == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert array[:-1].tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[],[]]" assert array[:-2].tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[]]" assert array[:-3].tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[]]" assert array[:-4].tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6]]" assert array[:-5].tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5]]" assert array[:-6].tojson() == "[[1.1,2.2,3.3],[]]" assert array[:-7].tojson() == "[[1.1,2.2,3.3]]" assert array[:-8].tojson() == "[]" assert array[:0].tojson() == "[]" assert array[:1].tojson() == "[[1.1,2.2,3.3]]" assert array[:2].tojson() == "[[1.1,2.2,3.3],[]]" assert array[:3].tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5]]" assert array[:4].tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6]]" assert array[:5].tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[]]" assert array[:6].tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[]]" assert array[:7].tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[],[]]" assert array[:8].tojson( ) == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" one = awkward1.from_iter([{ "x": 0.0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": [2, 2] }], highlevel=False) two = awkward1.from_iter([{ "x": 3.3, "y": [3, 3, 3] }, { "x": 4.4, "y": [4, 4, 4, 4] }], highlevel=False) array = awkward1.partition.IrregularlyPartitionedArray([one, two]) if not awkward1._util.py27 and not awkward1._util.py35: assert array.tojson( ) == '[{"x":0.0,"y":[]},{"x":1.1,"y":[1]},{"x":2.2,"y":[2,2]},{"x":3.3,"y":[3,3,3]},{"x":4.4,"y":[4,4,4,4]}]' assert array["x"].tojson() == "[0.0,1.1,2.2,3.3,4.4]" assert array["y"].tojson() == "[[],[1],[2,2],[3,3,3],[4,4,4,4]]" assert array[[ "x" ]].tojson() == '[{"x":0.0},{"x":1.1},{"x":2.2},{"x":3.3},{"x":4.4}]'
def test_basic(): one = awkward1.from_iter([[1.1, 2.2, 3.3], [], [4.4, 5.5]], highlevel=False) two = awkward1.from_iter([[6.6], [], [], [], [7.7, 8.8, 9.9]], highlevel=False) array = awkward1.partition.IrregularlyPartitionedArray([one, two]) assert len(array) == 8 assert [awkward1.to_list(x) for x in array.partitions] == [[[1.1, 2.2, 3.3], [], [4.4, 5.5]], [[6.6], [], [], [], [7.7, 8.8, 9.9]]] assert awkward1.to_list(array.partition(0)) == [[1.1, 2.2, 3.3], [], [4.4, 5.5]] assert awkward1.to_list(array.partition(1)) == [[6.6], [], [], [], [7.7, 8.8, 9.9]] assert array.start(0) == 0 assert array.start(1) == 3 assert array.stop(0) == 3 assert array.stop(1) == 8 assert array.partitionid_index_at(0) == (0, 0) assert array.partitionid_index_at(1) == (0, 1) assert array.partitionid_index_at(2) == (0, 2) assert array.partitionid_index_at(3) == (1, 0) assert array.partitionid_index_at(4) == (1, 1) assert array.partitionid_index_at(5) == (1, 2) assert array.partitionid_index_at(6) == (1, 3) assert array.partitionid_index_at(7) == (1, 4) assert array.tojson( ) == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert awkward1.to_list(array._ext.getitem_at(0)) == [1.1, 2.2, 3.3] assert awkward1.to_list(array._ext.getitem_at(1)) == [] assert awkward1.to_list(array._ext.getitem_at(2)) == [4.4, 5.5] assert awkward1.to_list(array._ext.getitem_at(3)) == [6.6] assert awkward1.to_list(array._ext.getitem_at(4)) == [] assert awkward1.to_list(array._ext.getitem_at(5)) == [] assert awkward1.to_list(array._ext.getitem_at(6)) == [] assert awkward1.to_list(array._ext.getitem_at(7)) == [7.7, 8.8, 9.9] assert awkward1.to_list(array._ext.getitem_at(-1)) == [7.7, 8.8, 9.9] assert awkward1.to_list(array._ext.getitem_at(-2)) == [] assert awkward1.to_list(array._ext.getitem_at(-3)) == [] assert awkward1.to_list(array._ext.getitem_at(-4)) == [] assert awkward1.to_list(array._ext.getitem_at(-5)) == [6.6] assert awkward1.to_list(array._ext.getitem_at(-6)) == [4.4, 5.5] assert awkward1.to_list(array._ext.getitem_at(-7)) == [] assert awkward1.to_list(array._ext.getitem_at(-8)) == [1.1, 2.2, 3.3] assert array._ext.getitem_range(0, 8, 1).tojson( ) == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range( 1, 8, 1).tojson() == "[[],[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range( 2, 8, 1).tojson() == "[[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range( 3, 8, 1).tojson() == "[[6.6],[],[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range(4, 8, 1).tojson() == "[[],[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range(5, 8, 1).tojson() == "[[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range(6, 8, 1).tojson() == "[[],[7.7,8.8,9.9]]" assert array._ext.getitem_range(7, 8, 1).tojson() == "[[7.7,8.8,9.9]]" assert array._ext.getitem_range(8, 8, 1).tojson() == "[]" assert array._ext.getitem_range(-1, 8, 1).tojson() == "[[7.7,8.8,9.9]]" assert array._ext.getitem_range(-2, 8, 1).tojson() == "[[],[7.7,8.8,9.9]]" assert array._ext.getitem_range(-3, 8, 1).tojson() == "[[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range(-4, 8, 1).tojson() == "[[],[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range( -5, 8, 1).tojson() == "[[6.6],[],[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range( -6, 8, 1).tojson() == "[[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range( -7, 8, 1).tojson() == "[[],[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range(-8, 8, 1).tojson( ) == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range(0, 8, 1).tojson( ) == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range( 0, 7, 1).tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[],[]]" assert array._ext.getitem_range( 0, 6, 1).tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[]]" assert array._ext.getitem_range( 0, 5, 1).tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[]]" assert array._ext.getitem_range( 0, 4, 1).tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6]]" assert array._ext.getitem_range( 0, 3, 1).tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5]]" assert array._ext.getitem_range(0, 2, 1).tojson() == "[[1.1,2.2,3.3],[]]" assert array._ext.getitem_range(0, 1, 1).tojson() == "[[1.1,2.2,3.3]]" assert array._ext.getitem_range(0, 0, 1).tojson() == "[]" assert array._ext.getitem_range(0, -8, 1).tojson() == "[]" assert array._ext.getitem_range(0, -7, 1).tojson() == "[[1.1,2.2,3.3]]" assert array._ext.getitem_range(0, -6, 1).tojson() == "[[1.1,2.2,3.3],[]]" assert array._ext.getitem_range( 0, -5, 1).tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5]]" assert array._ext.getitem_range( 0, -4, 1).tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6]]" assert array._ext.getitem_range( 0, -3, 1).tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[]]" assert array._ext.getitem_range( 0, -2, 1).tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[]]" assert array._ext.getitem_range( 0, -1, 1).tojson() == "[[1.1,2.2,3.3],[],[4.4,5.5],[6.6],[],[],[]]" assert array._ext.getitem_range( 0, 8, 2).tojson() == "[[1.1,2.2,3.3],[4.4,5.5],[],[]]" assert array._ext.getitem_range( 1, 8, 2).tojson() == "[[],[6.6],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range(1, 7, 2).tojson() == "[[],[6.6],[]]" assert array._ext.getitem_range(2, 8, 2).tojson() == "[[4.4,5.5],[],[]]" assert array._ext.getitem_range(3, 8, 2).tojson() == "[[6.6],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range(0, 8, 3).tojson() == "[[1.1,2.2,3.3],[6.6],[]]" assert array._ext.getitem_range(1, 8, 3).tojson() == "[[],[],[7.7,8.8,9.9]]" assert array._ext.getitem_range(1, 7, 3).tojson() == "[[],[]]" assert array._ext.getitem_range(2, 8, 3).tojson() == "[[4.4,5.5],[]]" assert array._ext.getitem_range(3, 8, 3).tojson() == "[[6.6],[]]" assert array._ext.getitem_range(2, 0, -1).tojson() == "[[4.4,5.5],[]]" assert array._ext.getitem_range( 2, None, -1).tojson() == "[[4.4,5.5],[],[1.1,2.2,3.3]]" assert array._ext.getitem_range(1, None, -1).tojson() == "[[],[1.1,2.2,3.3]]" assert array._ext.getitem_range( 3, None, -1).tojson() == "[[6.6],[4.4,5.5],[],[1.1,2.2,3.3]]" assert array._ext.getitem_range(None, None, -1).tojson( ) == "[[7.7,8.8,9.9],[],[],[],[6.6],[4.4,5.5],[],[1.1,2.2,3.3]]" assert array._ext.getitem_range(-1, None, -1).tojson( ) == "[[7.7,8.8,9.9],[],[],[],[6.6],[4.4,5.5],[],[1.1,2.2,3.3]]" assert array._ext.getitem_range( -2, None, -1).tojson() == "[[],[],[],[6.6],[4.4,5.5],[],[1.1,2.2,3.3]]" assert array._ext.getitem_range( -2, 0, -1).tojson() == "[[],[],[],[6.6],[4.4,5.5],[]]" assert array._ext.getitem_range( -2, 1, -1).tojson() == "[[],[],[],[6.6],[4.4,5.5]]" assert array._ext.getitem_range(-2, 2, -1).tojson() == "[[],[],[],[6.6]]" assert array._ext.getitem_range(-1, 3, -1).tojson() == "[[7.7,8.8,9.9],[],[],[]]" assert array._ext.getitem_range( -1, None, -2).tojson() == "[[7.7,8.8,9.9],[],[6.6],[]]" assert array._ext.getitem_range( -2, None, -2).tojson() == "[[],[],[4.4,5.5],[1.1,2.2,3.3]]" assert array._ext.getitem_range(-1, None, -3).tojson() == "[[7.7,8.8,9.9],[],[]]" assert array._ext.getitem_range(-2, None, -3).tojson() == "[[],[6.6],[1.1,2.2,3.3]]" assert array._ext.getitem_range(-3, None, -3).tojson() == "[[],[4.4,5.5]]" assert array._ext.getitem_range(-4, None, -3).tojson() == "[[],[]]" assert array._ext.getitem_range(-5, None, -3).tojson() == "[[6.6],[1.1,2.2,3.3]]" assert array._ext.getitem_range(-2, 0, -2).tojson() == "[[],[],[4.4,5.5]]" assert [awkward1.to_list(x) for x in array] == [[1.1, 2.2, 3.3], [], [4.4, 5.5], [6.6], [], [], [], [7.7, 8.8, 9.9]] assert awkward1.to_list(array.toContent()) == [[1.1, 2.2, 3.3], [], [4.4, 5.5], [6.6], [], [], [], [7.7, 8.8, 9.9]]
def test_0167_strings(): array = awkward1.repartition( awkward1.Array(["one", "two", "three", "two", "two", "one", "three"]), 3) assert awkward1.to_list(array == "two") == [ False, True, False, True, True, False, False ] assert awkward1.to_list("two" == array) == [ False, True, False, True, True, False, False ] assert awkward1.to_list(array == ["two"]) == [ False, True, False, True, True, False, False ] assert awkward1.to_list(["two"] == array) == [ False, True, False, True, True, False, False ] assert awkward1.to_list(array == awkward1.Array(["two"])) == [ False, True, False, True, True, False, False ] assert awkward1.to_list(awkward1.Array(["two"]) == array) == [ False, True, False, True, True, False, False ] assert awkward1.to_list(array < array) == [ [False, False, False], [False, False, False], [False, False, False, False, False], [False, False, False], [False, False, False], [False, False, False], [False, False, False, False, False] ] assert awkward1.to_list(array <= array) == [[True, True, True], [True, True, True], [True, True, True, True, True], [True, True, True], [True, True, True], [True, True, True], [True, True, True, True, True]] array = awkward1.Array([["one", "two", "three"], [], ["two"], ["two", "one"], ["three"]]) assert awkward1.to_list(array == "two") == [[False, True, False], [], [True], [True, False], [False]] assert awkward1.to_list("two" == array) == [[False, True, False], [], [True], [True, False], [False]] assert awkward1.to_list(array == ["two"]) == [[False, True, False], [], [True], [True, False], [False]] assert awkward1.to_list(["two"] == array) == [[False, True, False], [], [True], [True, False], [False]] assert awkward1.to_list(array == awkward1.Array(["two"])) == [[ False, True, False ], [], [True], [True, False], [False]] assert awkward1.to_list(awkward1.Array(["two"]) == array) == [[ False, True, False ], [], [True], [True, False], [False]] assert awkward1.to_list(array < array) == [ [[False, False, False], [False, False, False], [False, False, False, False, False]], [], [[False, False, False]], [[False, False, False], [False, False, False]], [[False, False, False, False, False]] ] assert awkward1.to_list(array <= array) == [ [[True, True, True], [True, True, True], [True, True, True, True, True]], [], [[True, True, True]], [[True, True, True], [True, True, True]], [[True, True, True, True, True]] ] array = awkward1.Array([["one", "two", "three"], [], ["two"], ["two", "one"], ["three"]]) assert awkward1.to_list( array == ["three", "two", "one", "one", "three"]) == [[ False, False, True ], [], [False], [False, True], [True]] assert awkward1.to_list( ["three", "two", "one", "one", "three"] == array) == [[ False, False, True ], [], [False], [False, True], [True]] assert awkward1.to_list( array == awkward1.Array(["three", "two", "one", "one", "three"])) == [[ False, False, True ], [], [False], [False, True], [True]] assert awkward1.to_list( awkward1.Array(["three", "two", "one", "one", "three"]) == array) == [[ False, False, True ], [], [False], [False, True], [True]]
def test_EmptyArray(): array = awkward1.layout.EmptyArray() assert awkward1.to_list(array.sort(0, True, False)) == [] assert awkward1.to_list(array.argsort(0, True, False)) == []
def test_argcombinations(): array = awkward1.Array([[0.0, 1.1, 2.2, 3.3], [], [4.4, 5.5, 6.6], [7.7], [8.8, 9.9, 10.0, 11.1, 12.2]]) assert awkward1.to_list(awkward1.argcombinations(array, 2, replacement=False)) == [[(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)], [], [(0, 1), (0, 2), (1, 2)], [], [(0, 1), (0, 2), (0, 3), (0, 4), (1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]]
def test_3d(): array = awkward1.layout.NumpyArray( numpy.array([ # axis 2: 0 1 2 3 4 # axis 1: [ [1.1, 2.2, 3.3, 4.4, 5.5], # 0 [6.6, 7.7, 8.8, 9.9, 10.10], # 1 [11.11, 12.12, 13.13, 14.14, 15.15] ], # 2 [ [-1.1, -2.2, -3.3, -4.4, -5.5], # 3 [-6.6, -7.7, -8.8, -9.9, -10.1], # 4 [-11.11, -12.12, -13.13, -14.14, -15.15] ] ])) # 5 assert awkward1.to_list(array.argsort(2, True, False)) == [[[0, 1, 2, 3, 4], [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], [[4, 3, 2, 1, 0], [4, 3, 2, 1, 0], [4, 3, 2, 1, 0]]] # np.sort(array, axis=2) # array([[[ 1.1 , 2.2 , 3.3 , 4.4 , 5.5 ], # [ 6.6 , 7.7 , 8.8 , 9.9 , 10.1 ], # [ 11.11, 12.12, 13.13, 14.14, 15.15]], # # [[ -5.5 , -4.4 , -3.3 , -2.2 , -1.1 ], # [-10.1 , -9.9 , -8.8 , -7.7 , -6.6 ], # [-15.15, -14.14, -13.13, -12.12, -11.11]]]) assert awkward1.to_list(array.sort( 2, True, False)) == [[[1.1, 2.2, 3.3, 4.4, 5.5], [6.6, 7.7, 8.8, 9.9, 10.10], [11.11, 12.12, 13.13, 14.14, 15.15]], [[-5.5, -4.4, -3.3, -2.2, -1.1], [-10.1, -9.9, -8.8, -7.7, -6.6], [-15.15, -14.14, -13.13, -12.12, -11.11]]] # np.argsort(array, 2) # array([[[0, 1, 2, 3, 4], # [0, 1, 2, 3, 4], # [0, 1, 2, 3, 4]], # # [[4, 3, 2, 1, 0], # [4, 3, 2, 1, 0], # [4, 3, 2, 1, 0]]]) assert awkward1.to_list(array.argsort(2, True, False)) == [[[0, 1, 2, 3, 4], [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], [[4, 3, 2, 1, 0], [4, 3, 2, 1, 0], [4, 3, 2, 1, 0]]] # np.sort(array, axis=1) # array([[[ 1.1 , 2.2 , 3.3 , 4.4 , 5.5 ], # [ 6.6 , 7.7 , 8.8 , 9.9 , 10.1 ], # [ 11.11, 12.12, 13.13, 14.14, 15.15]], # # [[-11.11, -12.12, -13.13, -14.14, -15.15], # [ -6.6 , -7.7 , -8.8 , -9.9 , -10.1 ], # [ -1.1 , -2.2 , -3.3 , -4.4 , -5.5 ]]]) assert awkward1.to_list(array.sort( 1, True, False)) == [[[1.1, 2.2, 3.3, 4.4, 5.5], [6.6, 7.7, 8.8, 9.9, 10.1], [11.11, 12.12, 13.13, 14.14, 15.15]], [[-11.11, -12.12, -13.13, -14.14, -15.15], [-6.6, -7.7, -8.8, -9.9, -10.1], [-1.1, -2.2, -3.3, -4.4, -5.5]]] assert awkward1.to_list(array.sort( 1, False, False)) == [[[11.11, 12.12, 13.13, 14.14, 15.15], [6.6, 7.7, 8.8, 9.9, 10.1], [1.1, 2.2, 3.3, 4.4, 5.5]], [[-1.1, -2.2, -3.3, -4.4, -5.5], [-6.6, -7.7, -8.8, -9.9, -10.1], [-11.11, -12.12, -13.13, -14.14, -15.15]]] # np.argsort(array, 1) # array([[[0, 0, 0, 0, 0], # [1, 1, 1, 1, 1], # [2, 2, 2, 2, 2]], # # [[2, 2, 2, 2, 2], # [1, 1, 1, 1, 1], # [0, 0, 0, 0, 0]]]) assert awkward1.to_list(array.argsort(1, True, False)) == [[[0, 0, 0, 0, 0], [1, 1, 1, 1, 1], [2, 2, 2, 2, 2]], [[2, 2, 2, 2, 2], [1, 1, 1, 1, 1], [0, 0, 0, 0, 0]]] # np.sort(array, axis=0) # array([[[ -1.1 , -2.2 , -3.3 , -4.4 , -5.5 ], # [ -6.6 , -7.7 , -8.8 , -9.9 , -10.1 ], # [-11.11, -12.12, -13.13, -14.14, -15.15]], # # [[ 1.1 , 2.2 , 3.3 , 4.4 , 5.5 ], # [ 6.6 , 7.7 , 8.8 , 9.9 , 10.1 ], # [ 11.11, 12.12, 13.13, 14.14, 15.15]]]) assert awkward1.to_list(array.sort( 0, True, False)) == [[[-1.1, -2.2, -3.3, -4.4, -5.5], [-6.6, -7.7, -8.8, -9.9, -10.1], [-11.11, -12.12, -13.13, -14.14, -15.15]], [[1.1, 2.2, 3.3, 4.4, 5.5], [6.6, 7.7, 8.8, 9.9, 10.1], [11.11, 12.12, 13.13, 14.14, 15.15]]] # np.argsort(array, 0) # array([[[1, 1, 1, 1, 1], # [1, 1, 1, 1, 1], # [1, 1, 1, 1, 1]], # # [[0, 0, 0, 0, 0], # [0, 0, 0, 0, 0], # [0, 0, 0, 0, 0]]]) assert awkward1.to_list(array.argsort(0, True, False)) == [[[1, 1, 1, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1]], [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]]
def test(): out = awkward1.ArrayBuilder() out.begin_record() if True: out.field("x") out.integer(3) out.field("extreme") out.begin_record() if True: out.field("pt") out.real(3.3) out.field("charge") out.integer(-1) out.field("iso") out.integer(100) out.end_record() out.end_record() out.begin_record() if True: out.field("x") out.integer(3) out.end_record() ss = out.snapshot() assert awkward1.to_list(ss) == [{ "x": 3, "extreme": { "pt": 3.3, "charge": -1, "iso": 100 } }, { "x": 3, "extreme": None }] assert awkward1.to_list( awkward1.Array([{ "x": 3, "extreme": { "pt": 3.3, "charge": -1, "iso": 100 } }, { "x": 3 }], check_valid=True)) == [{ "x": 3, "extreme": { "pt": 3.3, "charge": -1, "iso": 100 } }, { "x": 3, "extreme": None }] assert awkward1.to_list( awkward1.Array([{ "x": 3, "extreme": { "pt": 3.3, "charge": -1, "iso": 100 } }, { "x": 3, "what": 3 }], check_valid=True)) == [{ "x": 3, "extreme": { "pt": 3.3, "charge": -1, "iso": 100 }, "what": None }, { "x": 3, "extreme": None, "what": 3 }]
def test_IndexedOffsetArray(): array = awkward1.Array([[2.2, 1.1, 3.3], [None, None, None], [4.4, None, 5.5], [5.5, None, None], [-4.4, -5.5, -6.6]]).layout # np.sort(array, 1) # array([[ 1.1, 2.2, 3.3], # [1000. , 1000. , 1000. ], # [ 4.4, 5.5, 1000. ], # [ 5.5, 1000. , 1000. ], # [ -6.6, -5.5, -4.4]]) assert awkward1.to_list(array.sort(0, True, False)) == [[-4.4, -5.5, -6.6], [2.2, 1.1, 3.3], [4.4, None, 5.5], [5.5, None, None], [None, None, None]] assert awkward1.to_list(array.sort(1, True, False)) == [[1.1, 2.2, 3.3], [None, None, None], [4.4, 5.5, None], [5.5, None, None], [-6.6, -5.5, -4.4]] assert awkward1.to_list(array.sort(1, False, False)) == [[3.3, 2.2, 1.1], [None, None, None], [5.5, 4.4, None], [5.5, None, None], [-4.4, -5.5, -6.6]] # np.argsort(array, 1) # array([[1, 0, 2], # [0, 1, 2], # [0, 2, 1], # [0, 1, 2], # [2, 1, 0]]) assert awkward1.to_list(array.argsort(1, True, False)) == [[1, 0, 2], [None, None, None], [0, 1, None], [0, None, None], [2, 1, 0]] # np.sort(array, 0) # array([[ -4.4, -5.5, -6.6], # [ 2.2, 1.1, 3.3], # [ 4.4, 1000. , 5.5], # [ 5.5, 1000. , 1000. ], # [1000. , 1000. , 1000. ]]) assert awkward1.to_list(array.sort(0, True, False)) == [[-4.4, -5.5, -6.6], [2.2, 1.1, 3.3], [4.4, None, 5.5], [5.5, None, None], [None, None, None]] # np.argsort(array, 0) # array([[4, 4, 4], # [0, 0, 0], # [2, 1, 2], # [3, 2, 1], # [1, 3, 3]]) assert awkward1.to_list(array.argsort(0, True, False)) == [[3, 1, 2], [0, 0, 0], [1, None, 1], [2, None, None], [None, None, None]] array3 = awkward1.Array([[2.2, 1.1, 3.3], [], [4.4, 5.5], [5.5], [-4.4, -5.5, -6.6]]).layout assert awkward1.to_list(array3.sort(1, False, False)) == [[3.3, 2.2, 1.1], [], [5.5, 4.4], [5.5], [-4.4, -5.5, -6.6]] assert awkward1.to_list(array3.sort(0, True, False)) == [[-4.4, -5.5, -6.6], [], [2.2, 1.1], [4.4], [5.5, 5.5, 3.3]] #FIXME: Based on Numpy list sorting: # # array([list([2.2, 1.1, 3.3]), list([]), list([4.4, 5.5]), list([5.5]), # list([-4.4, -5.5, -6.6])], dtype=object) # np.sort(array, axis=0) # array([list([]), list([-4.4, -5.5, -6.6]), list([2.2, 1.1, 3.3]), # list([4.4, 5.5]), list([5.5])], dtype=object) # # the result should be: # # [[ -4.4, -5.5, -6.6 ], # [ 2.2, 1.1, 3.3 ], # [ 4.4, 5.5 ], # [ 5.5 ], # []] # This can be done following the steps: pad, sort, # and dropna to strip off the None's # array4 = array3.rpad(3, 1) assert awkward1.to_list(array4) == [[2.2, 1.1, 3.3], [None, None, None], [4.4, 5.5, None], [5.5, None, None], [-4.4, -5.5, -6.6]] array5 = array4.sort(0, True, False) assert awkward1.to_list(array5) == [[-4.4, -5.5, -6.6], [2.2, 1.1, 3.3], [4.4, 5.5, None], [5.5, None, None], [None, None, None]] # FIXME: implement dropna to strip off the None's # # array6 = array5.dropna(0) # assert awkward1.to_list(array6) == [ # [ -4.4, -5.5, -6.6 ], # [ 2.2, 1.1, 3.3 ], # [ 4.4, 5.5 ], # [ 5.5 ], # []] # FIXME: assert awkward1.to_list(array.argsort(1, True, False)) == [[1, 0, 2], [], [0, 1], [0], [2, 1, 0]] # [[1, 0, 2], [None, None, None], [-3, -2, None], [-4, None, None], [-4, -5, -6]] # FIXME: assert awkward1.to_list(array.argsort(0, True, False)) == [[3, 0, 1, 2], [2, 0, 1], [1, 0]] # [[3, 0, -2], [0, -1, -4], [1, None, -3], [2, None, None], [None, None, None]] content = awkward1.layout.NumpyArray(numpy.array([1.1, 2.2, 3.3, 4.4, 5.5])) index1 = awkward1.layout.Index32( numpy.array([1, 2, 3, 4], dtype=numpy.int32)) indexedarray1 = awkward1.layout.IndexedArray32(index1, content) assert awkward1.to_list(indexedarray1.argsort(0, True, False)) == [0, 1, 2, 3] index2 = awkward1.layout.Index64(numpy.array([1, 2, 3], dtype=numpy.int64)) indexedarray2 = awkward1.layout.IndexedArray64(index2, indexedarray1) assert awkward1.to_list(indexedarray2.sort(0, False, False)) == [5.5, 4.4, 3.3] index3 = awkward1.layout.Index32(numpy.array([1, 2], dtype=numpy.int32)) indexedarray3 = awkward1.layout.IndexedArray32(index3, indexedarray2) assert awkward1.to_list(indexedarray3.sort(0, True, False)) == [4.4, 5.5]
def test_write_read(tmp_path): array1 = awkward1.Array([[1, 2, 3], [], [4, 5], [], [], [6, 7, 8, 9]]) array2 = awkward1.repartition(array1, 2) array3 = awkward1.Array([ { "x": 1, "y": 1.1 }, { "x": 2, "y": 2.2 }, { "x": 3, "y": 3.3 }, { "x": 4, "y": 4.4 }, { "x": 5, "y": 5.5 }, { "x": 6, "y": 6.6 }, { "x": 7, "y": 7.7 }, { "x": 8, "y": 8.8 }, { "x": 9, "y": 9.9 }, ]) array4 = awkward1.repartition(array3, 2) awkward1.to_parquet(array1, os.path.join(tmp_path, "array1.parquet")) awkward1.to_parquet(array2, os.path.join(tmp_path, "array2.parquet")) awkward1.to_parquet(array3, os.path.join(tmp_path, "array3.parquet")) awkward1.to_parquet(array4, os.path.join(tmp_path, "array4.parquet")) assert awkward1.to_list( awkward1.from_parquet(os.path.join( tmp_path, "array1.parquet"))) == awkward1.to_list(array1) assert awkward1.to_list( awkward1.from_parquet(os.path.join( tmp_path, "array2.parquet"))) == awkward1.to_list(array2) assert awkward1.to_list( awkward1.from_parquet(os.path.join( tmp_path, "array3.parquet"))) == awkward1.to_list(array3) assert awkward1.to_list( awkward1.from_parquet(os.path.join( tmp_path, "array4.parquet"))) == awkward1.to_list(array4) assert awkward1.to_list( awkward1.from_parquet(os.path.join(tmp_path, "array1.parquet"), lazy=True)) == awkward1.to_list(array1) assert awkward1.to_list( awkward1.from_parquet(os.path.join(tmp_path, "array2.parquet"), lazy=True)) == awkward1.to_list(array2) assert awkward1.to_list( awkward1.from_parquet(os.path.join(tmp_path, "array3.parquet"), lazy=True)) == awkward1.to_list(array3) assert awkward1.to_list( awkward1.from_parquet(os.path.join(tmp_path, "array4.parquet"), lazy=True)) == awkward1.to_list(array4)
def test_explicit_broadcasting(): nparray = numpy.arange(2*3*5).reshape(2, 3, 5) lsarray = awkward1.Array(nparray.tolist(), check_valid=True) rgarray = awkward1.Array(nparray, check_valid=True) # explicit left-broadcasting assert awkward1.to_list(rgarray + numpy.array([[[100]], [[200]]])) == awkward1.to_list(nparray + numpy.array([[[100]], [[200]]])) assert awkward1.to_list(lsarray + numpy.array([[[100]], [[200]]])) == awkward1.to_list(nparray + numpy.array([[[100]], [[200]]])) assert awkward1.to_list(numpy.array([[[100]], [[200]]]) + rgarray) == awkward1.to_list(numpy.array([[[100]], [[200]]]) + nparray) assert awkward1.to_list(numpy.array([[[100]], [[200]]]) + lsarray) == awkward1.to_list(numpy.array([[[100]], [[200]]]) + nparray) # explicit right-broadcasting assert awkward1.to_list(rgarray + numpy.array([[[100, 200, 300, 400, 500]]])) == awkward1.to_list(nparray + numpy.array([[[100, 200, 300, 400, 500]]])) assert awkward1.to_list(lsarray + numpy.array([[[100, 200, 300, 400, 500]]])) == awkward1.to_list(nparray + numpy.array([[[100, 200, 300, 400, 500]]])) assert awkward1.to_list(numpy.array([[[100, 200, 300, 400, 500]]]) + rgarray) == awkward1.to_list(numpy.array([[[100, 200, 300, 400, 500]]]) + nparray) assert awkward1.to_list(numpy.array([[[100, 200, 300, 400, 500]]]) + lsarray) == awkward1.to_list(numpy.array([[[100, 200, 300, 400, 500]]]) + nparray)
def test_oamap_samples(): assert awkward1.to_list( awkward1.from_parquet("tests/samples/list-depths-simple.parquet")) == [ { "list0": 1, "list1": [1] }, { "list0": 2, "list1": [1, 2] }, { "list0": 3, "list1": [1, 2, 3] }, { "list0": 4, "list1": [1, 2, 3, 4] }, { "list0": 5, "list1": [1, 2, 3, 4, 5] }, ] assert awkward1.to_list( awkward1.from_parquet( "tests/samples/nullable-record-primitives.parquet")) == [ { "u1": None, "u4": 1, "u8": None, "f4": 1.100000023841858, "f8": None, "raw": b"one", "utf8": "one", }, { "u1": 1, "u4": None, "u8": 2, "f4": 2.200000047683716, "f8": None, "raw": None, "utf8": None, }, { "u1": None, "u4": None, "u8": 3, "f4": None, "f8": None, "raw": b"three", "utf8": None, }, { "u1": 0, "u4": None, "u8": 4, "f4": None, "f8": 4.4, "raw": None, "utf8": None, }, { "u1": None, "u4": 5, "u8": None, "f4": None, "f8": 5.5, "raw": None, "utf8": "five", }, ] assert awkward1.to_list( awkward1.from_parquet( "tests/samples/nullable-record-primitives-simple.parquet")) == [ { "u4": None, "u8": 1 }, { "u4": None, "u8": 2 }, { "u4": None, "u8": 3 }, { "u4": None, "u8": 4 }, { "u4": None, "u8": 5 }, ] assert awkward1.to_list( awkward1.from_parquet("tests/samples/record-primitives.parquet")) == [ { "u1": 0, "u4": 1, "u8": 1, "f4": 1.100000023841858, "f8": 1.1, "raw": b"one", "utf8": "one", }, { "u1": 1, "u4": 2, "u8": 2, "f4": 2.200000047683716, "f8": 2.2, "raw": b"two", "utf8": "two", }, { "u1": 1, "u4": 3, "u8": 3, "f4": 3.299999952316284, "f8": 3.3, "raw": b"three", "utf8": "three", }, { "u1": 0, "u4": 4, "u8": 4, "f4": 4.400000095367432, "f8": 4.4, "raw": b"four", "utf8": "four", }, { "u1": 0, "u4": 5, "u8": 5, "f4": 5.5, "f8": 5.5, "raw": b"five", "utf8": "five", }, ]
def test_ArrayBuilder_append(): array = awkward1.Array( [[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], check_valid=True) builder = awkward1.ArrayBuilder() builder.append(array, 3) builder.append(array, 2) builder.append(array, 2) builder.append(array, 0) builder.append(array, 1) builder.append(array, -1) assert awkward1.to_list(builder.snapshot()) == [[5.5], [3.3, 4.4], [3.3, 4.4], [0.0, 1.1, 2.2], [], [6.6, 7.7, 8.8, 9.9]] builder.extend(array) assert awkward1.to_list(builder.snapshot()) == [[5.5], [3.3, 4.4], [3.3, 4.4], [0.0, 1.1, 2.2], [], [6.6, 7.7, 8.8, 9.9], [0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]] builder = awkward1.ArrayBuilder() builder.null() builder.null() builder.null() builder.append(array, 3) builder.append(array, 2) builder.append(array, 2) builder.append(array, -1) assert awkward1.to_list(builder.snapshot()) == [ None, None, None, [5.5], [3.3, 4.4], [3.3, 4.4], [6.6, 7.7, 8.8, 9.9] ] builder.null() assert awkward1.to_list(builder.snapshot()) == [ None, None, None, [5.5], [3.3, 4.4], [3.3, 4.4], [6.6, 7.7, 8.8, 9.9], None ] one = awkward1.Array( [[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], check_valid=True) two = awkward1.Array([[3.3, 2.2, 1.1, 0.0], [5.5, 4.4], [], [6.6]], check_valid=True) builder = awkward1.ArrayBuilder() builder.append(one, 2) builder.append(two, 1) builder.append(one, 0) builder.append(two, -1) builder.append(one, -1) assert awkward1.to_list(builder.snapshot()) == [[3.3, 4.4], [5.5, 4.4], [0.0, 1.1, 2.2], [6.6], [6.6, 7.7, 8.8, 9.9]] builder = awkward1.ArrayBuilder() builder.null() builder.append(one, 2) builder.null() builder.append(two, 1) builder.null() assert awkward1.to_list( builder.snapshot()) == [None, [3.3, 4.4], None, [5.5, 4.4], None] builder = awkward1.ArrayBuilder() builder.string("hello") builder.append(one, 2) builder.string("there") builder.append(one, 0) assert awkward1.to_list( builder.snapshot()) == ["hello", [3.3, 4.4], "there", [0.0, 1.1, 2.2]] builder = awkward1.ArrayBuilder() builder.null() builder.string("hello") builder.null() builder.append(one, 2) builder.null() builder.string("there") builder.append(one, 0) assert awkward1.to_list(builder.snapshot()) == [ None, "hello", None, [3.3, 4.4], None, "there", [0.0, 1.1, 2.2] ] builder = awkward1.ArrayBuilder() builder.append(one, 2) builder.string("there") builder.append(one, 0) assert awkward1.to_list(builder.snapshot()) == [[3.3, 4.4], "there", [0.0, 1.1, 2.2]] builder = awkward1.ArrayBuilder() builder.null() builder.append(one, 2) builder.null() builder.string("there") builder.null() builder.append(one, 0) assert awkward1.to_list(builder.snapshot()) == [ None, [3.3, 4.4], None, "there", None, [0.0, 1.1, 2.2] ] array = awkward1.Array([ "zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine" ], check_valid=True) builder = awkward1.ArrayBuilder() builder.begin_list() builder.append(array, 1) builder.append(array, 2) builder.append(array, 3) builder.end_list() builder.begin_list() builder.end_list() builder.begin_list() builder.append(array, 4) builder.append(array, 5) builder.end_list() assert awkward1.to_list(builder.snapshot()) == [["one", "two", "three"], [], ["four", "five"]] builder.append(array, -1) assert awkward1.to_list(builder.snapshot()) == [["one", "two", "three"], [], ["four", "five"], "nine"] array = awkward1.Array([{ "x": 0.0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": [2, 2] }, { "x": 3.3, "y": [3, 3, 3] }], check_valid=True) builder = awkward1.ArrayBuilder() builder.append(array[2]) builder.append(array[2]) builder.append(array[1]) builder.append(array[-1]) tmp = builder.snapshot() assert awkward1.to_list(tmp) == [{ "x": 2.2, "y": [2, 2] }, { "x": 2.2, "y": [2, 2] }, { "x": 1.1, "y": [1] }, { "x": 3.3, "y": [3, 3, 3] }] assert isinstance(tmp.layout, awkward1.layout.IndexedArray64) assert isinstance(tmp.layout.content, awkward1.layout.RecordArray) builder.extend(array) tmp = builder.snapshot() assert awkward1.to_list(tmp) == [{ "x": 2.2, "y": [2, 2] }, { "x": 2.2, "y": [2, 2] }, { "x": 1.1, "y": [1] }, { "x": 3.3, "y": [3, 3, 3] }, { "x": 0.0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": [2, 2] }, { "x": 3.3, "y": [3, 3, 3] }] assert isinstance(tmp.layout, awkward1.layout.IndexedArray64) assert isinstance(tmp.layout.content, awkward1.layout.RecordArray) builder.append(999) tmp = builder.snapshot() assert awkward1.to_list(tmp) == [{ "x": 2.2, "y": [2, 2] }, { "x": 2.2, "y": [2, 2] }, { "x": 1.1, "y": [1] }, { "x": 3.3, "y": [3, 3, 3] }, { "x": 0.0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": [2, 2] }, { "x": 3.3, "y": [3, 3, 3] }, 999] assert isinstance(tmp.layout, awkward1.layout.UnionArray8_64) builder.append([1, 2, 3, 4, 5]) tmp = builder.snapshot() assert awkward1.to_list(tmp) == [{ "x": 2.2, "y": [2, 2] }, { "x": 2.2, "y": [2, 2] }, { "x": 1.1, "y": [1] }, { "x": 3.3, "y": [3, 3, 3] }, { "x": 0.0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": [2, 2] }, { "x": 3.3, "y": [3, 3, 3] }, 999, [1, 2, 3, 4, 5]] assert isinstance(tmp.layout, awkward1.layout.UnionArray8_64) array1 = awkward1.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]], check_valid=True) builder = awkward1.ArrayBuilder() builder.append(array1, 2) builder.append(array1, 1) builder.append(array1, 0) array2 = builder.snapshot() assert isinstance(array2.layout.content, awkward1.layout.ListOffsetArray64) builder = awkward1.ArrayBuilder() builder.append(array2, 2) builder.append(array2, 1) builder.append(array2, 0) array3 = builder.snapshot() assert isinstance(array3.layout.content, awkward1.layout.ListOffsetArray64) builder = awkward1.ArrayBuilder() builder.append(array3, 2) builder.append(array3, 1) builder.append(array3, 0) array4 = builder.snapshot() assert isinstance(array4.layout.content, awkward1.layout.ListOffsetArray64)
def test_recordarray_merge(): emptyarray = awkward1.layout.EmptyArray() arrayr1 = awkward1.from_iter([{ "x": 0, "y": [] }, { "x": 1, "y": [1, 1] }, { "x": 2, "y": [2, 2] }], highlevel=False) arrayr2 = awkward1.from_iter([{ "x": 2.2, "y": [2.2, 2.2] }, { "x": 1.1, "y": [1.1, 1.1] }, { "x": 0.0, "y": [0.0, 0.0] }], highlevel=False) arrayr3 = awkward1.from_iter([{ "x": 0, "y": 0.0 }, { "x": 1, "y": 1.1 }, { "x": 2, "y": 2.2 }], highlevel=False) arrayr4 = awkward1.from_iter([{ "y": [], "x": 0 }, { "y": [1, 1], "x": 1 }, { "y": [2, 2], "x": 2 }], highlevel=False) arrayr5 = awkward1.from_iter([{ "x": 0, "y": [], "z": 0 }, { "x": 1, "y": [1, 1], "z": 1 }, { "x": 2, "y": [2, 2], "z": 2 }], highlevel=False) arrayr6 = awkward1.from_iter([{ "z": 0, "x": 0, "y": [] }, { "z": 1, "x": 1, "y": [1, 1] }, { "z": 2, "x": 2, "y": [2, 2] }], highlevel=False) arrayr7 = awkward1.from_iter([{ "x": 0 }, { "x": 1 }, { "x": 2 }], highlevel=False) arrayt1 = awkward1.from_iter([(0, []), (1, [1.1]), (2, [2, 2])], highlevel=False) arrayt2 = awkward1.from_iter([(2.2, [2.2, 2.2]), (1.1, [1.1, 1.1]), (0.0, [0.0, 0.0])], highlevel=False) arrayt3 = awkward1.from_iter([(0, 0.0), (1, 1.1), (2, 2.2)], highlevel=False) arrayt4 = awkward1.from_iter([([], 0), ([1.1], 1), ([2.2, 2.2], 2)], highlevel=False) arrayt5 = awkward1.from_iter([(0, [], 0), (1, [1], 1), (2, [2, 2], 2)], highlevel=False) arrayt6 = awkward1.from_iter([(0, 0, []), (1, 1, [1]), (2, 2, [2, 2])], highlevel=False) arrayt7 = awkward1.from_iter([(0, ), (1, ), (2, )], highlevel=False) assert arrayr1.mergeable(arrayr2) assert arrayr2.mergeable(arrayr1) assert not arrayr1.mergeable(arrayr3) assert arrayr1.mergeable(arrayr4) assert arrayr4.mergeable(arrayr1) assert not arrayr1.mergeable(arrayr5) assert not arrayr1.mergeable(arrayr6) assert arrayr5.mergeable(arrayr6) assert arrayr6.mergeable(arrayr5) assert not arrayr1.mergeable(arrayr7) assert arrayt1.mergeable(arrayt2) assert arrayt2.mergeable(arrayt1) assert not arrayt1.mergeable(arrayt3) assert not arrayt1.mergeable(arrayt4) assert not arrayt1.mergeable(arrayt5) assert not arrayt1.mergeable(arrayt6) assert not arrayt5.mergeable(arrayt6) assert not arrayt1.mergeable(arrayt7) assert awkward1.to_list(arrayr1.merge(arrayr2)) == [{ "x": 0.0, "y": [] }, { "x": 1.0, "y": [1.0, 1.0] }, { "x": 2.0, "y": [2.0, 2.0] }, { "x": 2.2, "y": [2.2, 2.2] }, { "x": 1.1, "y": [1.1, 1.1] }, { "x": 0.0, "y": [0.0, 0.0] }] assert awkward1.to_list(arrayr2.merge(arrayr1)) == [{ "x": 2.2, "y": [2.2, 2.2] }, { "x": 1.1, "y": [1.1, 1.1] }, { "x": 0.0, "y": [0.0, 0.0] }, { "x": 0.0, "y": [] }, { "x": 1.0, "y": [1.0, 1.0] }, { "x": 2.0, "y": [2.0, 2.0] }] assert awkward1.to_list(arrayr1.merge(arrayr4)) == [{ "x": 0, "y": [] }, { "x": 1, "y": [1, 1] }, { "x": 2, "y": [2, 2] }, { "x": 0, "y": [] }, { "x": 1, "y": [1, 1] }, { "x": 2, "y": [2, 2] }] assert awkward1.to_list(arrayr4.merge(arrayr1)) == [{ "x": 0, "y": [] }, { "x": 1, "y": [1, 1] }, { "x": 2, "y": [2, 2] }, { "x": 0, "y": [] }, { "x": 1, "y": [1, 1] }, { "x": 2, "y": [2, 2] }] assert awkward1.to_list(arrayr5.merge(arrayr6)) == [{ "x": 0, "y": [], "z": 0 }, { "x": 1, "y": [1, 1], "z": 1 }, { "x": 2, "y": [2, 2], "z": 2 }, { "x": 0, "y": [], "z": 0 }, { "x": 1, "y": [1, 1], "z": 1 }, { "x": 2, "y": [2, 2], "z": 2 }] assert awkward1.to_list(arrayr6.merge(arrayr5)) == [{ "x": 0, "y": [], "z": 0 }, { "x": 1, "y": [1, 1], "z": 1 }, { "x": 2, "y": [2, 2], "z": 2 }, { "x": 0, "y": [], "z": 0 }, { "x": 1, "y": [1, 1], "z": 1 }, { "x": 2, "y": [2, 2], "z": 2 }] assert awkward1.to_list(arrayt1.merge(arrayt2)) == [(0.0, []), (1.0, [1.1]), (2.0, [2.0, 2.0]), (2.2, [2.2, 2.2]), (1.1, [1.1, 1.1]), (0.0, [0.0, 0.0])] assert awkward1.to_list(arrayt2.merge(arrayt1)) == [(2.2, [2.2, 2.2]), (1.1, [1.1, 1.1]), (0.0, [0.0, 0.0]), (0.0, []), (1.0, [1.1]), (2.0, [2.0, 2.0])] assert awkward1.to_list( arrayr1.merge(emptyarray)) == awkward1.to_list(arrayr1) assert awkward1.to_list( arrayr2.merge(emptyarray)) == awkward1.to_list(arrayr2) assert awkward1.to_list( arrayr3.merge(emptyarray)) == awkward1.to_list(arrayr3) assert awkward1.to_list( arrayr4.merge(emptyarray)) == awkward1.to_list(arrayr4) assert awkward1.to_list( arrayr5.merge(emptyarray)) == awkward1.to_list(arrayr5) assert awkward1.to_list( arrayr6.merge(emptyarray)) == awkward1.to_list(arrayr6) assert awkward1.to_list( arrayr7.merge(emptyarray)) == awkward1.to_list(arrayr7) assert awkward1.to_list( emptyarray.merge(arrayr1)) == awkward1.to_list(arrayr1) assert awkward1.to_list( emptyarray.merge(arrayr2)) == awkward1.to_list(arrayr2) assert awkward1.to_list( emptyarray.merge(arrayr3)) == awkward1.to_list(arrayr3) assert awkward1.to_list( emptyarray.merge(arrayr4)) == awkward1.to_list(arrayr4) assert awkward1.to_list( emptyarray.merge(arrayr5)) == awkward1.to_list(arrayr5) assert awkward1.to_list( emptyarray.merge(arrayr6)) == awkward1.to_list(arrayr6) assert awkward1.to_list( emptyarray.merge(arrayr7)) == awkward1.to_list(arrayr7) assert awkward1.to_list( arrayt1.merge(emptyarray)) == awkward1.to_list(arrayt1) assert awkward1.to_list( arrayt2.merge(emptyarray)) == awkward1.to_list(arrayt2) assert awkward1.to_list( arrayt3.merge(emptyarray)) == awkward1.to_list(arrayt3) assert awkward1.to_list( arrayt4.merge(emptyarray)) == awkward1.to_list(arrayt4) assert awkward1.to_list( arrayt5.merge(emptyarray)) == awkward1.to_list(arrayt5) assert awkward1.to_list( arrayt6.merge(emptyarray)) == awkward1.to_list(arrayt6) assert awkward1.to_list( arrayt7.merge(emptyarray)) == awkward1.to_list(arrayt7) assert awkward1.to_list( emptyarray.merge(arrayt1)) == awkward1.to_list(arrayt1) assert awkward1.to_list( emptyarray.merge(arrayt2)) == awkward1.to_list(arrayt2) assert awkward1.to_list( emptyarray.merge(arrayt3)) == awkward1.to_list(arrayt3) assert awkward1.to_list( emptyarray.merge(arrayt4)) == awkward1.to_list(arrayt4) assert awkward1.to_list( emptyarray.merge(arrayt5)) == awkward1.to_list(arrayt5) assert awkward1.to_list( emptyarray.merge(arrayt6)) == awkward1.to_list(arrayt6) assert awkward1.to_list( emptyarray.merge(arrayt7)) == awkward1.to_list(arrayt7)
def test_NumpyArray_getitem(): array = awkward1.Array([1.1, 2.2, 3.3, 4.4, 5.5], check_valid=True) @numba.njit def f1(x, i): return x[i] assert f1(array, 0) == 1.1 assert f1(array, 1) == 2.2 assert f1(array, 2) == 3.3 assert f1(array, 3) == 4.4 assert f1(array, 4) == 5.5 assert f1(array, -1) == 5.5 assert f1(array, -2) == 4.4 assert f1(array, -3) == 3.3 assert f1(array, -4) == 2.2 assert f1(array, -5) == 1.1 with pytest.raises(ValueError) as err: assert f1(array, 5) assert str(err.value) == "slice index out of bounds" with pytest.raises(ValueError) as err: assert f1(array, -6) assert str(err.value) == "slice index out of bounds" @numba.njit def f2(x, i1, i2): return x[i1:i2] assert awkward1.to_list(f2(array, 0, 5)) == [1.1, 2.2, 3.3, 4.4, 5.5] assert awkward1.to_list(f2(array, 1, 5)) == [2.2, 3.3, 4.4, 5.5] assert awkward1.to_list(f2(array, 2, 5)) == [3.3, 4.4, 5.5] assert awkward1.to_list(f2(array, 3, 5)) == [4.4, 5.5] assert awkward1.to_list(f2(array, 4, 5)) == [5.5] assert awkward1.to_list(f2(array, 5, 5)) == [] assert awkward1.to_list(f2(array, 6, 5)) == [] assert awkward1.to_list(f2(array, -1, 5)) == [5.5] assert awkward1.to_list(f2(array, -2, 5)) == [4.4, 5.5] assert awkward1.to_list(f2(array, -3, 5)) == [3.3, 4.4, 5.5] assert awkward1.to_list(f2(array, -4, 5)) == [2.2, 3.3, 4.4, 5.5] assert awkward1.to_list(f2(array, -5, 5)) == [1.1, 2.2, 3.3, 4.4, 5.5] assert awkward1.to_list(f2(array, -6, 5)) == [1.1, 2.2, 3.3, 4.4, 5.5] assert awkward1.to_list(f2(array, 0, -6)) == [] assert awkward1.to_list(f2(array, 0, -5)) == [] assert awkward1.to_list(f2(array, 0, -4)) == [1.1] assert awkward1.to_list(f2(array, 0, -3)) == [1.1, 2.2] assert awkward1.to_list(f2(array, 0, -2)) == [1.1, 2.2, 3.3] assert awkward1.to_list(f2(array, 0, -1)) == [1.1, 2.2, 3.3, 4.4] assert awkward1.to_list(f2(array, 0, 6)) == [1.1, 2.2, 3.3, 4.4, 5.5] assert awkward1.to_list(f2(array, 0, 5)) == [1.1, 2.2, 3.3, 4.4, 5.5] assert awkward1.to_list(f2(array, 0, 4)) == [1.1, 2.2, 3.3, 4.4] assert awkward1.to_list(f2(array, 0, 3)) == [1.1, 2.2, 3.3] assert awkward1.to_list(f2(array, 0, 2)) == [1.1, 2.2] assert awkward1.to_list(f2(array, 0, 1)) == [1.1] assert awkward1.to_list(f2(array, 0, 0)) == [] aslist = [1.1, 2.2, 3.3, 4.4, 5.5] for i1 in range(-6, 7): for i2 in range(-6, 7): assert awkward1.to_list(f2(array, i1, i2)) == aslist[i1:i2] @numba.njit def f3(x, i1, i2): return x[1:4][i1:i2] assert awkward1.to_list(f3(array, -1, 3)) == [4.4] assert awkward1.to_list(f3(array, -2, 3)) == [3.3, 4.4] assert awkward1.to_list(f3(array, -3, 3)) == [2.2, 3.3, 4.4] assert awkward1.to_list(f3(array, 0, 3)) == [2.2, 3.3, 4.4] assert awkward1.to_list(f3(array, 1, 3)) == [3.3, 4.4] assert awkward1.to_list(f3(array, 2, 3)) == [4.4] assert awkward1.to_list(f3(array, 3, 3)) == [] assert awkward1.to_list(f3(array, 0, -4)) == [] assert awkward1.to_list(f3(array, 0, -3)) == [] assert awkward1.to_list(f3(array, 0, -2)) == [2.2] assert awkward1.to_list(f3(array, 0, -1)) == [2.2, 3.3] assert awkward1.to_list(f3(array, 0, 3)) == [2.2, 3.3, 4.4] assert awkward1.to_list(f3(array, 0, 2)) == [2.2, 3.3] assert awkward1.to_list(f3(array, 0, 1)) == [2.2] assert awkward1.to_list(f3(array, 0, 0)) == []
def test_indexedarray_simplify_more(): content = awkward1.layout.NumpyArray( numpy.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) index1_32 = awkward1.layout.Index32( numpy.array([6, 5, 4, 3, 2, 1, 0], dtype=numpy.int32)) index1_U32 = awkward1.layout.IndexU32( numpy.array([6, 5, 4, 3, 2, 1, 0], dtype=numpy.uint32)) index1_64 = awkward1.layout.Index64( numpy.array([6, 5, 4, 3, 2, 1, 0], dtype=numpy.int64)) index2_32 = awkward1.layout.Index32( numpy.array([0, 2, 4, 6], dtype=numpy.int32)) index2_U32 = awkward1.layout.IndexU32( numpy.array([0, 2, 4, 6], dtype=numpy.uint32)) index2_64 = awkward1.layout.Index64( numpy.array([0, 2, 4, 6], dtype=numpy.int64)) array = awkward1.layout.IndexedArray32( index2_32, awkward1.layout.IndexedArray32(index1_32, content)) assert awkward1.to_list(array) == [6.6, 4.4, 2.2, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, 4.4, 2.2, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedArray32( index2_32, awkward1.layout.IndexedArrayU32(index1_U32, content)) assert awkward1.to_list(array) == [6.6, 4.4, 2.2, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, 4.4, 2.2, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedArray32( index2_32, awkward1.layout.IndexedArray64(index1_64, content)) assert awkward1.to_list(array) == [6.6, 4.4, 2.2, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, 4.4, 2.2, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedArrayU32( index2_U32, awkward1.layout.IndexedArray32(index1_32, content)) assert awkward1.to_list(array) == [6.6, 4.4, 2.2, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, 4.4, 2.2, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedArrayU32( index2_U32, awkward1.layout.IndexedArrayU32(index1_U32, content)) assert awkward1.to_list(array) == [6.6, 4.4, 2.2, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, 4.4, 2.2, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedArrayU32( index2_U32, awkward1.layout.IndexedArray64(index1_64, content)) assert awkward1.to_list(array) == [6.6, 4.4, 2.2, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, 4.4, 2.2, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedArray64( index2_64, awkward1.layout.IndexedArray32(index1_32, content)) assert awkward1.to_list(array) == [6.6, 4.4, 2.2, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, 4.4, 2.2, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedArray64( index2_64, awkward1.layout.IndexedArrayU32(index1_U32, content)) assert awkward1.to_list(array) == [6.6, 4.4, 2.2, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, 4.4, 2.2, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedArray64( index2_64, awkward1.layout.IndexedArray64(index1_64, content)) assert awkward1.to_list(array) == [6.6, 4.4, 2.2, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, 4.4, 2.2, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) index1_32 = awkward1.layout.Index32( numpy.array([6, 5, -1, 3, -1, 1, 0], dtype=numpy.int32)) index1_64 = awkward1.layout.Index64( numpy.array([6, 5, -1, 3, -1, 1, 0], dtype=numpy.int64)) index2_32 = awkward1.layout.Index32( numpy.array([0, 2, 4, 6], dtype=numpy.int32)) index2_U32 = awkward1.layout.IndexU32( numpy.array([0, 2, 4, 6], dtype=numpy.uint32)) index2_64 = awkward1.layout.Index64( numpy.array([0, 2, 4, 6], dtype=numpy.int64)) array = awkward1.layout.IndexedArray32( index2_32, awkward1.layout.IndexedOptionArray32(index1_32, content)) assert awkward1.to_list(array) == [6.6, None, None, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, None, None, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedArray32( index2_32, awkward1.layout.IndexedOptionArray64(index1_64, content)) assert awkward1.to_list(array) == [6.6, None, None, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, None, None, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedArrayU32( index2_U32, awkward1.layout.IndexedOptionArray32(index1_32, content)) assert awkward1.to_list(array) == [6.6, None, None, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, None, None, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedArrayU32( index2_U32, awkward1.layout.IndexedOptionArray64(index1_64, content)) assert awkward1.to_list(array) == [6.6, None, None, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, None, None, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedArray64( index2_64, awkward1.layout.IndexedOptionArray32(index1_32, content)) assert awkward1.to_list(array) == [6.6, None, None, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, None, None, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedArray64( index2_64, awkward1.layout.IndexedOptionArray64(index1_64, content)) assert awkward1.to_list(array) == [6.6, None, None, 0.0] assert awkward1.to_list(array.simplify()) == [6.6, None, None, 0.0] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) index1_32 = awkward1.layout.Index32( numpy.array([6, 5, 4, 3, 2, 1, 0], dtype=numpy.int32)) index1_U32 = awkward1.layout.IndexU32( numpy.array([6, 5, 4, 3, 2, 1, 0], dtype=numpy.uint32)) index1_64 = awkward1.layout.Index64( numpy.array([6, 5, 4, 3, 2, 1, 0], dtype=numpy.int64)) index2_32 = awkward1.layout.Index32( numpy.array([0, -1, 4, -1], dtype=numpy.int32)) index2_64 = awkward1.layout.Index64( numpy.array([0, -1, 4, -1], dtype=numpy.int64)) array = awkward1.layout.IndexedOptionArray32( index2_32, awkward1.layout.IndexedArray32(index1_32, content)) assert awkward1.to_list(array) == [6.6, None, 2.2, None] assert awkward1.to_list(array.simplify()) == [6.6, None, 2.2, None] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedOptionArray32( index2_32, awkward1.layout.IndexedArrayU32(index1_U32, content)) assert awkward1.to_list(array) == [6.6, None, 2.2, None] assert awkward1.to_list(array.simplify()) == [6.6, None, 2.2, None] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedOptionArray32( index2_32, awkward1.layout.IndexedArray64(index1_64, content)) assert awkward1.to_list(array) == [6.6, None, 2.2, None] assert awkward1.to_list(array.simplify()) == [6.6, None, 2.2, None] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedOptionArray64( index2_64, awkward1.layout.IndexedArray32(index1_32, content)) assert awkward1.to_list(array) == [6.6, None, 2.2, None] assert awkward1.to_list(array.simplify()) == [6.6, None, 2.2, None] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedOptionArray64( index2_64, awkward1.layout.IndexedArrayU32(index1_U32, content)) assert awkward1.to_list(array) == [6.6, None, 2.2, None] assert awkward1.to_list(array.simplify()) == [6.6, None, 2.2, None] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedOptionArray64( index2_64, awkward1.layout.IndexedArray64(index1_64, content)) assert awkward1.to_list(array) == [6.6, None, 2.2, None] assert awkward1.to_list(array.simplify()) == [6.6, None, 2.2, None] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) index1_32 = awkward1.layout.Index32( numpy.array([6, 5, -1, 3, -1, 1, 0], dtype=numpy.int32)) index1_64 = awkward1.layout.Index64( numpy.array([6, 5, -1, 3, -1, 1, 0], dtype=numpy.int64)) index2_32 = awkward1.layout.Index32( numpy.array([0, -1, 4, -1], dtype=numpy.int32)) index2_64 = awkward1.layout.Index64( numpy.array([0, -1, 4, -1], dtype=numpy.int64)) array = awkward1.layout.IndexedOptionArray32( index2_32, awkward1.layout.IndexedOptionArray32(index1_32, content)) assert awkward1.to_list(array) == [6.6, None, None, None] assert awkward1.to_list(array.simplify()) == [6.6, None, None, None] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedOptionArray32( index2_32, awkward1.layout.IndexedOptionArray64(index1_64, content)) assert awkward1.to_list(array) == [6.6, None, None, None] assert awkward1.to_list(array.simplify()) == [6.6, None, None, None] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedOptionArray64( index2_64, awkward1.layout.IndexedOptionArray32(index1_32, content)) assert awkward1.to_list(array) == [6.6, None, None, None] assert awkward1.to_list(array.simplify()) == [6.6, None, None, None] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray) array = awkward1.layout.IndexedOptionArray64( index2_64, awkward1.layout.IndexedOptionArray64(index1_64, content)) assert awkward1.to_list(array) == [6.6, None, None, None] assert awkward1.to_list(array.simplify()) == [6.6, None, None, None] assert isinstance(array.simplify(), awkward1.layout.IndexedOptionArray64) assert isinstance(array.simplify().content, awkward1.layout.NumpyArray)
def test_RecordArray_getitem(): array = awkward1.Array([{ "x": 0.0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": [2, 2] }, { "x": 3.3, "y": [3, 3, 3] }, { "x": 4.4, "y": [4, 4, 4, 4] }], check_valid=True) @numba.njit def f1(x, i): return x[i] assert awkward1.to_list(f1(array, 3)) == {"x": 3.3, "y": [3, 3, 3]} assert awkward1.to_list(f1(array, 2)) == {"x": 2.2, "y": [2, 2]} assert awkward1.to_list(f1(array, 1)) == {"x": 1.1, "y": [1]} @numba.njit def f2(x, i1, i2): return x[i1:i2] assert awkward1.to_list(f2(array, 1, 4)) == [{ "x": 1.1, "y": [1] }, { "x": 2.2, "y": [2, 2] }, { "x": 3.3, "y": [3, 3, 3] }] array = awkward1.Array([[{ "x": 0.0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": [2, 2] }], [], [{ "x": 3.3, "y": [3, 3, 3] }, { "x": 4.4, "y": [4, 4, 4, 4] }]], check_valid=True) @numba.njit def f3(x, i, j): return x[i][j] assert awkward1.to_list(f3(array, 2, -2)) == {"x": 3.3, "y": [3, 3, 3]}
def test_listarray_merge(): emptyarray = awkward1.layout.EmptyArray() content1 = awkward1.layout.NumpyArray( numpy.array([1.1, 2.2, 3.3, 4.4, 5.5])) content2 = awkward1.layout.NumpyArray(numpy.array([1, 2, 3, 4, 5, 6, 7])) for (dtype1, Index1, ListArray1), (dtype2, Index2, ListArray2) in [ ((numpy.int32, awkward1.layout.Index32, awkward1.layout.ListArray32), (numpy.int32, awkward1.layout.Index32, awkward1.layout.ListArray32)), ((numpy.int32, awkward1.layout.Index32, awkward1.layout.ListArray32), (numpy.uint32, awkward1.layout.IndexU32, awkward1.layout.ListArrayU32)), ((numpy.int32, awkward1.layout.Index32, awkward1.layout.ListArray32), (numpy.int64, awkward1.layout.Index64, awkward1.layout.ListArray64)), ((numpy.uint32, awkward1.layout.IndexU32, awkward1.layout.ListArrayU32), (numpy.int32, awkward1.layout.Index32, awkward1.layout.ListArray32)), ((numpy.uint32, awkward1.layout.IndexU32, awkward1.layout.ListArrayU32), (numpy.uint32, awkward1.layout.IndexU32, awkward1.layout.ListArrayU32)), ((numpy.uint32, awkward1.layout.IndexU32, awkward1.layout.ListArrayU32), (numpy.int64, awkward1.layout.Index64, awkward1.layout.ListArray64)), ((numpy.int64, awkward1.layout.Index64, awkward1.layout.ListArray64), (numpy.int32, awkward1.layout.Index32, awkward1.layout.ListArray32)), ((numpy.int64, awkward1.layout.Index64, awkward1.layout.ListArray64), (numpy.uint32, awkward1.layout.IndexU32, awkward1.layout.ListArrayU32)), ((numpy.int64, awkward1.layout.Index64, awkward1.layout.ListArray64), (numpy.int64, awkward1.layout.Index64, awkward1.layout.ListArray64)) ]: starts1 = Index1(numpy.array([0, 3, 3], dtype=dtype1)) stops1 = Index1(numpy.array([3, 3, 5], dtype=dtype1)) starts2 = Index2(numpy.array([2, 99, 0], dtype=dtype2)) stops2 = Index2(numpy.array([6, 99, 3], dtype=dtype2)) array1 = ListArray1(starts1, stops1, content1) array2 = ListArray2(starts2, stops2, content2) assert awkward1.to_list(array1) == [[1.1, 2.2, 3.3], [], [4.4, 5.5]] assert awkward1.to_list(array2) == [[3, 4, 5, 6], [], [1, 2, 3]] assert awkward1.to_list(array1.merge(array2)) == [[1.1, 2.2, 3.3], [], [4.4, 5.5], [3, 4, 5, 6], [], [1, 2, 3]] assert awkward1.to_list(array2.merge(array1)) == [[3, 4, 5, 6], [], [1, 2, 3], [1.1, 2.2, 3.3], [], [4.4, 5.5]] assert awkward1.to_list( array1.merge(emptyarray)) == awkward1.to_list(array1) assert awkward1.to_list( emptyarray.merge(array1)) == awkward1.to_list(array1) regulararray = awkward1.layout.RegularArray(content2, 2) assert awkward1.to_list(regulararray) == [[1, 2], [3, 4], [5, 6]] assert awkward1.to_list( regulararray.merge(emptyarray)) == awkward1.to_list(regulararray) assert awkward1.to_list( emptyarray.merge(regulararray)) == awkward1.to_list(regulararray) for (dtype1, Index1, ListArray1) in [ (numpy.int32, awkward1.layout.Index32, awkward1.layout.ListArray32), (numpy.uint32, awkward1.layout.IndexU32, awkward1.layout.ListArrayU32), (numpy.int64, awkward1.layout.Index64, awkward1.layout.ListArray64) ]: starts1 = Index1(numpy.array([0, 3, 3], dtype=dtype1)) stops1 = Index1(numpy.array([3, 3, 5], dtype=dtype1)) array1 = ListArray1(starts1, stops1, content1) assert awkward1.to_list( array1.merge(regulararray)) == [[1.1, 2.2, 3.3], [], [4.4, 5.5], [1, 2], [3, 4], [5, 6]] assert awkward1.to_list( regulararray.merge(array1)) == [[1, 2], [3, 4], [5, 6], [1.1, 2.2, 3.3], [], [4.4, 5.5]]
def test_numpyarray_merge(): emptyarray = awkward1.layout.EmptyArray() np1 = numpy.arange(2*7*5).reshape(2, 7, 5) np2 = numpy.arange(3*7*5).reshape(3, 7, 5) ak1 = awkward1.layout.NumpyArray(np1) ak2 = awkward1.layout.NumpyArray(np2) assert awkward1.to_list(ak1.merge(ak2)) == awkward1.to_list(numpy.concatenate([np1, np2])) assert awkward1.to_list(ak1[1:, :-1, ::-1].merge(ak2[1:, :-1, ::-1])) == awkward1.to_list(numpy.concatenate([np1[1:, :-1, ::-1], np2[1:, :-1, ::-1]])) for x, y, z in [(numpy.double, numpy.double, numpy.double), (numpy.double, numpy.float, numpy.double), (numpy.double, numpy.int64, numpy.double), (numpy.double, numpy.uint64, numpy.double), (numpy.double, numpy.int32, numpy.double), (numpy.double, numpy.uint32, numpy.double), (numpy.double, numpy.int16, numpy.double), (numpy.double, numpy.uint16, numpy.double), (numpy.double, numpy.int8, numpy.double), (numpy.double, numpy.uint8, numpy.double), (numpy.double, numpy.bool, numpy.double), (numpy.float, numpy.double, numpy.double), (numpy.float, numpy.float, numpy.double), (numpy.float, numpy.int64, numpy.double), (numpy.float, numpy.uint64, numpy.double), (numpy.float, numpy.int32, numpy.double), (numpy.float, numpy.uint32, numpy.double), (numpy.float, numpy.int16, numpy.double), (numpy.float, numpy.uint16, numpy.double), (numpy.float, numpy.int8, numpy.double), (numpy.float, numpy.uint8, numpy.double), (numpy.float, numpy.bool, numpy.double), (numpy.int64, numpy.double, numpy.double), (numpy.int64, numpy.float, numpy.double), (numpy.int64, numpy.int64, numpy.int64), (numpy.int64, numpy.uint64, numpy.int64), (numpy.int64, numpy.int32, numpy.int64), (numpy.int64, numpy.uint32, numpy.int64), (numpy.int64, numpy.int16, numpy.int64), (numpy.int64, numpy.uint16, numpy.int64), (numpy.int64, numpy.int8, numpy.int64), (numpy.int64, numpy.uint8, numpy.int64), (numpy.int64, numpy.bool, numpy.int64), (numpy.uint64, numpy.double, numpy.double), (numpy.uint64, numpy.float, numpy.double), (numpy.uint64, numpy.int64, numpy.int64), (numpy.uint64, numpy.uint64, numpy.uint64), (numpy.uint64, numpy.int32, numpy.int64), (numpy.uint64, numpy.uint32, numpy.int64), (numpy.uint64, numpy.int16, numpy.int64), (numpy.uint64, numpy.uint16, numpy.int64), (numpy.uint64, numpy.int8, numpy.int64), (numpy.uint64, numpy.uint8, numpy.int64), (numpy.uint64, numpy.bool, numpy.int64), (numpy.int32, numpy.double, numpy.double), (numpy.int32, numpy.float, numpy.double), (numpy.int32, numpy.int64, numpy.int64), (numpy.int32, numpy.uint64, numpy.int64), (numpy.int32, numpy.int32, numpy.int64), (numpy.int32, numpy.uint32, numpy.int64), (numpy.int32, numpy.int16, numpy.int64), (numpy.int32, numpy.uint16, numpy.int64), (numpy.int32, numpy.int8, numpy.int64), (numpy.int32, numpy.uint8, numpy.int64), (numpy.int32, numpy.bool, numpy.int64), (numpy.uint32, numpy.double, numpy.double), (numpy.uint32, numpy.float, numpy.double), (numpy.uint32, numpy.int64, numpy.int64), (numpy.uint32, numpy.uint64, numpy.int64), (numpy.uint32, numpy.int32, numpy.int64), (numpy.uint32, numpy.uint32, numpy.int64), (numpy.uint32, numpy.int16, numpy.int64), (numpy.uint32, numpy.uint16, numpy.int64), (numpy.uint32, numpy.int8, numpy.int64), (numpy.uint32, numpy.uint8, numpy.int64), (numpy.uint32, numpy.bool, numpy.int64), (numpy.int16, numpy.double, numpy.double), (numpy.int16, numpy.float, numpy.double), (numpy.int16, numpy.int64, numpy.int64), (numpy.int16, numpy.uint64, numpy.int64), (numpy.int16, numpy.int32, numpy.int64), (numpy.int16, numpy.uint32, numpy.int64), (numpy.int16, numpy.int16, numpy.int64), (numpy.int16, numpy.uint16, numpy.int64), (numpy.int16, numpy.int8, numpy.int64), (numpy.int16, numpy.uint8, numpy.int64), (numpy.int16, numpy.bool, numpy.int64), (numpy.uint16, numpy.double, numpy.double), (numpy.uint16, numpy.float, numpy.double), (numpy.uint16, numpy.int64, numpy.int64), (numpy.uint16, numpy.uint64, numpy.int64), (numpy.uint16, numpy.int32, numpy.int64), (numpy.uint16, numpy.uint32, numpy.int64), (numpy.uint16, numpy.int16, numpy.int64), (numpy.uint16, numpy.uint16, numpy.int64), (numpy.uint16, numpy.int8, numpy.int64), (numpy.uint16, numpy.uint8, numpy.int64), (numpy.uint16, numpy.bool, numpy.int64), (numpy.int8, numpy.double, numpy.double), (numpy.int8, numpy.float, numpy.double), (numpy.int8, numpy.int64, numpy.int64), (numpy.int8, numpy.uint64, numpy.int64), (numpy.int8, numpy.int32, numpy.int64), (numpy.int8, numpy.uint32, numpy.int64), (numpy.int8, numpy.int16, numpy.int64), (numpy.int8, numpy.uint16, numpy.int64), (numpy.int8, numpy.int8, numpy.int64), (numpy.int8, numpy.uint8, numpy.int64), (numpy.int8, numpy.bool, numpy.int64), (numpy.uint8, numpy.double, numpy.double), (numpy.uint8, numpy.float, numpy.double), (numpy.uint8, numpy.int64, numpy.int64), (numpy.uint8, numpy.uint64, numpy.int64), (numpy.uint8, numpy.int32, numpy.int64), (numpy.uint8, numpy.uint32, numpy.int64), (numpy.uint8, numpy.int16, numpy.int64), (numpy.uint8, numpy.uint16, numpy.int64), (numpy.uint8, numpy.int8, numpy.int64), (numpy.uint8, numpy.uint8, numpy.int64), (numpy.uint8, numpy.bool, numpy.int64), (numpy.bool, numpy.double, numpy.double), (numpy.bool, numpy.float, numpy.double), (numpy.bool, numpy.int64, numpy.int64), (numpy.bool, numpy.uint64, numpy.int64), (numpy.bool, numpy.int32, numpy.int64), (numpy.bool, numpy.uint32, numpy.int64), (numpy.bool, numpy.int16, numpy.int64), (numpy.bool, numpy.uint16, numpy.int64), (numpy.bool, numpy.int8, numpy.int64), (numpy.bool, numpy.uint8, numpy.int64), (numpy.bool, numpy.bool, numpy.bool)]: one = awkward1.layout.NumpyArray(numpy.array([1, 2, 3], dtype=x)) two = awkward1.layout.NumpyArray(numpy.array([4, 5], dtype=y)) three = one.merge(two) assert numpy.asarray(three).dtype == numpy.dtype(z) assert awkward1.to_list(three) == awkward1.to_list(numpy.concatenate([numpy.asarray(one), numpy.asarray(two)])) assert awkward1.to_list(one.merge(emptyarray)) == awkward1.to_list(one) assert awkward1.to_list(emptyarray.merge(one)) == awkward1.to_list(one)
def test_ByteMaskedArray_reduce(): content = awkward1.layout.NumpyArray( numpy.array([ 2, 3, 5, 7, 11, 0, 0, 0, 0, 0, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 0, 0, 0, 0, 0, 101, 103, 107, 109, 113 ], dtype=numpy.int64)) offsets1 = awkward1.layout.Index64( numpy.array([0, 5, 10, 15, 20, 25, 30], dtype=numpy.int64)) listoffsetarray = awkward1.layout.ListOffsetArray64(offsets1, content) # index = awkward1.layout.Index64(numpy.array([0, -1, 2, 3, -1, 5], dtype=numpy.int64)) # optionarray = awkward1.layout.IndexedOptionArray64(index, listoffsetarray) mask = awkward1.layout.Index8( numpy.array([0, 1, 0, 0, 1, 0], dtype=numpy.int8)) optionarray = awkward1.layout.ByteMaskedArray(mask, listoffsetarray, valid_when=False) offsets2 = awkward1.layout.Index64( numpy.array([0, 3, 6], dtype=numpy.int64)) depth2 = awkward1.layout.ListOffsetArray64(offsets2, optionarray) assert awkward1.to_list(depth2) == [[[2, 3, 5, 7, 11], None, [31, 37, 41, 43, 47]], [[53, 59, 61, 67, 71], None, [101, 103, 107, 109, 113]]] assert awkward1.to_list(depth2.prod(-1)) == [[ 2 * 3 * 5 * 7 * 11, None, 31 * 37 * 41 * 43 * 47 ], [53 * 59 * 61 * 67 * 71, None, 101 * 103 * 107 * 109 * 113]] assert awkward1.to_list(depth2.prod(-2)) == [[ 31 * 2, 37 * 3, 41 * 5, 43 * 7, 47 * 11 ], [101 * 53, 103 * 59, 107 * 61, 109 * 67, 113 * 71]] assert awkward1.to_list(depth2.prod(-3)) == [[ 53 * 2, 59 * 3, 61 * 5, 67 * 7, 71 * 11 ], [], [101 * 31, 103 * 37, 107 * 41, 109 * 43, 113 * 47]] content = awkward1.layout.NumpyArray( numpy.array([ 2, 3, 5, 7, 11, 0, 0, 0, 0, 0, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 0, 0, 0, 0, 0, 101, 103, 107, 109, 113 ], dtype=numpy.int64)) # index = awkward1.layout.Index64(numpy.array([0, 1, 2, 3, 4, -1, -1, -1, -1, -1, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, -1, -1, -1, -1, -1, 25, 26, 27, 28, 29], dtype=numpy.int64)) # optionarray = awkward1.layout.IndexedOptionArray64(index, content) mask = awkward1.layout.Index8( numpy.array([ 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0 ], dtype=numpy.int8)) optionarray = awkward1.layout.ByteMaskedArray(mask, content, valid_when=False) offsets1 = awkward1.layout.Index64( numpy.array([0, 5, 10, 15, 20, 25, 30], dtype=numpy.int64)) listoffsetarray = awkward1.layout.ListOffsetArray64(offsets1, optionarray) offsets2 = awkward1.layout.Index64( numpy.array([0, 3, 6], dtype=numpy.int64)) depth2 = awkward1.layout.ListOffsetArray64(offsets2, listoffsetarray) assert awkward1.to_list(depth2) == [[[2, 3, 5, 7, 11], [None, None, None, None, None], [31, 37, 41, 43, 47]], [[53, 59, 61, 67, 71], [None, None, None, None, None], [101, 103, 107, 109, 113]]] assert awkward1.to_list(depth2.prod(-1)) == [[ 2 * 3 * 5 * 7 * 11, 1 * 1 * 1 * 1 * 1, 31 * 37 * 41 * 43 * 47 ], [ 53 * 59 * 61 * 67 * 71, 1 * 1 * 1 * 1 * 1, 101 * 103 * 107 * 109 * 113 ]] assert awkward1.to_list(depth2.prod(-2)) == [[ 31 * 2, 37 * 3, 41 * 5, 43 * 7, 47 * 11 ], [101 * 53, 103 * 59, 107 * 61, 109 * 67, 113 * 71]] assert awkward1.to_list(depth2.prod(-3)) == [[ 53 * 2, 59 * 3, 61 * 5, 67 * 7, 71 * 11 ], [1, 1, 1, 1, 1], [101 * 31, 103 * 37, 107 * 41, 109 * 43, 113 * 47]]