~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/btree_index.py

  • Committer: Jelmer Vernooij
  • Date: 2011-09-26 15:21:01 UTC
  • mfrom: (6165.2.3 avoid-revision-history)
  • mto: This revision was merged to the branch mainline in revision 6216.
  • Revision ID: jelmer@samba.org-20110926152101-afcxw1hikybyivfd
merge avoid-revision-history.

Show diffs side-by-side

added added

removed removed

Lines of Context:
1
 
# Copyright (C) 2008 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
18
18
"""B+Tree indices"""
19
19
 
20
20
import cStringIO
21
 
from bisect import bisect_right
 
21
 
 
22
from bzrlib.lazy_import import lazy_import
 
23
lazy_import(globals(), """
 
24
import bisect
22
25
import math
23
26
import tempfile
24
27
import zlib
 
28
""")
25
29
 
26
30
from bzrlib import (
27
31
    chunk_writer,
33
37
    osutils,
34
38
    static_tuple,
35
39
    trace,
 
40
    transport,
36
41
    )
37
42
from bzrlib.index import _OPTION_NODE_REFS, _OPTION_KEY_ELEMENTS, _OPTION_LEN
38
 
from bzrlib.transport import get_transport
39
43
 
40
44
 
41
45
_BTSIGNATURE = "B+Tree Graph Index 2\n"
158
162
        :param references: An iterable of iterables of keys. Each is a
159
163
            reference to another key.
160
164
        :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.
 
165
            bytes as long as it does not contain \\0 or \\n.
162
166
        """
163
167
        # Ensure that 'key' is a StaticTuple
164
168
        key = static_tuple.StaticTuple.from_sequence(key).intern()
193
197
            new_backing_file, size = self._spill_mem_keys_without_combining()
194
198
        # Note: The transport here isn't strictly needed, because we will use
195
199
        #       direct access to the new_backing._file object
196
 
        new_backing = BTreeGraphIndex(get_transport('.'), '<temp>', size)
 
200
        new_backing = BTreeGraphIndex(transport.get_transport_from_path('.'),
 
201
                                      '<temp>', size)
197
202
        # GC will clean up the file
198
203
        new_backing._file = new_backing_file
199
204
        if self._combine_backing_indices:
567
572
                    else:
568
573
                        # yield keys
569
574
                        for value in key_dict.itervalues():
570
 
                            yield (self, ) + value
 
575
                            yield (self, ) + tuple(value)
571
576
            else:
572
577
                yield (self, ) + key_dict
573
578
 
601
606
        """In memory index's have no known corruption at the moment."""
602
607
 
603
608
 
604
 
class _LeafNode(object):
 
609
class _LeafNode(dict):
605
610
    """A leaf node for a serialised B+Tree index."""
606
611
 
607
 
    __slots__ = ('keys', 'min_key', 'max_key')
 
612
    __slots__ = ('min_key', 'max_key', '_keys')
608
613
 
609
614
    def __init__(self, bytes, key_length, ref_list_length):
610
615
        """Parse bytes to create a leaf node object."""
616
621
            self.max_key = key_list[-1][0]
617
622
        else:
618
623
            self.min_key = self.max_key = None
619
 
        self.keys = dict(key_list)
 
624
        super(_LeafNode, self).__init__(key_list)
 
625
        self._keys = dict(self)
 
626
 
 
627
    def all_items(self):
 
628
        """Return a sorted list of (key, (value, refs)) items"""
 
629
        items = self.items()
 
630
        items.sort()
 
631
        return items
 
632
 
 
633
    def all_keys(self):
 
634
        """Return a sorted list of all keys."""
 
635
        keys = self.keys()
 
636
        keys.sort()
 
637
        return keys
620
638
 
621
639
 
622
640
class _InternalNode(object):
647
665
    memory except when very large walks are done.
648
666
    """
649
667
 
650
 
    def __init__(self, transport, name, size, unlimited_cache=False):
 
668
    def __init__(self, transport, name, size, unlimited_cache=False,
 
669
                 offset=0):
651
670
        """Create a B+Tree index object on the index name.
652
671
 
653
672
        :param transport: The transport to read data for the index from.
