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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
|
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
from __future__ import annotations
import numpy as np # noqa: F401
import pytest # noqa: F401
import awkward as ak
def test():
array = ak.Array([3, 3, 3, 5, 5, 9, 9, 9, 9, 1, 3, 3])
assert ak.operations.run_lengths(array).to_list() == [3, 2, 4, 1, 2]
array = ak.Array([[3, 3, 3, 5], [5], [], [9, 9], [9, 9], [1, 3, 3]])
assert ak.operations.run_lengths(array).to_list() == [
[3, 1],
[1],
[],
[2],
[2],
[1, 2],
]
def test_all_same():
array = ak.Array([[3, 3, 3, 3], [3], [], [3, 3, 3], [3]])
assert ak.operations.run_lengths(array).to_list() == [
[4],
[1],
[],
[3],
[1],
]
def test_groupby():
array = ak.Array(
[
{"x": 1, "y": 1.1},
{"x": 2, "y": 2.2},
{"x": 1, "y": 1.1},
{"x": 3, "y": 3.3},
{"x": 1, "y": 1.1},
{"x": 2, "y": 2.2},
]
)
sorted = array[ak.operations.argsort(array.x)]
assert sorted.x.to_list() == [1, 1, 1, 2, 2, 3]
assert ak.operations.run_lengths(sorted.x).to_list() == [3, 2, 1]
assert ak.operations.unflatten(
sorted, ak.operations.run_lengths(sorted.x)
).to_list() == [
[{"x": 1, "y": 1.1}, {"x": 1, "y": 1.1}, {"x": 1, "y": 1.1}],
[{"x": 2, "y": 2.2}, {"x": 2, "y": 2.2}],
[{"x": 3, "y": 3.3}],
]
array = ak.Array(
[
[{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}, {"x": 1, "y": 1.1}],
[{"x": 3, "y": 3.3}, {"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}],
]
)
sorted = array[ak.operations.argsort(array.x)]
assert sorted.x.to_list() == [[1, 1, 2], [1, 2, 3]]
assert ak.operations.run_lengths(sorted.x).to_list() == [
[2, 1],
[1, 1, 1],
]
counts = ak.operations.flatten(ak.operations.run_lengths(sorted.x), axis=None)
assert ak.operations.unflatten(sorted, counts, axis=-1).to_list() == [
[[{"x": 1, "y": 1.1}, {"x": 1, "y": 1.1}], [{"x": 2, "y": 2.2}]],
[[{"x": 1, "y": 1.1}], [{"x": 2, "y": 2.2}], [{"x": 3, "y": 3.3}]],
]
def test_onstrings1():
data = ak.Array(["one", "one", "one", "two", "two", "three", "two", "two"])
assert ak.operations.run_lengths(data).to_list() == [3, 2, 1, 2]
def test_onstrings2():
data = ak.Array([["one", "one"], ["one", "two", "two"], ["three", "two", "two"]])
assert ak.operations.run_lengths(data).to_list() == [
[2],
[1, 2],
[1, 2],
]
|
