~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"
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('.'),
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