660
679
        :param unlimited_cache: If set to True, then instead of using an
661
680
            LRUCache with size _NODE_CACHE_SIZE, we will use a dict and always
662
681
            cache all leaf nodes.
 
682
        :param offset: The start of the btree index data isn't byte 0 of the
 
683
            file. Instead it starts at some point later.
663
684
        """
664
685
        self._transport = transport
665
686
        self._name = name
667
688
        self._file = None
668
689
        self._recommended_pages = self._compute_recommended_pages()
669
690
        self._root_node = None
 
691
        self._base_offset = offset
 
692
        self._leaf_factory = _LeafNode
670
693
        # Default max size is 100,000 leave values
671
694
        self._leaf_value_cache = None # lru_cache.LRUCache(100*1000)
672
695
        if unlimited_cache:
945
968
        """Cache directly from key => value, skipping the btree."""
946
969
        if self._leaf_value_cache is not None:
947
970
            for node in nodes.itervalues():
948
 
                for key, value in node.keys.iteritems():
 
971
                for key, value in node.all_items():
949
972
                    if key in self._leaf_value_cache:
950
973
                        # Don't add the rest of the keys, we've seen this node
951
974
                        # before.
975
998
        if self._row_offsets[-1] == 1:
976
999
            # There is only the root node, and we read that via key_count()
977
1000
            if self.node_ref_lists:
978
 
                for key, (value, refs) in sorted(self._root_node.keys.items()):
 
1001
                for key, (value, refs) in self._root_node.all_items():
979
1002
                    yield (self, key, value, refs)
980
1003
            else:
981
 
                for key, (value, refs) in sorted(self._root_node.keys.items()):
 
1004
                for key, (value, refs) in self._root_node.all_items():
982
1005
                    yield (self, key, value)
983
1006
            return
984
1007
        start_of_leaves = self._row_offsets[-2]
994
1017
        # for spilling index builds to disk.
995
1018
        if self.node_ref_lists:
996
1019
            for _, node in nodes:
997
 
                for key, (value, refs) in sorted(node.keys.items()):
 
1020
                for key, (value, refs) in node.all_items():
998
1021
                    yield (self, key, value, refs)
999
1022
        else:
1000
1023
            for _, node in nodes:
1001
 
                for key, (value, refs) in sorted(node.keys.items()):
 
1024
                for key, (value, refs) in node.all_items():
1002
1025
                    yield (self, key, value)
1003
1026
 
1004
1027
    @staticmethod
1028
1051
        # iter_steps = len(in_keys) + len(fixed_keys)
1029
1052
        # bisect_steps = len(in_keys) * math.log(len(fixed_keys), 2)
1030
1053
        if len(in_keys) == 1: # Bisect will always be faster for M = 1
1031
 
            return [(bisect_right(fixed_keys, in_keys[0]), in_keys)]
 
1054
            return [(bisect.bisect_right(fixed_keys, in_keys[0]), in_keys)]
1032
1055
        # elif bisect_steps < iter_steps:
1033
1056
        #     offsets = {}
1034
1057
        #     for key in in_keys:
1165
1188
                continue
1166
1189
            node = nodes[node_index]
1167
1190
            for next_sub_key in sub_keys:
1168
 
                if next_sub_key in node.keys:
1169
 
                    value, refs = node.keys[next_sub_key]
 
1191
                if next_sub_key in node:
 
1192
                    value, refs = node[next_sub_key]
1170
1193
                    if self.node_ref_lists:
1171
1194
                        yield (self, next_sub_key, value, refs)
1172
1195
                    else:
1240
1263
            # sub_keys is all of the keys we are looking for that should exist
1241
1264
            # on this page, if they aren't here, then they won't be found
1242
1265
            node = nodes[node_index]
1243
 
            node_keys = node.keys
1244
1266
            parents_to_check = set()
1245
1267
            for next_sub_key in sub_keys:
1246
 
                if next_sub_key not in node_keys:
 
1268
                if next_sub_key not in node:
1247
1269
                    # This one is just not present in the index at all
1248
1270
                    missing_keys.add(next_sub_key)
1249
1271
                else:
1250
 
                    value, refs = node_keys[next_sub_key]
 
1272
                    value, refs = node[next_sub_key]
1251
1273
                    parent_keys = refs[ref_list_num]
1252
1274
                    parent_map[next_sub_key] = parent_keys
1253
1275
                    parents_to_check.update(parent_keys)
1260
1282
            while parents_to_check:
1261
1283
                next_parents_to_check = set()
1262
1284
                for key in parents_to_check:
1263
 
                    if key in node_keys:
1264
 
                        value, refs = node_keys[key]
 
1285
                    if key in node:
 
1286
                        value, refs = node[key]
1265
1287
                        parent_keys = refs[ref_list_num]
1266
1288
                        parent_map[key] = parent_keys
1267
1289
                        next_parents_to_check.update(parent_keys)
1494
1516
        # list of (offset, length) regions of the file that should, evenually
1495
1517
        # be read in to data_ranges, either from 'bytes' or from the transport
1496
1518
        ranges = []
 
1519
        base_offset = self._base_offset
1497
1520
        for index in nodes:
1498
 
            offset = index * _PAGE_SIZE
 
1521
            offset = (index * _PAGE_SIZE)
1499
1522
            size = _PAGE_SIZE
1500
1523
            if index == 0:
1501
1524
                # Root node - special case
1505
1528
                    # The only case where we don't know the size, is for very
1506
1529
                    # small indexes. So we read the whole thing
1507
1530
                    bytes = self._transport.get_bytes(self._name)
1508
 
                    self._size = len(bytes)
 
1531
                    num_bytes = len(bytes)
 
1532
                    self._size = num_bytes - base_offset
1509
1533
                    # the whole thing should be parsed out of 'bytes'
1510
 
                    ranges.append((0, len(bytes)))
 
1534
                    ranges = [(start, min(_PAGE_SIZE, num_bytes - start))
 
1535
                        for start in xrange(base_offset, num_bytes, _PAGE_SIZE)]
1511
1536
                    break
1512
1537
            else:
1513
1538
                if offset > self._size:
1515
1540
                                         ' of the file %s > %s'
1516
1541
                                         % (offset, self._size))
1517
1542
                size = min(size, self._size - offset)
1518
 
            ranges.append((offset, size))
 
1543
            ranges.append((base_offset + offset, size))
1519
1544
        if not ranges:
1520
1545
            return
1521
1546
        elif bytes is not None:
1522
1547
            # already have the whole file
1523
 
            data_ranges = [(start, bytes[start:start+_PAGE_SIZE])
1524
 
                           for start in xrange(0, len(bytes), _PAGE_SIZE)]
 
1548
            data_ranges = [(start, bytes[start:start+size])
 
1549
                           for start, size in ranges]
1525
1550
        elif self._file is None:
1526
1551
            data_ranges = self._transport.readv(self._name, ranges)
1527
1552
        else:
1530
1555
                self._file.seek(offset)
1531
1556
                data_ranges.append((offset, self._file.read(size)))
1532
1557
        for offset, data in data_ranges:
 
1558
            offset -= base_offset
1533
1559
            if offset == 0:
1534
1560
                # extract the header
1535
1561
                offset, data = self._parse_header_from_bytes(data)
1537
1563
                    continue
1538
1564
            bytes = zlib.decompress(data)
1539
1565
            if bytes.startswith(_LEAF_FLAG):
1540
 
                node = _LeafNode(bytes, self._key_length, self.node_ref_lists)
 
1566
                node = self._leaf_factory(bytes, self._key_length,
 
1567
                                          self.node_ref_lists)
1541
1568
            elif bytes.startswith(_INTERNAL_FLAG):
1542
1569
                node = _InternalNode(bytes)
1543
1570
            else:
1562
1589
            pass
1563
1590
 
1564
1591
 
 
1592
_gcchk_factory = _LeafNode
 
1593
 
1565
1594
try:
1566
1595
    from bzrlib import _btree_serializer_pyx as _btree_serializer
 
1596
    _gcchk_factory = _btree_serializer._parse_into_chk
1567
1597
except ImportError, e:
1568
1598
    osutils.failed_to_load_extension(e)
1569
1599
    from bzrlib import _btree_serializer_py as _btree_serializer