Files
zeekay f5aaeb712e Hanzo Graph — the engine, on zapdb, ZAP-native
Derived from Dgraph (github.com/dgraph-io/dgraph) at commit 64804bd2, Apache-2.0.
See NOTICE.

WHAT THIS IS. The graph engine as a Go library: posting lists, the DQL parser,
the type and tokenizer system, the index and schema layers, and the wire types.
39,170 lines of the upstream's 166,337.

THE STORE is github.com/luxfi/zapdb, a Badger fork that kept the package name,
the subpackages (options, pb, y) and the API — Open, OpenManaged, DB, Options —
so the 74 import sites are aliased rather than rewritten.

THE CODEC is ZAP. zapdb serializes with ZAP, not protobuf: its pb package builds
the ZAP types by default and the protobuf ones only under `-tags grpc`, a path
that does not compile (table/builder.go calls pb.Marshal, which that tag
removes). One call site in posting needed converting; every other type here is
Hanzo's own protobuf and is untouched.

NO TRANSPORT. An embedded engine dials nothing, so nothing here imports
google.golang.org/grpc. Removed: 20 transport functions in x, the generated
service stubs for Raft, Zero and Worker (1,657 lines), and the JWT helper —
identity belongs to IAM. The core called none of them.

NOT KEPT: worker, conn, raftwal, the cluster commands, the GraphQL layer and the
test harnesses. `query`, the DQL execution engine, went with them because it
depends on worker; parsing DQL is available here, executing it is not.

The history of the upstream project lives at github.com/dgraph-io/dgraph. This
repository begins at the derivation because three quarters of that history
describes code this fork does not carry.
2026-08-20 21:26:27 -07:00

878 lines
25 KiB
Go

