github.com/weaviate/weaviate@v1.24.6/adapters/repos/db/priorityqueue/queue.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package priorityqueue
    13  
    14  type supportedValueType interface {
    15  	any | uint64
    16  }
    17  
    18  // Item represents a queue item supporting an optional additional Value
    19  type Item[T supportedValueType] struct {
    20  	ID       uint64
    21  	Dist     float32
    22  	Rescored bool
    23  	Value    T
    24  }
    25  
    26  // Queue is a priority queue supporting generic item values
    27  type Queue[T supportedValueType] struct {
    28  	items []Item[T]
    29  	less  func(items []Item[T], i, j int) bool
    30  }
    31  
    32  // NewMin constructs a priority queue which prioritizes items with smaller distance
    33  func NewMin[T supportedValueType](capacity int) *Queue[T] {
    34  	return &Queue[T]{
    35  		items: make([]Item[T], 0, capacity),
    36  		less: func(items []Item[T], i, j int) bool {
    37  			return items[i].Dist < items[j].Dist
    38  		},
    39  	}
    40  }
    41  
    42  // NewMax constructs a priority queue which prioritizes items with greater distance
    43  func NewMax[T supportedValueType](capacity int) *Queue[T] {
    44  	return &Queue[T]{
    45  		items: make([]Item[T], 0, capacity),
    46  		less: func(items []Item[T], i, j int) bool {
    47  			return items[i].Dist > items[j].Dist
    48  		},
    49  	}
    50  }
    51  
    52  // Pop removes the next item in the queue and returns it
    53  func (q *Queue[T]) Pop() Item[T] {
    54  	out := q.items[0]
    55  	q.items[0] = q.items[len(q.items)-1]
    56  	q.items = q.items[:len(q.items)-1]
    57  	q.heapify(0)
    58  	return out
    59  }
    60  
    61  // Top peeks at the next item in the queue
    62  func (q *Queue[T]) Top() Item[T] {
    63  	return q.items[0]
    64  }
    65  
    66  // Len returns the length of the queue
    67  func (q *Queue[T]) Len() int {
    68  	return len(q.items)
    69  }
    70  
    71  // Cap returns the remaining capacity of the queue
    72  func (q *Queue[T]) Cap() int {
    73  	return cap(q.items)
    74  }
    75  
    76  // Reset clears all items from the queue
    77  func (q *Queue[T]) Reset() {
    78  	q.items = q.items[:0]
    79  }
    80  
    81  // ResetCap drops existing queue items, and allocates a new queue with the given capacity
    82  func (q *Queue[T]) ResetCap(capacity int) {
    83  	q.items = make([]Item[T], 0, capacity)
    84  }
    85  
    86  // Insert creates a valueless item and adds it to the queue
    87  func (q *Queue[T]) Insert(id uint64, distance float32) int {
    88  	item := Item[T]{
    89  		ID:   id,
    90  		Dist: distance,
    91  	}
    92  	return q.insert(item)
    93  }
    94  
    95  // InsertWithValue creates an item with a T type value and adds it to the queue
    96  func (q *Queue[T]) InsertWithValue(id uint64, distance float32, val T) int {
    97  	item := Item[T]{
    98  		ID:    id,
    99  		Dist:  distance,
   100  		Value: val,
   101  	}
   102  	return q.insert(item)
   103  }
   104  
   105  func (q *Queue[T]) insert(item Item[T]) int {
   106  	q.items = append(q.items, item)
   107  	i := len(q.items) - 1
   108  	for i != 0 && q.less(q.items, i, q.parent(i)) {
   109  		q.swap(i, q.parent(i))
   110  		i = q.parent(i)
   111  	}
   112  	return i
   113  }
   114  
   115  func (q *Queue[T]) left(i int) int {
   116  	return 2*i + 1
   117  }
   118  
   119  func (q *Queue[T]) right(i int) int {
   120  	return 2*i + 2
   121  }
   122  
   123  func (q *Queue[T]) parent(i int) int {
   124  	return (i - 1) / 2
   125  }
   126  
   127  func (q *Queue[T]) swap(i, j int) {
   128  	q.items[i], q.items[j] = q.items[j], q.items[i]
   129  }
   130  
   131  func (q *Queue[T]) heapify(i int) {
   132  	left := q.left(i)
   133  	right := q.right(i)
   134  	smallest := i
   135  	if left < len(q.items) && q.less(q.items, left, i) {
   136  		smallest = left
   137  	}
   138  
   139  	if right < len(q.items) && q.less(q.items, right, smallest) {
   140  		smallest = right
   141  	}
   142  
   143  	if smallest != i {
   144  		q.swap(i, smallest)
   145  		q.heapify(smallest)
   146  	}
   147  }