def test_numpyarray_getitem_bystrides(): a = np.arange(10) b = ak.layout.NumpyArray(a) assert b[3] == a[3] assert b[-3] == a[-3] assert ak.to_list(b[()]) == ak.to_list(a[()]) assert ak.to_list(b[slice(None)]) == ak.to_list(a[slice(None)]) assert ak.to_list(b[slice(3, 7)]) == ak.to_list(a[slice(3, 7)]) assert ak.to_list(b[slice(3, 100)]) == ak.to_list(a[slice(3, 100)]) assert ak.to_list(b[slice(-100, 7)]) == ak.to_list(a[slice(-100, 7)]) assert ak.to_list(b[slice(3, -3)]) == ak.to_list(a[slice(3, -3)]) assert ak.to_list(b[slice(1, 7, 2)]) == ak.to_list(a[slice(1, 7, 2)]) assert ak.to_list(b[slice(-8, 7, 2)]) == ak.to_list(a[slice(-8, 7, 2)]) assert ak.to_list(b[slice(None, 7, 2)]) == ak.to_list(a[slice(None, 7, 2)]) assert ak.to_list(b[slice(None, -3, 2)]) == ak.to_list(a[slice(None, -3, 2)]) assert ak.to_list(b[slice(8, None, -1)]) == ak.to_list(a[slice(8, None, -1)]) assert ak.to_list(b[slice(8, None, -2)]) == ak.to_list(a[slice(8, None, -2)]) assert ak.to_list(b[slice(-2, None, -2)]) == ak.to_list(a[slice(-2, None, -2)]) a = np.arange(7 * 5).reshape(7, 5) b = ak.layout.NumpyArray(a) assert ak.to_list(b[()]) == ak.to_list(a[()]) assert ak.to_list(b[3]) == ak.to_list(a[3]) assert ak.to_list(b[(3, 2)]) == ak.to_list(a[3, 2]) assert ak.to_list(b[slice(1, 4)]) == ak.to_list(a[slice(1, 4)]) assert ak.to_list(b[(3, slice(1, 4))]) == ak.to_list(a[3, slice(1, 4)]) assert ak.to_list(b[(slice(1, 4), 3)]) == ak.to_list(a[slice(1, 4), 3]) assert ak.to_list(b[(slice(1, 4), slice(2, None))]) == ak.to_list(a[slice(1, 4), slice(2, None)]) assert ak.to_list(b[(slice(None, None, -1), slice(2, None))]) == ak.to_list(a[slice(None, None, -1), slice(2, None)]) assert ak.to_list(b[:, np.newaxis, :]) == ak.to_list(a[:, np.newaxis, :]) assert ak.to_list(b[np.newaxis, :, np.newaxis, :, np.newaxis]) == ak.to_list(a[np.newaxis, :, np.newaxis, :, np.newaxis]) if not ak._util.py27: assert ak.to_list(b[Ellipsis, 3]) == ak.to_list(a[Ellipsis, 3]) assert ak.to_list(b[Ellipsis, 3, 2]) == ak.to_list(a[Ellipsis, 3, 2]) assert ak.to_list(b[3, Ellipsis]) == ak.to_list(a[3, Ellipsis]) assert ak.to_list(b[3, 2, Ellipsis]) == ak.to_list(a[3, 2, Ellipsis])
def test_rpad_and_clip_indexed_array(): listoffsetarray = ak.from_iter( [[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False) backward = ak.from_iter( [[6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], [], [0.0, 1.1, 2.2]], highlevel=False) index = ak.layout.Index64(np.array([4, 3, 2, 1, 0], dtype=np.int64)) indexedarray = ak.layout.IndexedArray64(index, listoffsetarray) assert ak.to_list(indexedarray) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], [], [0.0, 1.1, 2.2], ] assert ak.to_list(backward.rpad_and_clip(4, 1)) == ak.to_list( indexedarray.rpad_and_clip(4, 1)) assert ak.to_list(indexedarray.rpad_and_clip(1, 0)) == [[6.6, 7.7, 8.8, 9.9]] assert ak.to_list(indexedarray.rpad_and_clip(2, 0)) == [[6.6, 7.7, 8.8, 9.9], [5.5]] assert ak.to_list(indexedarray.rpad_and_clip(3, 0)) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], ] assert ak.to_list(indexedarray.rpad_and_clip(4, 0)) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], [], ] assert ak.to_list(indexedarray.rpad_and_clip(5, 0)) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], [], [0.0, 1.1, 2.2], ] assert ak.to_list(indexedarray.rpad_and_clip(6, 0)) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], [], [0.0, 1.1, 2.2], None, ] assert ak.to_list(indexedarray.rpad_and_clip(7, 0)) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], [], [0.0, 1.1, 2.2], None, None, ] assert ak.to_list(indexedarray.rpad_and_clip(8, 0)) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], [], [0.0, 1.1, 2.2], None, None, None, ] assert ak.to_list(indexedarray.rpad_and_clip(1, 1)) == [ [6.6], [5.5], [3.3], [None], [0.0], ] assert ak.to_list(indexedarray.rpad_and_clip(2, 1)) == [ [6.6, 7.7], [5.5, None], [3.3, 4.4], [None, None], [0.0, 1.1], ] assert ak.to_list(indexedarray.rpad_and_clip(3, 1)) == [ [6.6, 7.7, 8.8], [5.5, None, None], [3.3, 4.4, None], [None, None, None], [0.0, 1.1, 2.2], ] assert ak.to_list(indexedarray.rpad_and_clip(4, 1)) == [ [6.6, 7.7, 8.8, 9.9], [5.5, None, None, None], [3.3, 4.4, None, None], [None, None, None, None], [0.0, 1.1, 2.2, None], ] assert ak.to_list(indexedarray.rpad_and_clip(5, 1)) == [ [6.6, 7.7, 8.8, 9.9, None], [5.5, None, None, None, None], [3.3, 4.4, None, None, None], [None, None, None, None, None], [0.0, 1.1, 2.2, None, None], ]
def test_listoffsetarray_at(): array1 = ak.layout.ListOffsetArrayU32(offsets1, content) assert (ak.to_list(array1[2, ]) == [4.4, 5.5]) assert ak.to_list(array1[2, 1:]) == [5.5] assert ak.to_list(array1[2:, 0]) == [4.4, 6.6, 7.7] assert ak.to_list(array1[2:, -1]) == [5.5, 6.6, 9.9]
def test_rpad_regular_array(): content = ak.layout.NumpyArray( np.array([ 2.1, 8.4, 7.4, 1.6, 2.2, 3.4, 6.2, 5.4, 1.5, 3.9, 3.8, 3.0, 8.5, 6.9, 4.3, 3.6, 6.7, 1.8, 3.2, ])) index = ak.layout.Index64( np.array([13, 9, 13, 4, 8, 3, 15, -1, 16, 2, 8], dtype=np.int64)) indexedarray = ak.layout.IndexedOptionArray64(index, content) array = ak.layout.RegularArray(indexedarray, 3, zeros_length=0) assert ak.to_list(array.rpad(5, 0)) == [ [6.9, 3.9, 6.9], [2.2, 1.5, 1.6], [3.6, None, 6.7], None, None, ] assert ak.to_list(array.rpad(4, 0)) == [ [6.9, 3.9, 6.9], [2.2, 1.5, 1.6], [3.6, None, 6.7], None, ] assert ak.to_list(array.rpad(3, 0)) == [ [6.9, 3.9, 6.9], [2.2, 1.5, 1.6], [3.6, None, 6.7], ] assert ak.to_list(array.rpad(1, 0)) == [ [6.9, 3.9, 6.9], [2.2, 1.5, 1.6], [3.6, None, 6.7], ] assert ak.to_list(array.rpad(5, 1)) == [ [6.9, 3.9, 6.9, None, None], [2.2, 1.5, 1.6, None, None], [3.6, None, 6.7, None, None], ] assert ak.to_list(array.rpad(4, 1)) == [ [6.9, 3.9, 6.9, None], [2.2, 1.5, 1.6, None], [3.6, None, 6.7, None], ] assert ak.to_list(array.rpad(3, 1)) == [ [6.9, 3.9, 6.9], [2.2, 1.5, 1.6], [3.6, None, 6.7], ] assert ak.to_list(array.rpad(1, 1)) == [ [6.9, 3.9, 6.9], [2.2, 1.5, 1.6], [3.6, None, 6.7], ] content = ak.layout.NumpyArray( np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) offsets = ak.layout.Index64(np.array([0, 3, 3, 5, 6, 10, 10])) listoffsetarray = ak.layout.ListOffsetArray64(offsets, content) regulararray = ak.layout.RegularArray(listoffsetarray, 2, zeros_length=0) assert ak.to_list(regulararray.rpad(1, 0)) == [ [[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], []], ] assert ak.to_list(regulararray.rpad(3, 0)) == [ [[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], []], ] assert ak.to_list(regulararray.rpad(4, 0)) == [ [[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], []], None, ] assert ak.to_list(regulararray.rpad(7, 0)) == [ [[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], []], None, None, None, None, ] assert ak.to_list(regulararray.rpad(1, 1)) == [ [[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], []], ] assert ak.to_list(regulararray.rpad(2, 1)) == [ [[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], []], ] assert ak.to_list(regulararray.rpad(3, 1)) == [ [[0.0, 1.1, 2.2], [], None], [[3.3, 4.4], [5.5], None], [[6.6, 7.7, 8.8, 9.9], [], None], ] assert ak.to_list(regulararray.rpad(5, 1)) == [ [[0.0, 1.1, 2.2], [], None, None, None], [[3.3, 4.4], [5.5], None, None, None], [[6.6, 7.7, 8.8, 9.9], [], None, None, None], ] assert ak.to_list(regulararray.rpad(7, 1)) == [ [[0.0, 1.1, 2.2], [], None, None, None, None, None], [[3.3, 4.4], [5.5], None, None, None, None, None], [[6.6, 7.7, 8.8, 9.9], [], None, None, None, None, None], ] assert ak.to_list(regulararray.rpad(1, 2)) == [ [[0.0, 1.1, 2.2], [None]], [[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], [None]], ] assert ak.to_list(regulararray.rpad(2, 2)) == [ [[0.0, 1.1, 2.2], [None, None]], [[3.3, 4.4], [5.5, None]], [[6.6, 7.7, 8.8, 9.9], [None, None]], ] assert ak.to_list(regulararray.rpad(3, 2)) == [ [[0.0, 1.1, 2.2], [None, None, None]], [[3.3, 4.4, None], [5.5, None, None]], [[6.6, 7.7, 8.8, 9.9], [None, None, None]], ] assert ak.to_list(regulararray.rpad(4, 2)) == [ [[0.0, 1.1, 2.2, None], [None, None, None, None]], [[3.3, 4.4, None, None], [5.5, None, None, None]], [[6.6, 7.7, 8.8, 9.9], [None, None, None, None]], ]
def test_rpad_listoffset_array(): content = ak.layout.NumpyArray( np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) offsets = ak.layout.Index64(np.array([0, 3, 3, 5, 6, 10, 10])) listoffsetarray = ak.layout.ListOffsetArray64(offsets, content) assert ak.to_list(listoffsetarray) == [ [0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9], [], ] assert ak.to_list(listoffsetarray.rpad(3, 0)) == [ [0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9], [], ] assert str(ak.type(listoffsetarray)) == str( ak.type(listoffsetarray.rpad(3, 0))) assert ak.to_list(listoffsetarray.rpad(7, 0)) == [ [0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9], [], None, ] assert "option[" + str(ak.type(listoffsetarray)) + "]" == str( ak.type(listoffsetarray.rpad(7, 0))) assert ak.to_list(listoffsetarray.rpad(5, 1)) == [ [0.0, 1.1, 2.2, None, None], [None, None, None, None, None], [3.3, 4.4, None, None, None], [5.5, None, None, None, None], [6.6, 7.7, 8.8, 9.9, None], [None, None, None, None, None], ] assert ak.to_list(listoffsetarray.rpad(1, 1)) == [ [0.0, 1.1, 2.2], [None], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9], [None], ] content = ak.layout.NumpyArray(np.array([1.5, 3.3])) index = ak.layout.Index64( np.array([ 0, -3, 1, -2, 1, 0, 0, -3, -13, 0, 1, 1, 0, 1, 1, 1, 1, -10, 0, -1, 0, 0, 0, 1, -1, 1, 1, ])) indexedarray = ak.layout.IndexedOptionArray64(index, content) offsets = ak.layout.Index64(np.array([14, 15, 15, 15, 26, 26, 26])) listoffsetarray = ak.layout.ListOffsetArray64(offsets, indexedarray) assert ak.to_list(listoffsetarray) == [ [3.3], [], [], [3.3, 3.3, None, 1.5, None, 1.5, 1.5, 1.5, 3.3, None, 3.3], [], [], ] assert ak.to_list(listoffsetarray.rpad(1, 0)) == [ [3.3], [], [], [3.3, 3.3, None, 1.5, None, 1.5, 1.5, 1.5, 3.3, None, 3.3], [], [], ] assert str(ak.type(listoffsetarray)) == str( ak.type(listoffsetarray.rpad(1, 0))) assert ak.to_list(listoffsetarray.rpad(6, 0)) == [ [3.3], [], [], [3.3, 3.3, None, 1.5, None, 1.5, 1.5, 1.5, 3.3, None, 3.3], [], [], ] assert "option[" + str(ak.type(listoffsetarray)) + "]" == str( ak.type(listoffsetarray.rpad(6, 0))) assert ak.to_list(listoffsetarray.rpad(7, 0)) == [ [3.3], [], [], [3.3, 3.3, None, 1.5, None, 1.5, 1.5, 1.5, 3.3, None, 3.3], [], [], None, ] assert "option[" + str(ak.type(listoffsetarray)) + "]" == str( ak.type(listoffsetarray.rpad(7, 0))) assert ak.to_list(listoffsetarray.rpad(9, 0)) == [ [3.3], [], [], [3.3, 3.3, None, 1.5, None, 1.5, 1.5, 1.5, 3.3, None, 3.3], [], [], None, None, None, ] assert "option[" + str(ak.type(listoffsetarray)) + "]" == str( ak.type(listoffsetarray.rpad(9, 0))) assert ak.to_list(listoffsetarray.rpad(1, 1)) == [ [3.3], [None], [None], [3.3, 3.3, None, 1.5, None, 1.5, 1.5, 1.5, 3.3, None, 3.3], [None], [None], ] assert str(ak.type(listoffsetarray)) == str( ak.type(listoffsetarray.rpad(1, 1))) assert ak.to_list(listoffsetarray.rpad(4, 1)) == [ [3.3, None, None, None], [None, None, None, None], [None, None, None, None], [3.3, 3.3, None, 1.5, None, 1.5, 1.5, 1.5, 3.3, None, 3.3], [None, None, None, None], [None, None, None, None], ] assert str(ak.type(listoffsetarray)) == str( ak.type(listoffsetarray.rpad(4, 1)))
def test_listarray(): one = ak.Array([[{ "x": 1 }, { "x": 2 }, { "x": 3 }], [], [{ "x": 4 }, { "x": 5 }]], check_valid=True) two = ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]], check_valid=True) assert ak.to_list(ak.with_field(one, two, "y")) == [ [{ "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 }], ] three = ak.Array([100, 200, 300], check_valid=True) assert ak.to_list(ak.with_field(one, three, "y")) == [ [{ "x": 1, "y": 100 }, { "x": 2, "y": 100 }, { "x": 3, "y": 100 }], [], [{ "x": 4, "y": 300 }, { "x": 5, "y": 300 }], ] assert ak.to_list(ak.with_field(one, [100, 200, 300], "y")) == [ [{ "x": 1, "y": 100 }, { "x": 2, "y": 100 }, { "x": 3, "y": 100 }], [], [{ "x": 4, "y": 300 }, { "x": 5, "y": 300 }], ]
def test_rpad_indexed_option_array(): listoffsetarray = ak.from_iter( [[0.0, None, None], None, [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False, ) backward = ak.from_iter( [[6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], None, [0.0, None, None]], highlevel=False, ) index = ak.layout.Index64(np.array([4, 3, 2, -1, 0], dtype=np.int64)) indexedarray = ak.layout.IndexedOptionArray64(index, listoffsetarray) assert ak.to_list(indexedarray) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], None, [0.0, None, None], ] assert ak.to_list(backward.rpad(4, 1)) == ak.to_list(indexedarray.rpad(4, 1)) assert ak.to_list(indexedarray.rpad(1, 0)) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], None, [0.0, None, None], ] assert ak.to_list(indexedarray.rpad(1, 1)) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], None, [0.0, None, None], ] assert ak.to_list(indexedarray.rpad(3, 1)) == [ [6.6, 7.7, 8.8, 9.9], [5.5, None, None], [3.3, 4.4, None], None, [0.0, None, None], ] assert ak.to_list(indexedarray.rpad(4, 0)) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], None, [0.0, None, None], ] assert ak.to_list(indexedarray.rpad(5, 0)) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], None, [0.0, None, None], ] assert ak.to_list(indexedarray.rpad(6, 0)) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], None, [0.0, None, None], None, ] assert ak.to_list(indexedarray.rpad(7, 0)) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], None, [0.0, None, None], None, None, ] assert ak.to_list(indexedarray.rpad(8, 0)) == [ [6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], None, [0.0, None, None], None, None, None, ] assert ak.to_list(indexedarray.rpad_and_clip(1, 0)) == [[6.6, 7.7, 8.8, 9.9]] assert ak.to_list(indexedarray.rpad_and_clip(1, 1)) == [ [6.6], [5.5], [3.3], None, [0.0], ]
def test_union_to_record(): recordarray1 = ak.Array([{"x": 1, "y": 1.1}, {"x": 3, "y": 3.3}]).layout recordarray2 = ak.Array([{"y": 2.2, "z": 999}]).layout tags = ak.layout.Index8(np.array([0, 1, 0], dtype=np.int8)) index = ak.layout.Index64(np.array([0, 0, 1], dtype=np.int64)) unionarray = ak.layout.UnionArray8_64(tags, index, [recordarray1, recordarray2]) assert ak.to_list(unionarray) == [ { "x": 1, "y": 1.1 }, { "y": 2.2, "z": 999 }, { "x": 3, "y": 3.3 }, ] converted = ak._util.union_to_record(unionarray, "values") assert isinstance(converted, ak.layout.RecordArray) assert ak.to_list(converted) == [ { "x": 1, "y": 1.1, "z": None }, { "x": None, "y": 2.2, "z": 999 }, { "x": 3, "y": 3.3, "z": None }, ] otherarray = ak.Array(["one", "two"]).layout tags2 = ak.layout.Index8(np.array([0, 2, 1, 2, 0], dtype=np.int8)) index2 = ak.layout.Index64(np.array([0, 0, 0, 1, 1], dtype=np.int64)) unionarray2 = ak.layout.UnionArray8_64( tags2, index2, [recordarray1, recordarray2, otherarray]) assert ak.to_list(unionarray2) == [ { "x": 1, "y": 1.1 }, "one", { "y": 2.2, "z": 999 }, "two", { "x": 3, "y": 3.3 }, ] converted2 = ak._util.union_to_record(unionarray2, "values") assert isinstance(converted2, ak.layout.RecordArray) assert ak.to_list(converted2) == [ { "x": 1, "y": 1.1, "z": None, "values": None }, { "x": None, "y": None, "z": None, "values": "one" }, { "x": None, "y": 2.2, "z": 999, "values": None }, { "x": None, "y": None, "z": None, "values": "two" }, { "x": 3, "y": 3.3, "z": None, "values": None }, ] df_unionarray = ak.to_pandas(unionarray) np.testing.assert_array_equal(df_unionarray["x"].values, np.array([1, np.nan, 3])) np.testing.assert_array_equal(df_unionarray["y"].values, np.array([1.1, 2.2, 3.3])) np.testing.assert_array_equal(df_unionarray["z"].values, np.array([np.nan, 999, np.nan])) df_unionarray2 = ak.to_pandas(unionarray2) np.testing.assert_array_equal(df_unionarray2["x"].values, [1, np.nan, np.nan, np.nan, 3]) np.testing.assert_array_equal(df_unionarray2["y"].values, [1.1, np.nan, 2.2, np.nan, 3.3]) np.testing.assert_array_equal(df_unionarray2["z"].values, [np.nan, np.nan, 999, np.nan, np.nan]) np.testing.assert_array_equal(df_unionarray2["values"].values, ["nan", "one", "nan", "two", "nan"])
def test_fromiter(): dataset = [ [(1, 1.1), (2, 2.2), (3, 3.3)], [(1, [1.1, 2.2, 3.3]), (2, []), (3, [4.4, 5.5])], [[(1, 1.1), (2, 2.2), (3, 3.3)], [], [(4, 4.4), (5, 5.5)]], [((1, 1), 1.1), ((2, 2), 2.2), ((3, 3), 3.3)], [({ "x": 1, "y": 1 }, 1.1), ({ "x": 2, "y": 2 }, 2.2), ({ "x": 3, "y": 3 }, 3.3)], [{ "one": 1, "two": 1.1 }, { "one": 2, "two": 2.2 }, { "one": 3, "two": 3.3 }], [ { "one": 1, "two": [1.1, 2.2, 3.3] }, { "one": 2, "two": [] }, { "one": 3, "two": [4.4, 5.5] }, ], [ [{ "one": 1, "two": 1.1 }, { "one": 2, "two": 2.2 }, { "one": 3, "two": 3.3 }], [], [{ "one": 4, "two": 4.4 }, { "one": 5, "two": 5.5 }], ], [ { "one": { "x": 1, "y": 1 }, "two": 1.1 }, { "one": { "x": 2, "y": 2 }, "two": 2.2 }, { "one": { "x": 3, "y": 3 }, "two": 3.3 }, ], [ { "one": (1, 1), "two": 1.1 }, { "one": (2, 2), "two": 2.2 }, { "one": (3, 3), "two": 3.3 }, ], ] for datum in dataset: assert ak.to_list(ak.from_iter(datum)) == datum
def test_json(): assert ak.to_list(ak.Array("[1.1, 2.2, 3.3, null, 4.4]", check_valid=True)) == [ 1.1, 2.2, 3.3, None, 4.4, ] assert ak.to_list(ak.Array("[null, 2.2, 3.3, null, 4.4]", check_valid=True)) == [ None, 2.2, 3.3, None, 4.4, ] assert ak.to_list( ak.Array("[[1.1, 2.2, 3.3], [], [null, 4.4]]", check_valid=True)) == [[1.1, 2.2, 3.3], [], [None, 4.4]] assert ak.to_list( ak.Array("[[1.1, 2.2, 3.3], [], null, [null, 4.4]]", check_valid=True)) == [[1.1, 2.2, 3.3], [], None, [None, 4.4]] assert ak.to_list( ak.Array("[[1.1, 2.2, 3.3], null, [], [null, 4.4]]", check_valid=True)) == [[1.1, 2.2, 3.3], None, [], [None, 4.4]] assert ak.to_list( ak.Array("[[1.1, 2.2, 3.3], null, [], [null, 4.4]]", check_valid=True)) != [[1.1, 2.2, 3.3], [], None, [None, 4.4]] assert ak.to_list( ak.Array("[[null, 1.1, 2.2, 3.3], [], [null, 4.4]]", check_valid=True)) == [[None, 1.1, 2.2, 3.3], [], [None, 4.4]] assert ak.to_list( ak.Array("[[null, 1.1, 2.2, 3.3], [], null, [null, 4.4]]", check_valid=True)) == [[None, 1.1, 2.2, 3.3], [], None, [None, 4.4]] assert ak.to_list( ak.Array("[[null, 1.1, 2.2, 3.3], null, [], [null, 4.4]]", check_valid=True)) == [[None, 1.1, 2.2, 3.3], None, [], [None, 4.4]] assert ak.to_list( ak.Array("[[null, 1.1, 2.2, 3.3], null, [], [null, 4.4]]", check_valid=True)) != [[None, 1.1, 2.2, 3.3], [], None, [None, 4.4]] assert ak.to_list( ak.Array( "[null, [1.1, 2.2, 3.3], [], [null, 4.4]]", check_valid=True)) == [None, [1.1, 2.2, 3.3], [], [None, 4.4]] assert ak.to_list( ak.Array("[null, [1.1, 2.2, 3.3], [], null, [null, 4.4]]", check_valid=True)) == [ None, [1.1, 2.2, 3.3], [], None, [None, 4.4] ] assert ak.to_list( ak.Array("[null, [1.1, 2.2, 3.3], null, [], [null, 4.4]]", check_valid=True)) == [ None, [1.1, 2.2, 3.3], None, [], [None, 4.4] ] assert ak.to_list( ak.Array("[null, [1.1, 2.2, 3.3], null, [], [null, 4.4]]", check_valid=True)) != [ None, [1.1, 2.2, 3.3], [], None, [None, 4.4] ] assert ak.to_list( ak.Array("[null, null, null, null, null]", check_valid=True)) == [ None, None, None, None, None, ] assert ak.to_list( ak.Array("[[null, null, null], [], [null, null]]", check_valid=True)) == [[None, None, None], [], [None, None]]
def test_null(): content = ak.layout.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3, 4.4])) index = ak.layout.Index64(np.array([2, 2, 0, -1, 4], dtype=np.int64)) array = ak.layout.IndexedOptionArray64(index, content) assert ak.to_list(array) == [2.2, 2.2, 0.0, None, 4.4]
def test_builder(): assert ak.to_list(ak.Array([1.1, 2.2, 3.3, None, 4.4], check_valid=True)) == [ 1.1, 2.2, 3.3, None, 4.4, ] assert ak.to_list(ak.Array([None, 2.2, 3.3, None, 4.4], check_valid=True)) == [ None, 2.2, 3.3, None, 4.4, ] assert ak.to_list( ak.Array([[1.1, 2.2, 3.3], [], [None, 4.4]], check_valid=True)) == [[1.1, 2.2, 3.3], [], [None, 4.4]] assert ak.to_list( ak.Array([[1.1, 2.2, 3.3], [], None, [None, 4.4]], check_valid=True)) == [[1.1, 2.2, 3.3], [], None, [None, 4.4]] assert ak.to_list( ak.Array([[1.1, 2.2, 3.3], None, [], [None, 4.4]], check_valid=True)) == [[1.1, 2.2, 3.3], None, [], [None, 4.4]] assert ak.to_list( ak.Array([[1.1, 2.2, 3.3], None, [], [None, 4.4]], check_valid=True)) != [[1.1, 2.2, 3.3], [], None, [None, 4.4]] assert ak.to_list( ak.Array([[None, 1.1, 2.2, 3.3], [], [None, 4.4]], check_valid=True)) == [[None, 1.1, 2.2, 3.3], [], [None, 4.4]] assert ak.to_list( ak.Array([[None, 1.1, 2.2, 3.3], [], None, [None, 4.4]], check_valid=True)) == [[None, 1.1, 2.2, 3.3], [], None, [None, 4.4]] assert ak.to_list( ak.Array([[None, 1.1, 2.2, 3.3], None, [], [None, 4.4]], check_valid=True)) == [[None, 1.1, 2.2, 3.3], None, [], [None, 4.4]] assert ak.to_list( ak.Array([[None, 1.1, 2.2, 3.3], None, [], [None, 4.4]], check_valid=True)) != [[None, 1.1, 2.2, 3.3], [], None, [None, 4.4]] assert ak.to_list( ak.Array( [None, [1.1, 2.2, 3.3], [], [None, 4.4]], check_valid=True)) == [None, [1.1, 2.2, 3.3], [], [None, 4.4]] assert ak.to_list( ak.Array([None, [1.1, 2.2, 3.3], [], None, [None, 4.4]], check_valid=True)) == [ None, [1.1, 2.2, 3.3], [], None, [None, 4.4] ] assert ak.to_list( ak.Array([None, [1.1, 2.2, 3.3], None, [], [None, 4.4]], check_valid=True)) == [ None, [1.1, 2.2, 3.3], None, [], [None, 4.4] ] assert ak.to_list( ak.Array([None, [1.1, 2.2, 3.3], None, [], [None, 4.4]], check_valid=True)) != [ None, [1.1, 2.2, 3.3], [], None, [None, 4.4] ] assert ak.to_list( ak.Array([None, None, None, None, None], check_valid=True)) == [ None, None, None, None, None, ] assert ak.to_list( ak.Array([[None, None, None], [], [None, None]], check_valid=True)) == [[None, None, None], [], [None, None]]
def test_numpyarray_getitem_next(): a = np.arange(10) b = ak.layout.NumpyArray(a) c = np.array([7, 3, 3, 5]) assert ak.to_list(b[c]) == ak.to_list(a[c]) c = np.array( [False, False, False, True, True, True, False, True, False, True]) assert ak.to_list(b[c]) == ak.to_list(a[c]) c = np.array([], dtype=int) assert ak.to_list(b[c]) == ak.to_list(a[c]) a = np.arange(10 * 3).reshape(10, 3) b = ak.layout.NumpyArray(a) c = np.array([7, 3, 3, 5]) assert ak.to_list(b[c]) == ak.to_list(a[c]) c = np.array( [False, False, False, True, True, True, False, True, False, True]) assert ak.to_list(b[c]) == ak.to_list(a[c]) c = np.array([], dtype=int) assert ak.to_list(b[c]) == ak.to_list(a[c]) a = np.arange(7 * 5).reshape(7, 5) b = ak.layout.NumpyArray(a) c1 = np.array([4, 1, 1, 3]) c2 = np.array([2, 2, 0, 1]) assert ak.to_list(b[c1, c2]) == ak.to_list(a[c1, c2]) c1 = np.array([[4, 1], [1, 3], [0, 4]]) c2 = np.array([[2, 2], [0, 1], [1, 3]]) assert ak.to_list(b[c1, c2]) == ak.to_list(a[c1, c2]) c = np.array([False, False, True, True, False, True, True]) assert ak.to_list(b[c]) == ak.to_list(a[c]) c = a % 2 == 0 # two dimensional assert ak.to_list(b[c]) == ak.to_list(a[c]) c = np.array([], dtype=int) assert ak.to_list(b[c]) == ak.to_list(a[c]) c1 = np.array([], dtype=int) c2 = np.array([], dtype=int) assert ak.to_list(b[c1, c2]) == ak.to_list(a[c1, c2]) a = np.arange(7 * 5).reshape(7, 5) b = ak.layout.NumpyArray(a) c = np.array([2, 0, 0, 1]) assert ak.to_list(b[1:4, c]) == ak.to_list(a[1:4, c]) assert ak.to_list(b[c, 1:4]) == ak.to_list(a[c, 1:4]) assert ak.to_list(b[1:4, np.newaxis, c]) == ak.to_list(a[1:4, np.newaxis, c]) assert ak.to_list(b[c, np.newaxis, 1:4]) == ak.to_list(a[c, np.newaxis, 1:4]) assert ak.to_list(b[1:4, np.newaxis, np.newaxis, c]) == ak.to_list(a[1:4, np.newaxis, np.newaxis, c]) assert ak.to_list(b[c, np.newaxis, np.newaxis, 1:4]) == ak.to_list(a[c, np.newaxis, np.newaxis, 1:4]) if not ak._util.py27: assert ak.to_list(b[Ellipsis, c]) == ak.to_list(a[Ellipsis, c]) assert ak.to_list(b[c, Ellipsis]) == ak.to_list(a[c, Ellipsis])
def test_numpyarray_contiguous(): a = np.arange(10)[8::-2] b = ak.layout.NumpyArray(a) assert ak.to_list(b) == ak.to_list(a) assert ak.to_list(b.contiguous()) == ak.to_list(a) b = b.contiguous() assert ak.to_list(b) == ak.to_list(a) a = np.arange(7 * 5).reshape(7, 5)[::-1, ::2] b = ak.layout.NumpyArray(a) assert ak.to_list(b) == ak.to_list(a) assert ak.to_list(b.contiguous()) b = b.contiguous() assert ak.to_list(b) == ak.to_list(a)
def test_withfield(): base = ak.Array([{"x": 1}, {"x": 2}, {"x": 3}], check_valid=True) what = ak.Array([1.1, 2.2, 3.3], check_valid=True) assert ak.to_list(ak.with_field(base, what)) == [ { "x": 1, "1": 1.1 }, { "x": 2, "1": 2.2 }, { "x": 3, "1": 3.3 }, ] assert ak.to_list(ak.with_field(base, what, where="y")) == [ { "x": 1, "y": 1.1 }, { "x": 2, "y": 2.2 }, { "x": 3, "y": 3.3 }, ] base["z"] = what assert ak.to_list(base) == [ { "x": 1, "z": 1.1 }, { "x": 2, "z": 2.2 }, { "x": 3, "z": 3.3 }, ] base["q"] = 123 assert ak.to_list(base) == [ { "x": 1, "z": 1.1, "q": 123 }, { "x": 2, "z": 2.2, "q": 123 }, { "x": 3, "z": 3.3, "q": 123 }, ] base = ak.Array([{"x": 1}, {"x": 2}, {"x": 3}], check_valid=True)[2] assert ak.to_list(ak.with_field(base, 100, "y")) == {"x": 3, "y": 100}
def test_ListOffsetArray(): array = ak.from_iter(["one", "two", "three", "four", "five"], highlevel=False) assert ak.to_list(array.sort(0, True, True)) == [ "five", "four", "one", "three", "two", ] assert array.is_unique() is True array2 = ak.from_iter(["one", "two", "one", "four", "two"], highlevel=False) assert ak.to_list(array2.sort(0, True, True)) == [ "four", "one", "one", "two", "two", ] assert array2.is_unique() is False content = ak.layout.NumpyArray( np.array([3.3, 1.1, 2.2, 0.0, 4.4, 9.9, 6.6, 7.7, 8.8, 5.5])) offsets = ak.layout.Index64(np.array([0, 3, 3, 5, 6, 10, 10])) listoffsetarray = ak.layout.ListOffsetArray64(offsets, content) assert ak.to_list(listoffsetarray) == [ [3.3, 1.1, 2.2], [], [0.0, 4.4], [9.9], [6.6, 7.7, 8.8, 5.5], [], ] assert listoffsetarray.is_unique() is True content = ak.layout.NumpyArray( np.array([3.3, 1.1, 2.2, 0.0, 4.4, 9.9, 2.2, 3.3, 1.1, 5.5])) offsets = ak.layout.Index64(np.array([0, 3, 3, 5, 6, 9, 10])) listoffsetarray = ak.layout.ListOffsetArray64(offsets, content) assert ak.to_list(listoffsetarray) == [ [3.3, 1.1, 2.2], [], [0.0, 4.4], [9.9], [2.2, 3.3, 1.1], [5.5], ] assert listoffsetarray.is_unique() is False content2 = ak.layout.NumpyArray( np.array([0.0, 1.1, 2.2, 3.3, 1.1, 5.5, 6.6, 7.7, 2.2, 9.9])) offsets2 = ak.layout.Index64(np.array([0, 3, 3, 5, 6, 10, 10])) listoffsetarray2 = ak.layout.ListOffsetArray64(offsets2, content2) assert ak.to_list(listoffsetarray2) == [ [0.0, 1.1, 2.2], [], [3.3, 1.1], [5.5], [6.6, 7.7, 2.2, 9.9], [], ] assert listoffsetarray2.is_unique() is True
def test_regulararray(): content = ak.layout.NumpyArray( np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) recordarray = ak.layout.RecordArray({"x": content}) regulararray = ak.Array(ak.layout.RegularArray(recordarray, 3, zeros_length=0), check_valid=True) content2 = ak.layout.NumpyArray(np.array([100, 200, 300])) regulararray2 = ak.Array(ak.layout.RegularArray(content2, 1, zeros_length=0), check_valid=True) assert ak.to_list(ak.with_field(regulararray, regulararray2, "y")) == [ [{ "x": 0.0, "y": 100 }, { "x": 1.1, "y": 100 }, { "x": 2.2, "y": 100 }], [{ "x": 3.3, "y": 200 }, { "x": 4.4, "y": 200 }, { "x": 5.5, "y": 200 }], [{ "x": 6.6, "y": 300 }, { "x": 7.7, "y": 300 }, { "x": 8.8, "y": 300 }], ] content2 = ak.layout.NumpyArray( np.array([100, 200, 300, 400, 500, 600, 700, 800, 900])) regulararray2 = ak.Array(ak.layout.RegularArray(content2, 3, zeros_length=0), check_valid=True) assert ak.to_list(ak.with_field(regulararray, regulararray2, "y")) == [ [{ "x": 0.0, "y": 100 }, { "x": 1.1, "y": 200 }, { "x": 2.2, "y": 300 }], [{ "x": 3.3, "y": 400 }, { "x": 4.4, "y": 500 }, { "x": 5.5, "y": 600 }], [{ "x": 6.6, "y": 700 }, { "x": 7.7, "y": 800 }, { "x": 8.8, "y": 900 }], ] content2 = ak.Array(ak.layout.NumpyArray(np.array([[100], [200], [300]])), check_valid=True) assert ak.to_list(ak.with_field(regulararray, content2, "y")) == [ [{ "x": 0.0, "y": 100 }, { "x": 1.1, "y": 100 }, { "x": 2.2, "y": 100 }], [{ "x": 3.3, "y": 200 }, { "x": 4.4, "y": 200 }, { "x": 5.5, "y": 200 }], [{ "x": 6.6, "y": 300 }, { "x": 7.7, "y": 300 }, { "x": 8.8, "y": 300 }], ] content2 = ak.Array( ak.layout.NumpyArray( np.array([[100, 200, 300], [400, 500, 600], [700, 800, 900]])), check_valid=True, ) assert ak.to_list(ak.with_field(regulararray, content2, "y")) == [ [{ "x": 0.0, "y": 100 }, { "x": 1.1, "y": 200 }, { "x": 2.2, "y": 300 }], [{ "x": 3.3, "y": 400 }, { "x": 4.4, "y": 500 }, { "x": 5.5, "y": 600 }], [{ "x": 6.6, "y": 700 }, { "x": 7.7, "y": 800 }, { "x": 8.8, "y": 900 }], ]
def test_string_equal2(): assert ak.to_list( ak.Array(["one", "two", "three", "two", "two", "one"], check_valid=True) == "two") == [ False, True, False, True, True, False ]
def test_record(): array1 = ak.from_iter([{ "x": 1, "y": 1.1 }, { "x": 2, "y": 2.2 }, { "x": 3, "y": 3.3 }], highlevel=False) assert ak.to_list(array1) == [ { "x": 1, "y": 1.1 }, { "x": 2, "y": 2.2 }, { "x": 3, "y": 3.3 }, ] array2 = array1.setitem_field( "z", ak.from_iter([[], [1], [2, 2]], highlevel=False)) assert ak.to_list(array2) == [ { "x": 1, "y": 1.1, "z": [] }, { "x": 2, "y": 2.2, "z": [1] }, { "x": 3, "y": 3.3, "z": [2, 2] }, ] array3 = array1.setitem_field( None, ak.from_iter([[], [1], [2, 2]], highlevel=False)) assert ak.to_list(array3) == [ { "x": 1, "y": 1.1, "2": [] }, { "x": 2, "y": 2.2, "2": [1] }, { "x": 3, "y": 3.3, "2": [2, 2] }, ] array3 = array1.setitem_field( 0, ak.from_iter([[], [1], [2, 2]], highlevel=False)) assert ak.to_list(array3) == [ { "x": 1, "y": 1.1, "0": [] }, { "x": 2, "y": 2.2, "0": [1] }, { "x": 3, "y": 3.3, "0": [2, 2] }, ] array1 = ak.from_iter([(1, 1.1), (2, 2.2), (3, 3.3)], highlevel=False) assert ak.to_list(array1) == [(1, 1.1), (2, 2.2), (3, 3.3)] array2 = array1.setitem_field( "z", ak.from_iter([[], [1], [2, 2]], highlevel=False)) assert ak.to_list(array2) == [ { "0": 1, "1": 1.1, "z": [] }, { "0": 2, "1": 2.2, "z": [1] }, { "0": 3, "1": 3.3, "z": [2, 2] }, ] array3 = array1.setitem_field( None, ak.from_iter([[], [1], [2, 2]], highlevel=False)) assert ak.to_list(array3) == [(1, 1.1, []), (2, 2.2, [1]), (3, 3.3, [2, 2])] array3 = array1.setitem_field( 0, ak.from_iter([[], [1], [2, 2]], highlevel=False)) assert ak.to_list(array3) == [([], 1, 1.1), ([1], 2, 2.2), ([2, 2], 3, 3.3)] array3 = array1.setitem_field( 1, ak.from_iter([[], [1], [2, 2]], highlevel=False)) assert ak.to_list(array3) == [(1, [], 1.1), (2, [1], 2.2), (3, [2, 2], 3.3)] array3 = array1.setitem_field( 100, ak.from_iter([[], [1], [2, 2]], highlevel=False)) assert ak.to_list(array3) == [(1, 1.1, []), (2, 2.2, [1]), (3, 3.3, [2, 2])]
def test_flatten(): assert ak.to_list( ak.flatten(ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]], check_valid=True), axis=1)) == [1.1, 2.2, 3.3, 4.4, 5.5]
def test_rpad_and_clip_regular_array(): content = ak.layout.NumpyArray( np.array([ 2.1, 8.4, 7.4, 1.6, 2.2, 3.4, 6.2, 5.4, 1.5, 3.9, 3.8, 3.0, 8.5, 6.9, 4.3, 3.6, 6.7, 1.8, 3.2, ])) index = ak.layout.Index64( np.array([13, 9, 13, 4, 8, 3, 15, -1, 16, 2, 8], dtype=np.int64)) indexedarray = ak.layout.IndexedOptionArray64(index, content) array = ak.layout.RegularArray(indexedarray, 3, zeros_length=0) assert ak.to_list(array.rpad_and_clip(5, 0)) == [ [6.9, 3.9, 6.9], [2.2, 1.5, 1.6], [3.6, None, 6.7], None, None, ] assert ak.to_list(array.rpad_and_clip(4, 0)) == [ [6.9, 3.9, 6.9], [2.2, 1.5, 1.6], [3.6, None, 6.7], None, ] assert ak.to_list(array.rpad_and_clip(3, 0)) == [ [6.9, 3.9, 6.9], [2.2, 1.5, 1.6], [3.6, None, 6.7], ] assert ak.to_list(array.rpad_and_clip(2, 0)) == [[6.9, 3.9, 6.9], [2.2, 1.5, 1.6]] assert ak.to_list(array.rpad_and_clip(1, 0)) == [[6.9, 3.9, 6.9]] assert ak.to_list(array.rpad_and_clip(5, 1)) == [ [6.9, 3.9, 6.9, None, None], [2.2, 1.5, 1.6, None, None], [3.6, None, 6.7, None, None], ] assert ak.to_list(array.rpad_and_clip(4, 1)) == [ [6.9, 3.9, 6.9, None], [2.2, 1.5, 1.6, None], [3.6, None, 6.7, None], ] assert ak.to_list(array.rpad_and_clip(3, 1)) == [ [6.9, 3.9, 6.9], [2.2, 1.5, 1.6], [3.6, None, 6.7], ] assert ak.to_list(array.rpad_and_clip(2, 1)) == [ [6.9, 3.9], [2.2, 1.5], [3.6, None], ] assert ak.to_list(array.rpad_and_clip(1, 1)) == [[6.9], [2.2], [3.6]] array = ak.layout.NumpyArray(np.arange(2 * 3 * 5).reshape(2, 3, 5)) assert ak.to_list(array) == [ [[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(array.rpad_and_clip(7, 2)) == [ [ [0, 1, 2, 3, 4, None, None], [5, 6, 7, 8, 9, None, None], [10, 11, 12, 13, 14, None, None], ], [ [15, 16, 17, 18, 19, None, None], [20, 21, 22, 23, 24, None, None], [25, 26, 27, 28, 29, None, None], ], ] content = ak.layout.NumpyArray( np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) offsets = ak.layout.Index64(np.array([0, 3, 3, 5, 6, 10, 10])) listoffsetarray = ak.layout.ListOffsetArray64(offsets, content) regulararray = ak.layout.RegularArray(listoffsetarray, 2, zeros_length=0) assert ak.to_list(regulararray.rpad_and_clip(1, 0)) == [[[0.0, 1.1, 2.2], []]] assert ak.to_list(regulararray.rpad_and_clip(2, 0)) == [ [[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]], ] assert ak.to_list(regulararray.rpad_and_clip(3, 0)) == [ [[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], []], ] assert ak.to_list(regulararray.rpad_and_clip(4, 0)) == [ [[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], []], None, ] assert ak.to_list(regulararray.rpad_and_clip(5, 0)) == [ [[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], []], None, None, ] assert ak.to_list(regulararray.rpad_and_clip(1, 1)) == [ [[0.0, 1.1, 2.2]], [[3.3, 4.4]], [[6.6, 7.7, 8.8, 9.9]], ] assert ak.to_list(regulararray.rpad_and_clip(2, 1)) == [ [[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], []], ] assert ak.to_list(regulararray.rpad_and_clip(3, 1)) == [ [[0.0, 1.1, 2.2], [], None], [[3.3, 4.4], [5.5], None], [[6.6, 7.7, 8.8, 9.9], [], None], ] assert ak.to_list(regulararray.rpad_and_clip(7, 1)) == [ [[0.0, 1.1, 2.2], [], None, None, None, None, None], [[3.3, 4.4], [5.5], None, None, None, None, None], [[6.6, 7.7, 8.8, 9.9], [], None, None, None, None, None], ] assert ak.to_list(regulararray.rpad_and_clip(1, 2)) == [ [[0.0], [None]], [[3.3], [5.5]], [[6.6], [None]], ] assert ak.to_list(regulararray.rpad_and_clip(2, 2)) == [ [[0.0, 1.1], [None, None]], [[3.3, 4.4], [5.5, None]], [[6.6, 7.7], [None, None]], ] assert ak.to_list(regulararray.rpad_and_clip(3, 2)) == [ [[0.0, 1.1, 2.2], [None, None, None]], [[3.3, 4.4, None], [5.5, None, None]], [[6.6, 7.7, 8.8], [None, None, None]], ] assert ak.to_list(regulararray.rpad_and_clip(4, 2)) == [ [[0.0, 1.1, 2.2, None], [None, None, None, None]], [[3.3, 4.4, None, None], [5.5, None, None, None]], [[6.6, 7.7, 8.8, 9.9], [None, None, None, None]], ] assert ak.to_list(regulararray.rpad_and_clip(5, 2)) == [ [[0.0, 1.1, 2.2, None, None], [None, None, None, None, None]], [[3.3, 4.4, None, None, None], [5.5, None, None, None, None]], [[6.6, 7.7, 8.8, 9.9, None], [None, None, None, None, None]], ]
def test_records(): one = ak.Array([{ "x": 1, "y": [1] }, { "x": 2, "y": [1, 2] }, { "x": 3, "y": [1, 2, 3] }]).layout two = ak.Array([{"y": [], "x": 4}, {"y": [3, 2, 1], "x": 5}]).layout three = two[0:0] four = ak.Array([{"x": 6, "y": [1]}, {"x": 7, "y": [1, 2]}]).layout assert ak.to_list(one.mergemany([two, three, four])) == [ { "x": 1, "y": [1] }, { "x": 2, "y": [1, 2] }, { "x": 3, "y": [1, 2, 3] }, { "y": [], "x": 4 }, { "y": [3, 2, 1], "x": 5 }, { "x": 6, "y": [1] }, { "x": 7, "y": [1, 2] }, ] three = ak.layout.EmptyArray() assert ak.to_list(one.mergemany([two, three, four])) == [ { "x": 1, "y": [1] }, { "x": 2, "y": [1, 2] }, { "x": 3, "y": [1, 2, 3] }, { "y": [], "x": 4 }, { "y": [3, 2, 1], "x": 5 }, { "x": 6, "y": [1] }, { "x": 7, "y": [1, 2] }, ]
def test_rpad_numpy_array(): array = ak.layout.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5])) assert ak.to_list(array.rpad(10, 0)) == [ 1.1, 2.2, 3.3, 4.4, 5.5, None, None, None, None, None, ] array = ak.layout.NumpyArray(np.array([[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]])) assert ak.to_list(array.rpad(5, 0)) == [ [1.1, 2.2, 3.3], [4.4, 5.5, 6.6], None, None, None, ] assert ak.to_list(array.rpad(5, 1)) == [ [1.1, 2.2, 3.3, None, None], [4.4, 5.5, 6.6, None, None], ] array = ak.layout.NumpyArray( np.arange(2 * 3 * 5, dtype=np.int64).reshape(2, 3, 5)) assert ak.to_list(array) == [ [[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(array.rpad(1, 0)) == [ [[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(array.rpad(2, 0)) == [ [[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(array.rpad(3, 0)) == [ [[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]], None, ] assert ak.to_list(array.rpad(4, 0)) == [ [[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]], None, None, ] assert ak.to_list(array.rpad(5, 0)) == [ [[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]], None, None, None, ] assert ak.to_list(array.rpad(2, 1)) == [ [[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(array.rpad(3, 1)) == [ [[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(array.rpad(4, 1)) == [ [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, 14], None], [[15, 16, 17, 18, 19], [20, 21, 22, 23, 24], [25, 26, 27, 28, 29], None], ] assert ak.to_list(array.rpad(5, 1)) == [ [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, 14], None, None], [[15, 16, 17, 18, 19], [20, 21, 22, 23, 24], [25, 26, 27, 28, 29], None, None], ] assert ak.to_list(array.rpad(3, 2)) == [ [[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(array.rpad(7, 2)) == [ [ [0, 1, 2, 3, 4, None, None], [5, 6, 7, 8, 9, None, None], [10, 11, 12, 13, 14, None, None], ], [ [15, 16, 17, 18, 19, None, None], [20, 21, 22, 23, 24, None, None], [25, 26, 27, 28, 29, None, None], ], ] assert ak.to_list(array.rpad(2, 2)) == [ [[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]], ]
def test_bytemasked(): one = (ak.Array([1, 2, 3, 4, 5, 6]).mask[[True, True, False, True, False, True]].layout) two = ak.Array([7, 99, 999, 8, 9]).mask[[True, False, False, True, True]].layout three = ak.Array([100, 200, 300]).layout four = ak.Array([None, None, 123, None]).layout assert ak.to_list(one.mergemany([two, three, four])) == [ 1, 2, None, 4, None, 6, 7, None, None, 8, 9, 100, 200, 300, None, None, 123, None, ] assert ak.to_list(four.mergemany([three, two, one])) == [ None, None, 123, None, 100, 200, 300, 7, None, None, 8, 9, 1, 2, None, 4, None, 6, ] assert ak.to_list(three.mergemany([four, one])) == [ 100, 200, 300, None, None, 123, None, 1, 2, None, 4, None, 6, ] assert ak.to_list(three.mergemany([four, one, two])) == [ 100, 200, 300, None, None, 123, None, 1, 2, None, 4, None, 6, 7, None, None, 8, 9, ] assert ak.to_list(three.mergemany([two, one])) == [ 100, 200, 300, 7, None, None, 8, 9, 1, 2, None, 4, None, 6, ] assert ak.to_list(three.mergemany([two, one, four])) == [ 100, 200, 300, 7, None, None, 8, 9, 1, 2, None, 4, None, 6, None, None, 123, None, ]
def test_rpad_list_array(): content = ak.layout.NumpyArray( np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) starts = ak.layout.Index64(np.array([0, 3, 4, 5, 8])) stops = ak.layout.Index64(np.array([3, 3, 6, 8, 9])) array = ak.layout.ListArray64(starts, stops, content) assert ak.to_list(array) == [ [0.0, 1.1, 2.2], [], [4.4, 5.5], [5.5, 6.6, 7.7], [8.8], ] assert ak.to_list(array.rpad(1, 0)) == [ [0.0, 1.1, 2.2], [], [4.4, 5.5], [5.5, 6.6, 7.7], [8.8], ] assert str(ak.type(array)) == str(ak.type(array.rpad(1, 0))) assert ak.to_list(array.rpad(2, 0)) == [ [0.0, 1.1, 2.2], [], [4.4, 5.5], [5.5, 6.6, 7.7], [8.8], ] assert str(ak.type(array)) == str(ak.type(array.rpad(2, 0))) assert ak.to_list(array.rpad(7, 0)) == [ [0.0, 1.1, 2.2], [], [4.4, 5.5], [5.5, 6.6, 7.7], [8.8], None, None, ] assert "option[" + str(ak.type(array)) + "]" == str( ak.type(array.rpad(7, 0))) assert ak.to_list(array.rpad(1, 1)) == [ [0.0, 1.1, 2.2], [None], [4.4, 5.5], [5.5, 6.6, 7.7], [8.8], ] assert ak.to_list(array.rpad(2, 1)) == [ [0.0, 1.1, 2.2], [None, None], [4.4, 5.5], [5.5, 6.6, 7.7], [8.8, None], ] assert ak.to_list(array.rpad(3, 1)) == [ [0.0, 1.1, 2.2], [None, None, None], [4.4, 5.5, None], [5.5, 6.6, 7.7], [8.8, None, None], ] assert ak.to_list(array.rpad(4, 1)) == [ [0.0, 1.1, 2.2, None], [None, None, None, None], [4.4, 5.5, None, None], [5.5, 6.6, 7.7, None], [8.8, None, None, None], ]
def test_union_option(): one = ak.Array([1, 2, [], [3, 4]]).layout two = ak.Array([100, None, 300]).layout three = ak.Array([{"x": 1}, {"x": 2}, 5, 6, 7]).layout assert ak.to_list(one.mergemany([two, three])) == [ 1, 2, [], [3, 4], 100, None, 300, { "x": 1 }, { "x": 2 }, 5, 6, 7, ] assert ak.to_list(one.mergemany([three, two])) == [ 1, 2, [], [3, 4], { "x": 1 }, { "x": 2 }, 5, 6, 7, 100, None, 300, ] assert ak.to_list(two.mergemany([one, three])) == [ 100, None, 300, 1, 2, [], [3, 4], { "x": 1 }, { "x": 2 }, 5, 6, 7, ] assert ak.to_list(two.mergemany([three, one])) == [ 100, None, 300, { "x": 1 }, { "x": 2 }, 5, 6, 7, 1, 2, [], [3, 4], ] assert ak.to_list(three.mergemany([one, two])) == [ { "x": 1 }, { "x": 2 }, 5, 6, 7, 1, 2, [], [3, 4], 100, None, 300, ] assert ak.to_list(three.mergemany([two, one])) == [ { "x": 1 }, { "x": 2 }, 5, 6, 7, 100, None, 300, 1, 2, [], [3, 4], ] one = ak.Array([1, 2, [], [3, 4]]).layout two = ak.Array([100, None, 300]).layout three = ak.Array([{"x": 1}, {"x": 2}, 5, None, 7]).layout assert ak.to_list(one.mergemany([two, three])) == [ 1, 2, [], [3, 4], 100, None, 300, { "x": 1 }, { "x": 2 }, 5, None, 7, ] assert ak.to_list(one.mergemany([three, two])) == [ 1, 2, [], [3, 4], { "x": 1 }, { "x": 2 }, 5, None, 7, 100, None, 300, ] assert ak.to_list(two.mergemany([one, three])) == [ 100, None, 300, 1, 2, [], [3, 4], { "x": 1 }, { "x": 2 }, 5, None, 7, ] assert ak.to_list(two.mergemany([three, one])) == [ 100, None, 300, { "x": 1 }, { "x": 2 }, 5, None, 7, 1, 2, [], [3, 4], ] assert ak.to_list(three.mergemany([one, two])) == [ { "x": 1 }, { "x": 2 }, 5, None, 7, 1, 2, [], [3, 4], 100, None, 300, ] assert ak.to_list(three.mergemany([two, one])) == [ { "x": 1 }, { "x": 2 }, 5, None, 7, 100, None, 300, 1, 2, [], [3, 4], ] one = ak.Array([1, 2, [], [3, 4]]).layout two = ak.Array([100, 200, 300]).layout three = ak.Array([{"x": 1}, {"x": 2}, 5, None, 7]).layout assert ak.to_list(one.mergemany([two, three])) == [ 1, 2, [], [3, 4], 100, 200, 300, { "x": 1 }, { "x": 2 }, 5, None, 7, ] assert ak.to_list(one.mergemany([three, two])) == [ 1, 2, [], [3, 4], { "x": 1 }, { "x": 2 }, 5, None, 7, 100, 200, 300, ] assert ak.to_list(two.mergemany([one, three])) == [ 100, 200, 300, 1, 2, [], [3, 4], { "x": 1 }, { "x": 2 }, 5, None, 7, ] assert ak.to_list(two.mergemany([three, one])) == [ 100, 200, 300, { "x": 1 }, { "x": 2 }, 5, None, 7, 1, 2, [], [3, 4], ] assert ak.to_list(three.mergemany([one, two])) == [ { "x": 1 }, { "x": 2 }, 5, None, 7, 1, 2, [], [3, 4], 100, 200, 300, ] assert ak.to_list(three.mergemany([two, one])) == [ { "x": 1 }, { "x": 2 }, 5, None, 7, 100, 200, 300, 1, 2, [], [3, 4], ]
def test_pyarrow_constructor(): a = ak.Array(pyarrow.array([[1.1, 2.2, 3.3], [], [4.4, 5.5]])) assert ak.to_list(a) == [[1.1, 2.2, 3.3], [], [4.4, 5.5]]
def test_unmasked_array(): content = ak.layout.NumpyArray( np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) layout = ak.layout.UnmaskedArray(content) packed = ak.packed(layout, highlevel=False) assert ak.to_list(packed) == ak.to_list(layout)
def test_listoffsetarray_slice_slice(): array1 = ak.layout.ListOffsetArrayU32(offsets1, content) assert ak.to_list(array1[2:]) == [[4.4, 5.5], [6.6], [7.7, 8.8, 9.9]] assert ak.to_list(array1[2:, 1:]) == [[5.5], [], [8.8, 9.9]] assert ak.to_list(array1[2:, :-1]) == [[4.4], [], [7.7, 8.8]]
def test_getitem_deeper(): assert ak.to_list(listarray) == [ [[[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]]], [[[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], []]], ] assert ak.to_list(listarray[0, 0, 0]) == [0.0, 1.1, 2.2] assert ak.to_list(listarray[0, 0, 1]) == [] assert ak.to_list(listarray[0, 1, 0]) == [3.3, 4.4] assert ak.to_list(listarray[0, 1, 1]) == [5.5] assert ak.to_list(listarray[1, 0, 0]) == [3.3, 4.4] assert ak.to_list(listarray[1, 0, 1]) == [5.5] assert ak.to_list(listarray[1, 1, 0]) == [6.6, 7.7, 8.8, 9.9] assert ak.to_list(listarray[1, 1, 1]) == [] assert ak.to_list(listarray[0, 0, 0:]) == [[0.0, 1.1, 2.2], []] assert ak.to_list(listarray[0, 0, 1:]) == [[]] assert ak.to_list(listarray[0, 1, 0:]) == [[3.3, 4.4], [5.5]] assert ak.to_list(listarray[0, 1, 1:]) == [[5.5]] assert ak.to_list(listarray[1, 0, 0:]) == [[3.3, 4.4], [5.5]] assert ak.to_list(listarray[1, 0, 1:]) == [[5.5]] assert ak.to_list(listarray[1, 1, 0:]) == [[6.6, 7.7, 8.8, 9.9], []] assert ak.to_list(listarray[1, 1, 1:]) == [[]] assert ak.to_list(listarray[[1], 0, 0:]) == [[[3.3, 4.4], [5.5]]] assert ak.to_list(listarray[[1, 0], 0, 0:]) == [ [[3.3, 4.4], [5.5]], [[0.0, 1.1, 2.2], []], ] assert ak.to_list(listarray[:, :, [0, 1]]) == [ [[[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]]], [[[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], []]], ] assert ak.to_list(listarray[:, :, [1, 0]]) == [ [[[], [0.0, 1.1, 2.2]], [[5.5], [3.3, 4.4]]], [[[5.5], [3.3, 4.4]], [[], [6.6, 7.7, 8.8, 9.9]]], ] assert ak.to_list(listarray[:, :, [1, 0, 1]]) == [ [[[], [0.0, 1.1, 2.2], []], [[5.5], [3.3, 4.4], [5.5]]], [[[5.5], [3.3, 4.4], [5.5]], [[], [6.6, 7.7, 8.8, 9.9], []]], ] assert ak.to_list(listarray[:, :2, [0, 1]]) == [ [[[0.0, 1.1, 2.2], []], [[3.3, 4.4], [5.5]]], [[[3.3, 4.4], [5.5]], [[6.6, 7.7, 8.8, 9.9], []]], ] assert ak.to_list(listarray[:1, [0, 0, 1, 1], [0, 1, 0, 1]]) == [[[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5]]] assert ak.to_list(listarray[:1, [1, 1, 0, 0], [1, 0, 1, 0]]) == [[[5.5], [3.3, 4.4], [], [0.0, 1.1, 2.2]]]