~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/btree_index.py

  • Committer: Joe Julian
  • Date: 2010-01-10 02:25:31 UTC
  • mto: (4634.119.7 2.0)
  • mto: This revision was merged to the branch mainline in revision 4959.
  • Revision ID: joe@julianfamily.org-20100110022531-wqk61rsagz8xsiga
Added MANIFEST.in to allow bdist_rpm to have all the required include files and tools. bdist_rpm will still fail to build correctly on some distributions due to a disttools bug http://bugs.python.org/issue644744

Show diffs side-by-side

added added

removed removed

Lines of Context:
1
 
# Copyright (C) 2008-2011 Canonical Ltd
 
1
# Copyright (C) 2008 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
 
 
22
 
import cStringIO
23
 
 
24
 
from bzrlib.lazy_import import lazy_import
25
 
lazy_import(globals(), """
26
 
import bisect
 
20
from bisect import bisect_right
27
21
import math
28
22
import tempfile
29
23
import zlib
30
 
""")
31
24
 
32
25
from bzrlib import (
33
26
    chunk_writer,
37
30
    index,
38
31
    lru_cache,
39
32
    osutils,
40
 
    static_tuple,
41
33
    trace,
42
 
    transport,
43
34
    )
44
35
from bzrlib.index import _OPTION_NODE_REFS, _OPTION_KEY_ELEMENTS, _OPTION_LEN
 
36
from bzrlib.transport import get_transport
45
37
 
46
38
 
47
39
_BTSIGNATURE = "B+Tree Graph Index 2\n"
68
60
    def __init__(self):
69
61
        """Create a _BuilderRow."""
70
62
        self.nodes = 0
71
 
        self.spool = None# tempfile.TemporaryFile(prefix='bzr-index-row-')
 
63
        self.spool = tempfile.TemporaryFile()
72
64
        self.writer = None
73
65
 
74
66
    def finish_node(self, pad=True):
75
67
        byte_lines, _, padding = self.writer.finish()
76
68
        if self.nodes == 0:
77
 
            self.spool = cStringIO.StringIO()
78
69
            # padded note:
79
70
            self.spool.write("\x00" * _RESERVED_HEADER_BYTES)
80
 
        elif self.nodes == 1:
81
 
            # We got bigger than 1 node, switch to a temp file
82
 
            spool = tempfile.TemporaryFile(prefix='bzr-index-row-')
83
 
            spool.write(self.spool.getvalue())
84
 
            self.spool = spool
85
71
        skipped_bytes = 0
86
72
        if not pad and padding:
87
73
            del byte_lines[-1]
164
150
        :param references: An iterable of iterables of keys. Each is a
165
151
            reference to another key.
166
152
        :param value: The value to associate with the key. It may be any
167
 
            bytes as long as it does not contain \\0 or \\n.
 
153
            bytes as long as it does not contain \0 or \n.
168
154
        """
169
 
        # Ensure that 'key' is a StaticTuple
170
 
        key = static_tuple.StaticTuple.from_sequence(key).intern()
171
155
        # we don't care about absent_references
172
156
        node_refs, _ = self._check_key_ref_value(key, references, value)
173
157
        if key in self._nodes:
174
158
            raise errors.BadIndexDuplicateKey(key, self)
175
 
        self._nodes[key] = static_tuple.StaticTuple(node_refs, value)
 
159
        self._nodes[key] = (node_refs, value)
 
160
        self._keys.add(key)
176
161
        if self._nodes_by_key is not None and self._key_length > 1:
177
162
            self._update_nodes_by_key(key, value, node_refs)
178
 
        if len(self._nodes) < self._spill_at:
 
163
        if len(self._keys) < self._spill_at:
179
164
            return
180
165
        self._spill_mem_keys_to_disk()
181
166
 
197
182
             backing_pos) = self._spill_mem_keys_and_combine()
198
183
        else:
199
184
            new_backing_file, size = self._spill_mem_keys_without_combining()
 
185
        dir_path, base_name = osutils.split(new_backing_file.name)
200
186
        # Note: The transport here isn't strictly needed, because we will use
201
187
        #       direct access to the new_backing._file object
202
 
        new_backing = BTreeGraphIndex(transport.get_transport_from_path('.'),
203
 
                                      '<temp>', size)
 
188
        new_backing = BTreeGraphIndex(get_transport(dir_path),
 
189
                                      base_name, size)
204
190
        # GC will clean up the file
205
191
        new_backing._file = new_backing_file
206
192
        if self._combine_backing_indices:
211
197
                self._backing_indices[backing_pos] = None
212
198
        else:
213
199
            self._backing_indices.append(new_backing)
 
200
        self._keys = set()
214
201
        self._nodes = {}
215
202
        self._nodes_by_key = None
216
203
 
296
283
            flag when writing out. This is used by the _spill_mem_keys_to_disk
297
284
            functionality.
298
285
        """
