github.com/weaviate/weaviate@v1.24.6/adapters/repos/db/vector/compressionhelpers/quantizer.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 compressionhelpers
    13  
    14  import "encoding/binary"
    15  
    16  type quantizerDistancer[T byte | uint64] interface {
    17  	Distance(x []T) (float32, bool, error)
    18  	DistanceToFloat(x []float32) (float32, bool, error)
    19  }
    20  
    21  type quantizer[T byte | uint64] interface {
    22  	DistanceBetweenCompressedVectors(x, y []T) (float32, error)
    23  	DistanceBetweenCompressedAndUncompressedVectors(x []float32, encoded []T) (float32, error)
    24  	Encode(vec []float32) []T
    25  	NewQuantizerDistancer(a []float32) quantizerDistancer[T]
    26  	NewCompressedQuantizerDistancer(a []T) quantizerDistancer[T]
    27  	ReturnQuantizerDistancer(distancer quantizerDistancer[T])
    28  	CompressedBytes(compressed []T) []byte
    29  	FromCompressedBytes(compressed []byte) []T
    30  	ExposeFields() PQData
    31  }
    32  
    33  func (bq *BinaryQuantizer) ExposeFields() PQData {
    34  	return PQData{}
    35  }
    36  
    37  func (bq *BinaryQuantizer) DistanceBetweenCompressedAndUncompressedVectors(x []float32, y []uint64) (float32, error) {
    38  	encoded := bq.Encode(x)
    39  	return bq.DistanceBetweenCompressedVectors(encoded, y)
    40  }
    41  
    42  func (pq *ProductQuantizer) NewQuantizerDistancer(vec []float32) quantizerDistancer[byte] {
    43  	return pq.NewDistancer(vec)
    44  }
    45  
    46  func (pq *ProductQuantizer) ReturnQuantizerDistancer(distancer quantizerDistancer[byte]) {
    47  	concreteDistancer := distancer.(*PQDistancer)
    48  	if concreteDistancer == nil {
    49  		return
    50  	}
    51  	pq.ReturnDistancer(concreteDistancer)
    52  }
    53  
    54  func (bq *BinaryQuantizer) CompressedBytes(compressed []uint64) []byte {
    55  	slice := make([]byte, len(compressed)*8)
    56  	for i := range compressed {
    57  		binary.LittleEndian.PutUint64(slice[i*8:], compressed[i])
    58  	}
    59  	return slice
    60  }
    61  
    62  func (bq *BinaryQuantizer) FromCompressedBytes(compressed []byte) []uint64 {
    63  	l := len(compressed) / 8
    64  	if len(compressed)%8 != 0 {
    65  		l++
    66  	}
    67  	slice := make([]uint64, l)
    68  
    69  	for i := range slice {
    70  		slice[i] = binary.LittleEndian.Uint64(compressed[i*8:])
    71  	}
    72  	return slice
    73  }
    74  
    75  func (pq *ProductQuantizer) CompressedBytes(compressed []byte) []byte {
    76  	return compressed
    77  }
    78  
    79  func (pq *ProductQuantizer) FromCompressedBytes(compressed []byte) []byte {
    80  	return compressed
    81  }
    82  
    83  type BQDistancer struct {
    84  	x          []float32
    85  	bq         *BinaryQuantizer
    86  	compressed []uint64
    87  }
    88  
    89  func (bq *BinaryQuantizer) NewDistancer(a []float32) *BQDistancer {
    90  	return &BQDistancer{
    91  		x:          a,
    92  		bq:         bq,
    93  		compressed: bq.Encode(a),
    94  	}
    95  }
    96  
    97  func (bq *BinaryQuantizer) NewCompressedQuantizerDistancer(a []uint64) quantizerDistancer[uint64] {
    98  	return &BQDistancer{
    99  		x:          nil,
   100  		bq:         bq,
   101  		compressed: a,
   102  	}
   103  }
   104  
   105  func (d *BQDistancer) Distance(x []uint64) (float32, bool, error) {
   106  	dist, err := d.bq.DistanceBetweenCompressedVectors(d.compressed, x)
   107  	return dist, err == nil, err
   108  }
   109  
   110  func (d *BQDistancer) DistanceToFloat(x []float32) (float32, bool, error) {
   111  	if len(d.x) > 0 {
   112  		return d.bq.distancer.SingleDist(d.x, x)
   113  	}
   114  	xComp := d.bq.Encode(x)
   115  	dist, err := d.bq.DistanceBetweenCompressedVectors(d.compressed, xComp)
   116  	return dist, err == nil, err
   117  }
   118  
   119  func (bq *BinaryQuantizer) NewQuantizerDistancer(vec []float32) quantizerDistancer[uint64] {
   120  	return bq.NewDistancer(vec)
   121  }
   122  
   123  func (bq *BinaryQuantizer) ReturnQuantizerDistancer(distancer quantizerDistancer[uint64]) {}