-
Notifications
You must be signed in to change notification settings - Fork 0
/
brin_overlap_test.py
109 lines (89 loc) · 4.2 KB
/
brin_overlap_test.py
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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
from __future__ import annotations
from datetime import datetime
from brin_lib import BlockRange
from brin_overlap import compute_overlap, find_position, try_insert
def blknums(brs: list[BlockRange]) -> list[int]:
return [br.blknum for br in brs]
def test_find_position() -> None:
a = BlockRange(1, datetime.fromtimestamp(1), datetime.fromtimestamp(2))
b = BlockRange(2, datetime.fromtimestamp(3), datetime.fromtimestamp(4))
c = BlockRange(3, datetime.fromtimestamp(5), datetime.fromtimestamp(6))
assert find_position([b, c], a) == 0
assert find_position([a, c], b) == 1
assert find_position([a, b], c) == 2
def test_find_position_equal_endpoints() -> None:
a = BlockRange(1, datetime.fromtimestamp(1), datetime.fromtimestamp(2))
b = BlockRange(2, datetime.fromtimestamp(2), datetime.fromtimestamp(3))
c = BlockRange(3, datetime.fromtimestamp(3), datetime.fromtimestamp(4))
assert find_position([b, c], a) == 0
assert find_position([a, c], b) == 1
assert find_position([a, b], c) == 2
def test_find_position_overlap() -> None:
a = BlockRange(1, datetime.fromtimestamp(1), datetime.fromtimestamp(2))
b = BlockRange(2, datetime.fromtimestamp(3), datetime.fromtimestamp(4))
ab = BlockRange(3, datetime.fromtimestamp(2), datetime.fromtimestamp(4))
assert find_position([a, b], ab) is None
assert find_position([a, ab], b) is None
def test_try_insert() -> None:
a = BlockRange(1, datetime.fromtimestamp(1), datetime.fromtimestamp(2))
b = BlockRange(2, datetime.fromtimestamp(3), datetime.fromtimestamp(4))
c = BlockRange(3, datetime.fromtimestamp(5), datetime.fromtimestamp(6))
level = [a, c]
assert try_insert(level, b)
assert blknums(level) == [1, 2, 3]
level = [a, b]
assert try_insert(level, c)
assert blknums(level) == [1, 2, 3]
level = [b, c]
assert try_insert(level, a)
assert blknums(level) == [1, 2, 3]
def test_try_insert_fail() -> None:
a = BlockRange(1, datetime.fromtimestamp(1), datetime.fromtimestamp(2))
b = BlockRange(2, datetime.fromtimestamp(3), datetime.fromtimestamp(4))
ab1 = BlockRange(3, datetime.fromtimestamp(1), datetime.fromtimestamp(4))
ab2 = BlockRange(3, datetime.fromtimestamp(1), datetime.fromtimestamp(3))
ab3 = BlockRange(3, datetime.fromtimestamp(2), datetime.fromtimestamp(4))
levels = [a, b]
assert not try_insert(levels, ab1)
assert not try_insert(levels, ab2)
assert not try_insert(levels, ab3)
def test_try_insert_edge() -> None:
a = BlockRange(1, datetime.fromtimestamp(1), datetime.fromtimestamp(2))
b = BlockRange(2, datetime.fromtimestamp(3), datetime.fromtimestamp(4))
# empty list
level: list[BlockRange] = []
assert try_insert(level, a)
assert blknums(level) == [1]
# insert before first element
level = [b]
assert try_insert(level, a)
assert blknums(level) == [1, 2]
# insert after last element
level = [a]
assert try_insert(level, b)
assert blknums(level) == [1, 2]
def test_compute_overlap_ordered() -> None:
a = BlockRange(1, datetime.fromtimestamp(1), datetime.fromtimestamp(2))
b = BlockRange(2, datetime.fromtimestamp(3), datetime.fromtimestamp(4))
c = BlockRange(3, datetime.fromtimestamp(5), datetime.fromtimestamp(6))
d = BlockRange(4, datetime.fromtimestamp(7), datetime.fromtimestamp(8))
bro = compute_overlap([a, b, c, d])
assert bro.min_val.timestamp() == 1
assert bro.max_val.timestamp() == 8
assert bro.min_blknum == 1
assert bro.max_blknum == 4
assert len(bro.levels) == 1
assert blknums(bro.levels[0]) == [1, 2, 3, 4]
def test_compute_overlap_overlap() -> None:
a = BlockRange(1, datetime.fromtimestamp(1), datetime.fromtimestamp(2))
b = BlockRange(2, datetime.fromtimestamp(3), datetime.fromtimestamp(4))
c = BlockRange(3, datetime.fromtimestamp(3), datetime.fromtimestamp(4))
d = BlockRange(4, datetime.fromtimestamp(1), datetime.fromtimestamp(2))
bro = compute_overlap([a, b, c, d])
assert bro.min_val.timestamp() == 1
assert bro.max_val.timestamp() == 4
assert bro.min_blknum == 1
assert bro.max_blknum == 4
assert len(bro.levels) == 2
assert blknums(bro.levels[0]) == [1, 2]
assert blknums(bro.levels[1]) == [4, 3]