← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/btree_index.py

  • Committer: Canonical.com Patch Queue Manager
  • Date: 2010-09-01 08:02:42 UTC
  • mfrom: (5390.3.3 faster-revert-593560)
  • Revision ID: pqm@pqm.ubuntu.com-20100901080242-esg62ody4frwmy66
(spiv) Avoid repeatedly calling self.target.all_file_ids() in
 InterTree.iter_changes. (Andrew Bennetts)

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/btree_index.py
1
 
# Copyright (C) 2008 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
17
17
 
18
18
"""B+Tree indices"""
19
19
 
 
20
import cStringIO
20
21
from bisect import bisect_right
21
22
import math
22
23
import tempfile
30
31
    index,
31
32
    lru_cache,
32
33
    osutils,
 
34
    static_tuple,
33
35
    trace,
 
36
    transport,
34
37
    )
35
38
from bzrlib.index import _OPTION_NODE_REFS, _OPTION_KEY_ELEMENTS, _OPTION_LEN
36
 
from bzrlib.transport import get_transport
37
39
 
38
40
 
39
41
_BTSIGNATURE = "B+Tree Graph Index 2\n"
60
62
    def __init__(self):
61
63
        """Create a _BuilderRow."""
62
64
        self.nodes = 0
63
 
        self.spool = tempfile.TemporaryFile()
 
65
        self.spool = None# tempfile.TemporaryFile(prefix='bzr-index-row-')
64
66
        self.writer = None
65
67
 
66
68
    def finish_node(self, pad=True):
67
69
        byte_lines, _, padding = self.writer.finish()
68
70
        if self.nodes == 0:
 
71
            self.spool = cStringIO.StringIO()
69
72
            # padded note:
70
73
            self.spool.write("\x00" * _RESERVED_HEADER_BYTES)
 
74
        elif self.nodes == 1:
 
75
            # We got bigger than 1 node, switch to a temp file
 
76
            spool = tempfile.TemporaryFile(prefix='bzr-index-row-')
 
77
            spool.write(self.spool.getvalue())
 
78
            self.spool = spool
71
79
        skipped_bytes = 0
72
80
        if not pad and padding:
73
81
            del byte_lines[-1]
152
160
        :param value: The value to associate with the key. It may be any
153
161
            bytes as long as it does not contain \0 or \n.
154
162
        """
 
163
        # Ensure that 'key' is a StaticTuple
 
164
        key = static_tuple.StaticTuple.from_sequence(key).intern()
155
165
        # we don't care about absent_references
156
166
        node_refs, _ = self._check_key_ref_value(key, references, value)
157
167
        if key in self._nodes:
158
168
            raise errors.BadIndexDuplicateKey(key, self)
159
 
        self._nodes[key] = (node_refs, value)
160
 
        self._keys.add(key)
 
169
        self._nodes[key] = static_tuple.StaticTuple(node_refs, value)
161
170
        if self._nodes_by_key is not None and self._key_length > 1:
162
171
            self._update_nodes_by_key(key, value, node_refs)
163
 
        if len(self._keys) < self._spill_at:
 
172
        if len(self._nodes) < self._spill_at:
164
173
            return
165
174
        self._spill_mem_keys_to_disk()
166
175
 
182
191
             backing_pos) = self._spill_mem_keys_and_combine()
183
192
        else:
184
193
            new_backing_file, size = self._spill_mem_keys_without_combining()
185
 
        dir_path, base_name = osutils.split(new_backing_file.name)
186
194
        # Note: The transport here isn't strictly needed, because we will use
187
195
        #       direct access to the new_backing._file object
188
 
        new_backing = BTreeGraphIndex(get_transport(dir_path),
189
 
                                      base_name, size)
 
196
        new_backing = BTreeGraphIndex(transport.get_transport('.'),
 
197
                                      '<temp>', size)
190
198
        # GC will clean up the file
191
199
        new_backing._file = new_backing_file
192
200
        if self._combine_backing_indices:
197
205
                self._backing_indices[backing_pos] = None
198
206
        else:
199
207
            self._backing_indices.append(new_backing)
200
 
        self._keys = set()
201
208
        self._nodes = {}
202
209
        self._nodes_by_key = None
203
210
 
379
386
        for row in reversed(rows):
380
387
            pad = (type(row) != _LeafBuilderRow)
381
388
            row.finish_node(pad=pad)
382
 
        result = tempfile.NamedTemporaryFile(prefix='bzr-index-')
383
389
        lines = [_BTSIGNATURE]
384
390
        lines.append(_OPTION_NODE_REFS + str(self.reference_lists) + '\n')
385
391
        lines.append(_OPTION_KEY_ELEMENTS + str(self._key_length) + '\n')
386
392
        lines.append(_OPTION_LEN + str(key_count) + '\n')
387
393
        row_lengths = [row.nodes for row in rows]
388
394
        lines.append(_OPTION_ROW_LENGTHS + ','.join(map(str, row_lengths)) + '\n')
 
395
        if row_lengths and row_lengths[-1] > 1:
 
396
            result = tempfile.NamedTemporaryFile(prefix='bzr-index-')
 
397
        else:
 
398
            result = cStringIO.StringIO()
389
399
        result.writelines(lines)
390
400
        position = sum(map(len, lines))
391
401
        root_row = True
402
412
            # Special case the first node as it may be prefixed
403
413
            node = row.spool.read(_PAGE_SIZE)
404
414
            result.write(node[reserved:])
405
 
            result.write("\x00" * (reserved - position))
 
415
            if len(node) == _PAGE_SIZE:
 
416
                result.write("\x00" * (reserved - position))
406
417
            position = 0 # Only the root row actually has an offset
407
418
            copied_len = osutils.pumpfile(row.spool, result)
408
419
            if copied_len != (row.nodes - 1) * _PAGE_SIZE:
453
464
            efficient order for the index (keys iteration order in this case).
454
465
        """
