forked from DataIntelligenceCrew/go-faiss
-
Notifications
You must be signed in to change notification settings - Fork 3
/
index_flat.go
56 lines (49 loc) · 1.39 KB
/
index_flat.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
package faiss
/*
#include <faiss/c_api/IndexFlat_c.h>
#include <faiss/c_api/Index_c.h>
*/
import "C"
import "unsafe"
// IndexFlat is an index that stores the full vectors and performs exhaustive
// search.
type IndexFlat struct {
Index
}
// NewIndexFlat creates a new flat index.
func NewIndexFlat(d int, metric int) (*IndexFlat, error) {
var idx faissIndex
if c := C.faiss_IndexFlat_new_with(
&idx.idx,
C.idx_t(d),
C.FaissMetricType(metric),
); c != 0 {
return nil, getLastError()
}
return &IndexFlat{&idx}, nil
}
// NewIndexFlatIP creates a new flat index with the inner product metric type.
func NewIndexFlatIP(d int) (*IndexFlat, error) {
return NewIndexFlat(d, MetricInnerProduct)
}
// NewIndexFlatL2 creates a new flat index with the L2 metric type.
func NewIndexFlatL2(d int) (*IndexFlat, error) {
return NewIndexFlat(d, MetricL2)
}
// Xb returns the index's vectors.
// The returned slice becomes invalid after any add or remove operation.
func (idx *IndexFlat) Xb() []float32 {
var size C.size_t
var ptr *C.float
C.faiss_IndexFlat_xb(idx.cPtr(), &ptr, &size)
return (*[1 << 30]float32)(unsafe.Pointer(ptr))[:size:size]
}
// AsFlat casts idx to a flat index.
// AsFlat panics if idx is not a flat index.
func (idx *IndexImpl) AsFlat() *IndexFlat {
ptr := C.faiss_IndexFlat_cast(idx.cPtr())
if ptr == nil {
panic("index is not a flat index")
}
return &IndexFlat{&faissIndex{ptr}}
}