/brz/remove-bazaar

To get this branch, use:
bzr branch http://gegoxaren.bato24.eu/bzr/brz/remove-bazaar

« back to all changes in this revision

Viewing changes to breezy/btree_index.py

  • Committer: Breezy landing bot
  • Author(s): Martin
  • Date: 2017-06-05 14:25:10 UTC
  • mfrom: (6654.1.1 index_prefix_dedupe)
  • Revision ID: breezy.the.bot@gmail.com-20170605142510-v8g9z86d4yw78fvm
Add index._iter_entries_prefix helper to remove copied code

Merged from https://code.launchpad.net/~gz/brz/index_prefix_dedupe/+merge/325052

Show diffs side-by-side

added added

removed removed

Lines of Context:
breezy/btree_index.py
1
 
# Copyright (C) 2008, 2009, 2010 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
17
17
 
18
18
"""B+Tree indices"""
19
19
 
20
 
import cStringIO
21
 
from bisect import bisect_right
 
20
from __future__ import absolute_import
 
21
 
 
22
from .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
 
from bzrlib import (
 
30
from . import (
27
31
    chunk_writer,
28
32
    debug,
29
33
    errors,
33
37
    osutils,
34
38
    static_tuple,
35
39
    trace,
36
 
    )
37
 
from bzrlib.index import _OPTION_NODE_REFS, _OPTION_KEY_ELEMENTS, _OPTION_LEN
38
 
from bzrlib.transport import get_transport
 
40
    transport,
 
41
    )
 
42
from .index import _OPTION_NODE_REFS, _OPTION_KEY_ELEMENTS, _OPTION_LEN
 
43
from .sixish import (
 
44
    BytesIO,
 
45
    map,
 
46
    range,
 
47
    )
39
48
 
40
49
 
41
50
_BTSIGNATURE = "B+Tree Graph Index 2\n"
68
77
    def finish_node(self, pad=True):
69
78
        byte_lines, _, padding = self.writer.finish()
70
79
        if self.nodes == 0:
71
 
            self.spool = cStringIO.StringIO()
 
80
            self.spool = BytesIO()
72
81
            # padded note:
73
82
            self.spool.write("\x00" * _RESERVED_HEADER_BYTES)
74
83
        elif self.nodes == 1:
158
167
        :param references: An iterable of iterables of keys. Each is a
159
168
            reference to another key.
160
169
        :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.
 
170
            bytes as long as it does not contain \\0 or \\n.
162
171
        """
163
172
        # Ensure that 'key' is a StaticTuple
164
173
        key = static_tuple.StaticTuple.from_sequence(key).intern()
193
202
            new_backing_file, size = self._spill_mem_keys_without_combining()
194
203
        # Note: The transport here isn't strictly needed, because we will use
195
204
        #       direct access to the new_backing._file object
196
 
        new_backing = BTreeGraphIndex(get_transport('.'), '<temp>', size)
 
205
        new_backing = BTreeGraphIndex(transport.get_transport_from_path('.'),
 
206
                                      '<temp>', size)
197
207
        # GC will clean up the file
198
208
        new_backing._file = new_backing_file
199
209
        if self._combine_backing_indices:
256
266
        current_values = []
257
267
        for iterator in iterators_to_combine:
258
268
            try:
259
 
                current_values.append(iterator.next())
 
269
                current_values.append(next(iterator))
260
270
            except StopIteration:
261
271
                current_values.append(None)
262
272
        last = None
276
286
            yield (self,) + selected[1][1:]
277
287
            pos = selected[0]
278
288
            try:
279
 
                current_values[pos] = iterators_to_combine[pos].next()
 
289
                current_values[pos] = next(iterators_to_combine[pos])
280
290
            except StopIteration:
281
291
                current_values[pos] = None
282
292
 
289
299
            flag when writing out. This is used by the _spill_mem_keys_to_disk
290
300
            functionality.
291
301
        """
 
302
        new_leaf = False
292
303
        if rows[-1].writer is None:
293
304
            # opening a new leaf chunk;
 
305
            new_leaf = True
294
306
            for pos, internal_row in enumerate(rows[:-1]):
295
307
                # flesh out any internal nodes that are needed to
296
308
                # preserve the height of the tree
315
327
                optimize_for_size=self._optimize_for_size)
316
328
            rows[-1].writer.write(_LEAF_FLAG)
317
329
        if rows[-1].writer.write(line):
 
330
            # if we failed to write, despite having an empty page to write to,
 
331
            # then line is too big. raising the error avoids infinite recursion
 
332
            # searching for a suitably large page that will not be found.
 