299
 
        new_leaf = False
300
286
        if rows[-1].writer is None:
301
287
            # opening a new leaf chunk;
302
 
            new_leaf = True
303
288
            for pos, internal_row in enumerate(rows[:-1]):
304
289
                # flesh out any internal nodes that are needed to
305
290
                # preserve the height of the tree
324
309
                optimize_for_size=self._optimize_for_size)
325
310
            rows[-1].writer.write(_LEAF_FLAG)
326
311
        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)
332
312
            # this key did not fit in the node:
333
313
            rows[-1].finish_node()
334
314
            key_line = string_key + "\n"
399
379
        for row in reversed(rows):
400
380
            pad = (type(row) != _LeafBuilderRow)
401
381
            row.finish_node(pad=pad)
 
382
        result = tempfile.NamedTemporaryFile(prefix='bzr-index-')
402
383
        lines = [_BTSIGNATURE]
403
384
        lines.append(_OPTION_NODE_REFS + str(self.reference_lists) + '\n')
404
385
        lines.append(_OPTION_KEY_ELEMENTS + str(self._key_length) + '\n')
405
386
        lines.append(_OPTION_LEN + str(key_count) + '\n')
406
387
        row_lengths = [row.nodes for row in rows]
407
388
        lines.append(_OPTION_ROW_LENGTHS + ','.join(map(str, row_lengths)) + '\n')
408
 
        if row_lengths and row_lengths[-1] > 1:
409
 
            result = tempfile.NamedTemporaryFile(prefix='bzr-index-')
410
 
        else:
411
 
            result = cStringIO.StringIO()
412
389
        result.writelines(lines)
413
390
        position = sum(map(len, lines))
414
391
        root_row = True
425
402
            # Special case the first node as it may be prefixed
426
403
            node = row.spool.read(_PAGE_SIZE)
427
404
            result.write(node[reserved:])
428
 
            if len(node) == _PAGE_SIZE:
429
 
                result.write("\x00" * (reserved - position))
 
405
            result.write("\x00" * (reserved - position))
430
406
            position = 0 # Only the root row actually has an offset
431
407
            copied_len = osutils.pumpfile(row.spool, result)
432
408
            if copied_len != (row.nodes - 1) * _PAGE_SIZE:
477
453
            efficient order for the index (keys iteration order in this case).