455
466
        keys = set(keys)
456
 
        local_keys = keys.intersection(self._keys)
 
467
        # Note: We don't use keys.intersection() here. If you read the C api,
 
468
        #       set.intersection(other) special cases when other is a set and
 
469
        #       will iterate the smaller of the two and lookup in the other.
 
470
        #       It does *not* do this for any other type (even dict, unlike
 
471
        #       some other set functions.) Since we expect keys is generally <<
 
472
        #       self._nodes, it is faster to iterate over it in a list
 
473
        #       comprehension
 
474
        nodes = self._nodes
 
475
        local_keys = [key for key in keys if key in nodes]
457
476
        if self.reference_lists:
458
477
            for key in local_keys:
459
 
                node = self._nodes[key]
 
478
                node = nodes[key]
460
479
                yield self, key, node[1], node[0]
461
480
        else:
462
481
            for key in local_keys:
463
 
                node = self._nodes[key]
 
482
                node = nodes[key]
464
483
                yield self, key, node[1]
465
484
        # Find things that are in backing indices that have not been handled
466
485
        # yet.
549
568
                    else:
550
569
                        # yield keys
551
570
                        for value in key_dict.itervalues():
552
 
                            yield (self, ) + value
 
571
                            yield (self, ) + tuple(value)
553
572
            else:
554
573
                yield (self, ) + key_dict
555
574
 
576
595
 
577
596
        For InMemoryGraphIndex the estimate is exact.
578
597
        """
579
 
        return len(self._keys) + sum(backing.key_count() for backing in
 
598
        return len(self._nodes) + sum(backing.key_count() for backing in
580
599
            self._backing_indices if backing is not None)
581
600
 
582
601
    def validate(self):
583
602
        """In memory index's have no known corruption at the moment."""
584
603
 
585
604
 
586
 
class _LeafNode(object):
 
605
class _LeafNode(dict):
587
606
    """A leaf node for a serialised B+Tree index."""
588
607
 
589
 
    __slots__ = ('keys', 'min_key', 'max_key')
 
608
    __slots__ = ('min_key', 'max_key', '_keys')
590
609
 
591
610
    def __init__(self, bytes, key_length, ref_list_length):
592
611
        """Parse bytes to create a leaf node object."""
598
617
            self.max_key = key_list[-1][0]
599
618
        else:
600
619
            self.min_key = self.max_key = None
601
 
        self.keys = dict(key_list)
 
620
        super(_LeafNode, self).__init__(key_list)
 
621
        self._keys = dict(self)
 
622
 
 
623
    def all_items(self):
 
624
        """Return a sorted list of (key, (value, refs)) items"""
 
625
        items = self.items()
 
626
        items.sort()
 
627
        return items
 
628
 
 
629
    def all_keys(self):
 
630
        """Return a sorted list of all keys."""
 
631
        keys = self.keys()
 
632
        keys.sort()
 
633
        return keys
602
634
 
603
635
 
604
636
class _InternalNode(object):
614
646
    def _parse_lines(self, lines):
615
647
        nodes = []
616
648
        self.offset = int(lines[1][7:])
 
649
        as_st = static_tuple.StaticTuple.from_sequence
617
650
        for line in lines[2:]:
618
651
            if line == '':
