github.com/weaviate/weaviate@v1.24.6/adapters/repos/db/vector/hnsw/pools.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 hnsw
    13  
    14  import (
    15  	"sync"
    16  
    17  	"github.com/weaviate/weaviate/adapters/repos/db/priorityqueue"
    18  	"github.com/weaviate/weaviate/adapters/repos/db/vector/cache"
    19  	"github.com/weaviate/weaviate/adapters/repos/db/vector/common"
    20  	"github.com/weaviate/weaviate/adapters/repos/db/vector/hnsw/visited"
    21  )
    22  
    23  type pools struct {
    24  	visitedLists     *visited.Pool
    25  	visitedListsLock *sync.RWMutex
    26  
    27  	pqItemSlice  *sync.Pool
    28  	pqHeuristic  *pqMinWithIndexPool
    29  	pqResults    *common.PqMaxPool
    30  	pqCandidates *pqMinPool
    31  
    32  	tempVectors *common.TempVectorsPool
    33  }
    34  
    35  func newPools(maxConnectionsLayerZero int) *pools {
    36  	return &pools{
    37  		visitedLists:     visited.NewPool(1, cache.InitialSize+500),
    38  		visitedListsLock: &sync.RWMutex{},
    39  		pqItemSlice: &sync.Pool{
    40  			New: func() interface{} {
    41  				return make([]priorityqueue.Item[uint64], 0, maxConnectionsLayerZero)
    42  			},
    43  		},
    44  		pqHeuristic:  newPqMinWithIndexPool(maxConnectionsLayerZero),
    45  		pqResults:    common.NewPqMaxPool(maxConnectionsLayerZero),
    46  		pqCandidates: newPqMinPool(maxConnectionsLayerZero),
    47  		tempVectors:  common.NewTempVectorsPool(),
    48  	}
    49  }
    50  
    51  type pqMinPool struct {
    52  	pool *sync.Pool
    53  }
    54  
    55  func newPqMinPool(defaultCap int) *pqMinPool {
    56  	return &pqMinPool{
    57  		pool: &sync.Pool{
    58  			New: func() interface{} {
    59  				return priorityqueue.NewMin[any](defaultCap)
    60  			},
    61  		},
    62  	}
    63  }
    64  
    65  func (pqh *pqMinPool) GetMin(capacity int) *priorityqueue.Queue[any] {
    66  	pq := pqh.pool.Get().(*priorityqueue.Queue[any])
    67  	if pq.Cap() < capacity {
    68  		pq.ResetCap(capacity)
    69  	} else {
    70  		pq.Reset()
    71  	}
    72  
    73  	return pq
    74  }
    75  
    76  func (pqh *pqMinPool) Put(pq *priorityqueue.Queue[any]) {
    77  	pqh.pool.Put(pq)
    78  }
    79  
    80  type pqMinWithIndexPool struct {
    81  	pool *sync.Pool
    82  }
    83  
    84  func newPqMinWithIndexPool(defaultCap int) *pqMinWithIndexPool {
    85  	return &pqMinWithIndexPool{
    86  		pool: &sync.Pool{
    87  			New: func() interface{} {
    88  				return priorityqueue.NewMin[uint64](defaultCap)
    89  			},
    90  		},
    91  	}
    92  }
    93  
    94  func (pqh *pqMinWithIndexPool) GetMin(capacity int) *priorityqueue.Queue[uint64] {
    95  	pq := pqh.pool.Get().(*priorityqueue.Queue[uint64])
    96  	if pq.Cap() < capacity {
    97  		pq.ResetCap(capacity)
    98  	} else {
    99  		pq.Reset()
   100  	}
   101  
   102  	return pq
   103  }
   104  
   105  func (pqh *pqMinWithIndexPool) Put(pq *priorityqueue.Queue[uint64]) {
   106  	pqh.pool.Put(pq)
   107  }