3641.3.29
by John Arbash Meinel
Cleanup the copyright headers |
1 |
# Copyright (C) 2008 Canonical Ltd
|
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
2 |
#
|
3 |
# This program is free software; you can redistribute it and/or modify
|
|
3641.3.29
by John Arbash Meinel
Cleanup the copyright headers |
4 |
# it under the terms of the GNU General Public License as published by
|
5 |
# the Free Software Foundation; either version 2 of the License, or
|
|
6 |
# (at your option) any later version.
|
|
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
7 |
#
|
8 |
# This program is distributed in the hope that it will be useful,
|
|
9 |
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
10 |
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
11 |
# GNU General Public License for more details.
|
|
12 |
#
|
|
13 |
# You should have received a copy of the GNU General Public License
|
|
14 |
# along with this program; if not, write to the Free Software
|
|
3641.3.29
by John Arbash Meinel
Cleanup the copyright headers |
15 |
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
|
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
16 |
#
|
17 |
||
18 |
"""B+Tree indices"""
|
|
19 |
||
20 |
import array |
|
21 |
import bisect |
|
22 |
from bisect import bisect_right |
|
23 |
from copy import deepcopy |
|
24 |
import math |
|
25 |
import sha |
|
26 |
import struct |
|
27 |
import tempfile |
|
28 |
import zlib |
|
29 |
||
30 |
from bzrlib import ( |
|
31 |
chunk_writer, |
|
32 |
debug, |
|
33 |
errors, |
|
34 |
index, |
|
35 |
lru_cache, |
|
36 |
osutils, |
|
37 |
trace, |
|
38 |
)
|
|
39 |
from bzrlib.index import _OPTION_NODE_REFS, _OPTION_KEY_ELEMENTS, _OPTION_LEN |
|
40 |
from bzrlib.transport import get_transport |
|
41 |
||
42 |
||
3641.3.3
by John Arbash Meinel
Change the header to indicate these indexes are |
43 |
_BTSIGNATURE = "B+Tree Graph Index 2\n" |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
44 |
_OPTION_ROW_LENGTHS = "row_lengths=" |
45 |
_LEAF_FLAG = "type=leaf\n" |
|
46 |
_INTERNAL_FLAG = "type=internal\n" |
|
47 |
_INTERNAL_OFFSET = "offset=" |
|
48 |
||
49 |
_RESERVED_HEADER_BYTES = 120 |
|
50 |
_PAGE_SIZE = 4096 |
|
51 |
||
52 |
# 4K per page: 4MB - 1000 entries
|
|
53 |
_NODE_CACHE_SIZE = 1000 |
|
54 |
||
55 |
||
56 |
class _BuilderRow(object): |
|
57 |
"""The stored state accumulated while writing out a row in the index.
|
|
58 |
||
59 |
:ivar spool: A temporary file used to accumulate nodes for this row
|
|
60 |
in the tree.
|
|
61 |
:ivar nodes: The count of nodes emitted so far.
|
|
62 |
"""
|
|
63 |
||
64 |
def __init__(self): |
|
65 |
"""Create a _BuilderRow."""
|
|
66 |
self.nodes = 0 |
|
67 |
self.spool = tempfile.TemporaryFile() |
|
68 |
self.writer = None |
|
69 |
||
70 |
def finish_node(self, pad=True): |
|
71 |
byte_lines, _, padding = self.writer.finish() |
|
72 |
if self.nodes == 0: |
|
73 |
# padded note:
|
|
74 |
self.spool.write("\x00" * _RESERVED_HEADER_BYTES) |
|
75 |
skipped_bytes = 0 |
|
76 |
if not pad and padding: |
|
77 |
del byte_lines[-1] |
|
78 |
skipped_bytes = padding |
|
79 |
self.spool.writelines(byte_lines) |
|
3644.2.3
by John Arbash Meinel
Do a bit more work to get all the tests to pass. |
80 |
remainder = (self.spool.tell() + skipped_bytes) % _PAGE_SIZE |
81 |
if remainder != 0: |
|
82 |
raise AssertionError("incorrect node length: %d, %d" |
|
83 |
% (self.spool.tell(), remainder)) |
|
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
84 |
self.nodes += 1 |
85 |
self.writer = None |
|
86 |
||
87 |
||
88 |
class _InternalBuilderRow(_BuilderRow): |
|
89 |
"""The stored state accumulated while writing out internal rows."""
|
|
90 |
||
91 |
def finish_node(self, pad=True): |
|
92 |
if not pad: |
|
93 |
raise AssertionError("Must pad internal nodes only.") |
|
94 |
_BuilderRow.finish_node(self) |
|
95 |
||
96 |
||
97 |
class _LeafBuilderRow(_BuilderRow): |
|
98 |
"""The stored state accumulated while writing out a leaf rows."""
|
|
99 |
||
100 |
||
101 |
class BTreeBuilder(index.GraphIndexBuilder): |
|
102 |
"""A Builder for B+Tree based Graph indices.
|
|
103 |
||
104 |
The resulting graph has the structure:
|
|
105 |
||
106 |
_SIGNATURE OPTIONS NODES
|
|
107 |
_SIGNATURE := 'B+Tree Graph Index 1' NEWLINE
|
|
108 |
OPTIONS := REF_LISTS KEY_ELEMENTS LENGTH
|
|
109 |
REF_LISTS := 'node_ref_lists=' DIGITS NEWLINE
|
|
110 |
KEY_ELEMENTS := 'key_elements=' DIGITS NEWLINE
|
|
111 |
LENGTH := 'len=' DIGITS NEWLINE
|
|
112 |
ROW_LENGTHS := 'row_lengths' DIGITS (COMMA DIGITS)*
|
|
113 |
NODES := NODE_COMPRESSED*
|
|
114 |
NODE_COMPRESSED:= COMPRESSED_BYTES{4096}
|
|
115 |
NODE_RAW := INTERNAL | LEAF
|
|
116 |
INTERNAL := INTERNAL_FLAG POINTERS
|
|
117 |
LEAF := LEAF_FLAG ROWS
|
|
118 |
KEY_ELEMENT := Not-whitespace-utf8
|
|
119 |
KEY := KEY_ELEMENT (NULL KEY_ELEMENT)*
|
|
120 |
ROWS := ROW*
|
|
121 |
ROW := KEY NULL ABSENT? NULL REFERENCES NULL VALUE NEWLINE
|
|
122 |
ABSENT := 'a'
|
|
123 |
REFERENCES := REFERENCE_LIST (TAB REFERENCE_LIST){node_ref_lists - 1}
|
|
124 |
REFERENCE_LIST := (REFERENCE (CR REFERENCE)*)?
|
|
125 |
REFERENCE := KEY
|
|
126 |
VALUE := no-newline-no-null-bytes
|
|
127 |
"""
|
|
128 |
||
129 |
def __init__(self, reference_lists=0, key_elements=1, spill_at=100000): |
|
130 |
"""See GraphIndexBuilder.__init__.
|
|
131 |
||
132 |
:param spill_at: Optional parameter controlling the maximum number
|
|
133 |
of nodes that BTreeBuilder will hold in memory.
|
|
134 |
"""
|
|
135 |
index.GraphIndexBuilder.__init__(self, reference_lists=reference_lists, |
|
136 |
key_elements=key_elements) |
|
137 |
self._spill_at = spill_at |
|
138 |
self._backing_indices = [] |
|
3644.2.11
by John Arbash Meinel
Document the new form of _nodes and remove an unnecessary cast. |
139 |
# A map of {key: (node_refs, value)}
|
140 |
self._nodes = {} |
|
3644.2.1
by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed. |
141 |
# Indicate it hasn't been built yet
|
142 |
self._nodes_by_key = None |
|
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
143 |
|
144 |
def add_node(self, key, value, references=()): |
|
145 |
"""Add a node to the index.
|
|
146 |
||
147 |
If adding the node causes the builder to reach its spill_at threshold,
|
|
148 |
disk spilling will be triggered.
|
|
149 |
||
150 |
:param key: The key. keys are non-empty tuples containing
|
|
151 |
as many whitespace-free utf8 bytestrings as the key length
|
|
152 |
defined for this index.
|
|
153 |
:param references: An iterable of iterables of keys. Each is a
|
|
154 |
reference to another key.
|
|
155 |
:param value: The value to associate with the key. It may be any
|
|
156 |
bytes as long as it does not contain \0 or \n.
|
|
157 |
"""
|
|
3644.2.9
by John Arbash Meinel
Refactor some code. |
158 |
# we don't care about absent_references
|
159 |
node_refs, _ = self._check_key_ref_value(key, references, value) |
|
3644.2.2
by John Arbash Meinel
the new btree index doesn't have 'absent' keys in its _nodes |
160 |
if key in self._nodes: |
3644.2.1
by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed. |
161 |
raise errors.BadIndexDuplicateKey(key, self) |
3644.2.11
by John Arbash Meinel
Document the new form of _nodes and remove an unnecessary cast. |
162 |
self._nodes[key] = (node_refs, value) |
3644.2.1
by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed. |
163 |
self._keys.add(key) |
3644.2.9
by John Arbash Meinel
Refactor some code. |
164 |
if self._nodes_by_key is not None and self._key_length > 1: |
165 |
self._update_nodes_by_key(key, value, node_refs) |
|
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
166 |
if len(self._keys) < self._spill_at: |
167 |
return
|
|
3644.2.9
by John Arbash Meinel
Refactor some code. |
168 |
self._spill_mem_keys_to_disk() |
169 |
||
170 |
def _spill_mem_keys_to_disk(self): |
|
171 |
"""Write the in memory keys down to disk to cap memory consumption.
|
|
172 |
||
173 |
If we already have some keys written to disk, we will combine them so
|
|
174 |
as to preserve the sorted order. The algorithm for combining uses
|
|
175 |
powers of two. So on the first spill, write all mem nodes into a
|
|
176 |
single index. On the second spill, combine the mem nodes with the nodes
|
|
177 |
on disk to create a 2x sized disk index and get rid of the first index.
|
|
178 |
On the third spill, create a single new disk index, which will contain
|
|
179 |
the mem nodes, and preserve the existing 2x sized index. On the fourth,
|
|
180 |
combine mem with the first and second indexes, creating a new one of
|
|
181 |
size 4x. On the fifth create a single new one, etc.
|
|
182 |
"""
|
|
3644.2.8
by John Arbash Meinel
Two quick tweaks. |
183 |
iterators_to_combine = [self._iter_mem_nodes()] |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
184 |
pos = -1 |
185 |
for pos, backing in enumerate(self._backing_indices): |
|
186 |
if backing is None: |
|
187 |
pos -= 1 |
|
188 |
break
|
|
189 |
iterators_to_combine.append(backing.iter_all_entries()) |
|
190 |
backing_pos = pos + 1 |
|
191 |
new_backing_file, size = \ |
|
192 |
self._write_nodes(self._iter_smallest(iterators_to_combine)) |
|
3644.2.9
by John Arbash Meinel
Refactor some code. |
193 |
dir_path, base_name = osutils.split(new_backing_file.name) |
194 |
# Note: The transport here isn't strictly needed, because we will use
|
|
195 |
# direct access to the new_backing._file object
|
|
196 |
new_backing = BTreeGraphIndex(get_transport(dir_path), |
|
197 |
base_name, size) |
|
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
198 |
# GC will clean up the file
|
199 |
new_backing._file = new_backing_file |
|
200 |
if len(self._backing_indices) == backing_pos: |
|
201 |
self._backing_indices.append(None) |
|
202 |
self._backing_indices[backing_pos] = new_backing |
|
203 |
for pos in range(backing_pos): |
|
204 |
self._backing_indices[pos] = None |
|
205 |
self._keys = set() |
|
206 |
self._nodes = {} |
|
3644.2.1
by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed. |
207 |
self._nodes_by_key = None |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
208 |
|
209 |
def add_nodes(self, nodes): |
|
210 |
"""Add nodes to the index.
|
|
211 |
||
212 |
:param nodes: An iterable of (key, node_refs, value) entries to add.
|
|
213 |
"""
|
|
214 |
if self.reference_lists: |
|
215 |
for (key, value, node_refs) in nodes: |
|
216 |
self.add_node(key, value, node_refs) |
|
217 |
else: |
|
218 |
for (key, value) in nodes: |
|
219 |
self.add_node(key, value) |
|
220 |
||
221 |
def _iter_mem_nodes(self): |
|
222 |
"""Iterate over the nodes held in memory."""
|
|
3644.2.8
by John Arbash Meinel
Two quick tweaks. |
223 |
nodes = self._nodes |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
224 |
if self.reference_lists: |
3644.2.8
by John Arbash Meinel
Two quick tweaks. |
225 |
for key in sorted(nodes): |
226 |
references, value = nodes[key] |
|
227 |
yield self, key, value, references |
|
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
228 |
else: |
3644.2.8
by John Arbash Meinel
Two quick tweaks. |
229 |
for key in sorted(nodes): |
230 |
references, value = nodes[key] |
|
231 |
yield self, key, value |
|
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
232 |
|
233 |
def _iter_smallest(self, iterators_to_combine): |
|
3641.3.9
by John Arbash Meinel
Special case around _iter_smallest when we have only |
234 |
if len(iterators_to_combine) == 1: |
235 |
for value in iterators_to_combine[0]: |
|
236 |
yield value |
|
237 |
return
|
|
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
238 |
current_values = [] |
239 |
for iterator in iterators_to_combine: |
|
240 |
try: |
|
241 |
current_values.append(iterator.next()) |
|
242 |
except StopIteration: |
|
243 |
current_values.append(None) |
|
244 |
last = None |
|
245 |
while True: |
|
246 |
# Decorate candidates with the value to allow 2.4's min to be used.
|
|
247 |
candidates = [(item[1][1], item) for item |
|
248 |
in enumerate(current_values) if item[1] is not None] |
|
249 |
if not len(candidates): |
|
250 |
return
|
|
251 |
selected = min(candidates) |
|
252 |
# undecorate back to (pos, node)
|
|
253 |
selected = selected[1] |
|
254 |
if last == selected[1][1]: |
|
255 |
raise errors.BadIndexDuplicateKey(last, self) |
|
256 |
last = selected[1][1] |
|
257 |
# Yield, with self as the index
|
|
258 |
yield (self,) + selected[1][1:] |
|
259 |
pos = selected[0] |
|
260 |
try: |
|
261 |
current_values[pos] = iterators_to_combine[pos].next() |
|
262 |
except StopIteration: |
|
263 |
current_values[pos] = None |
|
264 |
||
3641.3.11
by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state. |
265 |
def _add_key(self, string_key, line, rows): |
3641.3.8
by John Arbash Meinel
Move the add_key helper function into a separate func |
266 |
"""Add a key to the current chunk.
|
267 |
||
268 |
:param string_key: The key to add.
|
|
3641.3.11
by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state. |
269 |
:param line: The fully serialised key and value.
|
3641.3.8
by John Arbash Meinel
Move the add_key helper function into a separate func |
270 |
"""
|
271 |
if rows[-1].writer is None: |
|
272 |
# opening a new leaf chunk;
|
|
273 |
for pos, internal_row in enumerate(rows[:-1]): |
|
274 |
# flesh out any internal nodes that are needed to
|
|
3641.3.11
by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state. |
275 |
# preserve the height of the tree
|
3641.3.8
by John Arbash Meinel
Move the add_key helper function into a separate func |
276 |
if internal_row.writer is None: |
277 |
length = _PAGE_SIZE |
|
278 |
if internal_row.nodes == 0: |
|
279 |
length -= _RESERVED_HEADER_BYTES # padded |
|
3641.3.11
by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state. |
280 |
internal_row.writer = chunk_writer.ChunkWriter(length, 0) |
3641.3.8
by John Arbash Meinel
Move the add_key helper function into a separate func |
281 |
internal_row.writer.write(_INTERNAL_FLAG) |
282 |
internal_row.writer.write(_INTERNAL_OFFSET + |
|
283 |
str(rows[pos + 1].nodes) + "\n") |
|
284 |
# add a new leaf
|
|
285 |
length = _PAGE_SIZE |
|
286 |
if rows[-1].nodes == 0: |
|
287 |
length -= _RESERVED_HEADER_BYTES # padded |
|
288 |
rows[-1].writer = chunk_writer.ChunkWriter(length) |
|
289 |
rows[-1].writer.write(_LEAF_FLAG) |
|
3641.3.11
by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state. |
290 |
if rows[-1].writer.write(line): |
3641.3.8
by John Arbash Meinel
Move the add_key helper function into a separate func |
291 |
# this key did not fit in the node:
|
292 |
rows[-1].finish_node() |
|
3641.3.11
by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state. |
293 |
key_line = string_key + "\n" |
3641.3.8
by John Arbash Meinel
Move the add_key helper function into a separate func |
294 |
new_row = True |
295 |
for row in reversed(rows[:-1]): |
|
296 |
# Mark the start of the next node in the node above. If it
|
|
297 |
# doesn't fit then propogate upwards until we find one that
|
|
298 |
# it does fit into.
|
|
3641.3.11
by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state. |
299 |
if row.writer.write(key_line): |
3641.3.8
by John Arbash Meinel
Move the add_key helper function into a separate func |
300 |
row.finish_node() |
301 |
else: |
|
302 |
# We've found a node that can handle the pointer.
|
|
303 |
new_row = False |
|
304 |
break
|
|
305 |
# If we reached the current root without being able to mark the
|
|
306 |
# division point, then we need a new root:
|
|
307 |
if new_row: |
|
308 |
# We need a new row
|
|
309 |
if 'index' in debug.debug_flags: |
|
310 |
trace.mutter('Inserting new global row.') |
|
311 |
new_row = _InternalBuilderRow() |
|
312 |
reserved_bytes = 0 |
|
313 |
rows.insert(0, new_row) |
|
314 |
# This will be padded, hence the -100
|
|
315 |
new_row.writer = chunk_writer.ChunkWriter( |
|
316 |
_PAGE_SIZE - _RESERVED_HEADER_BYTES, |
|
317 |
reserved_bytes) |
|
318 |
new_row.writer.write(_INTERNAL_FLAG) |
|
319 |
new_row.writer.write(_INTERNAL_OFFSET + |
|
320 |
str(rows[1].nodes - 1) + "\n") |
|
3641.3.11
by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state. |
321 |
new_row.writer.write(key_line) |
322 |
self._add_key(string_key, line, rows) |
|
3641.3.8
by John Arbash Meinel
Move the add_key helper function into a separate func |
323 |
|
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
324 |
def _write_nodes(self, node_iterator): |
325 |
"""Write node_iterator out as a B+Tree.
|
|
326 |
||
327 |
:param node_iterator: An iterator of sorted nodes. Each node should
|
|
328 |
match the output given by iter_all_entries.
|
|
329 |
:return: A file handle for a temporary file containing a B+Tree for
|
|
330 |
the nodes.
|
|
331 |
"""
|
|
332 |
# The index rows - rows[0] is the root, rows[1] is the layer under it
|
|
333 |
# etc.
|
|
334 |
rows = [] |
|
335 |
# forward sorted by key. In future we may consider topological sorting,
|
|
336 |
# at the cost of table scans for direct lookup, or a second index for
|
|
337 |
# direct lookup
|
|
338 |
key_count = 0 |
|
339 |
# A stack with the number of nodes of each size. 0 is the root node
|
|
340 |
# and must always be 1 (if there are any nodes in the tree).
|
|
341 |
self.row_lengths = [] |
|
342 |
# Loop over all nodes adding them to the bottom row
|
|
343 |
# (rows[-1]). When we finish a chunk in a row,
|
|
344 |
# propogate the key that didn't fit (comes after the chunk) to the
|
|
345 |
# row above, transitively.
|
|
346 |
for node in node_iterator: |
|
347 |
if key_count == 0: |
|
348 |
# First key triggers the first row
|
|
349 |
rows.append(_LeafBuilderRow()) |
|
350 |
key_count += 1 |
|
3641.3.11
by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state. |
351 |
# TODO: Flattening the node into a string key and a line should
|
352 |
# probably be put into a pyrex function. We can do a quick
|
|
353 |
# iter over all the entries to determine the final length,
|
|
354 |
# and then do a single malloc() rather than lots of
|
|
355 |
# intermediate mallocs as we build everything up.
|
|
356 |
# ATM 3 / 13s are spent flattening nodes (10s is compressing)
|
|
3641.3.30
by John Arbash Meinel
Rename _parse_btree to _btree_serializer |
357 |
string_key, line = _btree_serializer._flatten_node(node, |
358 |
self.reference_lists) |
|
3641.3.8
by John Arbash Meinel
Move the add_key helper function into a separate func |
359 |
self._add_key(string_key, line, rows) |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
360 |
for row in reversed(rows): |
361 |
pad = (type(row) != _LeafBuilderRow) |
|
362 |
row.finish_node(pad=pad) |
|
363 |
result = tempfile.NamedTemporaryFile() |
|
364 |
lines = [_BTSIGNATURE] |
|
365 |
lines.append(_OPTION_NODE_REFS + str(self.reference_lists) + '\n') |
|
366 |
lines.append(_OPTION_KEY_ELEMENTS + str(self._key_length) + '\n') |
|
367 |
lines.append(_OPTION_LEN + str(key_count) + '\n') |
|
368 |
row_lengths = [row.nodes for row in rows] |
|
369 |
lines.append(_OPTION_ROW_LENGTHS + ','.join(map(str, row_lengths)) + '\n') |
|
370 |
result.writelines(lines) |
|
371 |
position = sum(map(len, lines)) |
|
372 |
root_row = True |
|
373 |
if position > _RESERVED_HEADER_BYTES: |
|
374 |
raise AssertionError("Could not fit the header in the" |
|
375 |
" reserved space: %d > %d" |
|
376 |
% (position, _RESERVED_HEADER_BYTES)) |
|
377 |
# write the rows out:
|
|
378 |
for row in rows: |
|
379 |
reserved = _RESERVED_HEADER_BYTES # reserved space for first node |
|
380 |
row.spool.flush() |
|
381 |
row.spool.seek(0) |
|
382 |
# copy nodes to the finalised file.
|
|
383 |
# Special case the first node as it may be prefixed
|
|
384 |
node = row.spool.read(_PAGE_SIZE) |
|
385 |
result.write(node[reserved:]) |
|
386 |
result.write("\x00" * (reserved - position)) |
|
387 |
position = 0 # Only the root row actually has an offset |
|
388 |
copied_len = osutils.pumpfile(row.spool, result) |
|
389 |
if copied_len != (row.nodes - 1) * _PAGE_SIZE: |
|
390 |
if type(row) != _LeafBuilderRow: |
|
3644.2.3
by John Arbash Meinel
Do a bit more work to get all the tests to pass. |
391 |
raise AssertionError("Incorrect amount of data copied" |
392 |
" expected: %d, got: %d" |
|
393 |
% ((row.nodes - 1) * _PAGE_SIZE, |
|
394 |
copied_len)) |
|
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
395 |
result.flush() |
396 |
size = result.tell() |
|
397 |
result.seek(0) |
|
398 |
return result, size |
|
399 |
||
400 |
def finish(self): |
|
401 |
"""Finalise the index.
|
|
402 |
||
403 |
:return: A file handle for a temporary file containing the nodes added
|
|
404 |
to the index.
|
|
405 |
"""
|
|
406 |
return self._write_nodes(self.iter_all_entries())[0] |
|
407 |
||
408 |
def iter_all_entries(self): |
|
409 |
"""Iterate over all keys within the index
|
|
410 |
||
411 |
:return: An iterable of (index, key, reference_lists, value). There is no
|
|
412 |
defined order for the result iteration - it will be in the most
|
|
413 |
efficient order for the index (in this case dictionary hash order).
|
|
414 |
"""
|
|
415 |
if 'evil' in debug.debug_flags: |
|
416 |
trace.mutter_callsite(3, |
|
417 |
"iter_all_entries scales with size of history.") |
|
418 |
# Doing serial rather than ordered would be faster; but this shouldn't
|
|
419 |
# be getting called routinely anyway.
|
|
3644.2.8
by John Arbash Meinel
Two quick tweaks. |
420 |
iterators = [self._iter_mem_nodes()] |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
421 |
for backing in self._backing_indices: |
422 |
if backing is not None: |
|
423 |
iterators.append(backing.iter_all_entries()) |
|
3641.3.9
by John Arbash Meinel
Special case around _iter_smallest when we have only |
424 |
if len(iterators) == 1: |
425 |
return iterators[0] |
|
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
426 |
return self._iter_smallest(iterators) |
427 |
||
428 |
def iter_entries(self, keys): |
|
429 |
"""Iterate over keys within the index.
|
|
430 |
||
431 |
:param keys: An iterable providing the keys to be retrieved.
|
|
432 |
:return: An iterable of (index, key, value, reference_lists). There is no
|
|
433 |
defined order for the result iteration - it will be in the most
|
|
434 |
efficient order for the index (keys iteration order in this case).
|
|
435 |
"""
|
|
436 |
keys = set(keys) |
|
437 |
if self.reference_lists: |
|
438 |
for key in keys.intersection(self._keys): |
|
439 |
node = self._nodes[key] |
|
3644.2.1
by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed. |
440 |
yield self, key, node[1], node[0] |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
441 |
else: |
442 |
for key in keys.intersection(self._keys): |
|
443 |
node = self._nodes[key] |
|
3644.2.1
by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed. |
444 |
yield self, key, node[1] |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
445 |
keys.difference_update(self._keys) |
446 |
for backing in self._backing_indices: |
|
447 |
if backing is None: |
|
448 |
continue
|
|
449 |
if not keys: |
|
450 |
return
|
|
451 |
for node in backing.iter_entries(keys): |
|
452 |
keys.remove(node[1]) |
|
453 |
yield (self,) + node[1:] |
|
454 |
||
455 |
def iter_entries_prefix(self, keys): |
|
456 |
"""Iterate over keys within the index using prefix matching.
|
|
457 |
||
458 |
Prefix matching is applied within the tuple of a key, not to within
|
|
459 |
the bytestring of each key element. e.g. if you have the keys ('foo',
|
|
460 |
'bar'), ('foobar', 'gam') and do a prefix search for ('foo', None) then
|
|
461 |
only the former key is returned.
|
|
462 |
||
463 |
:param keys: An iterable providing the key prefixes to be retrieved.
|
|
464 |
Each key prefix takes the form of a tuple the length of a key, but
|
|
465 |
with the last N elements 'None' rather than a regular bytestring.
|
|
466 |
The first element cannot be 'None'.
|
|
467 |
:return: An iterable as per iter_all_entries, but restricted to the
|
|
468 |
keys with a matching prefix to those supplied. No additional keys
|
|
469 |
will be returned, and every match that is in the index will be
|
|
470 |
returned.
|
|
471 |
"""
|
|
472 |
# XXX: To much duplication with the GraphIndex class; consider finding
|
|
473 |
# a good place to pull out the actual common logic.
|
|
474 |
keys = set(keys) |
|
475 |
if not keys: |
|
476 |
return
|
|
477 |
for backing in self._backing_indices: |
|
478 |
if backing is None: |
|
479 |
continue
|
|
480 |
for node in backing.iter_entries_prefix(keys): |
|
481 |
yield (self,) + node[1:] |
|
482 |
if self._key_length == 1: |
|
483 |
for key in keys: |
|
484 |
# sanity check
|
|
485 |
if key[0] is None: |
|
486 |
raise errors.BadIndexKey(key) |
|
487 |
if len(key) != self._key_length: |
|
488 |
raise errors.BadIndexKey(key) |
|
489 |
try: |
|
490 |
node = self._nodes[key] |
|
491 |
except KeyError: |
|
492 |
continue
|
|
493 |
if self.reference_lists: |
|
3644.2.1
by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed. |
494 |
yield self, key, node[1], node[0] |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
495 |
else: |
3644.2.1
by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed. |
496 |
yield self, key, node[1] |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
497 |
return
|
498 |
for key in keys: |
|
499 |
# sanity check
|
|
500 |
if key[0] is None: |
|
501 |
raise errors.BadIndexKey(key) |
|
502 |
if len(key) != self._key_length: |
|
503 |
raise errors.BadIndexKey(key) |
|
504 |
# find what it refers to:
|
|
3644.2.1
by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed. |
505 |
key_dict = self._get_nodes_by_key() |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
506 |
elements = list(key) |
507 |
# find the subdict to return
|
|
508 |
try: |
|
509 |
while len(elements) and elements[0] is not None: |
|
510 |
key_dict = key_dict[elements[0]] |
|
511 |
elements.pop(0) |
|
512 |
except KeyError: |
|
513 |
# a non-existant lookup.
|
|
514 |
continue
|
|
515 |
if len(elements): |
|
516 |
dicts = [key_dict] |
|
517 |
while dicts: |
|
518 |
key_dict = dicts.pop(-1) |
|
519 |
# can't be empty or would not exist
|
|
520 |
item, value = key_dict.iteritems().next() |
|
521 |
if type(value) == dict: |
|
522 |
# push keys
|
|
523 |
dicts.extend(key_dict.itervalues()) |
|
524 |
else: |
|
525 |
# yield keys
|
|
526 |
for value in key_dict.itervalues(): |
|
527 |
yield (self, ) + value |
|
528 |
else: |
|
529 |
yield (self, ) + key_dict |
|
530 |
||
3644.2.1
by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed. |
531 |
def _get_nodes_by_key(self): |
532 |
if self._nodes_by_key is None: |
|
533 |
nodes_by_key = {} |
|
534 |
if self.reference_lists: |
|
535 |
for key, (references, value) in self._nodes.iteritems(): |
|
536 |
key_dict = nodes_by_key |
|
537 |
for subkey in key[:-1]: |
|
538 |
key_dict = key_dict.setdefault(subkey, {}) |
|
539 |
key_dict[key[-1]] = key, value, references |
|
540 |
else: |
|
541 |
for key, (references, value) in self._nodes.iteritems(): |
|
542 |
key_dict = nodes_by_key |
|
543 |
for subkey in key[:-1]: |
|
544 |
key_dict = key_dict.setdefault(subkey, {}) |
|
545 |
key_dict[key[-1]] = key, value |
|
546 |
self._nodes_by_key = nodes_by_key |
|
547 |
return self._nodes_by_key |
|
548 |
||
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
549 |
def key_count(self): |
550 |
"""Return an estimate of the number of keys in this index.
|
|
551 |
||
552 |
For InMemoryGraphIndex the estimate is exact.
|
|
553 |
"""
|
|
554 |
return len(self._keys) + sum(backing.key_count() for backing in |
|
555 |
self._backing_indices if backing is not None) |
|
556 |
||
557 |
def validate(self): |
|
558 |
"""In memory index's have no known corruption at the moment."""
|
|
559 |
||
560 |
||
561 |
class _LeafNode(object): |
|
562 |
"""A leaf node for a serialised B+Tree index."""
|
|
563 |
||
564 |
def __init__(self, bytes, key_length, ref_list_length): |
|
565 |
"""Parse bytes to create a leaf node object."""
|
|
566 |
# splitlines mangles the \r delimiters.. don't use it.
|
|
3641.3.30
by John Arbash Meinel
Rename _parse_btree to _btree_serializer |
567 |
self.keys = dict(_btree_serializer._parse_leaf_lines(bytes, |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
568 |
key_length, ref_list_length)) |
569 |
||
570 |
||
571 |
class _InternalNode(object): |
|
572 |
"""An internal node for a serialised B+Tree index."""
|
|
573 |
||
574 |
def __init__(self, bytes): |
|
575 |
"""Parse bytes to create an internal node object."""
|
|
576 |
# splitlines mangles the \r delimiters.. don't use it.
|
|
577 |
self.keys = self._parse_lines(bytes.split('\n')) |
|
578 |
||
579 |
def _parse_lines(self, lines): |
|
580 |
nodes = [] |
|
581 |
self.offset = int(lines[1][7:]) |
|
582 |
for line in lines[2:]: |
|
583 |
if line == '': |
|
584 |
break
|
|
585 |
nodes.append(tuple(line.split('\0'))) |
|
586 |
return nodes |
|
587 |
||
588 |
||
589 |
class BTreeGraphIndex(object): |
|
590 |
"""Access to nodes via the standard GraphIndex interface for B+Tree's.
|
|
591 |
||
592 |
Individual nodes are held in a LRU cache. This holds the root node in
|
|
593 |
memory except when very large walks are done.
|
|
594 |
"""
|
|
595 |
||
596 |
def __init__(self, transport, name, size): |
|
597 |
"""Create a B+Tree index object on the index name.
|
|
598 |
||
599 |
:param transport: The transport to read data for the index from.
|
|
600 |
:param name: The file name of the index on transport.
|
|
601 |
:param size: Optional size of the index in bytes. This allows
|
|
602 |
compatibility with the GraphIndex API, as well as ensuring that
|
|
603 |
the initial read (to read the root node header) can be done
|
|
604 |
without over-reading even on empty indices, and on small indices
|
|
605 |
allows single-IO to read the entire index.
|
|
606 |
"""
|
|
607 |
self._transport = transport |
|
608 |
self._name = name |
|
609 |
self._size = size |
|
610 |
self._file = None |
|
611 |
self._page_size = transport.recommended_page_size() |
|
612 |
self._root_node = None |
|
613 |
# Default max size is 100,000 leave values
|
|
614 |
self._leaf_value_cache = None # lru_cache.LRUCache(100*1000) |
|
615 |
self._leaf_node_cache = lru_cache.LRUCache(_NODE_CACHE_SIZE) |
|
616 |
self._internal_node_cache = lru_cache.LRUCache() |
|
617 |
self._key_count = None |
|
618 |
self._row_lengths = None |
|
619 |
self._row_offsets = None # Start of each row, [-1] is the end |
|
620 |
||
621 |
def __eq__(self, other): |
|
622 |
"""Equal when self and other were created with the same parameters."""
|
|
623 |
return ( |
|
624 |
type(self) == type(other) and |
|
625 |
self._transport == other._transport and |
|
626 |
self._name == other._name and |
|
627 |
self._size == other._size) |
|
628 |
||
629 |
def __ne__(self, other): |
|
630 |
return not self.__eq__(other) |
|
631 |
||
632 |
def _get_root_node(self): |
|
633 |
if self._root_node is None: |
|
634 |
# We may not have a root node yet
|
|
635 |
nodes = list(self._read_nodes([0])) |
|
636 |
if len(nodes): |
|
637 |
self._root_node = nodes[0][1] |
|
638 |
return self._root_node |
|
639 |
||
640 |
def _cache_nodes(self, nodes, cache): |
|
641 |
"""Read nodes and cache them in the lru.
|
|
642 |
||
643 |
The nodes list supplied is sorted and then read from disk, each node
|
|
644 |
being inserted it into the _node_cache.
|
|
645 |
||
646 |
Note: Asking for more nodes than the _node_cache can contain will
|
|
647 |
result in some of the results being immediately discarded, to prevent
|
|
648 |
this an assertion is raised if more nodes are asked for than are
|
|
649 |
cachable.
|
|
650 |
||
651 |
:return: A dict of {node_pos: node}
|
|
652 |
"""
|
|
653 |
if len(nodes) > cache._max_cache: |
|
654 |
trace.mutter('Requesting %s > %s nodes, not all will be cached', |
|
655 |
len(nodes), cache._max_cache) |
|
656 |
found = {} |
|
657 |
for node_pos, node in self._read_nodes(sorted(nodes)): |
|
658 |
if node_pos == 0: # Special case |
|
659 |
self._root_node = node |
|
660 |
else: |
|
661 |
cache.add(node_pos, node) |
|
662 |
found[node_pos] = node |
|
663 |
return found |
|
664 |
||
3641.5.18
by John Arbash Meinel
Clean out the global state, good for prototyping and tuning, bad for production code. |
665 |
def _get_nodes(self, cache, node_indexes): |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
666 |
found = {} |
667 |
needed = [] |
|
668 |
for idx in node_indexes: |
|
669 |
if idx == 0 and self._root_node is not None: |
|
670 |
found[0] = self._root_node |
|
671 |
continue
|
|
672 |
try: |
|
673 |
found[idx] = cache[idx] |
|
674 |
except KeyError: |
|
675 |
needed.append(idx) |
|
676 |
found.update(self._cache_nodes(needed, cache)) |
|
677 |
return found |
|
678 |
||
679 |
def _get_internal_nodes(self, node_indexes): |
|
680 |
"""Get a node, from cache or disk.
|
|
681 |
||
682 |
After getting it, the node will be cached.
|
|
683 |
"""
|
|
3641.5.18
by John Arbash Meinel
Clean out the global state, good for prototyping and tuning, bad for production code. |
684 |
return self._get_nodes(self._internal_node_cache, node_indexes) |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
685 |
|
686 |
def _get_leaf_nodes(self, node_indexes): |
|
687 |
"""Get a bunch of nodes, from cache or disk."""
|
|
3641.5.18
by John Arbash Meinel
Clean out the global state, good for prototyping and tuning, bad for production code. |
688 |
found = self._get_nodes(self._leaf_node_cache, node_indexes) |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
689 |
if self._leaf_value_cache is not None: |
690 |
for node in found.itervalues(): |
|
691 |
for key, value in node.keys.iteritems(): |
|
692 |
if key in self._leaf_value_cache: |
|
693 |
# Don't add the rest of the keys, we've seen this node
|
|
694 |
# before.
|
|
695 |
break
|
|
696 |
self._leaf_value_cache[key] = value |
|
697 |
return found |
|
698 |
||
699 |
def iter_all_entries(self): |
|
700 |
"""Iterate over all keys within the index.
|
|
701 |
||
702 |
:return: An iterable of (index, key, value) or (index, key, value, reference_lists).
|
|
703 |
The former tuple is used when there are no reference lists in the
|
|
704 |
index, making the API compatible with simple key:value index types.
|
|
705 |
There is no defined order for the result iteration - it will be in
|
|
706 |
the most efficient order for the index.
|
|
707 |
"""
|
|
708 |
if 'evil' in debug.debug_flags: |
|
709 |
trace.mutter_callsite(3, |
|
710 |
"iter_all_entries scales with size of history.") |
|
711 |
if not self.key_count(): |
|
712 |
return
|
|
713 |
start_of_leaves = self._row_offsets[-2] |
|
714 |
end_of_leaves = self._row_offsets[-1] |
|
715 |
needed_nodes = range(start_of_leaves, end_of_leaves) |
|
716 |
# We iterate strictly in-order so that we can use this function
|
|
717 |
# for spilling index builds to disk.
|
|
718 |
if self.node_ref_lists: |
|
719 |
for _, node in self._read_nodes(needed_nodes): |
|
720 |
for key, (value, refs) in sorted(node.keys.items()): |
|
721 |
yield (self, key, value, refs) |
|
722 |
else: |
|
723 |
for _, node in self._read_nodes(needed_nodes): |
|
724 |
for key, (value, refs) in sorted(node.keys.items()): |
|
725 |
yield (self, key, value) |
|
726 |
||
727 |
@staticmethod
|
|
728 |
def _multi_bisect_right(in_keys, fixed_keys): |
|
729 |
"""Find the positions where each 'in_key' would fit in fixed_keys.
|
|
730 |
||
731 |
This is equivalent to doing "bisect_right" on each in_key into
|
|
732 |
fixed_keys
|
|
733 |
||
734 |
:param in_keys: A sorted list of keys to match with fixed_keys
|
|
735 |
:param fixed_keys: A sorted list of keys to match against
|
|
736 |
:return: A list of (integer position, [key list]) tuples.
|
|
737 |
"""
|
|
738 |
if not in_keys: |
|
739 |
return [] |
|
740 |
if not fixed_keys: |
|
741 |
# no pointers in the fixed_keys list, which means everything must
|
|
742 |
# fall to the left.
|
|
743 |
return [(0, in_keys)] |
|
744 |
||
745 |
# TODO: Iterating both lists will generally take M + N steps
|
|
746 |
# Bisecting each key will generally take M * log2 N steps.
|
|
747 |
# If we had an efficient way to compare, we could pick the method
|
|
748 |
# based on which has the fewer number of steps.
|
|
749 |
# There is also the argument that bisect_right is a compiled
|
|
750 |
# function, so there is even more to be gained.
|
|
751 |
# iter_steps = len(in_keys) + len(fixed_keys)
|
|
752 |
# bisect_steps = len(in_keys) * math.log(len(fixed_keys), 2)
|
|
753 |
if len(in_keys) == 1: # Bisect will always be faster for M = 1 |
|
754 |
return [(bisect_right(fixed_keys, in_keys[0]), in_keys)] |
|
755 |
# elif bisect_steps < iter_steps:
|
|
756 |
# offsets = {}
|
|
757 |
# for key in in_keys:
|
|
758 |
# offsets.setdefault(bisect_right(fixed_keys, key),
|
|
759 |
# []).append(key)
|
|
760 |
# return [(o, offsets[o]) for o in sorted(offsets)]
|
|
761 |
in_keys_iter = iter(in_keys) |
|
762 |
fixed_keys_iter = enumerate(fixed_keys) |
|
763 |
cur_in_key = in_keys_iter.next() |
|
764 |
cur_fixed_offset, cur_fixed_key = fixed_keys_iter.next() |
|
765 |
||
766 |
class InputDone(Exception): pass |
|
767 |
class FixedDone(Exception): pass |
|
768 |
||
769 |
output = [] |
|
770 |
cur_out = [] |
|
771 |
||
772 |
# TODO: Another possibility is that rather than iterating on each side,
|
|
773 |
# we could use a combination of bisecting and iterating. For
|
|
774 |
# example, while cur_in_key < fixed_key, bisect to find its
|
|
775 |
# point, then iterate all matching keys, then bisect (restricted
|
|
776 |
# to only the remainder) for the next one, etc.
|
|
777 |
try: |
|
778 |
while True: |
|
779 |
if cur_in_key < cur_fixed_key: |
|
780 |
cur_keys = [] |
|
781 |
cur_out = (cur_fixed_offset, cur_keys) |
|
782 |
output.append(cur_out) |
|
783 |
while cur_in_key < cur_fixed_key: |
|
784 |
cur_keys.append(cur_in_key) |
|
785 |
try: |
|
786 |
cur_in_key = in_keys_iter.next() |
|
787 |
except StopIteration: |
|
788 |
raise InputDone |
|
789 |
# At this point cur_in_key must be >= cur_fixed_key
|
|
790 |
# step the cur_fixed_key until we pass the cur key, or walk off
|
|
791 |
# the end
|
|
792 |
while cur_in_key >= cur_fixed_key: |
|
793 |
try: |
|
794 |
cur_fixed_offset, cur_fixed_key = fixed_keys_iter.next() |
|
795 |
except StopIteration: |
|
796 |
raise FixedDone |
|
797 |
except InputDone: |
|
798 |
# We consumed all of the input, nothing more to do
|
|
799 |
pass
|
|
800 |
except FixedDone: |
|
801 |
# There was some input left, but we consumed all of fixed, so we
|
|
802 |
# have to add one more for the tail
|
|
803 |
cur_keys = [cur_in_key] |
|
804 |
cur_keys.extend(in_keys_iter) |
|
805 |
cur_out = (len(fixed_keys), cur_keys) |
|
806 |
output.append(cur_out) |
|
807 |
return output |
|
808 |
||
809 |
def iter_entries(self, keys): |
|
810 |
"""Iterate over keys within the index.
|
|
811 |
||
812 |
:param keys: An iterable providing the keys to be retrieved.
|
|
813 |
:return: An iterable as per iter_all_entries, but restricted to the
|
|
814 |
keys supplied. No additional keys will be returned, and every
|
|
815 |
key supplied that is in the index will be returned.
|
|
816 |
"""
|
|
817 |
# 6 seconds spent in miss_torture using the sorted() line.
|
|
818 |
# Even with out of order disk IO it seems faster not to sort it when
|
|
819 |
# large queries are being made.
|
|
820 |
# However, now that we are doing multi-way bisecting, we need the keys
|
|
821 |
# in sorted order anyway. We could change the multi-way code to not
|
|
822 |
# require sorted order. (For example, it bisects for the first node,
|
|
823 |
# does an in-order search until a key comes before the current point,
|
|
824 |
# which it then bisects for, etc.)
|
|
825 |
keys = frozenset(keys) |
|
826 |
if not keys: |
|
827 |
return
|
|
828 |
||
829 |
if not self.key_count(): |
|
830 |
return
|
|
831 |
||
832 |
needed_keys = [] |
|
833 |
if self._leaf_value_cache is None: |
|
834 |
needed_keys = keys |
|
835 |
else: |
|
836 |
for key in keys: |
|
837 |
value = self._leaf_value_cache.get(key, None) |
|
838 |
if value is not None: |
|
839 |
# This key is known not to be here, skip it
|
|
840 |
value, refs = value |
|
841 |
if self.node_ref_lists: |
|
842 |
yield (self, key, value, refs) |
|
843 |
else: |
|
844 |
yield (self, key, value) |
|
845 |
else: |
|
846 |
needed_keys.append(key) |
|
847 |
||
848 |
last_key = None |
|
849 |
needed_keys = keys |
|
850 |
if not needed_keys: |
|
851 |
return
|
|
852 |
# 6 seconds spent in miss_torture using the sorted() line.
|
|
853 |
# Even with out of order disk IO it seems faster not to sort it when
|
|
854 |
# large queries are being made.
|
|
855 |
needed_keys = sorted(needed_keys) |
|
856 |
||
857 |
nodes_and_keys = [(0, needed_keys)] |
|
858 |
||
859 |
for row_pos, next_row_start in enumerate(self._row_offsets[1:-1]): |
|
860 |
node_indexes = [idx for idx, s_keys in nodes_and_keys] |
|
861 |
nodes = self._get_internal_nodes(node_indexes) |
|
862 |
||
863 |
next_nodes_and_keys = [] |
|
864 |
for node_index, sub_keys in nodes_and_keys: |
|
865 |
node = nodes[node_index] |
|
866 |
positions = self._multi_bisect_right(sub_keys, node.keys) |
|
867 |
node_offset = next_row_start + node.offset |
|
868 |
next_nodes_and_keys.extend([(node_offset + pos, s_keys) |
|
869 |
for pos, s_keys in positions]) |
|
870 |
nodes_and_keys = next_nodes_and_keys |
|
871 |
# We should now be at the _LeafNodes
|
|
872 |
node_indexes = [idx for idx, s_keys in nodes_and_keys] |
|
873 |
||
874 |
# TODO: We may *not* want to always read all the nodes in one
|
|
875 |
# big go. Consider setting a max size on this.
|
|
876 |
||
877 |
nodes = self._get_leaf_nodes(node_indexes) |
|
878 |
for node_index, sub_keys in nodes_and_keys: |
|
879 |
if not sub_keys: |
|
880 |
continue
|
|
881 |
node = nodes[node_index] |
|
882 |
for next_sub_key in sub_keys: |
|
883 |
if next_sub_key in node.keys: |
|
884 |
value, refs = node.keys[next_sub_key] |
|
885 |
if self.node_ref_lists: |
|
886 |
yield (self, next_sub_key, value, refs) |
|
887 |
else: |
|
888 |
yield (self, next_sub_key, value) |
|
889 |
||
890 |
def iter_entries_prefix(self, keys): |
|
891 |
"""Iterate over keys within the index using prefix matching.
|
|
892 |
||
893 |
Prefix matching is applied within the tuple of a key, not to within
|
|
894 |
the bytestring of each key element. e.g. if you have the keys ('foo',
|
|
895 |
'bar'), ('foobar', 'gam') and do a prefix search for ('foo', None) then
|
|
896 |
only the former key is returned.
|
|
897 |
||
898 |
WARNING: Note that this method currently causes a full index parse
|
|
899 |
unconditionally (which is reasonably appropriate as it is a means for
|
|
900 |
thunking many small indices into one larger one and still supplies
|
|
901 |
iter_all_entries at the thunk layer).
|
|
902 |
||
903 |
:param keys: An iterable providing the key prefixes to be retrieved.
|
|
904 |
Each key prefix takes the form of a tuple the length of a key, but
|
|
905 |
with the last N elements 'None' rather than a regular bytestring.
|
|
906 |
The first element cannot be 'None'.
|
|
907 |
:return: An iterable as per iter_all_entries, but restricted to the
|
|
908 |
keys with a matching prefix to those supplied. No additional keys
|
|
909 |
will be returned, and every match that is in the index will be
|
|
910 |
returned.
|
|
911 |
"""
|
|
912 |
keys = sorted(set(keys)) |
|
913 |
if not keys: |
|
914 |
return
|
|
915 |
# Load if needed to check key lengths
|
|
916 |
if self._key_count is None: |
|
917 |
self._get_root_node() |
|
918 |
# TODO: only access nodes that can satisfy the prefixes we are looking
|
|
919 |
# for. For now, to meet API usage (as this function is not used by
|
|
920 |
# current bzrlib) just suck the entire index and iterate in memory.
|
|
921 |
nodes = {} |
|
922 |
if self.node_ref_lists: |
|
923 |
if self._key_length == 1: |
|
924 |
for _1, key, value, refs in self.iter_all_entries(): |
|
925 |
nodes[key] = value, refs |
|
926 |
else: |
|
927 |
nodes_by_key = {} |
|
928 |
for _1, key, value, refs in self.iter_all_entries(): |
|
929 |
key_value = key, value, refs |
|
930 |
# For a key of (foo, bar, baz) create
|
|
931 |
# _nodes_by_key[foo][bar][baz] = key_value
|
|
932 |
key_dict = nodes_by_key |
|
933 |
for subkey in key[:-1]: |
|
934 |
key_dict = key_dict.setdefault(subkey, {}) |
|
935 |
key_dict[key[-1]] = key_value |
|
936 |
else: |
|
937 |
if self._key_length == 1: |
|
938 |
for _1, key, value in self.iter_all_entries(): |
|
939 |
nodes[key] = value |
|
940 |
else: |
|
941 |
nodes_by_key = {} |
|
942 |
for _1, key, value in self.iter_all_entries(): |
|
943 |
key_value = key, value |
|
944 |
# For a key of (foo, bar, baz) create
|
|
945 |
# _nodes_by_key[foo][bar][baz] = key_value
|
|
946 |
key_dict = nodes_by_key |
|
947 |
for subkey in key[:-1]: |
|
948 |
key_dict = key_dict.setdefault(subkey, {}) |
|
949 |
key_dict[key[-1]] = key_value |
|
950 |
if self._key_length == 1: |
|
951 |
for key in keys: |
|
952 |
# sanity check
|
|
953 |
if key[0] is None: |
|
954 |
raise errors.BadIndexKey(key) |
|
955 |
if len(key) != self._key_length: |
|
956 |
raise errors.BadIndexKey(key) |
|
957 |
try: |
|
958 |
if self.node_ref_lists: |
|
959 |
value, node_refs = nodes[key] |
|
960 |
yield self, key, value, node_refs |
|
961 |
else: |
|
962 |
yield self, key, nodes[key] |
|
963 |
except KeyError: |
|
964 |
pass
|
|
965 |
return
|
|
966 |
for key in keys: |
|
967 |
# sanity check
|
|
968 |
if key[0] is None: |
|
969 |
raise errors.BadIndexKey(key) |
|
970 |
if len(key) != self._key_length: |
|
971 |
raise errors.BadIndexKey(key) |
|
972 |
# find what it refers to:
|
|
973 |
key_dict = nodes_by_key |
|
974 |
elements = list(key) |
|
975 |
# find the subdict whose contents should be returned.
|
|
976 |
try: |
|
977 |
while len(elements) and elements[0] is not None: |
|
978 |
key_dict = key_dict[elements[0]] |
|
979 |
elements.pop(0) |
|
980 |
except KeyError: |
|
981 |
# a non-existant lookup.
|
|
982 |
continue
|
|
983 |
if len(elements): |
|
984 |
dicts = [key_dict] |
|
985 |
while dicts: |
|
986 |
key_dict = dicts.pop(-1) |
|
987 |
# can't be empty or would not exist
|
|
988 |
item, value = key_dict.iteritems().next() |
|
989 |
if type(value) == dict: |
|
990 |
# push keys
|
|
991 |
dicts.extend(key_dict.itervalues()) |
|
992 |
else: |
|
993 |
# yield keys
|
|
994 |
for value in key_dict.itervalues(): |
|
995 |
# each value is the key:value:node refs tuple
|
|
996 |
# ready to yield.
|
|
997 |
yield (self, ) + value |
|
998 |
else: |
|
999 |
# the last thing looked up was a terminal element
|
|
1000 |
yield (self, ) + key_dict |
|
1001 |
||
1002 |
def key_count(self): |
|
1003 |
"""Return an estimate of the number of keys in this index.
|
|
1004 |
||
1005 |
For BTreeGraphIndex the estimate is exact as it is contained in the
|
|
1006 |
header.
|
|
1007 |
"""
|
|
1008 |
if self._key_count is None: |
|
1009 |
self._get_root_node() |
|
1010 |
return self._key_count |
|
1011 |
||
1012 |
def _parse_header_from_bytes(self, bytes): |
|
1013 |
"""Parse the header from a region of bytes.
|
|
1014 |
||
1015 |
:param bytes: The data to parse.
|
|
1016 |
:return: An offset, data tuple such as readv yields, for the unparsed
|
|
1017 |
data. (which may be of length 0).
|
|
1018 |
"""
|
|
1019 |
signature = bytes[0:len(self._signature())] |
|
1020 |
if not signature == self._signature(): |
|
1021 |
raise errors.BadIndexFormatSignature(self._name, BTreeGraphIndex) |
|
1022 |
lines = bytes[len(self._signature()):].splitlines() |
|
1023 |
options_line = lines[0] |
|
1024 |
if not options_line.startswith(_OPTION_NODE_REFS): |
|
1025 |
raise errors.BadIndexOptions(self) |
|
1026 |
try: |
|
1027 |
self.node_ref_lists = int(options_line[len(_OPTION_NODE_REFS):]) |
|
1028 |
except ValueError: |
|
1029 |
raise errors.BadIndexOptions(self) |
|
1030 |
options_line = lines[1] |
|
1031 |
if not options_line.startswith(_OPTION_KEY_ELEMENTS): |
|
1032 |
raise errors.BadIndexOptions(self) |
|
1033 |
try: |
|
1034 |
self._key_length = int(options_line[len(_OPTION_KEY_ELEMENTS):]) |
|
1035 |
except ValueError: |
|
1036 |
raise errors.BadIndexOptions(self) |
|
1037 |
options_line = lines[2] |
|
1038 |
if not options_line.startswith(_OPTION_LEN): |
|
1039 |
raise errors.BadIndexOptions(self) |
|
1040 |
try: |
|
1041 |
self._key_count = int(options_line[len(_OPTION_LEN):]) |
|
1042 |
except ValueError: |
|
1043 |
raise errors.BadIndexOptions(self) |
|
1044 |
options_line = lines[3] |
|
1045 |
if not options_line.startswith(_OPTION_ROW_LENGTHS): |
|
1046 |
raise errors.BadIndexOptions(self) |
|
1047 |
try: |
|
1048 |
self._row_lengths = map(int, [length for length in |
|
1049 |
options_line[len(_OPTION_ROW_LENGTHS):].split(',') |
|
1050 |
if len(length)]) |
|
1051 |
except ValueError: |
|
1052 |
raise errors.BadIndexOptions(self) |
|
1053 |
offsets = [] |
|
1054 |
row_offset = 0 |
|
1055 |
for row in self._row_lengths: |
|
1056 |
offsets.append(row_offset) |
|
1057 |
row_offset += row |
|
1058 |
offsets.append(row_offset) |
|
1059 |
self._row_offsets = offsets |
|
1060 |
||
1061 |
# calculate the bytes we have processed
|
|
1062 |
header_end = (len(signature) + sum(map(len, lines[0:4])) + 4) |
|
1063 |
return header_end, bytes[header_end:] |
|
1064 |
||
1065 |
def _read_nodes(self, nodes): |
|
1066 |
"""Read some nodes from disk into the LRU cache.
|
|
1067 |
||
1068 |
This performs a readv to get the node data into memory, and parses each
|
|
1069 |
node, the yields it to the caller. The nodes are requested in the
|
|
1070 |
supplied order. If possible doing sort() on the list before requesting
|
|
1071 |
a read may improve performance.
|
|
1072 |
||
1073 |
:param nodes: The nodes to read. 0 - first node, 1 - second node etc.
|
|
1074 |
:return: None
|
|
1075 |
"""
|
|
1076 |
ranges = [] |
|
1077 |
for index in nodes: |
|
1078 |
offset = index * _PAGE_SIZE |
|
1079 |
size = _PAGE_SIZE |
|
1080 |
if index == 0: |
|
1081 |
# Root node - special case
|
|
1082 |
if self._size: |
|
1083 |
size = min(_PAGE_SIZE, self._size) |
|
1084 |
else: |
|
1085 |
stream = self._transport.get(self._name) |
|
1086 |
start = stream.read(_PAGE_SIZE) |
|
1087 |
# Avoid doing this again
|
|
1088 |
self._size = len(start) |
|
1089 |
size = min(_PAGE_SIZE, self._size) |
|
1090 |
else: |
|
1091 |
size = min(size, self._size - offset) |
|
1092 |
ranges.append((offset, size)) |
|
1093 |
if not ranges: |
|
1094 |
return
|
|
1095 |
if self._file is None: |
|
1096 |
data_ranges = self._transport.readv(self._name, ranges) |
|
1097 |
else: |
|
1098 |
data_ranges = [] |
|
1099 |
for offset, size in ranges: |
|
1100 |
self._file.seek(offset) |
|
1101 |
data_ranges.append((offset, self._file.read(size))) |
|
1102 |
for offset, data in data_ranges: |
|
1103 |
if offset == 0: |
|
1104 |
# extract the header
|
|
1105 |
offset, data = self._parse_header_from_bytes(data) |
|
1106 |
if len(data) == 0: |
|
1107 |
continue
|
|
1108 |
bytes = zlib.decompress(data) |
|
1109 |
if bytes.startswith(_LEAF_FLAG): |
|
1110 |
node = _LeafNode(bytes, self._key_length, self.node_ref_lists) |
|
1111 |
elif bytes.startswith(_INTERNAL_FLAG): |
|
1112 |
node = _InternalNode(bytes) |
|
1113 |
else: |
|
1114 |
raise AssertionError("Unknown node type for %r" % bytes) |
|
1115 |
yield offset / _PAGE_SIZE, node |
|
1116 |
||
1117 |
def _signature(self): |
|
1118 |
"""The file signature for this index type."""
|
|
1119 |
return _BTSIGNATURE |
|
1120 |
||
1121 |
def validate(self): |
|
1122 |
"""Validate that everything in the index can be accessed."""
|
|
1123 |
# just read and parse every node.
|
|
1124 |
self._get_root_node() |
|
1125 |
if len(self._row_lengths) > 1: |
|
1126 |
start_node = self._row_offsets[1] |
|
1127 |
else: |
|
1128 |
# We shouldn't be reading anything anyway
|
|
1129 |
start_node = 1 |
|
1130 |
node_end = self._row_offsets[-1] |
|
1131 |
for node in self._read_nodes(range(start_node, node_end)): |
|
1132 |
pass
|
|
1133 |
||
1134 |
||
1135 |
try: |
|
3641.3.30
by John Arbash Meinel
Rename _parse_btree to _btree_serializer |
1136 |
from bzrlib import _btree_serializer_c as _btree_serializer |
3641.3.1
by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests. |
1137 |
except ImportError: |
3641.3.30
by John Arbash Meinel
Rename _parse_btree to _btree_serializer |
1138 |
from bzrlib import _btree_serializer_py as _btree_serializer |