619
652
                break
620
 
            nodes.append(tuple(map(intern, line.split('\0'))))
 
653
            nodes.append(as_st(map(intern, line.split('\0'))).intern())
621
654
        return nodes
622
655
 
623
656
 
628
661
    memory except when very large walks are done.
629
662
    """
630
663
 
631
 
    def __init__(self, transport, name, size, unlimited_cache=False):
 
664
    def __init__(self, transport, name, size, unlimited_cache=False,
 
665
                 offset=0):
632
666
        """Create a B+Tree index object on the index name.
633
667
 
634
668
        :param transport: The transport to read data for the index from.
641
675
        :param unlimited_cache: If set to True, then instead of using an
642
676
            LRUCache with size _NODE_CACHE_SIZE, we will use a dict and always
643
677
            cache all leaf nodes.
 
678
        :param offset: The start of the btree index data isn't byte 0 of the
 
679
            file. Instead it starts at some point later.
644
680
        """
645
681
        self._transport = transport
646
682
        self._name = name
648
684
        self._file = None
649
685
        self._recommended_pages = self._compute_recommended_pages()
650
686
        self._root_node = None
 
687
        self._base_offset = offset
 
688
        self._leaf_factory = _LeafNode
651
689
        # Default max size is 100,000 leave values
652
690
        self._leaf_value_cache = None # lru_cache.LRUCache(100*1000)
653
691
        if unlimited_cache:
843
881
            new_tips = next_tips
844
882
        return final_offsets
845
883
 
 
884
    def clear_cache(self):
 
885
        """Clear out any cached/memoized values.
 
886
 
 
887
        This can be called at any time, but generally it is used when we have
 
888
        extracted some information, but don't expect to be requesting any more
 
889
        from this index.
 