478
454
        """
479
455
        keys = set(keys)
480
 
        # Note: We don't use keys.intersection() here. If you read the C api,
481
 
        #       set.intersection(other) special cases when other is a set and
482
 
        #       will iterate the smaller of the two and lookup in the other.
483
 
        #       It does *not* do this for any other type (even dict, unlike
484
 
        #       some other set functions.) Since we expect keys is generally <<
485
 
        #       self._nodes, it is faster to iterate over it in a list
486
 
        #       comprehension
487
 
        nodes = self._nodes
488
 
        local_keys = [key for key in keys if key in nodes]
 
456
        local_keys = keys.intersection(self._keys)
489
457
        if self.reference_lists:
490
458
            for key in local_keys:
491
 
                node = nodes[key]
 
459
                node = self._nodes[key]
492
460
                yield self, key, node[1], node[0]
493
461
        else:
494
462
            for key in local_keys:
495
 
                node = nodes[key]
 
463
                node = self._nodes[key]
496
464
                yield self, key, node[1]
497
465
        # Find things that are in backing indices that have not been handled
498
466
        # yet.
581
549
                    else:
582
550
                        # yield keys
583
551
                        for value in key_dict.itervalues():
584
 
                            yield (self, ) + tuple(value)
 
552
                            yield (self, ) + value
585
553
            else:
586
554
                yield (self, ) + key_dict
587
555
 
608
576
 
609
577
        For InMemoryGraphIndex the estimate is exact.
610
578
        """
