~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/btree_index.py

  • Committer: Martin Pool
  • Date: 2010-02-25 06:17:27 UTC
  • mfrom: (5055 +trunk)
  • mto: This revision was merged to the branch mainline in revision 5057.
  • Revision ID: mbp@sourcefrog.net-20100225061727-4sd9lt0qmdc6087t
merge news

Show diffs side-by-side

added added

removed removed

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