← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/btree_index.py

(vila) Make all transport put_bytes() raises TypeError when given unicode
 strings rather than bytes (Vincent Ladeuil)

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/btree_index.py
1
 
# Copyright (C) 2008, 2009, 2010 Canonical Ltd
 
1
# Copyright (C) 2008-2011 Canonical Ltd
2
2
#
3
3
# This program is free software; you can redistribute it and/or modify
4
4
# it under the terms of the GNU General Public License as published by
17
17
 
18
18
"""B+Tree indices"""
19
19
 
 
20
from __future__ import absolute_import
 
21
 
20
22
import cStringIO
21
 
from bisect import bisect_right
 
23
 
 
24
from bzrlib.lazy_import import lazy_import
 
25
lazy_import(globals(), """
 
26
import bisect
22
27
import math
23
28
import tempfile
24
29
import zlib
 
30
""")
25
31
 
26
32
from bzrlib import (
27
33
    chunk_writer,
33
39
    osutils,
34
40
    static_tuple,
35
41
    trace,
 
42
    transport,
36
43
    )
37
44
from bzrlib.index import _OPTION_NODE_REFS, _OPTION_KEY_ELEMENTS, _OPTION_LEN
38
 
from bzrlib.transport import get_transport
39
45
 
40
46
 
41
47
_BTSIGNATURE = "B+Tree Graph Index 2\n"
158
164
        :param references: An iterable of iterables of keys. Each is a
159
165
            reference to another key.
160
166
        :param value: The value to associate with the key. It may be any
161
 
            bytes as long as it does not contain \0 or \n.
 
167
            bytes as long as it does not contain \\0 or \\n.
162
168
        """
163
169
        # Ensure that 'key' is a StaticTuple
164
170
        key = static_tuple.StaticTuple.from_sequence(key).intern()
193
199
            new_backing_file, size = self._spill_mem_keys_without_combining()
194
200
        # Note: The transport here isn't strictly needed, because we will use
195
201
        #       direct access to the new_backing._file object
196
 
        new_backing = BTreeGraphIndex(get_transport('.'), '<temp>', size)
 
202
        new_backing = BTreeGraphIndex(transport.get_transport_from_path('.'),
 
203
                                      '<temp>', size)
197
204
        # GC will clean up the file
198
205
        new_backing._file = new_backing_file
199
206
        if self._combine_backing_indices:
289
296
            flag when writing out. This is used by the _spill_mem_keys_to_disk
290
297
            functionality.
291
298
        """
 
299
        new_leaf = False
292
300
        if rows[-1].writer is None:
293
301
            # opening a new leaf chunk;
 
302
            new_leaf = True
294
303
            for pos, internal_row in enumerate(rows[:-1]):
295
304
                # flesh out any internal nodes that are needed to
296
305
                # preserve the height of the tree
315
324
                optimize_for_size=self._optimize_for_size)
316
325
            rows[-1].writer.write(_LEAF_FLAG)
317
326
        if rows[-1].writer.write(line):
 
327
            # if we failed to write, despite having an empty page to write to,
 
328
            # then line is too big. raising the error avoids infinite recursion
 
329
            # searching for a suitably large page that will not be found.
 
330
            if new_leaf:
 
331
                raise errors.BadIndexKey(string_key)
318
332
            # this key did not fit in the node:
319
333
            rows[-1].finish_node()
320
334
            key_line = string_key + "\n"
601
615
        """In memory index's have no known corruption at the moment."""
602
616
 
603
617
 
604
 
class _LeafNode(object):
 
618
class _LeafNode(dict):
605
619
    """A leaf node for a serialised B+Tree index."""
606
620
 
607
 
    __slots__ = ('keys', 'min_key', 'max_key')
 
621
    __slots__ = ('min_key', 'max_key', '_keys')
608
622
 
609
623
    def __init__(self, bytes, key_length, ref_list_length):
610
624
        """Parse bytes to create a leaf node object."""
616
630
            self.max_key = key_list[-1][0]
617
631
        else:
618
632
            self.min_key = self.max_key = None
619
 
        self.keys = dict(key_list)
 
633
        super(_LeafNode, self).__init__(key_list)
 
634
        self._keys = dict(self)
 
635
 
 
636
    def all_items(self):
 
637
        """Return a sorted list of (key, (value, refs)) items"""
 
638
        items = self.items()
 
639
        items.sort()
 
640
        return items
 
641
 
 
642
    def all_keys(self):
 
643
        """Return a sorted list of all keys."""
 
644
        keys = self.keys()
 
645
        keys.sort()
 
646
        return keys
620
647
 
621
648
 
622
649
class _InternalNode(object):
671
698
        self._recommended_pages = self._compute_recommended_pages()
672
699
        self._root_node = None
673
700
        self._base_offset = offset
 
701
        self._leaf_factory = _LeafNode
674
702
        # Default max size is 100,000 leave values
675
703
        self._leaf_value_cache = None # lru_cache.LRUCache(100*1000)
676
704
        if unlimited_cache:
949
977
        """Cache directly from key => value, skipping the btree."""