/*
* SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc.
* SPDX-License-Identifier: Apache-2.0
*/
package hnsw
import (
"context"
"encoding/binary"
"fmt"
"log"
"math"
"math/rand"
"sort"
"strconv"
"strings"
"unsafe"
c "github.com/hanzoai/graph/tok/constraints"
"github.com/hanzoai/graph/tok/index"
"github.com/pkg/errors"
"github.com/viterin/vek"
"github.com/viterin/vek/vek32"
)
const (
Euclidean = "euclidean"
Cosine = "cosine"
DotProd = "dotproduct"
EmptyHNSWTreeError = "HNSW tree has no elements"
VecKeyword = "__vector_"
visitedVectorsLevel = "visited_vectors_level_"
distanceComputations = "vector_distance_computations"
searchTime = "vector_search_time"
VecEntry = "__vector_entry"
VecDead = "__vector_dead"
VectorIndexMaxLevels = 5
EfConstruction = 16
EfSearch = 12
// ByteData indicates the key stores data.
ByteData = byte(0x00)
// DefaultPrefix is the prefix used for data, index and reverse keys so that relative
DefaultPrefix = byte(0x00)
// NsSeparator is the separator between the namespace and attribute.
NsSeparator = "-"
)
var (
errNilVector = errors.New("nil vector returned")
errFetchingPostingList = errors.New("error fetching posting list")
)
type SearchResult struct {
nnUids []uint64
traversalPath []uint64
extraMetrics map[string]uint64
}
func (s *SearchResult) GetNnUids() []uint64 {
return s.nnUids
}
func (s *SearchResult) GetTraversalPath() []uint64 {
return s.traversalPath
}
func (s *SearchResult) GetExtraMetrics() map[string]uint64 {
return s.extraMetrics
}
func applyDistanceFunction[T c.Float](a, b []T, floatBits int, funcName string,
applyFn32 func(a, b []float32) float32, applyFn64 func(a, b []float64) float64) (T, error) {
if len(a) != len(b) {
err := errors.New(fmt.Sprintf("can not compute %s on vectors of different lengths", funcName))
return T(0), err
}
switch floatBits {
case 32:
var a1, b1 []float32
a1 = *(*[]float32)(unsafe.Pointer(&a))
b1 = *(*[]float32)(unsafe.Pointer(&b))
return T(applyFn32(a1, b1)), nil
case 64:
var a1, b1 []float64
a1 = *(*[]float64)(unsafe.Pointer(&a))
b1 = *(*[]float64)(unsafe.Pointer(&b))
return T(applyFn64(a1, b1)), nil
default:
panic("While applying function on two floats, found an invalid number of float bits")
}
}
// This needs to implement signature of SimilarityType[T].distanceScore
// function, hence it takes in a floatBits parameter,
// but doesn't actually use it.
func dotProduct[T c.Float](a, b []T, floatBits int) (T, error) {
return applyDistanceFunction(a, b, floatBits, "dot product", vek32.Dot, vek.Dot)
}
// This needs to implement signature of SimilarityType[T].distanceScore
// function, hence it takes in a floatBits parameter.
func cosineSimilarity[T c.Float](a, b []T, floatBits int) (T, error) {
return applyDistanceFunction(a, b, floatBits, "cosine distance", vek32.CosineSimilarity, vek.CosineSimilarity)
}
// This needs to implement signature of SimilarityType[T].distanceScore
// function, hence it takes in a floatBits parameter,
// but doesn't actually use it.
func euclideanDistanceSq[T c.Float](a, b []T, floatBits int) (T, error) {
return applyDistanceFunction(a, b, floatBits, "euclidean distance", vek32.Distance, vek.Distance)
}
// Used for distance, since shorter distance is better
func insortPersistentHeapAscending[T c.Float](
slice []persistentHeapElement[T],
val persistentHeapElement[T]) []persistentHeapElement[T] {
i := sort.Search(len(slice), func(i int) bool { return slice[i].value > val.value })
var empty T
slice = append(slice, *initPersistentHeapElement(empty, notAUid, false))
copy(slice[i+1:], slice[i:])
slice[i] = val
return slice
}
// Used for cosine similarity, since higher similarity score is better
func insortPersistentHeapDescending[T c.Float](
slice []persistentHeapElement[T],
val persistentHeapElement[T]) []persistentHeapElement[T] {
i := sort.Search(len(slice), func(i int) bool { return slice[i].value < val.value })
var empty T
slice = append(slice, *initPersistentHeapElement(empty, notAUid, false))
copy(slice[i+1:], slice[i:])
slice[i] = val
return slice
}
func isBetterScoreForDistance[T c.Float](a, b T) bool {
return a < b
}
func isBetterScoreForSimilarity[T c.Float](a, b T) bool {
return a > b
}
func ParseEdges(s string) ([]uint64, error) {
s = strings.ReplaceAll(s, "\n", " ")
s = strings.ReplaceAll(s, "\t", " ")
s = strings.TrimSpace(s)
if len(s) == 0 {
return []uint64{}, nil
}
trimmedPre := strings.TrimPrefix(s, "[")
if len(trimmedPre) == len(s) {
return nil, cannotConvertToUintSlice(s)
}
trimmed := strings.TrimRight(trimmedPre, "]")
if len(trimmed) == len(trimmedPre) {
return nil, cannotConvertToUintSlice(s)
}
if len(trimmed) == 0 {
return []uint64{}, nil
}
if strings.Contains(trimmed, ",") {
// Splitting based on comma-separation.
values := strings.Split(trimmed, ",")
result := make([]uint64, len(values))
for i := range values {
trimmedVal := strings.TrimSpace(values[i])
val, err := strconv.ParseUint(trimmedVal, 10, 64)
if err != nil {
return nil, cannotConvertToUintSlice(s)
}
result[i] = val
}
return result, nil
}
values := strings.Split(trimmed, " ")
result := make([]uint64, 0, len(values))
for i := range values {
if len(values[i]) == 0 {
// skip if we have an empty string. This can naturally
// occur if input s was "[1.0 2.0]"
// notice the extra whitespace in separation!
continue
}
if len(values[i]) > 0 {
val, err := strconv.ParseUint(values[i], 10, 64)
if err != nil {
return nil, cannotConvertToUintSlice(s)
}
result = append(result, val)
}
}
return result, nil
}
func cannotConvertToUintSlice(s string) error {
return errors.Errorf("Cannot convert %s to uint slice", s)
}
// TODO: Move SimilarityType to index package.
//
// Remove "hnsw-isms".
type SimilarityType[T c.Float] struct {
indexType string
distanceScore func(v, w []T, floatBits int) (T, error)
insortHeap func(slice []persistentHeapElement[T], val persistentHeapElement[T]) []persistentHeapElement[T]
isBetterScore func(a, b T) bool
// isSimilarityMetric is true for metrics where higher values indicate better matches
// (e.g., cosine similarity, dot product). For distance metrics like euclidean,
// this is false because lower values indicate better matches.
isSimilarityMetric bool
}
func GetSimType[T c.Float](indexType string, floatBits int) SimilarityType[T] {
switch {
case indexType == Euclidean:
return SimilarityType[T]{indexType: Euclidean, distanceScore: euclideanDistanceSq[T],
insortHeap: insortPersistentHeapAscending[T], isBetterScore: isBetterScoreForDistance[T],
isSimilarityMetric: false}
case indexType == Cosine:
return SimilarityType[T]{indexType: Cosine, distanceScore: cosineSimilarity[T],
insortHeap: insortPersistentHeapDescending[T], isBetterScore: isBetterScoreForSimilarity[T],
isSimilarityMetric: true}
case indexType == DotProd:
return SimilarityType[T]{indexType: DotProd, distanceScore: dotProduct[T],
insortHeap: insortPersistentHeapDescending[T], isBetterScore: isBetterScoreForSimilarity[T],
isSimilarityMetric: true}
default:
return SimilarityType[T]{indexType: Euclidean, distanceScore: euclideanDistanceSq[T],
insortHeap: insortPersistentHeapAscending[T], isBetterScore: isBetterScoreForDistance[T],
isSimilarityMetric: false}
}
}
// TxnCache implements CacheType interface
type TxnCache struct {
txn index.Txn
startTs uint64
}
func (tc *TxnCache) Get(key []byte) (rval []byte, rerr error) {
return tc.txn.Get(key)
}
func (tc *TxnCache) Ts() uint64 {
return tc.startTs
}
func (tc *TxnCache) Find(prefix []byte, filter func([]byte) bool) (uint64, error) {
return tc.txn.Find(prefix, filter)
}
func NewTxnCache(txn index.Txn, startTs uint64) *TxnCache {
return &TxnCache{
txn: txn,
startTs: startTs,
}
}
// QueryCache implements index.CacheType interface
type QueryCache struct {
cache index.LocalCache
readTs uint64
}
func (qc *QueryCache) Find(prefix []byte, filter func([]byte) bool) (uint64, error) {
return qc.cache.Find(prefix, filter)
}
func (qc *QueryCache) Get(key []byte) (rval []byte, rerr error) {
return qc.cache.Get(key)
}
func (qc *QueryCache) Ts() uint64 {
return qc.readTs
}
func NewQueryCache(cache index.LocalCache, readTs uint64) *QueryCache {
return &QueryCache{
cache: cache,
readTs: readTs,
}
}
// getDataFromKeyWithCacheType(keyString, uid, c) looks up data in c
// associated with keyString and uid.
func getDataFromKeyWithCacheType(keyString string, uid uint64, c index.CacheType) ([]byte, error) {
key := DataKey(keyString, uid)
data, err := c.Get(key)
if err != nil {
return nil, fmt.Errorf("%w: %w; %s", err, errFetchingPostingList, keyString+" with uid "+strconv.FormatUint(uid, 10))
}
return data, nil
}
// populateEdgeDataFromStore(keyString, uid, c, edgeData)
// will fill edgeData with the contents of the neighboring edges for
// a given DataKey by looking into the given cache (which may result
// in a call to the underlying persistent storage).
// If data is found for the key, this returns true, otherwise, it
// returns false. If the data was found (and there were no errors),
// it populates edgeData with the found contents.
func populateEdgeDataFromKeyWithCacheType(
keyString string,
uid uint64,
c index.CacheType,
edgeData *[][]uint64) (bool, error) {
data, err := getDataFromKeyWithCacheType(keyString, uid, c)
// Note that posting list fetching errors are treated as just not having
// found the data -- no harm, no foul, as it is probably a
// dead reference that we can ignore.
if err != nil && !errors.Is(err, errFetchingPostingList) {
return false, err
}
if data == nil {
return false, nil
}
err = decodeUint64MatrixUnsafe(data, edgeData)
return true, err
}
// entryUuidInsert adds the entry uuid to the given key
func entryUuidInsert(
ctx context.Context,
key []byte,
txn index.Txn,
predEntryKey string,
entryUuid []byte) (*index.KeyValue, error) {
edge := &index.KeyValue{
Entity: 1,
Attr: predEntryKey,
Value: entryUuid,
}
err := txn.AddMutationWithLockHeld(ctx, key, edge)
return edge, err
}
func ConcatStrings(strs ...string) string {
total := ""
for _, s := range strs {
total += s
}
return total
}
func getInsertLayer(maxLevels int) int {
// multFactor is a multiplicative factor used to normalize the distribution
var level int
randFloat := rand.Float64()
for i := range maxLevels {
// calculate level based on section 3.1 here
if randFloat < math.Pow(1.0/float64(5), float64(maxLevels-1-i)) {
level = i
break
}
}
return level
}
var emptyVec = []byte{}
// adds the data corresponding to a uid to the given vec variable in the form of []T
// this does not allocate memory for vec, so it must be allocated before calling this function
func (ph *persistentHNSW[T]) getVecFromUid(uid uint64, c index.CacheType, vec *[]T) error {
data, err := getDataFromKeyWithCacheType(ph.pred, uid, c)
if err != nil {
if errors.Is(err, errFetchingPostingList) {
// no vector. Return empty array of floats
index.BytesAsFloatArray(emptyVec, vec, ph.floatBits)
return fmt.Errorf("%w; %w", errNilVector, err)
}
return err
}
if data != nil {
index.BytesAsFloatArray(data, vec, ph.floatBits)
return nil
} else {
index.BytesAsFloatArray(emptyVec, vec, ph.floatBits)
return errNilVector
}
}
// chooses whether to create the entry and start nodes based on if it already
// exists, and if it hasnt been created yet, it adds the startNode to all
// levels.
func (ph *persistentHNSW[T]) createEntryAndStartNodes(
ctx context.Context,
c *TxnCache,
inUuid uint64,
vec *[]T) (uint64, []*index.KeyValue, error) {
txn := c.txn
edges := []*index.KeyValue{}
entryKey := DataKey(ph.vecEntryKey, 1) // 0-profile_vector_entry
txn.LockKey(entryKey)
defer txn.UnlockKey(entryKey)
data, _ := txn.GetWithLockHeld(entryKey)
create_edges := func(inUuid uint64) (uint64, []*index.KeyValue, error) {
startEdges, err := ph.addStartNodeToAllLevels(ctx, entryKey, txn, inUuid)
if err != nil {
return 0, []*index.KeyValue{}, err
}
// return entry node at all levels
edges = append(edges, startEdges...)
return 0, edges, nil
}
if data == nil {
// no entries in vector index yet b/c no entry exists, so put in all levels
return create_edges(inUuid)
}
entry := BytesToUint64(data) // convert entry Uuid returned from Get to uint64
err := ph.getVecFromUid(entry, c, vec)
if err != nil || len(*vec) == 0 {
// The entry vector has been deleted. We have to create a new entry vector.
entry, err := ph.calculateNewEntryVec(ctx, c, vec)
if err != nil {
// No other node exists, go with the new node that has come
return create_edges(inUuid)
}
return create_edges(entry)
}
return entry, edges, nil
}
// Converts the matrix into linear array that looks like
// [0: Number of rows 1: Length of row1 2-n: Data of row1 3: Length of row2 ..]
func encodeUint64MatrixUnsafe(matrix [][]uint64) []byte {
if len(matrix) == 0 {
return nil
}
// Calculate the total size
var totalSize uint64
for _, row := range matrix {
totalSize += uint64(len(row))*uint64(unsafe.Sizeof(uint64(0))) + uint64(unsafe.Sizeof(uint64(0)))
}
totalSize += uint64(unsafe.Sizeof(uint64(0)))
// Create a byte slice with the appropriate size
data := make([]byte, totalSize)
offset := 0
// Write number of rows
rows := uint64(len(matrix))
copy(data[offset:offset+8], (*[8]byte)(unsafe.Pointer(&rows))[:])
offset += 8
// Write each row's length and data
for _, row := range matrix {
rowLen := uint64(len(row))
copy(data[offset:offset+8], (*[8]byte)(unsafe.Pointer(&rowLen))[:])
offset += 8
for i := range row {
copy(data[offset:offset+8], (*[8]byte)(unsafe.Pointer(&row[i]))[:])
offset += 8
}
}
return data
}
func decodeUint64MatrixUnsafe(data []byte, matrix *[][]uint64) error {
if len(data) == 0 {
return nil
}
offset := 0
// Read number of rows
rows := *(*uint64)(unsafe.Pointer(&data[offset]))
offset += 8
*matrix = make([][]uint64, rows)
for i := 0; i < int(rows); i++ {
// Read row length
rowLen := *(*uint64)(unsafe.Pointer(&data[offset]))
offset += 8
(*matrix)[i] = make([]uint64, rowLen)
for j := 0; j < int(rowLen); j++ {
(*matrix)[i][j] = *(*uint64)(unsafe.Pointer(&data[offset]))
offset += 8
}
}
return nil
}
// adds empty layers to all levels
func (ph *persistentHNSW[T]) addStartNodeToAllLevels(
ctx context.Context,
entryKey []byte,
txn index.Txn,
inUuid uint64) ([]*index.KeyValue, error) {
edges := []*index.KeyValue{}
key := DataKey(ph.vecKey, inUuid)
emptyEdgesBytes := encodeUint64MatrixUnsafe(make([][]uint64, ph.maxLevels))
// creates empty at all levels only for entry node
edge, err := ph.newPersistentEdgeKeyValueEntry(ctx, key, txn, inUuid, emptyEdgesBytes)
if err != nil {
return []*index.KeyValue{}, err
}
edges = append(edges, edge)
inUuidByte := Uint64ToBytes(inUuid)
// add inUuid as entry for this structure from now on
edge, err = entryUuidInsert(ctx, entryKey, txn, ph.vecEntryKey, inUuidByte)
if err != nil {
return []*index.KeyValue{}, err
}
edges = append(edges, edge)
return edges, nil
}
// creates a new edge with the given uuid and edges. Lock must be held before calling this function
func (ph *persistentHNSW[T]) newPersistentEdgeKeyValueEntry(ctx context.Context, key []byte,
txn index.Txn, uuid uint64, edges []byte) (*index.KeyValue, error) {
txn.LockKey(key)
defer txn.UnlockKey(key)
edge := &index.KeyValue{
Entity: uuid,
Attr: ph.vecKey,
Value: edges,
}
if err := txn.AddMutationWithLockHeld(ctx, key, edge); err != nil {
return nil, err
}
return edge, nil
}
func dedupeUidsPreserveOrder(uids []uint64) []uint64 {
if len(uids) <= 1 {
return uids
}
seen := make(map[uint64]struct{}, len(uids))
out := uids[:0]
for _, uid := range uids {
if uid == notAUid {
continue
}
if _, ok := seen[uid]; ok {
continue
}
seen[uid] = struct{}{}
out = append(out, uid)
}
return out
}
func worstScore[T c.Float](simType SimilarityType[T]) T {
// For distance metrics, lower is better so the worst score is +Inf
// For similarity metrics, higher is better so the worst score is -Inf
if simType.isSimilarityMetric {
return T(math.Inf(-1))
}
return T(math.Inf(1))
}
func (ph *persistentHNSW[T]) uidScoreForNode(
tc *TxnCache,
node uint64,
nodeVec []T,
uid uint64,
outVec *[]T,
) T {
score := worstScore(ph.simType)
if uid == notAUid || uid == node {
return score
}
if err := ph.getVecFromUid(uid, tc, outVec); err == nil && len(*outVec) != 0 {
if s, err := ph.simType.distanceScore(nodeVec, *outVec, ph.floatBits); err == nil {
score = s
}
}
return score
}
// insertUidsTopKByScore incrementally maintains a best-to-worst sorted uid list capped at `keep`.
// It assumes `edges` is already sorted best-to-worst according to simType semantics.
func insertUidsTopKByScore[T c.Float](
edges []uint64,
newUids []uint64,
keep int,
simType SimilarityType[T],
score func(uint64) T,
) []uint64 {
if keep <= 0 {
return []uint64{}
}
if len(edges) > keep {
edges = edges[:keep]
}
seen := make(map[uint64]struct{}, len(edges)+len(newUids))
filtered := edges[:0]
for _, uid := range edges {
if uid == notAUid {
continue
}
if _, ok := seen[uid]; ok {
continue
}
seen[uid] = struct{}{}
filtered = append(filtered, uid)
}
edges = filtered
insertSorted := func(edges []uint64, uid uint64) []uint64 {
scoreNew := score(uid)
pos := sort.Search(len(edges), func(i int) bool {
return simType.isBetterScore(scoreNew, score(edges[i]))
})
edges = append(edges, 0)
copy(edges[pos+1:], edges[pos:])
edges[pos] = uid
return edges
}
for _, uid := range newUids {
if uid == notAUid {
continue
}
if _, ok := seen[uid]; ok {
continue
}
seen[uid] = struct{}{}
edges = insertSorted(edges, uid)
if len(edges) > keep {
edges = edges[:keep]
}
}
return edges
}
// addNeighbors adds the neighbors of the given uuid to the given level.
// It returns the edge created and the error if any.
func (ph *persistentHNSW[T]) addNeighbors(ctx context.Context, tc *TxnCache,
uuid uint64, allLayerNeighbors [][]uint64) (*index.KeyValue, error) {
txn := tc.txn
keyPred := ph.vecKey
key := DataKey(keyPred, uuid)
txn.LockKey(key)
defer txn.UnlockKey(key)
var nnEdgesErr error
var allLayerEdges [][]uint64
var ok bool
allLayerEdges, ok = ph.nodeAllEdges[uuid]
if !ok {
data, _ := txn.GetWithLockHeld(key)
if data == nil {
allLayerEdges = allLayerNeighbors
} else {
// all edges of nearest neighbor
err := decodeUint64MatrixUnsafe(data, &allLayerEdges)
if err != nil {
return nil, err
}
}
}
var inVec []T
inVecReady := false
var outVec []T
for level := range ph.maxLevels {
allLayerEdges[level], nnEdgesErr = ph.removeDeadNodes(allLayerEdges[level], tc)
if nnEdgesErr != nil {
return nil, nnEdgesErr
}
// Fast-path: if we're not adding any neighbours at this level, don't do any extra work.
// This matters because addNeighbors() is called for many nodes where only one layer
// actually changes (e.g. inbound edge updates during construction)
if len(allLayerNeighbors[level]) == 0 {
continue
}
// We maintain the invariant that allLayerEdges[level] is sorted best-to-worst
// (according to ph.simType semantics) and bounded by efConstruction.
//
// This lets us do incremental updates cheaply: for each new neighbor, compute its
// score once, binary-search insertion into the sorted list, then truncate.
if !inVecReady {
if err := ph.getVecFromUid(uuid, tc, &inVec); err != nil || len(inVec) == 0 {
// Without the source vector we can't score edges reliably.
// Fall back to "append then truncate" after a cheap de-dupe.
allLayerEdges[level] = append(allLayerEdges[level], allLayerNeighbors[level]...)
allLayerEdges[level] = dedupeUidsPreserveOrder(allLayerEdges[level])
if len(allLayerEdges[level]) > ph.efConstruction {
allLayerEdges[level] = allLayerEdges[level][:ph.efConstruction]
}
continue
}
inVecReady = true
}
// Small local cache for scores computed during this call.
scoreCache := make(map[uint64]T, len(allLayerEdges[level])+len(allLayerNeighbors[level]))
getScore := func(uid uint64) T {
if s, ok := scoreCache[uid]; ok {
return s
}
s := ph.uidScoreForNode(tc, uuid, inVec, uid, &outVec)
scoreCache[uid] = s
return s
}
// Filter out self-loops here (helper only knows about notAUid).
newUids := allLayerNeighbors[level]
if uuid != notAUid {
dst := newUids[:0]
for _, uid := range newUids {
if uid != uuid {
dst = append(dst, uid)
}
}
newUids = dst
}
allLayerEdges[level] = insertUidsTopKByScore(allLayerEdges[level], newUids, ph.efConstruction, ph.simType, getScore)
}
// on every modification of the layer edges, add it to in mem map so you dont have to always be reading
// from persistent storage
ph.nodeAllEdges[uuid] = allLayerEdges
inboundEdgesBytes := encodeUint64MatrixUnsafe(allLayerEdges)
edge := &index.KeyValue{
Entity: uuid,
Attr: ph.vecKey,
Value: inboundEdgesBytes,
}
if err := txn.AddMutationWithLockHeld(ctx, key, edge); err != nil {
return nil, err
}
return edge, nil
}
// removeDeadNodes(nnEdges, tc) removes dead nodes from nnEdges and returns the new nnEdges
func (ph *persistentHNSW[T]) removeDeadNodes(nnEdges []uint64, tc *TxnCache) ([]uint64, error) {
// TODO add a path to delete deadNodes
if ph.deadNodes == nil {
data, err := getDataFromKeyWithCacheType(ph.vecDead, 1, tc)
if err != nil && !errors.Is(err, errFetchingPostingList) {
return []uint64{}, err
}
var deadNodes []uint64
if data != nil { // if dead nodes exist, convert to []uint64
deadNodes, err = ParseEdges(string(data))
if err != nil {
return []uint64{}, err
}
}
ph.deadNodes = make(map[uint64]struct{})
for _, n := range deadNodes {
ph.deadNodes[n] = struct{}{}
}
}
if len(ph.deadNodes) == 0 {
return nnEdges, nil
}
var diff []uint64
for _, s := range nnEdges {
if _, ok := ph.deadNodes[s]; !ok {
diff = append(diff, s)
continue
}
}
return diff, nil
}
func Uint64ToBytes(key uint64) []byte {
b := make([]byte, 8)
binary.BigEndian.PutUint64(b, key)
return b
}
func BytesToUint64(bytes []byte) uint64 {
return binary.BigEndian.Uint64(bytes)
}
func isEqual[T c.Float](a []T, b []T) bool {
if len(a) != len(b) {
return false
}
for i, val := range a {
if val != b[i] {
return false
}
}
return true
}
// DataKey generates a data key with the given attribute and UID.
// The structure of a data key is as follows:
//
// byte 0: key type prefix (set to DefaultPrefix or ByteSplit if part of a multi-part list)
// byte 1-2: length of attr
// next len(attr) bytes: value of attr
// next byte: data type prefix (set to ByteData)
// next eight bytes: value of uid
// next eight bytes (optional): if the key corresponds to a split list, the startUid of
// the split stored in this key and the first byte will be sets to ByteSplit.
func DataKey(attr string, uid uint64) []byte {
extra := 1 + 8 // ByteData + UID
buf, prefixLen := generateKey(DefaultPrefix, attr, extra)
rest := buf[prefixLen:]
rest[0] = ByteData
rest = rest[1:]
binary.BigEndian.PutUint64(rest, uid)
return buf
}
// genKey creates the key and writes the initial bytes (type byte, length of attribute,
// and the attribute itself). It leaves the rest of the key empty for further processing
// if necessary. It also returns next index from where further processing should be done.
func generateKey(typeByte byte, attr string, extra int) ([]byte, int) {
// Separate namespace and attribute from attr and write namespace in the first 8 bytes of key.
namespace, attr := ParseNamespaceBytes(attr)
prefixLen := 1 + 8 + 2 + len(attr) // byteType + ns + len(pred) + pred
buf := make([]byte, prefixLen+extra)
buf[0] = typeByte
AssertTrue(copy(buf[1:], namespace) == 8)
rest := buf[9:]
writeAttr(rest, attr)
return buf, prefixLen
}
func ParseNamespaceBytes(attr string) ([]byte, string) {
splits := strings.SplitN(attr, NsSeparator, 2)
ns := make([]byte, 8)
binary.BigEndian.PutUint64(ns, strToUint(splits[0]))
return ns, splits[1]
}
// AssertTrue asserts that b is true. Otherwise, it would log fatal.
func AssertTrue(b bool) {
if !b {
log.Fatalf("%+v", errors.Errorf("Assert failed"))
}
}
func writeAttr(buf []byte, attr string) []byte {
AssertTrue(len(attr) < math.MaxUint16)
binary.BigEndian.PutUint16(buf[:2], uint16(len(attr)))
rest := buf[2:]
AssertTrue(len(attr) == copy(rest, attr))
return rest[len(attr):]
}
// For consistency, use base16 to encode/decode the namespace.
func strToUint(s string) uint64 {
ns, err := strconv.ParseUint(s, 16, 64)
Check(err)
return ns
}
// Check logs fatal if err != nil.
func Check(err error) {
if err != nil {
err = errors.Wrap(err, "")
log.Fatalf("%+v", err)
}
}