3198: Apply StreamFileReader tests to ArvadosFileReader
[arvados.git] / sdk / python / arvados / ranges.py
1 class Range(object):
2     def __init__(self, locator, range_start, range_size, segment_offset=0):
3         self.locator = locator
4         self.range_start = range_start
5         self.range_size = range_size
6         self.segment_offset = segment_offset
7
8     def __repr__(self):
9         return "Range(\"%s\", %i, %i, %i)" % (self.locator, self.range_start, self.range_size, self.segment_offset)
10
11 def first_block(data_locators, range_start, range_size, debug=False):
12     block_start = 0L
13
14     # range_start/block_start is the inclusive lower bound
15     # range_end/block_end is the exclusive upper bound
16
17     hi = len(data_locators)
18     lo = 0
19     i = int((hi + lo) / 2)
20     block_size = data_locators[i].range_size
21     block_start = data_locators[i].range_start
22     block_end = block_start + block_size
23     if debug: print '---'
24
25     # perform a binary search for the first block
26     # assumes that all of the blocks are contigious, so range_start is guaranteed
27     # to either fall into the range of a block or be outside the block range entirely
28     while not (range_start >= block_start and range_start < block_end):
29         if lo == i:
30             # must be out of range, fail
31             return None
32         if range_start > block_start:
33             lo = i
34         else:
35             hi = i
36         i = int((hi + lo) / 2)
37         if debug: print lo, i, hi
38         block_size = data_locators[i].range_size
39         block_start = data_locators[i].range_start
40         block_end = block_start + block_size
41
42     return i
43
44 class LocatorAndRange(object):
45     def __init__(self, locator, block_size, segment_offset, segment_size):
46         self.locator = locator
47         self.block_size = block_size
48         self.segment_offset = segment_offset
49         self.segment_size = segment_size
50
51     def __eq__(self, other):
52         return  (self.locator == other.locator and
53                  self.block_size == other.block_size and
54                  self.segment_offset == other.segment_offset and
55                  self.segment_size == other.segment_size)
56
57     def __repr__(self):
58         return "LocatorAndRange(\"%s\", %i, %i, %i)" % (self.locator, self.block_size, self.segment_offset, self.segment_size)
59
60 def locators_and_ranges(data_locators, range_start, range_size, debug=False):
61     '''
62     Get blocks that are covered by the range
63     data_locators: list of Range objects, assumes that blocks are in order and contigous
64     range_start: start of range
65     range_size: size of range
66     returns list of LocatorAndRange objects
67     '''
68     if range_size == 0:
69         return []
70     resp = []
71     range_start = long(range_start)
72     range_size = long(range_size)
73     range_end = range_start + range_size
74
75     i = first_block(data_locators, range_start, range_size, debug)
76     if i is None:
77         return []
78
79     while i < len(data_locators):
80         dl = data_locators[i]
81         block_start = dl.range_start
82         block_size = dl.range_size
83         block_end = block_start + block_size
84         if debug:
85             print dl.locator, "range_start", range_start, "block_start", block_start, "range_end", range_end, "block_end", block_end
86         if range_end <= block_start:
87             # range ends before this block starts, so don't look at any more locators
88             break
89
90         #if range_start >= block_end:
91             # range starts after this block ends, so go to next block
92             # we should always start at the first block due to the binary above, so this test is redundant
93             #next
94
95         if range_start >= block_start and range_end <= block_end:
96             # range starts and ends in this block
97             resp.append(LocatorAndRange(dl.locator, block_size, dl.segment_offset + (range_start - block_start), range_size))
98         elif range_start >= block_start and range_end > block_end:
99             # range starts in this block
100             resp.append(LocatorAndRange(dl.locator, block_size, dl.segment_offset + (range_start - block_start), block_end - range_start))
101         elif range_start < block_start and range_end > block_end:
102             # range starts in a previous block and extends to further blocks
103             resp.append(LocatorAndRange(dl.locator, block_size, dl.segment_offset, block_size))
104         elif range_start < block_start and range_end <= block_end:
105             # range starts in a previous block and ends in this block
106             resp.append(LocatorAndRange(dl.locator, block_size, dl.segment_offset, range_end - block_start))
107         block_start = block_end
108         i += 1
109     return resp
110
111 def replace_range(data_locators, new_range_start, new_range_size, new_locator, new_segment_offset, debug=False):
112     '''
113     Replace a file segment range with a new segment.
114     data_locators: list of Range objects, assumes that segments are in order and contigous
115     range_start: start of range
116     range_size: size of range
117     new_locator: locator for new segment to be inserted
118     !!! data_locators will be updated in place !!!
119     '''
120     if new_range_size == 0:
121         return
122
123     new_range_start = long(new_range_start)
124     new_range_size = long(new_range_size)
125     new_range_end = new_range_start + new_range_size
126
127     if len(data_locators) == 0:
128         data_locators.append(Range(new_locator, new_range_start, new_range_size, new_segment_offset))
129         return
130
131     last = data_locators[-1]
132     if (last.range_start+last.range_size) == new_range_start:
133         if last.locator == new_locator:
134             # extend last segment
135             last.range_size += new_range_size
136         else:
137             data_locators.append(Range(new_locator, new_range_start, new_range_size, new_segment_offset))
138         return
139
140     i = first_block(data_locators, new_range_start, new_range_size, debug)
141     if i is None:
142         return
143
144     while i < len(data_locators):
145         dl = data_locators[i]
146         old_segment_start = dl.range_start
147         old_segment_end = old_segment_start + dl.range_size
148         if debug:
149             print locator, "range_start", new_range_start, "segment_start", old_segment_start, "range_end", new_range_end, "segment_end", old_segment_end
150         if new_range_end <= old_segment_start:
151             # range ends before this segment starts, so don't look at any more locators
152             break
153
154         #if range_start >= old_segment_end:
155             # range starts after this segment ends, so go to next segment
156             # we should always start at the first segment due to the binary above, so this test is redundant
157             #next
158
159         if  old_segment_start <= new_range_start and new_range_end <= old_segment_end:
160             # new range starts and ends in old segment
161             # split segment into up to 3 pieces
162             if (new_range_start-old_segment_start) > 0:
163                 data_locators[i] = Range(dl.locator, old_segment_start, (new_range_start-old_segment_start), dl.segment_offset)
164                 data_locators.insert(i+1, Range(new_locator, new_range_start, new_range_size, new_segment_offset))
165             else:
166                 data_locators[i] = Range(new_locator, new_range_start, new_range_size, new_segment_offset)
167                 i -= 1
168             if (old_segment_end-new_range_end) > 0:
169                 data_locators.insert(i+2, Range(dl.locator, new_range_end, (old_segment_end-new_range_end), dl.segment_offset + (new_range_start-old_segment_start) + new_range_size))
170             return
171         elif old_segment_start <= new_range_start and new_range_end > old_segment_end:
172             # range starts in this segment
173             # split segment into 2 pieces
174             data_locators[i] = Range(dl.locator, old_segment_start, (new_range_start-old_segment_start), dl.segment_offset)
175             data_locators.insert(i+1, Range(new_locator, new_range_start, new_range_size, new_segment_offset))
176             i += 1
177         elif new_range_start < old_segment_start and new_range_end >= old_segment_end:
178             # range starts in a previous segment and extends to further segments
179             # delete this segment
180             del data_locators[i]
181             i -= 1
182         elif new_range_start < old_segment_start and new_range_end < old_segment_end:
183             # range starts in a previous segment and ends in this segment
184             # move the starting point of this segment up, and shrink it.
185             data_locators[i] = Range(dl.locator, new_range_end, (old_segment_end-new_range_end), dl.segment_offset + (new_range_end-old_segment_start))
186             return
187         i += 1