950
978
        if self._leaf_value_cache is not None:
951
979
            for node in nodes.itervalues():
952
 
                for key, value in node.keys.iteritems():
 
980
                for key, value in node.all_items():
953
981
                    if key in self._leaf_value_cache:
954
982
                        # Don't add the rest of the keys, we've seen this node
955
983
                        # before.
979
1007
        if self._row_offsets[-1] == 1:
980
1008
            # There is only the root node, and we read that via key_count()
981
1009
            if self.node_ref_lists:
982
 
                for key, (value, refs) in sorted(self._root_node.keys.items()):
 
1010
                for key, (value, refs) in self._root_node.all_items():
983
1011
                    yield (self, key, value, refs)
984
1012
            else:
985
 
                for key, (value, refs) in sorted(self._root_node.keys.items()):
 
1013
                for key, (value, refs) in self._root_node.all_items():
986
1014
                    yield (self, key, value)
987
1015
            return
988
1016
        start_of_leaves = self._row_offsets[-2]
998
1026
        # for spilling index builds to disk.
999
1027
        if self.node_ref_lists:
1000
1028
            for _, node in nodes:
1001
 
                for key, (value, refs) in sorted(node.keys.items()):
 
1029
                for key, (value, refs) in node.all_items():
1002
1030
                    yield (self, key, value, refs)
1003
1031
        else:
1004
1032
            for _, node in nodes:
1005
 
                for key, (value, refs) in sorted(node.keys.items()):
 
1033
                for key, (value, refs) in node.all_items():
1006
1034
                    yield (self, key, value)
1007
1035
 
1008
1036
    @staticmethod
1032
1060
        # iter_steps = len(in_keys) + len(fixed_keys)
1033
1061
        # bisect_steps = len(in_keys) * math.log(len(fixed_keys), 2)
1034
1062
        if len(in_keys) == 1: # Bisect will always be faster for M = 1
1035
 
            return [(bisect_right(fixed_keys, in_keys[0]), in_keys)]
 
1063
            return [(bisect.bisect_right(fixed_keys, in_keys[0]), in_keys)]
1036
1064
        # elif bisect_steps < iter_steps:
1037
1065
        #     offsets = {}
1038
1066
        #     for key in in_keys:
1169
1197
                continue
1170
1198
            node = nodes[node_index]
1171
1199
            for next_sub_key in sub_keys:
1172
 
                if next_sub_key in node.keys:
1173
 
                    value, refs = node.keys[next_sub_key]
 
1200
                if next_sub_key in node:
 
1201
                    value, refs = node[next_sub_key]
1174
1202
                    if self.node_ref_lists:
1175
1203
                        yield (self, next_sub_key, value, refs)
1176
1204
                    else:
1244
1272
            # sub_keys is all of the keys we are looking for that should exist
1245
1273
            # on this page, if they aren't here, then they won't be found
1246
1274
            node = nodes[node_index]
1247
 
            node_keys = node.keys
1248
1275
            parents_to_check = set()
1249
1276
            for next_sub_key in sub_keys:
1250
 
                if next_sub_key not in node_keys:
 
1277
                if next_sub_key not in node:
1251
1278
                    # This one is just not present in the index at all
1252
1279
                    missing_keys.add(next_sub_key)
1253
1280
                else:
1254
 
                    value, refs = node_keys[next_sub_key]
 
1281
                    value, refs = node[next_sub_key]
1255
1282
                    parent_keys = refs[ref_list_num]
1256
1283
                    parent_map[next_sub_key] = parent_keys
1257
1284
                    parents_to_check.update(parent_keys)
1264
1291
            while parents_to_check:
1265
1292
                next_parents_to_check = set()
1266
1293
                for key in parents_to_check:
1267
 
                    if key in node_keys:
1268
 
                        value, refs = node_keys[key]
 
1294
                    if key in node:
 
1295
                        value, refs = node[key]
1269
1296
                        parent_keys = refs[ref_list_num]
1270
1297
                        parent_map[key] = parent_keys
1271
1298
                        next_parents_to_check.update(parent_keys)
1545
1572
                    continue
1546
1573
            bytes = zlib.decompress(data)
1547
1574
            if bytes.startswith(_LEAF_FLAG):
1548
 
                node = _LeafNode(bytes, self._key_length, self.node_ref_lists)
 
1575
                node = self._leaf_factory(bytes, self._key_length,
 
1576
                                          self.node_ref_lists)
1549
1577
            elif bytes.startswith(_INTERNAL_FLAG):
1550
1578
                node = _InternalNode(bytes)
1551
1579
            else:
1570
1598
            pass
1571
1599
 
1572
1600
 
 
1601
_gcchk_factory = _LeafNode
 
1602
 
1573
1603
try:
1574
1604
    from bzrlib import _btree_serializer_pyx as _btree_serializer
 
1605
    _gcchk_factory = _btree_serializer._parse_into_chk
1575
1606
except ImportError, e:
1576
1607
    osutils.failed_to_load_extension(e)
1577
1608
    from bzrlib import _btree_serializer_py as _btree_serializer