333
            if new_leaf:
 
334
                raise errors.BadIndexKey(string_key)
318
335
            # this key did not fit in the node:
319
336
            rows[-1].finish_node()
320
337
            key_line = string_key + "\n"
383
400
                                    self.reference_lists)
384
401
            self._add_key(string_key, line, rows, allow_optimize=allow_optimize)
385
402
        for row in reversed(rows):
386
 
            pad = (type(row) != _LeafBuilderRow)
 
403
            pad = (not isinstance(row, _LeafBuilderRow))
387
404
            row.finish_node(pad=pad)
388
405
        lines = [_BTSIGNATURE]
389
406
        lines.append(_OPTION_NODE_REFS + str(self.reference_lists) + '\n')
394
411
        if row_lengths and row_lengths[-1] > 1:
395
412
            result = tempfile.NamedTemporaryFile(prefix='bzr-index-')
396
413
        else:
397
 
            result = cStringIO.StringIO()
 
414
            result = BytesIO()
398
415
        result.writelines(lines)
399
416
        position = sum(map(len, lines))
400
417
        root_row = True
416
433
            position = 0 # Only the root row actually has an offset
417
434
            copied_len = osutils.pumpfile(row.spool, result)
418
435
            if copied_len != (row.nodes - 1) * _PAGE_SIZE:
419
 
                if type(row) != _LeafBuilderRow:
 
436
                if not isinstance(row, _LeafBuilderRow):
