github.com/grailbio/base@v0.0.11/mapio/merged.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package mapio
     6  
     7  import (
     8  	"bytes"
     9  	"container/heap"
    10  )
    11  
    12  var scanSentinel = new(MapScanner)
    13  
    14  // Merged represents the merged contents of multiple underlying maps.
    15  // Like Map, Merged presents a sorted, scannable map, but it does not
    16  // guarantee that the order of traversal is stable.
    17  type Merged []*Map
    18  
    19  // Seek returns a scanner for the merged map that starts at the first
    20  // entry where entryKey <= key.
    21  func (m Merged) Seek(key []byte) *MergedScanner {
    22  	merged := make(MergedScanner, 0, len(m)+1)
    23  	for i := range m {
    24  		s := m[i].Seek(key)
    25  		if !s.Scan() {
    26  			if err := s.Err(); err != nil {
    27  				return &MergedScanner{s}
    28  			}
    29  			// Otherwise it's just empty and we can skip it.
    30  			continue
    31  		}
    32  		merged = append(merged, s)
    33  	}
    34  	if len(merged) == 0 {
    35  		return &MergedScanner{}
    36  	}
    37  	heap.Init(&merged)
    38  	merged = append(merged, scanSentinel)
    39  	return &merged
    40  }
    41  
    42  // MergedScanner is a scanner for merged maps.
    43  type MergedScanner []*MapScanner
    44  
    45  // Len implements heap.Interface
    46  func (m MergedScanner) Len() int { return len(m) }
    47  
    48  // Less implements heap.Interface
    49  func (m MergedScanner) Less(i, j int) bool { return bytes.Compare(m[i].Key(), m[j].Key()) < 0 }
    50  
    51  // Swap implements heap.Interface
    52  func (m MergedScanner) Swap(i, j int) { m[i], m[j] = m[j], m[i] }
    53  
    54  // Push implements heap.Interface
    55  func (m *MergedScanner) Push(x interface{}) {
    56  	*m = append(*m, x.(*MapScanner))
    57  }
    58  
    59  // Pop implements heap.Interface
    60  func (m *MergedScanner) Pop() interface{} {
    61  	n := len(*m)
    62  	elem := (*m)[n-1]
    63  	*m = (*m)[:n-1]
    64  	return elem
    65  }
    66  
    67  // Scan scans the next entry in the merged map, returning true on
    68  // success. If Scan returns false, the caller should check Err to
    69  // distinguish between scan completion and scan error.
    70  func (m *MergedScanner) Scan() bool {
    71  	if len(*m) == 0 || (*m)[0].err != nil {
    72  		return false
    73  	}
    74  	if len(*m) > 0 && (*m)[len(*m)-1] == scanSentinel {
    75  		*m = (*m)[:len(*m)-1]
    76  		return true
    77  	}
    78  
    79  	if (*m)[0].Scan() {
    80  		heap.Fix(m, 0)
    81  	} else if (*m)[0].err == nil {
    82  		heap.Remove(m, 0)
    83  	}
    84  	ok := len(*m) > 0 && (*m)[0].err == nil
    85  	return ok
    86  
    87  }
    88  
    89  // Err returns the last error encountered while scanning, if any.
    90  func (m MergedScanner) Err() error {
    91  	if len(m) == 0 {
    92  		return nil
    93  	}
    94  	return m[0].err
    95  }
    96  
    97  // Key returns the last key scanned.
    98  func (m MergedScanner) Key() []byte {
    99  	return m[0].Key()
   100  }
   101  
   102  // Value returns the last value scanned.
   103  func (m MergedScanner) Value() []byte {
   104  	return m[0].Value()
   105  }