def test():
content = ak.repartition(range(10), 3)
assert ak.unflatten(content, [3, 0, 2, 1, 3, 1]).tolist() == [
[0, 1, 2],
[],
[3, 4],
[5],
[6, 7, 8],
[9],
]
original = ak.repartition(ak.Array([[1, 2, 3, 4], [5, 6, 7], [8, 9]]), 2)
assert ak.unflatten(original, [2, 2, 1, 2, 1, 1], axis=1).tolist() == [
[[1, 2], [3, 4]],
[[5], [6, 7]],
[[8], [9]],
]
original = ak.repartition(ak.Array([[1, 2, 3, 4], [], [5, 6, 7], [8, 9]]),
2)
assert ak.unflatten(original, [2, 2, 1, 2, 1, 1], axis=1).tolist() == [
[[1, 2], [3, 4]],
[],
[[5], [6, 7]],
[[8], [9]],
]
def test_view():
aslist = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
asarray = ak.repartition(ak.Array(aslist), 3)
asview = ak_numba_arrayview.ArrayView.fromarray(asarray)
for start in range(10):
for stop in range(start, 10):
asview.start = start
asview.stop = stop
assert ak.to_list(asview.toarray()) == aslist[start:stop]
asarray = ak.repartition(ak.Array(aslist), [3, 2, 0, 1, 4])
asview = ak_numba_arrayview.ArrayView.fromarray(asarray)
for start in range(10):
for stop in range(start, 10):
asview.start = start
asview.stop = stop
assert ak.to_list(asview.toarray()) == aslist[start:stop]
aslist = [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]]
asarray = ak.repartition(ak.Array(aslist), 3)
asview = ak_numba_arrayview.ArrayView.fromarray(asarray)
for start in range(5):
for stop in range(start, 5):
asview.start = start
asview.stop = stop
assert ak.to_list(asview.toarray()) == aslist[start:stop]
def test_with_name_field():
array = ak.Array(
[
{"x": 0.0, "y": []},
{"x": 1.1, "y": [1]},
{"x": 2.2, "y": [2, 2]},
{"x": 3.3, "y": [3, 3, 3]},
]
)
array2 = ak.repartition(array, 2)
z = ak.Array([100, 200, 300, 400])
z2 = ak.repartition(z, 3)
one = ak.with_name(array, "Wilbur")
assert isinstance(one.layout, ak.layout.Content)
assert one.layout.parameters["__record__"] == "Wilbur"
one = ak.with_name(array2, "Wilbur")
assert isinstance(one.layout, ak.partition.PartitionedArray)
assert one.layout.partition(0).parameters["__record__"] == "Wilbur"
assert one.layout.partition(1).parameters["__record__"] == "Wilbur"
two = ak.with_field(array, z, "z")
assert isinstance(two.layout, ak.layout.Content)
assert ak.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 = ak.with_field(array2, z, "z")
assert isinstance(two.layout, ak.partition.PartitionedArray)
assert ak.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 = ak.with_field(array2, z2, "z")
assert isinstance(two.layout, ak.partition.PartitionedArray)
assert ak.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 = ak.with_field(array2, z2, "z")
assert isinstance(two.layout, ak.partition.PartitionedArray)
assert ak.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_iter():
asnumpy = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9])
assert sys.getrefcount(asnumpy) == 2
array = ak.repartition(ak.Array(asnumpy), 3)
assert sys.getrefcount(asnumpy) == 3
@numba.njit
def f1(x):
out = 0
for xi in x:
out += xi
return out
for _ in range(10):
assert f1(array) == 45
assert sys.getrefcount(asnumpy) == 3
del array
assert sys.getrefcount(asnumpy) == 2
aslist = [
{"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": []},
]
asarray = ak.repartition(ak.Array(aslist), 2)
@numba.njit
def f2(x):
i = 0
for xi in x:
if i == 6:
return xi["y"]
i += 1
assert ak.to_list(f2(asarray)) == [6, 6]
@numba.njit
def f3(x):
i = 0
for xi in x:
if i == 6:
return xi
i += 1
assert ak.to_list(f3(asarray)) == {"x": 6.6, "y": [6, 6]}
def test():
array = ak.Array([3, 3, 3, 5, 5, 9, 9, 9, 9, 1, 3, 3])
assert ak.run_lengths(array).tolist() == [3, 2, 4, 1, 2]
array = ak.Array([[3, 3, 3, 5], [5], [], [9, 9], [9, 9], [1, 3, 3]])
assert ak.run_lengths(array).tolist() == [[3, 1], [1], [], [2], [2], [1, 2]]
array = ak.repartition(ak.Array([3, 3, 3, 5, 5, 9, 9, 9, 9, 1, 3, 3]), 7)
assert ak.run_lengths(array).tolist() == [3, 2, 4, 1, 2]
array = ak.repartition(
ak.Array([[3, 3, 3, 5], [5], [], [9, 9], [9, 9], [1, 3, 3]]), 4
)
assert ak.run_lengths(array).tolist() == [[3, 1], [1], [], [2], [2], [1, 2]]
def test_len():
array = ak.repartition(ak.Array([1.1, 2.2, 3.3, 4.4, 5.5]), 3)
@numba.njit
def f1(x):
return len(x)
assert f1(array) == 5
aslist = [
{
"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": []
},
]
asarray = ak.repartition(ak.Array(aslist), 2)
assert f1(asarray) == 9
def test_getitem_1a():
array = ak.repartition(
ak.Array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9]), 3
)
@numba.njit
def f1(x, i):
return x[i]
assert [f1(array, i) for i in range(10)] == [
0.0,
1.1,
2.2,
3.3,
4.4,
5.5,
6.6,
7.7,
8.8,
9.9,
]
assert [f1(array, -i) for i in range(1, 11)] == [
9.9,
8.8,
7.7,
6.6,
5.5,
4.4,
3.3,
2.2,
1.1,
0.0,
]
def test_boxing1():
asnumpy = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
assert sys.getrefcount(asnumpy) == 2
aslayout = ak.layout.NumpyArray(asnumpy)
aspart = ak.repartition(aslayout, 3, highlevel=False)
asarray = ak.Array(aspart)
aspart = asarray._layout
assert (
sys.getrefcount(asnumpy),
sys.getrefcount(aslayout),
sys.getrefcount(aspart),
) == (3, 2, 3)
@numba.njit
def f1(x):
return 3.14
for _ in range(5):
f1(asarray)
assert (
sys.getrefcount(asnumpy),
sys.getrefcount(aslayout),
sys.getrefcount(aspart),
) == (3, 2, 3)
del asarray
del aspart
del aslayout
import gc
gc.collect()
assert sys.getrefcount(asnumpy) == 2
def test(tmp_path):
filename = os.path.join(tmp_path, "test.parquet")
ak.to_parquet(ak.repartition(range(8), 2), filename)
assert ak.from_parquet(filename, row_groups=[1,
3]).tolist() == [2, 3, 6, 7]
assert ak.from_parquet(filename, row_groups=[1, 3],
lazy=True).tolist() == [
2,
3,
6,
7,
]
assert ak.from_parquet(tmp_path, row_groups=[1,
3]).tolist() == [2, 3, 6, 7]
assert ak.from_parquet(tmp_path, row_groups=[1, 3],
lazy=True).tolist() == [
2,
3,
6,
7,
]
ak.to_parquet.dataset(tmp_path)
assert ak.from_parquet(tmp_path, row_groups=[1,
3]).tolist() == [2, 3, 6, 7]
assert ak.from_parquet(tmp_path, row_groups=[1, 3],
lazy=True).tolist() == [
2,
3,
6,
7,
]
def test_getitem_1b():
asnumpy = np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])
array = ak.repartition(ak.Array(asnumpy), 3)
assert sys.getrefcount(asnumpy) == 3
@numba.njit
def f2(x, i1, i2):
out = x[i1:i2]
return out
assert isinstance(f2(array, 0, 10).layout, ak.partition.PartitionedArray)
assert isinstance(f2(array, 4, 5).layout, ak.partition.PartitionedArray)
assert isinstance(f2(array, 5, 5).layout, ak.partition.PartitionedArray)
for start in range(-10, 10):
for stop in range(-10, 10):
assert (
ak.to_list(f2(array, start, stop))
== [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9][start:stop]
)
assert sys.getrefcount(asnumpy) == 3
del array
assert sys.getrefcount(asnumpy) == 2
def test_enumerate_Partitioned():
@numba.njit
def f1(array):
for _, _ in enumerate(array):
pass
array = ak.repartition(ak.Array(range(100)), 10)
f1(array)
def test_atleast_1d():
array = ak.Array([1.1, 2.2, 3.3, 4.4, 5.5])
array2 = ak.repartition(array, 2)
one = ak.atleast_1d(array)
assert isinstance(one, np.ndarray)
assert ak.to_list(one) == [1.1, 2.2, 3.3, 4.4, 5.5]
one = ak.atleast_1d(array2)
assert isinstance(one, np.ndarray)
assert ak.to_list(one) == [1.1, 2.2, 3.3, 4.4, 5.5]
def test_mask2():
array = ak.Array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])
assert isinstance(array.layout, ak.layout.Content)
mask = ak.Array([False, False, True, True, False, True, True, False, True])
assert isinstance(mask.layout, ak.layout.Content)
one = array.mask[mask]
assert isinstance(one.layout, ak.layout.Content)
assert ak.to_list(one) == [None, None, 3.3, 4.4, None, 6.6, 7.7, None, 9.9]
array = ak.repartition(array, 4)
assert isinstance(array.layout, ak.partition.PartitionedArray)
one = array.mask[mask]
assert isinstance(one.layout, ak.partition.PartitionedArray)
assert ak.to_list(one) == [None, None, 3.3, 4.4, None, 6.6, 7.7, None, 9.9]
mask = ak.repartition(mask, 3)
assert isinstance(mask.layout, ak.partition.PartitionedArray)
one = array.mask[mask]
assert isinstance(one.layout, ak.partition.PartitionedArray)
assert ak.to_list(one) == [None, None, 3.3, 4.4, None, 6.6, 7.7, None, 9.9]
array = ak.repartition(array, None)
assert isinstance(array.layout, ak.layout.Content)
one = array.mask[mask]
assert isinstance(one.layout, ak.partition.PartitionedArray)
assert ak.to_list(one) == [None, None, 3.3, 4.4, None, 6.6, 7.7, None, 9.9]
mask = ak.repartition(mask, None)
assert isinstance(mask.layout, ak.layout.Content)
one = array.mask[mask]
assert isinstance(one.layout, ak.layout.Content)
assert ak.to_list(one) == [None, None, 3.3, 4.4, None, 6.6, 7.7, None, 9.9]
def test_firsts_singletons():
array = ak.Array([None, 1.1, 2.2, None, 3.3, None, None, 4.4, 5.5, None])
one = ak.singletons(array)
assert ak.to_list(one) == [[], [1.1], [2.2], [], [3.3], [], [], [4.4], [5.5], []]
two = ak.firsts(one)
assert ak.to_list(two) == [None, 1.1, 2.2, None, 3.3, None, None, 4.4, 5.5, None]
array = ak.repartition(array, 3)
assert isinstance(array.layout, ak.partition.PartitionedArray)
one = ak.singletons(array)
assert isinstance(one.layout, ak.partition.PartitionedArray)
assert ak.to_list(one) == [[], [1.1], [2.2], [], [3.3], [], [], [4.4], [5.5], []]
two = ak.firsts(one)
assert isinstance(two.layout, ak.partition.PartitionedArray)
assert ak.to_list(two) == [None, 1.1, 2.2, None, 3.3, None, None, 4.4, 5.5, None]
def test_lazy_partitioned():
array = ak.repartition(ak.Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]), 3)
form, length, container = ak.to_buffers(array)
assert length == [3, 3, 3, 1]
assert ak.from_buffers(form, length, container, lazy=True).tolist() == [
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
]
def test_boxing3():
asnumpy = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
assert sys.getrefcount(asnumpy) == 2
aslayout = ak.layout.NumpyArray(asnumpy)
aspart = ak.repartition(aslayout, 3, highlevel=False)
asarray = ak.Array(aspart)
aspart = asarray._layout
assert (
sys.getrefcount(asnumpy),
sys.getrefcount(aslayout),
sys.getrefcount(aspart),
) == (3, 2, 3)
@numba.njit
def f3(x):
return x, x
for _ in range(10):
out1, out2 = f3(asarray)
assert isinstance(out1.layout, ak.partition.PartitionedArray)
assert isinstance(out2.layout, ak.partition.PartitionedArray)
assert ak.to_list(out1) == [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
assert ak.to_list(out2) == [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
assert (
sys.getrefcount(asnumpy),
sys.getrefcount(aslayout),
sys.getrefcount(aspart),
) == (3, 2, 3)
del out1
del out2
del asarray
del aspart
del aslayout
import gc
gc.collect()
assert sys.getrefcount(asnumpy) == 2
def test():
array = ak.repartition(ak.Array(range(10)), [3, 3, 3, 1])
arrow = ak.to_arrow(array)
assert isinstance(arrow, pyarrow.ChunkedArray)
assert arrow.to_pylist() == list(range(10))
def check_awkward_accessible(col: ak.Array):
'Check to make sure we can look at every item in column'
ak.repartition(col, 3) # type: ignore
def test_partitioned():
array = ak.repartition(ak.Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]), 3)
form, length, container = ak.to_buffers(array)
assert ak.from_buffers(form, length, container).tolist() == [
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
]
form, length, container = ak.to_buffers(array)
assert ak.from_buffers(form, length, container).tolist() == [
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
]
one = ak.Array([1, 2, 3, 4, 5])
two = ak.Array([6, 7, 8, 9, 10])
container = {}
lengths = []
form1, length, _ = ak.to_buffers(one, container, 0)
lengths.append(length)
form2, length, _ = ak.to_buffers(two, container, 1)
lengths.append(length)
assert form1 == form2
assert ak.from_buffers(form1, lengths, container).tolist() == [
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
]
assert pickle.loads(pickle.dumps(array, -1)).tolist() == [
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
]
def test_0167_strings():
array = ak.repartition(
ak.Array(["one", "two", "three", "two", "two", "one", "three"]), 3)
assert ak.to_list(array == "two") == [
False, True, False, True, True, False, False
]
assert ak.to_list("two" == array) == [
False, True, False, True, True, False, False
]
assert ak.to_list(array == ["two"]) == [
False,
True,
False,
True,
True,
False,
False,
]
assert ak.to_list(["two"] == array) == [
False,
True,
False,
True,
True,
False,
False,
]
assert ak.to_list(array == ak.Array(["two"])) == [
False,
True,
False,
True,
True,
False,
False,
]
assert ak.to_list(ak.Array(["two"]) == array) == [
False,
True,
False,
True,
True,
False,
False,
]
array = ak.Array([["one", "two", "three"], [], ["two"], ["two", "one"],
["three"]])
assert ak.to_list(array == "two") == [
[False, True, False],
[],
[True],
[True, False],
[False],
]
assert ak.to_list("two" == array) == [
[False, True, False],
[],
[True],
[True, False],
[False],
]
assert ak.to_list(array == ["two"]) == [
[False, True, False],
[],
[True],
[True, False],
[False],
]
assert ak.to_list(["two"] == array) == [
[False, True, False],
[],
[True],
[True, False],
[False],
]
assert ak.to_list(array == ak.Array(["two"])) == [
[False, True, False],
[],
[True],
[True, False],
[False],
]
assert ak.to_list(ak.Array(["two"]) == array) == [
[False, True, False],
[],
[True],
[True, False],
[False],
]
array = ak.Array([["one", "two", "three"], [], ["two"], ["two", "one"],
["three"]])
assert ak.to_list(array == ["three", "two", "one", "one", "three"]) == [
[False, False, True],
[],
[False],
[False, True],
[True],
]
assert ak.to_list(["three", "two", "one", "one", "three"] == array) == [
[False, False, True],
[],
[False],
[False, True],
[True],
]
assert ak.to_list(
array == ak.Array(["three", "two", "one", "one", "three"])) == [
[False, False, True],
[],
[False],
[False, True],
[True],
]
assert ak.to_list(
ak.Array(["three", "two", "one", "one", "three"]) == array) == [
[False, False, True],
[],
[False],
[False, True],
[True],
]
def test_repartition_again():
array = ak.Array([[], [1.1, 2.2, 3.3], [], [], [4.4, 5.5], [], [6.6],
[7.7, 8.8, 9.9], []])
array2 = ak.repartition(array, 2)
array3 = ak.repartition(array, 3)
array4 = ak.repartition(array, [3, 2, 3, 1])
array5 = ak.repartition(array2, None)
assert isinstance(array.layout, ak.layout.Content)
assert isinstance(array2.layout, ak.partition.PartitionedArray)
assert isinstance(array3.layout, ak.partition.PartitionedArray)
assert isinstance(array4.layout, ak.partition.PartitionedArray)
assert isinstance(array5.layout, ak.layout.Content)
assert ak.to_list(array) == [
[],
[1.1, 2.2, 3.3],
[],
[],
[4.4, 5.5],
[],
[6.6],
[7.7, 8.8, 9.9],
[],
]
assert ak.to_list(array2) == [
[],
[1.1, 2.2, 3.3],
[],
[],
[4.4, 5.5],
[],
[6.6],
[7.7, 8.8, 9.9],
[],
]
assert ak.to_list(array3) == [
[],
[1.1, 2.2, 3.3],
[],
[],
[4.4, 5.5],
[],
[6.6],
[7.7, 8.8, 9.9],
[],
]
assert ak.to_list(array4) == [
[],
[1.1, 2.2, 3.3],
[],
[],
[4.4, 5.5],
[],
[6.6],
[7.7, 8.8, 9.9],
[],
]
assert ak.to_list(array5) == [
[],
[1.1, 2.2, 3.3],
[],
[],
[4.4, 5.5],
[],
[6.6],
[7.7, 8.8, 9.9],
[],
]
assert [ak.to_list(x) for x in array2.layout.partitions] == [
[[], [1.1, 2.2, 3.3]],
[[], []],
[[4.4, 5.5], []],
[[6.6], [7.7, 8.8, 9.9]],
[[]],
]
assert [ak.to_list(x) for x in array3.layout.partitions] == [
[[], [1.1, 2.2, 3.3], []],
[[], [4.4, 5.5], []],
[[6.6], [7.7, 8.8, 9.9], []],
]
assert [ak.to_list(x) for x in array4.layout.partitions] == [
[[], [1.1, 2.2, 3.3], []],
[[], [4.4, 5.5]],
[[], [6.6], [7.7, 8.8, 9.9]],
[[]],
]
def test():
assert ak.flatten(
ak.repartition(ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]]),
3)).tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
def test_explode(tmp_path):
array3 = ak.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 = ak.repartition(array3, 2)
ak.to_parquet(array3,
os.path.join(tmp_path, "array3.parquet"),
explode_records=True)
ak.to_parquet(array4,
os.path.join(tmp_path, "array4.parquet"),
explode_records=True)
assert ak.from_parquet(os.path.join(tmp_path,
"array3.parquet")).tolist() == [
{
"x": [1, 2, 3],
"y": [1.1, 2.2, 3.3]
},
{
"x": [],
"y": []
},
{
"x": [4, 5],
"y": [4.4, 5.5]
},
{
"x": [],
"y": []
},
{
"x": [],
"y": []
},
{
"x": [6, 7, 8, 9],
"y": [6.6, 7.7, 8.8, 9.9]
},
]
assert ak.from_parquet(os.path.join(tmp_path,
"array4.parquet")).tolist() == [
{
"x": [1, 2, 3],
"y": [1.1, 2.2, 3.3]
},
{
"x": [],
"y": []
},
{
"x": [4, 5],
"y": [4.4, 5.5]
},
{
"x": [],
"y": []
},
{
"x": [],
"y": []
},
{
"x": [6, 7, 8, 9],
"y": [6.6, 7.7, 8.8, 9.9]
},
]
def test_write_read(tmp_path):
array1 = ak.Array([[1, 2, 3], [], [4, 5], [], [], [6, 7, 8, 9]])
array2 = ak.repartition(array1, 2)
array3 = ak.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 = ak.repartition(array3, 2)
ak.to_parquet(array1, os.path.join(tmp_path, "array1.parquet"))
ak.to_parquet(array2, os.path.join(tmp_path, "array2.parquet"))
ak.to_parquet(array3, os.path.join(tmp_path, "array3.parquet"))
ak.to_parquet(array4, os.path.join(tmp_path, "array4.parquet"))
assert ak.to_list(ak.from_parquet(os.path.join(
tmp_path, "array1.parquet"))) == ak.to_list(array1)
assert ak.to_list(ak.from_parquet(os.path.join(
tmp_path, "array2.parquet"))) == ak.to_list(array2)
assert ak.to_list(ak.from_parquet(os.path.join(
tmp_path, "array3.parquet"))) == ak.to_list(array3)
assert ak.to_list(ak.from_parquet(os.path.join(
tmp_path, "array4.parquet"))) == ak.to_list(array4)
assert ak.to_list(
ak.from_parquet(os.path.join(tmp_path, "array1.parquet"),
lazy=True)) == ak.to_list(array1)
assert ak.to_list(
ak.from_parquet(os.path.join(tmp_path, "array2.parquet"),
lazy=True)) == ak.to_list(array2)
assert ak.to_list(
ak.from_parquet(os.path.join(tmp_path, "array3.parquet"),
lazy=True)) == ak.to_list(array3)
assert ak.to_list(
ak.from_parquet(os.path.join(tmp_path, "array4.parquet"),
lazy=True)) == ak.to_list(array4)
def test_zip():
x = ak.Array([[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]])
y = ak.Array([1.1, 2.2, 3.3, 4.4, 5.5])
one = ak.zip({"x": x, "y": y})
two = ak.zip({"x": x, "y": y}, depth_limit=1)
xx, yy = ak.unzip(two)
assert isinstance(one.layout, ak.layout.Content)
assert isinstance(two.layout, ak.layout.Content)
assert isinstance(xx.layout, ak.layout.Content)
assert isinstance(yy.layout, ak.layout.Content)
assert ak.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 ak.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 ak._util.py27 and not ak._util.py35:
assert ak.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]]
assert ak.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5]
x = ak.repartition(x, 3)
assert isinstance(x.layout, ak.partition.PartitionedArray)
assert ak.to_list(x) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]]
one = ak.zip({"x": x, "y": y})
two = ak.zip({"x": x, "y": y}, depth_limit=1)
xx, yy = ak.unzip(two)
assert isinstance(one.layout, ak.partition.PartitionedArray)
assert isinstance(two.layout, ak.partition.PartitionedArray)
assert isinstance(xx.layout, ak.partition.PartitionedArray)
assert isinstance(yy.layout, ak.partition.PartitionedArray)
assert ak.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 ak.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 ak._util.py27 and not ak._util.py35:
assert ak.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]]
assert ak.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5]
y = ak.repartition(y, 2)
assert isinstance(x.layout, ak.partition.PartitionedArray)
assert ak.to_list(y) == [1.1, 2.2, 3.3, 4.4, 5.5]
one = ak.zip({"x": x, "y": y})
two = ak.zip({"x": x, "y": y}, depth_limit=1)
xx, yy = ak.unzip(two)
assert isinstance(one.layout, ak.partition.PartitionedArray)
assert isinstance(two.layout, ak.partition.PartitionedArray)
assert isinstance(xx.layout, ak.partition.PartitionedArray)
assert isinstance(yy.layout, ak.partition.PartitionedArray)
assert ak.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 ak.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 ak._util.py27 and not ak._util.py35:
assert ak.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]]
assert ak.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5]
x = ak.repartition(x, None)
assert isinstance(x.layout, ak.layout.Content)
assert ak.to_list(x) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]]
one = ak.zip({"x": x, "y": y})
two = ak.zip({"x": x, "y": y}, depth_limit=1)
xx, yy = ak.unzip(two)
assert isinstance(one.layout, ak.partition.PartitionedArray)
assert isinstance(two.layout, ak.partition.PartitionedArray)
assert isinstance(xx.layout, ak.partition.PartitionedArray)
assert isinstance(yy.layout, ak.partition.PartitionedArray)
assert ak.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 ak.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 ak._util.py27 and not ak._util.py35:
assert ak.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]]
assert ak.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5]
y = ak.repartition(y, None)
assert isinstance(y.layout, ak.layout.Content)
assert ak.to_list(y) == [1.1, 2.2, 3.3, 4.4, 5.5]
one = ak.zip({"x": x, "y": y})
two = ak.zip({"x": x, "y": y}, depth_limit=1)
xx, yy = ak.unzip(two)
assert isinstance(one.layout, ak.layout.Content)
assert isinstance(two.layout, ak.layout.Content)
assert isinstance(xx.layout, ak.layout.Content)
assert isinstance(yy.layout, ak.layout.Content)
assert ak.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 ak.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 ak._util.py27 and not ak._util.py35:
assert ak.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]]
assert ak.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5]
def test():
array = ak.repartition([{"x": x, "y": x * 10} for x in range(10)], 2)
assert ak.to_list(ak.num(array.layout, axis=0)) == {"x": 10, "y": 10}
def test_getitem_2():
aslist = [
{"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": []},
]
asarray = ak.repartition(ak.Array(aslist), 2)
@numba.njit
def f3a(x):
return x["x"]
assert ak.to_list(f3a(asarray)) == [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]
@numba.njit
def f3b(x):
return x.x
assert ak.to_list(f3b(asarray)) == [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]
@numba.njit
def f4a(x):
return x["y"]
assert ak.to_list(f4a(asarray)) == [
[],
[1],
[2, 2],
[3, 3, 3],
[4, 4, 4, 4],
[5, 5, 5],
[6, 6],
[7],
[],
]
@numba.njit
def f4b(x):
return x.y
assert ak.to_list(f4b(asarray)) == [
[],
[1],
[2, 2],
[3, 3, 3],
[4, 4, 4, 4],
[5, 5, 5],
[6, 6],
[7],
[],
]
@numba.njit
def f5a(x, i):
return x["x"][i]
assert [f5a(asarray, i) for i in range(-9, 9)]
@numba.njit
def f5b(x, i):
return x[i]["x"]
assert [f5b(asarray, i) for i in range(-9, 9)]
@numba.njit
def f5c(x, i):
return x.x[i]
assert [f5c(asarray, i) for i in range(-9, 9)]
@numba.njit
def f5d(x, i):
return x[i].x
assert [f5d(asarray, i) for i in range(-9, 9)]
@numba.njit
def f6a(x, i):
return x["y"][i]
assert ak.to_list(f6a(asarray, 6)) == [6, 6]
assert ak.to_list(f6a(asarray, -3)) == [6, 6]
@numba.njit
def f6b(x, i):
return x[i]["y"]
assert ak.to_list(f6b(asarray, 6)) == [6, 6]
assert ak.to_list(f6b(asarray, -3)) == [6, 6]
@numba.njit
def f6c(x, i):
return x.y[i]
assert ak.to_list(f6c(asarray, 6)) == [6, 6]
assert ak.to_list(f6c(asarray, -3)) == [6, 6]
@numba.njit
def f6d(x, i):
return x[i].y
assert ak.to_list(f6d(asarray, 6)) == [6, 6]
assert ak.to_list(f6d(asarray, -3)) == [6, 6]