420
437
                    raise AssertionError("Incorrect amount of data copied"
421
438
                        " expected: %d, got: %d"
422
439
                        % ((row.nodes - 1) * _PAGE_SIZE,
512
529
            will be returned, and every match that is in the index will be
513
530
            returned.
514
531
        """
515
 
        # XXX: To much duplication with the GraphIndex class; consider finding
516
 
        # a good place to pull out the actual common logic.
517
532
        keys = set(keys)
518
533
        if not keys:
519
534
            return
524
539
                yield (self,) + node[1:]
525
540
        if self._key_length == 1:
526
541
            for key in keys:
527
 
                # sanity check
528
 
                if key[0] is None:
529
 
                    raise errors.BadIndexKey(key)
530
 
                if len(key) != self._key_length:
531
 
                    raise errors.BadIndexKey(key)
 
542
                index._sanity_check_key(self, key)
532
543
                try:
533
544
                    node = self._nodes[key]
534
545
                except KeyError:
538
549
                else:
539
550
                    yield self, key, node[1]
540
551
            return
541
 
        for key in keys:
542
 
            # sanity check
543
 
            if key[0] is None:
544
 
                raise errors.BadIndexKey(key)
545
 
            if len(key) != self._key_length:
546
 
                raise errors.BadIndexKey(key)
547
 
            # find what it refers to:
548
 
            key_dict = self._get_nodes_by_key()
549
 
            elements = list(key)
550
 
            # find the subdict to return
551
 
            try:
552
 
                while len(elements) and elements[0] is not None:
553
 
                    key_dict = key_dict[elements[0]]
554
 
                    elements.pop(0)
555
 
            except KeyError:
556
 
                # a non-existant lookup.
557
 
                continue
558
 
            if len(elements):
559
 
                dicts = [key_dict]
560
 
                while dicts:
561
 
                    key_dict = dicts.pop(-1)
562
 
                    # can't be empty or would not exist
563
 
                    item, value = key_dict.iteritems().next()
564
 
                    if type(value) == dict:
565
 
                        # push keys
566
 
                        dicts.extend(key_dict.itervalues())
567
 
                    else:
568
 
                        # yield keys
569
 
                        for value in key_dict.itervalues():
570
 
                            yield (self, ) + tuple(value)
571
 
            else:
572
 
                yield (self, ) + key_dict
 
552
        nodes_by_key = self._get_nodes_by_key()
 
553
        for entry in index._iter_entries_prefix(self, nodes_by_key, keys):
 
554
            yield entry
573
555
 
574
556
    def _get_nodes_by_key(self):
575
557
        if self._nodes_by_key is None:
601
583
        """In memory index's have no known corruption at the moment."""
602
584
 
603
585
 
604
 
class _LeafNode(object):
 
586
class _LeafNode(dict):
605
587
    """A leaf node for a serialised B+Tree index."""
606
588
 
607
 
    __slots__ = ('keys', 'min_key', 'max_key')
 
589
    __slots__ = ('min_key', 'max_key', '_keys')
608
590
 
609
591
    def __init__(self, bytes, key_length, ref_list_length):
610
592
        """Parse bytes to create a leaf node object."""
616
598
            self.max_key = key_list[-1][0]
617
599
        else:
618
600
            self.min_key = self.max_key = None
619
 
        self.keys = dict(key_list)
 
601
        super(_LeafNode, self).__init__(key_list)
 
602
        self._keys = dict(self)
 
603
 
 
604
    def all_items(self):
 
605
        """Return a sorted list of (key, (value, refs)) items"""
 
606
        items = sorted(self.items())
 
607
        return items
 
608
 
 
609
    def all_keys(self):
 
610
        """Return a sorted list of all keys."""
 
611
        keys = sorted(self.keys())
 
612
        return keys
620
613
 
621
614
 
622
615
class _InternalNode(object):
636
629
        for line in lines[2:]:
637
630
            if line == '':
638
631
                break
639
 
            nodes.append(as_st(map(intern, line.split('\0'))).intern())
 
632
            # GZ 2017-05-24: Used to intern() each chunk of line as well, need
 
633
            # to recheck performance and perhaps adapt StaticTuple to adjust.
 
634
            nodes.append(as_st(line.split(b'\0')).intern())
640
635
        return nodes
641
636
 
642
637
 
671
666
        self._recommended_pages = self._compute_recommended_pages()
672
667
        self._root_node = None
673
668
        self._base_offset = offset
 
669
        self._leaf_factory = _LeafNode
674
670
        # Default max size is 100,000 leave values
675
671
        self._leaf_value_cache = None # lru_cache.LRUCache(100*1000)
676
672
        if unlimited_cache:
689
685
    def __eq__(self, other):
690
686
        """Equal when self and other were created with the same parameters."""
691
687
        return (
692
 
            type(self) == type(other) and
 
688
            isinstance(self, type(other)) and
693
689
            self._transport == other._transport and
694
690
            self._name == other._name and
695
691
            self._size == other._size)
787
783
        if total_pages - len(cached_offsets) <= self._recommended_pages:
788
784
            # Read whatever is left
789
785
            if cached_offsets:
790
 
                expanded = [x for x in xrange(total_pages)
 
786
                expanded = [x for x in range(total_pages)
791
787
                               if x not in cached_offsets]
792
788
            else:
793
 
                expanded = range(total_pages)
 
789
                expanded = list(range(total_pages))
794
790
            if 'index' in debug.debug_flags:
795
791
                trace.mutter('  reading all unread pages: %s', expanded)
796
792
            return expanded
949
945
        """Cache directly from key => value, skipping the btree."""
950
946
        if self._leaf_value_cache is not None:
951
947
            for node in nodes.itervalues():
952
 
                for key, value in node.keys.iteritems():
 
948
                for key, value in node.all_items():
953
949
                    if key in self._leaf_value_cache:
954
950
                        # Don't add the rest of the keys, we've seen this node
955
951
                        # before.
979
975
        if self._row_offsets[-1] == 1:
980
976
            # There is only the root node, and we read that via key_count()
981
977
            if self.node_ref_lists:
982
 
                for key, (value, refs) in sorted(self._root_node.keys.items()):
 
978
                for key, (value, refs) in self._root_node.all_items():
983
979
                    yield (self, key, value, refs)
984
980
            else:
985
 
                for key, (value, refs) in sorted(self._root_node.keys.items()):
 
981
                for key, (value, refs) in self._root_node.all_items():
986
982
                    yield (self, key, value)
987
983
            return
988
984
        start_of_leaves = self._row_offsets[-2]
989
985
        end_of_leaves = self._row_offsets[-1]
990
 
        needed_offsets = range(start_of_leaves, end_of_leaves)
 
986
        needed_offsets = list(range(start_of_leaves, end_of_leaves))
991
987
        if needed_offsets == [0]:
992
988
            # Special case when we only have a root node, as we have already
993
989
            # read everything
998
994
        # for spilling index builds to disk.
999
995
        if self.node_ref_lists:
1000
996
            for _, node in nodes:
1001
 
                for key, (value, refs) in sorted(node.keys.items()):
 
997
                for key, (value, refs) in node.all_items():
1002
998
                    yield (self, key, value, refs)
1003
999
        else:
1004
1000
            for _, node in nodes:
1005
 
                for key, (value, refs) in sorted(node.keys.items()):
 
1001
                for key, (value, refs) in node.all_items():
1006
1002
                    yield (self, key, value)
1007
1003
 
1008
1004
    @staticmethod
1032
1028
        # iter_steps = len(in_keys) + len(fixed_keys)
1033
1029
        # bisect_steps = len(in_keys) * math.log(len(fixed_keys), 2)
1034
1030
        if len(in_keys) == 1: # Bisect will always be faster for M = 1
1035
 
            return [(bisect_right(fixed_keys, in_keys[0]), in_keys)]
 
1031
            return [(bisect.bisect_right(fixed_keys, in_keys[0]), in_keys)]
1036
1032
        # elif bisect_steps < iter_steps:
1037
1033
        #     offsets = {}
1038
1034
        #     for key in in_keys:
1041
1037
        #     return [(o, offsets[o]) for o in sorted(offsets)]
1042
1038
        in_keys_iter = iter(in_keys)
1043
1039
        fixed_keys_iter = enumerate(fixed_keys)
1044
 
        cur_in_key = in_keys_iter.next()
1045
 
        cur_fixed_offset, cur_fixed_key = fixed_keys_iter.next()
 
1040
        cur_in_key = next(in_keys_iter)
 
1041
        cur_fixed_offset, cur_fixed_key = next(fixed_keys_iter)
1046
1042
 
1047
1043
        class InputDone(Exception): pass
1048
1044
        class FixedDone(Exception): pass
1064
1060
                    while cur_in_key < cur_fixed_key:
1065
1061
                        cur_keys.append(cur_in_key)
1066
1062
                        try:
1067
 
                            cur_in_key = in_keys_iter.next()
 
1063
                            cur_in_key = next(in_keys_iter)
1068
1064
                        except StopIteration:
1069
1065
                            raise InputDone
1070
1066
                    # At this point cur_in_key must be >= cur_fixed_key
1072
1068
                # the end
1073
1069
                while cur_in_key >= cur_fixed_key:
1074
1070
                    try:
1075
 
                        cur_fixed_offset, cur_fixed_key = fixed_keys_iter.next()
 
1071
                        cur_fixed_offset, cur_fixed_key = next(fixed_keys_iter)
1076
1072
                    except StopIteration:
1077
1073
                        raise FixedDone
1078
1074
        except InputDone:
1169
1165
                continue
1170
1166
            node = nodes[node_index]
1171
1167
            for next_sub_key in sub_keys:
1172
 
                if next_sub_key in node.keys:
1173
 
                    value, refs = node.keys[next_sub_key]
 
1168
                if next_sub_key in node:
 
1169
                    value, refs = node[next_sub_key]
1174
1170
                    if self.node_ref_lists:
1175
1171
                        yield (self, next_sub_key, value, refs)
1176
1172
                    else:
1244
1240
            # sub_keys is all of the keys we are looking for that should exist
1245
1241
            # on this page, if they aren't here, then they won't be found
1246
1242
            node = nodes[node_index]
1247
 
            node_keys = node.keys
1248
1243
            parents_to_check = set()
1249
1244
            for next_sub_key in sub_keys:
1250
 
                if next_sub_key not in node_keys:
 
1245
                if next_sub_key not in node:
1251
1246
                    # This one is just not present in the index at all
1252
1247
                    missing_keys.add(next_sub_key)
1253
1248
                else:
1254
 
                    value, refs = node_keys[next_sub_key]
 
1249
                    value, refs = node[next_sub_key]
1255
1250
                    parent_keys = refs[ref_list_num]
1256
1251
                    parent_map[next_sub_key] = parent_keys
1257
1252
                    parents_to_check.update(parent_keys)
1264
1259
            while parents_to_check:
1265
1260
                next_parents_to_check = set()
1266
1261
                for key in parents_to_check:
1267
 
                    if key in node_keys:
1268
 
                        value, refs = node_keys[key]
 
1262
                    if key in node:
 
1263
                        value, refs = node[key]
1269
1264
                        parent_keys = refs[ref_list_num]
1270
1265
                        parent_map[key] = parent_keys
1271
1266
                        next_parents_to_check.update(parent_keys)
1333
1328
            self._get_root_node()
1334
1329
        # TODO: only access nodes that can satisfy the prefixes we are looking
1335
1330
        # for. For now, to meet API usage (as this function is not used by
1336
 
        # current bzrlib) just suck the entire index and iterate in memory.
 
1331
        # current breezy) just suck the entire index and iterate in memory.
1337
1332
        nodes = {}
1338
1333
        if self.node_ref_lists:
1339
1334
            if self._key_length == 1:
1365
1360
                    key_dict[key[-1]] = key_value
1366
1361
        if self._key_length == 1:
1367
1362
            for key in keys:
1368
 
                # sanity check
1369
 
                if key[0] is None:
1370
 
                    raise errors.BadIndexKey(key)
1371
 
                if len(key) != self._key_length:
1372
 
                    raise errors.BadIndexKey(key)
 
1363
                index._sanity_check_key(self, key)
1373
1364
                try:
1374
1365
                    if self.node_ref_lists:
1375
1366
                        value, node_refs = nodes[key]
1379
1370
                except KeyError:
1380
1371
                    pass
1381
1372
            return
1382
 
        for key in keys:
1383
 
            # sanity check
1384
 
            if key[0] is None:
1385
 
                raise errors.BadIndexKey(key)
1386
 
            if len(key) != self._key_length:
1387
 
                raise errors.BadIndexKey(key)
1388
 
            # find what it refers to:
1389
 
            key_dict = nodes_by_key
1390
 
            elements = list(key)
1391
 
            # find the subdict whose contents should be returned.
1392
 
            try:
1393
 
                while len(elements) and elements[0] is not None:
1394
 
                    key_dict = key_dict[elements[0]]
1395
 
                    elements.pop(0)
1396
 
            except KeyError:
1397
 
                # a non-existant lookup.
1398
 
                continue
1399
 
            if len(elements):
1400
 
                dicts = [key_dict]
1401
 
                while dicts:
1402
 
                    key_dict = dicts.pop(-1)
1403
 
                    # can't be empty or would not exist
1404
 
                    item, value = key_dict.iteritems().next()
1405
 
                    if type(value) == dict:
1406
 
                        # push keys
1407
 
                        dicts.extend(key_dict.itervalues())
1408
 
                    else:
1409
 
                        # yield keys
1410
 
                        for value in key_dict.itervalues():
1411
 
                            # each value is the key:value:node refs tuple
1412
 
                            # ready to yield.
1413
 
                            yield (self, ) + value
1414
 
            else:
1415
 
                # the last thing looked up was a terminal element
1416
 
                yield (self, ) + key_dict
 
1373
        for entry in index._iter_entries_prefix(self, nodes_by_key, keys):
 
1374
            yield entry
1417
1375
 
1418
1376
    def key_count(self):
1419
1377
        """Return an estimate of the number of keys in this index.
1471
1429
        if not options_line.startswith(_OPTION_ROW_LENGTHS):
1472
1430
            raise errors.BadIndexOptions(self)
1473
1431
        try:
1474
 
            self._row_lengths = map(int, [length for length in
 
1432
            self._row_lengths = [int(length) for length in
1475
1433
                options_line[len(_OPTION_ROW_LENGTHS):].split(',')
1476
 
                if len(length)])
 
1434
                if length]
1477
1435
        except ValueError:
1478
1436
            raise errors.BadIndexOptions(self)
1479
1437
        self._compute_row_offsets()
1514
1472
                    self._size = num_bytes - base_offset
1515
1473
                    # the whole thing should be parsed out of 'bytes'
1516
1474
                    ranges = [(start, min(_PAGE_SIZE, num_bytes - start))
1517
 
                        for start in xrange(base_offset, num_bytes, _PAGE_SIZE)]
 
1475
                        for start in range(base_offset, num_bytes, _PAGE_SIZE)]
1518
1476
                    break
1519
1477
            else:
1520
1478
                if offset > self._size:
1545
1503
                    continue
1546
1504
            bytes = zlib.decompress(data)
1547
1505
            if bytes.startswith(_LEAF_FLAG):
1548
 
                node = _LeafNode(bytes, self._key_length, self.node_ref_lists)
 
1506
                node = self._leaf_factory(bytes, self._key_length,
 
1507
                                          self.node_ref_lists)
1549
1508
            elif bytes.startswith(_INTERNAL_FLAG):
1550
1509
                node = _InternalNode(bytes)
1551
1510
            else:
1566
1525
            # We shouldn't be reading anything anyway
1567
1526
            start_node = 1
1568
1527
        node_end = self._row_offsets[-1]
1569
 
        for node in self._read_nodes(range(start_node, node_end)):
 
1528
        for node in self._read_nodes(list(range(start_node, node_end))):
1570
1529
            pass
1571
1530
 
1572
1531
 
 
1532
_gcchk_factory = _LeafNode
 
1533
 
1573
1534
try:
1574
 
    from bzrlib import _btree_serializer_pyx as _btree_serializer
1575
 
except ImportError, e:
 
1535
    from breezy import _btree_serializer_pyx as _btree_serializer
 
1536
    _gcchk_factory = _btree_serializer._parse_into_chk
 
1537
except ImportError as e:
1576
1538
    osutils.failed_to_load_extension(e)
1577
 
    from bzrlib import _btree_serializer_py as _btree_serializer
 
1539
    from breezy import _btree_serializer_py as _btree_serializer