890
        """
 
891
        # Note that we don't touch self._root_node or self._internal_node_cache
 
892
        # We don't expect either of those to be big, and it can save
 
893
        # round-trips in the future. We may re-evaluate this if InternalNode
 
894
        # memory starts to be an issue.
 
895
        self._leaf_node_cache.clear()
 
896
 
846
897
    def external_references(self, ref_list_num):
847
898
        if self._root_node is None:
848
899
            self._get_root_node()
913
964
        """Cache directly from key => value, skipping the btree."""
914
965
        if self._leaf_value_cache is not None:
915
966
            for node in nodes.itervalues():
916
 
                for key, value in node.keys.iteritems():
 
967
                for key, value in node.all_items():
917
968
                    if key in self._leaf_value_cache:
918
969
                        # Don't add the rest of the keys, we've seen this node
919
970
                        # before.
943
994
        if self._row_offsets[-1] == 1:
944
995
            # There is only the root node, and we read that via key_count()
945
996
            if self.node_ref_lists:
946
 
                for key, (value, refs) in sorted(self._root_node.keys.items()):
 
997
                for key, (value, refs) in self._root_node.all_items():
947
998
                    yield (self, key, value, refs)
948
999
            else:
949
 
                for key, (value, refs) in sorted(self._root_node.keys.items()):
 
1000
                for key, (value, refs) in self._root_node.all_items():
950
1001
                    yield (self, key, value)
951
1002
            return
952
1003
        start_of_leaves = self._row_offsets[-2]
962
1013
        # for spilling index builds to disk.
963
1014
        if self.node_ref_lists:
964
1015
            for _, node in nodes:
965
 
                for key, (value, refs) in sorted(node.keys.items()):
 
1016
                for key, (value, refs) in node.all_items():
966
1017
                    yield (self, key, value, refs)
967
1018
        else:
968
1019
            for _, node in nodes:
969
 
                for key, (value, refs) in sorted(node.keys.items()):
 
1020
                for key, (value, refs) in node.all_items():
970
1021
                    yield (self, key, value)
971
1022
 
972
1023
    @staticmethod
1133
1184
                continue
1134
1185
            node = nodes[node_index]
1135
1186
            for next_sub_key in sub_keys:
1136
 
                if next_sub_key in node.keys:
1137
 
                    value, refs = node.keys[next_sub_key]
 
1187
                if next_sub_key in node:
 
1188
                    value, refs = node[next_sub_key]
1138
1189
                    if self.node_ref_lists:
1139
1190
                        yield (self, next_sub_key, value, refs)
1140
1191
                    else:
1208
1259
            # sub_keys is all of the keys we are looking for that should exist
1209
1260
            # on this page, if they aren't here, then they won't be found
1210
1261
            node = nodes[node_index]
1211
 
            node_keys = node.keys
1212
1262
            parents_to_check = set()
1213
1263
            for next_sub_key in sub_keys:
1214
 
                if next_sub_key not in node_keys:
 
1264
                if next_sub_key not in node:
1215
1265
                    # This one is just not present in the index at all
1216
1266
                    missing_keys.add(next_sub_key)
1217
1267
                else:
1218
 
                    value, refs = node_keys[next_sub_key]
 
1268
                    value, refs = node[next_sub_key]
1219
1269
                    parent_keys = refs[ref_list_num]
1220
1270
                    parent_map[next_sub_key] = parent_keys
1221
1271
                    parents_to_check.update(parent_keys)
1228
1278
            while parents_to_check:
1229
1279
                next_parents_to_check = set()
1230
1280
                for key in parents_to_check:
1231
 
                    if key in node_keys:
1232
 
                        value, refs = node_keys[key]
 
1281
                    if key in node:
 
1282
                        value, refs = node[key]
1233
1283
                        parent_keys = refs[ref_list_num]
1234
1284
                        parent_map[key] = parent_keys
1235
1285
                        next_parents_to_check.update(parent_keys)
1462
1512
        # list of (offset, length) regions of the file that should, evenually
1463
1513
        # be read in to data_ranges, either from 'bytes' or from the transport
1464
1514
        ranges = []
 
1515
        base_offset = self._base_offset
1465
1516
        for index in nodes:
1466
 
            offset = index * _PAGE_SIZE
 
1517
            offset = (index * _PAGE_SIZE)
1467
1518
            size = _PAGE_SIZE
1468
1519
            if index == 0:
1469
1520
                # Root node - special case
1473
1524
                    # The only case where we don't know the size, is for very
1474
1525
                    # small indexes. So we read the whole thing
1475
1526
                    bytes = self._transport.get_bytes(self._name)
1476
 
                    self._size = len(bytes)
 
1527
                    num_bytes = len(bytes)
 
1528
                    self._size = num_bytes - base_offset
1477
1529
                    # the whole thing should be parsed out of 'bytes'
1478
 
                    ranges.append((0, len(bytes)))
 
1530
                    ranges = [(start, min(_PAGE_SIZE, num_bytes - start))
 
1531
                        for start in xrange(base_offset, num_bytes, _PAGE_SIZE)]
1479
1532
                    break
1480
1533
            else:
1481
1534
                if offset > self._size:
1483
1536
                                         ' of the file %s > %s'
1484
1537
                                         % (offset, self._size))
1485
1538
                size = min(size, self._size - offset)
1486
 
            ranges.append((offset, size))
 
1539
            ranges.append((base_offset + offset, size))
1487
1540
        if not ranges:
1488
1541
            return
1489
1542
        elif bytes is not None:
1490
1543
            # already have the whole file
1491
 
            data_ranges = [(start, bytes[start:start+_PAGE_SIZE])
1492
 
                           for start in xrange(0, len(bytes), _PAGE_SIZE)]
 
1544
            data_ranges = [(start, bytes[start:start+size])
 
1545
                           for start, size in ranges]
1493
1546
        elif self._file is None:
1494
1547
            data_ranges = self._transport.readv(self._name, ranges)
1495
1548
        else:
1498
1551
                self._file.seek(offset)
1499
1552
                data_ranges.append((offset, self._file.read(size)))
1500
1553
        for offset, data in data_ranges:
 
1554
            offset -= base_offset
1501
1555
            if offset == 0:
1502
1556
                # extract the header
1503
1557
                offset, data = self._parse_header_from_bytes(data)
1505
1559
                    continue
1506
1560
            bytes = zlib.decompress(data)
1507
1561
            if bytes.startswith(_LEAF_FLAG):
1508
 
                node = _LeafNode(bytes, self._key_length, self.node_ref_lists)
 
1562
                node = self._leaf_factory(bytes, self._key_length,
 
1563
                                          self.node_ref_lists)
1509
1564
            elif bytes.startswith(_INTERNAL_FLAG):
1510
1565
                node = _InternalNode(bytes)
1511
1566
            else:
1530
1585
            pass
1531
1586
 
1532
1587
 
 
1588
_gcchk_factory = _LeafNode
 
1589
 
1533
1590
try:
1534
1591
    from bzrlib import _btree_serializer_pyx as _btree_serializer
1535
 
except ImportError:
 
1592
    _gcchk_factory = _btree_serializer._parse_into_chk
 
1593
except ImportError, e:
 
1594
    osutils.failed_to_load_extension(e)
1536
1595
    from bzrlib import _btree_serializer_py as _btree_serializer