611
 
        return len(self._nodes) + sum(backing.key_count() for backing in
 
579
        return len(self._keys) + sum(backing.key_count() for backing in
612
580
            self._backing_indices if backing is not None)
613
581
 
614
582
    def validate(self):
615
583
        """In memory index's have no known corruption at the moment."""
616
584
 
617
585
 
618
 
class _LeafNode(dict):
 
586
class _LeafNode(object):
619
587
    """A leaf node for a serialised B+Tree index."""
620
588
 
621
 
    __slots__ = ('min_key', 'max_key', '_keys')
 
589
    __slots__ = ('keys', 'min_key', 'max_key')
622
590
 
623
591
    def __init__(self, bytes, key_length, ref_list_length):
624
592
        """Parse bytes to create a leaf node object."""
630
598
            self.max_key = key_list[-1][0]
631
599
        else:
632
600
            self.min_key = self.max_key = None
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
 
601
        self.keys = dict(key_list)
647
602
 
648
603
 
649
604
class _InternalNode(object):
659
614
    def _parse_lines(self, lines):
660
615
        nodes = []
661
616
        self.offset = int(lines[1][7:])
662
 
        as_st = static_tuple.StaticTuple.from_sequence
663
617
        for line in lines[2:]:
664
618
            if line == '':
665
619
                break
666
 
            nodes.append(as_st(map(intern, line.split('\0'))).intern())
 
620
            nodes.append(tuple(map(intern, line.split('\0'))))
667
621
        return nodes
668
622
 
669
623
 
674
628
    memory except when very large walks are done.
675
629
    """
676
630
 
677
 
    def __init__(self, transport, name, size, unlimited_cache=False,
678
 
                 offset=0):
 
631
    def __init__(self, transport, name, size, unlimited_cache=False):
679
632
        """Create a B+Tree index object on the index name.
680
633
 
681
634
        :param transport: The transport to read data for the index from.
688
641
        :param unlimited_cache: If set to True, then instead of using an
689
642
            LRUCache with size _NODE_CACHE_SIZE, we will use a dict and always
690
643
            cache all leaf nodes.
691
 
        :param offset: The start of the btree index data isn't byte 0 of the
692
 
            file. Instead it starts at some point later.
693
644
        """
694
645
        self._transport = transport
695
646
        self._name = name
697
648
        self._file = None
698
649
        self._recommended_pages = self._compute_recommended_pages()
699
650
        self._root_node = None
700
 
        self._base_offset = offset
701
 
        self._leaf_factory = _LeafNode
702
651
        # Default max size is 100,000 leave values
703
652
        self._leaf_value_cache = None # lru_cache.LRUCache(100*1000)
704
653
        if unlimited_cache:
894
843
            new_tips = next_tips
895
844
        return final_offsets
896
845
 
897
 
    def clear_cache(self):
898
 
        """Clear out any cached/memoized values.
899
 
 
900
 
        This can be called at any time, but generally it is used when we have
901
 
        extracted some information, but don't expect to be requesting any more
902
 
        from this index.
903
 
        """
904
 
        # Note that we don't touch self._root_node or self._internal_node_cache
905
 
        # We don't expect either of those to be big, and it can save
906
 
        # round-trips in the future. We may re-evaluate this if InternalNode
907
 
        # memory starts to be an issue.
908
 
        self._leaf_node_cache.clear()
909
 
 
910
846
    def external_references(self, ref_list_num):
911
847
        if self._root_node is None:
912
848
            self._get_root_node()
977
913
        """Cache directly from key => value, skipping the btree."""
978
914
        if self._leaf_value_cache is not None:
979
915
            for node in nodes.itervalues():
980
 
                for key, value in node.all_items():
 
916
                for key, value in node.keys.iteritems():
981
917
                    if key in self._leaf_value_cache:
982
918
                        # Don't add the rest of the keys, we've seen this node
983
919
                        # before.
1007
943
        if self._row_offsets[-1] == 1:
1008
944
            # There is only the root node, and we read that via key_count()
1009
945
            if self.node_ref_lists:
1010
 
                for key, (value, refs) in self._root_node.all_items():
 
946
                for key, (value, refs) in sorted(self._root_node.keys.items()):
1011
947
                    yield (self, key, value, refs)
1012
948
            else:
1013
 
                for key, (value, refs) in self._root_node.all_items():
 
949
                for key, (value, refs) in sorted(self._root_node.keys.items()):
1014
950
                    yield (self, key, value)
1015
951
            return
1016
952
        start_of_leaves = self._row_offsets[-2]
1026
962
        # for spilling index builds to disk.
1027
963
        if self.node_ref_lists:
1028
964
            for _, node in nodes:
1029
 
                for key, (value, refs) in node.all_items():
 
965
                for key, (value, refs) in sorted(node.keys.items()):
1030
966
                    yield (self, key, value, refs)
1031
967
        else:
1032
968
            for _, node in nodes:
1033
 
                for key, (value, refs) in node.all_items():
 
969
                for key, (value, refs) in sorted(node.keys.items()):
1034
970
                    yield (self, key, value)
1035
971
 
1036
972
    @staticmethod
1060
996
        # iter_steps = len(in_keys) + len(fixed_keys)
1061
997
        # bisect_steps = len(in_keys) * math.log(len(fixed_keys), 2)
1062
998
        if len(in_keys) == 1: # Bisect will always be faster for M = 1
1063
 
            return [(bisect.bisect_right(fixed_keys, in_keys[0]), in_keys)]
 
999
            return [(bisect_right(fixed_keys, in_keys[0]), in_keys)]
1064
1000
        # elif bisect_steps < iter_steps:
1065
1001
        #     offsets = {}
1066
1002
        #     for key in in_keys:
1197
1133
                continue
1198
1134
            node = nodes[node_index]
1199
1135
            for next_sub_key in sub_keys:
1200
 
                if next_sub_key in node:
1201
 
                    value, refs = node[next_sub_key]
 
1136
                if next_sub_key in node.keys:
 
1137
                    value, refs = node.keys[next_sub_key]
1202
1138
                    if self.node_ref_lists:
1203
1139
                        yield (self, next_sub_key, value, refs)
1204
1140
                    else:
1272
1208
            # sub_keys is all of the keys we are looking for that should exist
1273
1209
            # on this page, if they aren't here, then they won't be found
1274
1210
            node = nodes[node_index]
 
1211
            node_keys = node.keys
1275
1212
            parents_to_check = set()
1276
1213
            for next_sub_key in sub_keys:
1277
 
                if next_sub_key not in node:
 
1214
                if next_sub_key not in node_keys:
1278
1215
                    # This one is just not present in the index at all
1279
1216
                    missing_keys.add(next_sub_key)
1280
1217
                else:
1281
 
                    value, refs = node[next_sub_key]
 
1218
                    value, refs = node_keys[next_sub_key]
1282
1219
                    parent_keys = refs[ref_list_num]
1283
1220
                    parent_map[next_sub_key] = parent_keys
1284
1221
                    parents_to_check.update(parent_keys)
1291
1228
            while parents_to_check:
1292
1229
                next_parents_to_check = set()
1293
1230
                for key in parents_to_check:
1294
 
                    if key in node:
1295
 
                        value, refs = node[key]
 
1231
                    if key in node_keys:
 
1232
                        value, refs = node_keys[key]
1296
1233
                        parent_keys = refs[ref_list_num]
1297
1234
                        parent_map[key] = parent_keys
1298
1235
                        next_parents_to_check.update(parent_keys)
1525
1462
        # list of (offset, length) regions of the file that should, evenually
1526
1463
        # be read in to data_ranges, either from 'bytes' or from the transport
1527
1464
        ranges = []
1528
 
        base_offset = self._base_offset
1529
1465
        for index in nodes:
1530
 
            offset = (index * _PAGE_SIZE)
 
1466
            offset = index * _PAGE_SIZE
1531
1467
            size = _PAGE_SIZE
1532
1468
            if index == 0:
1533
1469
                # Root node - special case
1537
1473
                    # The only case where we don't know the size, is for very
1538
1474
                    # small indexes. So we read the whole thing
1539
1475
                    bytes = self._transport.get_bytes(self._name)
1540
 
                    num_bytes = len(bytes)
1541
 
                    self._size = num_bytes - base_offset
 
1476
                    self._size = len(bytes)
1542
1477
                    # the whole thing should be parsed out of 'bytes'
1543
 
                    ranges = [(start, min(_PAGE_SIZE, num_bytes - start))
1544
 
                        for start in xrange(base_offset, num_bytes, _PAGE_SIZE)]
 
1478
                    ranges.append((0, len(bytes)))
1545
1479
                    break
1546
1480
            else:
1547
1481
                if offset > self._size:
1549
1483
                                         ' of the file %s > %s'
1550
1484
                                         % (offset, self._size))
1551
1485
                size = min(size, self._size - offset)
1552
 
            ranges.append((base_offset + offset, size))
 
1486
            ranges.append((offset, size))
1553
1487
        if not ranges:
1554
1488
            return
1555
1489
        elif bytes is not None:
1556
1490
            # already have the whole file
1557
 
            data_ranges = [(start, bytes[start:start+size])
1558
 
                           for start, size in ranges]
 
1491
            data_ranges = [(start, bytes[start:start+_PAGE_SIZE])
 
1492
                           for start in xrange(0, len(bytes), _PAGE_SIZE)]
1559
1493
        elif self._file is None:
1560
1494
            data_ranges = self._transport.readv(self._name, ranges)
1561
1495
        else:
1564
1498
                self._file.seek(offset)
1565
1499
                data_ranges.append((offset, self._file.read(size)))
1566
1500
        for offset, data in data_ranges:
1567
 
            offset -= base_offset
1568
1501
            if offset == 0:
1569
1502
                # extract the header
1570
1503
                offset, data = self._parse_header_from_bytes(data)
1572
1505
                    continue
1573
1506
            bytes = zlib.decompress(data)
1574
1507
            if bytes.startswith(_LEAF_FLAG):
1575
 
                node = self._leaf_factory(bytes, self._key_length,
1576
 
                                          self.node_ref_lists)
 
1508
                node = _LeafNode(bytes, self._key_length, self.node_ref_lists)
1577
1509
            elif bytes.startswith(_INTERNAL_FLAG):
1578
1510
                node = _InternalNode(bytes)
1579
1511
            else:
1598
1530
            pass
1599
1531
 
1600
1532
 
1601
 
_gcchk_factory = _LeafNode
1602
 
 
1603
1533
try:
1604
1534
    from bzrlib import _btree_serializer_pyx as _btree_serializer
1605
 
    _gcchk_factory = _btree_serializer._parse_into_chk
1606
 
except ImportError, e:
1607
 
    osutils.failed_to_load_extension(e)
 
1535
except ImportError:
1608
1536
    from bzrlib import _btree_serializer_py as _btree_serializer