perf(vector): Add heap to neighbour edges

posting/index.go

@@ -25,8 +25,10 @@ import (
 	"fmt"
 	"math"
 	"os"
+	"strings"
 	"sync/atomic"
 	"time"
+	"unsafe"
 
 	"github.com/golang/glog"
 	"github.com/pkg/errors"
@@ -653,7 +655,7 @@ func (r *rebuilder) RunWithoutTemp(ctx context.Context) error {
 	stream := pstore.NewStreamAt(r.startTs)
 	stream.LogPrefix = fmt.Sprintf("Rebuilding index for predicate %s (1/2):", r.attr)
 	stream.Prefix = r.prefix
-	stream.NumGo = 128
+	stream.NumGo = 16
 	txn := NewTxn(r.startTs)
 	stream.KeyToList = func(key []byte, it *badger.Iterator) (*bpb.KVList, error) {
 		// We should return quickly if the context is no longer valid.
@@ -742,6 +744,10 @@ func (r *rebuilder) RunWithoutTemp(ctx context.Context) error {
 		return err
 	}
 
+	if os.Getenv("DEBUG_SHOW_HNSW_TREE") != "" {
+		printTreeStats(txn)
+	}
+
 	txn.Update()
 	writer := NewTxnWriter(pstore)
 
@@ -763,6 +769,84 @@ func (r *rebuilder) RunWithoutTemp(ctx context.Context) error {
 		})
 }
 
+func printTreeStats(txn *Txn) {
+	txn.cache.Lock()
+
+	numLevels := 20
+	numNodes := make([]int, numLevels)
+	numConnections := make([]int, numLevels)
+
+	var temp [][]uint64
+	for key, pl := range txn.cache.plists {
+		pk, _ := x.Parse([]byte(key))
+		if strings.HasSuffix(pk.Attr, "__vector_") {
+			data := pl.getPosting(txn.cache.startTs)
+			if data == nil || len(data.Postings) == 0 {
+				continue
+			}
+
+			err := decodeUint64MatrixUnsafe(data.Postings[0].Value, &temp)
+			if err != nil {
+				fmt.Println("Error while decoding", err)
+			}
+
+			for i := 0; i < len(temp); i++ {
+				if len(temp[i]) > 0 {
+					numNodes[i] += 1
+				}
+				numConnections[i] += len(temp[i])
+			}
+
+		}
+	}
+
+	for i := 0; i < numLevels; i++ {
+		fmt.Printf("%d, ", numNodes[i])
+	}
+	fmt.Println("")
+	for i := 0; i < numLevels; i++ {
+		fmt.Printf("%d, ", numConnections[i])
+	}
+	fmt.Println("")
+	for i := 0; i < numLevels; i++ {
+		if numNodes[i] == 0 {
+			fmt.Printf("0, ")
+			continue
+		}
+		fmt.Printf("%d, ", numConnections[i]/numNodes[i])
+	}
+	fmt.Println("")
+
+	txn.cache.Unlock()
+}
+
+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
+}
+
 func (r *rebuilder) Run(ctx context.Context) error {
 	if r.startTs == 0 {
 		glog.Infof("maxassigned is 0, no indexing work for predicate %s", r.attr)

posting/list.go

@@ -545,6 +545,16 @@ func (l *List) addMutationInternal(ctx context.Context, txn *Txn, t *pb.Directed
 	return nil
 }
 
+// getMutation returns a marshaled version of posting list mutation stored internally.
+func (l *List) getPosting(startTs uint64) *pb.PostingList {
+	l.RLock()
+	defer l.RUnlock()
+	if pl, ok := l.mutationMap[startTs]; ok {
+		return pl
+	}
+	return nil
+}
+
 // getMutation returns a marshaled version of posting list mutation stored internally.
 func (l *List) getMutation(startTs uint64) []byte {
 	l.RLock()

posting/lists.go

@@ -143,9 +143,17 @@ func (vc *viLocalCache) GetWithLockHeld(key []byte) (rval index.Value, rerr erro
 }
 
 func (vc *viLocalCache) GetValueFromPostingList(pl *List) (rval index.Value, rerr error) {
-	val, err := pl.ValueWithLockHeld(vc.delegate.startTs)
-	rval = val.Value
-	return rval, err
+	value := pl.findStaticValue(vc.delegate.startTs)
+
+	if value == nil {
+		return nil, ErrNoValue
+	}
+
+	if hasDeleteAll(value.Postings[0]) || value.Postings[0].Op == Del {
+		return nil, ErrNoValue
+	}
+
+	return value.Postings[0].Value, nil
 }
 
 func NewViLocalCache(delegate *LocalCache) *viLocalCache {
@@ -280,24 +288,23 @@ func (lc *LocalCache) SetIfAbsent(key string, updated *List) *List {
 }
 
 func (lc *LocalCache) getInternal(key []byte, readFromDisk bool) (*List, error) {
+	skey := string(key)
 	getNewPlistNil := func() (*List, error) {
 		lc.RLock()
 		defer lc.RUnlock()
 		if lc.plists == nil {
 			return getNew(key, pstore, lc.startTs)
 		}
+		if l, ok := lc.plists[skey]; ok {
+			return l, nil
+		}
 		return nil, nil
 	}
 
 	if l, err := getNewPlistNil(); l != nil || err != nil {
 		return l, err
 	}
 
-	skey := string(key)
-	if pl := lc.getNoStore(skey); pl != nil {
-		return pl, nil
-	}
-
 	var pl *List
 	if readFromDisk {
 		var err error

posting/oracle.go

@@ -103,9 +103,18 @@ func (vt *viTxn) GetWithLockHeld(key []byte) (rval index.Value, rerr error) {
 }
 
 func (vt *viTxn) GetValueFromPostingList(pl *List) (rval index.Value, rerr error) {
-	val, err := pl.ValueWithLockHeld(vt.delegate.StartTs)
-	rval = val.Value
-	return rval, err
+	value := pl.findStaticValue(vt.delegate.StartTs)
+
+	if value == nil {
+		//fmt.Println("DIFF", val, err, nil, badger.ErrKeyNotFound)
+		return nil, ErrNoValue
+	}
+
+	if hasDeleteAll(value.Postings[0]) || value.Postings[0].Op == Del {
+		return nil, ErrNoValue
+	}
+
+	return value.Postings[0].Value, nil
 }
 
 func (vt *viTxn) AddMutation(ctx context.Context, key []byte, t *index.KeyValue) error {

query/math.go

@@ -84,6 +84,25 @@ func processBinary(mNode *mathTree) error {
 		return nil
 	}
 
+	// If mpl or mpr have 0 and just 0 in it, that means it's an output of aggregation somewhere.
+	// This value would need to be applied to all.
+	checkAggrResult := func(value map[uint64]types.Val) (types.Val, bool) {
+		if len(value) != 1 {
+			return types.Val{}, false
+		}
+
+		val, ok := value[0]
+		return val, ok
+	}
+
+	if val, ok := checkAggrResult(mpl); ok {
+		cl = val
+		mpl = nil
+	} else if val, ok := checkAggrResult(mpr); ok {
+		cr = val
+		mpr = nil
+	}
+
 	if len(mpl) != 0 || len(mpr) != 0 {
 		for k := range mpr {
 			if err := f(k); err != nil {

query/math_test.go

@@ -26,6 +26,53 @@ import (
 	"github.com/dgraph-io/dgraph/types"
 )
 
+func TestVector(t *testing.T) {
+	tree := &mathTree{
+		Fn: "sqrt",
+		Child: []*mathTree{{
+			Fn: "dot",
+			Child: []*mathTree{
+				{
+					Fn: "-",
+					Child: []*mathTree{
+						{
+							Var: "v1",
+							Val: map[uint64]types.Val{
+								0: {Tid: 12, Value: []float32{1.0, 2.0}},
+							},
+						},
+						{
+							Var: "v2",
+							Val: map[uint64]types.Val{
+								123: {Tid: 12, Value: []float32{1.0, 2.0}},
+							},
+						},
+					},
+				},
+				{
+					Fn: "-",
+					Child: []*mathTree{
+						{
+							Var: "v1",
+							Val: map[uint64]types.Val{
+								0: {Tid: 12, Value: []float32{1.0, 2.0}},
+							},
+						},
+						{
+							Var: "v2",
+							Val: map[uint64]types.Val{
+								123: {Tid: 12, Value: []float32{1.0, 2.0}},
+							},
+						},
+					},
+				},
+			},
+		}},
+	}
+
+	require.NoError(t, evalMathTree(tree))
+}
+
 func TestProcessBinary(t *testing.T) {
 	tests := []struct {
 		in  *mathTree

query/query.go

@@ -1185,6 +1185,7 @@ func (sg *SubGraph) transformVars(doneVars map[string]varValue, path []*SubGraph
 			mt.Const = val
 			continue
 		}
+
 		mt.Val = newMap
 	}
 	return nil

query/vector/vector_test.go

@@ -598,10 +598,10 @@ func TestVectorDelete(t *testing.T) {
 		deleteTriplesInCluster(triple)
 		uid := strings.Split(triple, " ")[0]
 		query = fmt.Sprintf(`{
-		 vector(func: uid(%s)) {
-		  vtest
-		 }
-	}`, uid[1:len(uid)-1])
+		  vector(func: uid(%s)) {
+		   vtest
+		  }
+		}`, uid[1:len(uid)-1])
 
 		result = processQueryNoErr(t, query)
 		require.JSONEq(t, `{"data": {"vector":[]}}`, result)

tok/hnsw/heap.go

@@ -1,3 +1,21 @@
+/*
+ * Copyright 2016-2024 Dgraph Labs, Inc. and Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * Co-authored by: [email protected], [email protected], [email protected]
+ */
+
 package hnsw
 
 import (