perf(core): Add sharded map for post query speedup (#9431)

Author: harshil-goel

dql/math.go

@@ -49,7 +49,7 @@ type MathTree struct {
 	Fn    string
 	Var   string
 	Const types.Val // This will always be parsed as a float value
-	Val   map[uint64]types.Val
+	Val   *types.ShardedMap
 	Child []*MathTree
 }
 

edgraph/server.go

@@ -834,7 +834,7 @@ func findMutationVars(qc *queryContext) []string {
 // If val(variable) exists in a query, but the values are not there for the variable,
 // it will ignore the mutation silently.
 func updateValInNQuads(nquads []*api.NQuad, qc *queryContext, isSet bool) []*api.NQuad {
-	getNewVals := func(s string) (map[uint64]types.Val, bool) {
+	getNewVals := func(s string) (*types.ShardedMap, bool) {
 		if strings.HasPrefix(s, "val(") {
 			varName := s[4 : len(s)-1]
 			if v, ok := qc.valRes[varName]; ok && v != nil {
@@ -845,15 +845,15 @@ func updateValInNQuads(nquads []*api.NQuad, qc *queryContext, isSet bool) []*api
 		return nil, false
 	}
 
-	getValue := func(key uint64, uidToVal map[uint64]types.Val) (types.Val, bool) {
-		val, ok := uidToVal[key]
+	getValue := func(key uint64, uidToVal *types.ShardedMap) (types.Val, bool) {
+		val, ok := uidToVal.Get(key)
 		if ok {
 			return val, true
 		}
 
 		// Check if the variable is aggregate variable
 		// Only 0 key would exist for aggregate variable
-		val, ok = uidToVal[0]
+		val, ok = uidToVal.Get(0)
 		return val, ok
 	}
 
@@ -1026,7 +1026,7 @@ type queryContext struct {
 	uidRes map[string][]string
 	// valRes stores mapping from variable names to values for value
 	// variables used in the mutation block of incoming request.
-	valRes map[string]map[uint64]types.Val
+	valRes map[string]*types.ShardedMap
 	// l stores latency numbers
 	latency *query.Latency
 	// span stores a opencensus span used throughout the query processing
@@ -1494,10 +1494,11 @@ func processQuery(ctx context.Context, qc *queryContext) (*api.Response, error)
 		if v.Uids != nil && len(v.Uids.Uids) > 0 {
 			uidList = v.Uids.Uids
 		} else {
-			uidList = make([]uint64, 0, len(v.Vals))
-			for uid := range v.Vals {
+			uidList = make([]uint64, 0, v.Vals.Len())
+			v.Vals.Iterate(func(uid uint64, val types.Val) error {
 				uidList = append(uidList, uid)
-			}
+				return nil
+			})
 		}
 		if len(uidList) == 0 {
 			continue
@@ -1565,7 +1566,7 @@ func parseRequest(ctx context.Context, qc *queryContext) error {
 		}
 
 		qc.uidRes = make(map[string][]string)
-		qc.valRes = make(map[string]map[uint64]types.Val)
+		qc.valRes = make(map[string]*types.ShardedMap)
 		upsertQuery = buildUpsertQuery(qc)
 		needVars = findMutationVars(qc)
 		if upsertQuery == "" {
@@ -1627,19 +1628,24 @@ func verifyUnique(qc *queryContext, qr query.Request) error {
 		var predValue interface{}
 		if strings.HasPrefix(pred.ObjectId, "val(") {
 			varName := qr.Vars[pred.ObjectId[4:len(pred.ObjectId)-1]]
-			val, ok := varName.Vals[0]
+			val, ok := varName.Vals.Get(0)
 			if !ok {
-				_, isValueGoingtoSet := varName.Vals[subjectUid]
+				_, isValueGoingtoSet := varName.Vals.Get(subjectUid)
 				if !isValueGoingtoSet {
 					continue
 				}
 
 				results := qr.Vars[queryVar.valVar]
-				for uidOfv, v := range results.Vals {
-					if v.Value == varName.Vals[subjectUid].Value && uidOfv != subjectUid {
+				err := results.Vals.Iterate(func(uidOfv uint64, v types.Val) error {
+					varNameVal, _ := varName.Vals.Get(subjectUid)
+					if v.Value == varNameVal.Value && uidOfv != subjectUid {
 						return errors.Errorf("could not insert duplicate value [%v] for predicate [%v]",
 							v.Value, pred.Predicate)
 					}
+					return nil
+				})
+				if err != nil {
+					return err
 				}
 				continue
 			} else {

query/groupby.go

@@ -334,7 +334,7 @@ func (sg *SubGraph) fillGroupedVars(doneVars map[string]varValue, path []*SubGra
 		}
 		chVar := child.Params.Var
 
-		tempMap := make(map[uint64]types.Val)
+		tempMap := types.NewShardedMap()
 		for _, grp := range res.group {
 			if len(grp.keys) == 0 {
 				continue
@@ -349,7 +349,7 @@ func (sg *SubGraph) fillGroupedVars(doneVars map[string]varValue, path []*SubGra
 			}
 			// grp.aggregates could be empty if schema conversion failed during aggregation
 			if len(grp.aggregates) > 0 {
-				tempMap[uid] = grp.aggregates[len(grp.aggregates)-1].key
+				tempMap.Set(uid, grp.aggregates[len(grp.aggregates)-1].key)
 			}
 		}
 		doneVars[chVar] = varValue{

query/math.go

@@ -6,6 +6,8 @@
 package query
 
 import (
+	"sync"
+
 	"github.com/golang/glog"
 	"github.com/pkg/errors"
 
@@ -16,7 +18,7 @@ type mathTree struct {
 	Fn    string
 	Var   string
 	Const types.Val // If its a const value node.
-	Val   map[uint64]types.Val
+	Val   *types.ShardedMap
 	Child []*mathTree
 }
 
@@ -36,24 +38,23 @@ var (
 // processBinary handles the binary operands like
 // +, -, *, /, %, max, min, logbase, dot
 func processBinary(mNode *mathTree) error {
-	destMap := make(map[uint64]types.Val)
 	aggName := mNode.Fn
 
 	mpl := mNode.Child[0].Val
 	mpr := mNode.Child[1].Val
 	cl := mNode.Child[0].Const
 	cr := mNode.Child[1].Const
 
-	f := func(k uint64) error {
+	f := func(k uint64, lshard, rshard, destMapi *map[uint64]types.Val) error {
 		ag := aggregator{
 			name: aggName,
 		}
-		lVal := mpl[k]
+		lVal := (*lshard)[k]
 		if cl.Value != nil {
 			// Use the constant value that was supplied.
 			lVal = cl
 		}
-		rVal := mpr[k]
+		rVal := (*rshard)[k]
 		if cr.Value != nil {
 			// Use the constant value that was supplied.
 			rVal = cr
@@ -66,7 +67,7 @@ func processBinary(mNode *mathTree) error {
 		if err != nil {
 			return err
 		}
-		destMap[k], err = ag.Value()
+		(*destMapi)[k], err = ag.Value()
 		if err != nil {
 			return err
 		}
@@ -75,12 +76,12 @@ func processBinary(mNode *mathTree) error {
 
 	// 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 {
+	checkAggrResult := func(value *types.ShardedMap) (types.Val, bool) {
+		if value.Len() != 1 {
 			return types.Val{}, false
 		}
 
-		val, ok := value[0]
+		val, ok := value.Get(0)
 		return val, ok
 	}
 
@@ -92,21 +93,31 @@ func processBinary(mNode *mathTree) error {
 		mpr = nil
 	}
 
-	if len(mpl) != 0 || len(mpr) != 0 {
-		for k := range mpr {
-			if err := f(k); err != nil {
-				return err
-			}
-		}
-		for k := range mpl {
-			if _, ok := mpr[k]; ok {
-				continue
-			}
-			if err := f(k); err != nil {
-				return err
-			}
+	if mpl.Len() != 0 || mpr.Len() != 0 {
+		var wg sync.WaitGroup
+		returnMap := types.NewShardedMap()
+
+		for i := range types.NumShards {
+			wg.Add(1)
+			mlps := mpl.GetShardOrNil(i)
+			mprs := mpr.GetShardOrNil(i)
+			destMapi := returnMap.GetShardOrNil(i)
+			go func(i int) {
+				defer wg.Done()
+				for k := range mlps {
+					f(k, &mlps, &mprs, &destMapi)
+				}
+				for k := range mprs {
+					if _, ok := mlps[k]; ok {
+						continue
+					}
+					f(k, &mlps, &mprs, &destMapi)
+				}
+			}(i)
 		}
-		mNode.Val = destMap
+
+		wg.Wait()
+		mNode.Val = returnMap
 		return nil
 	}
 
@@ -132,8 +143,8 @@ func processBinary(mNode *mathTree) error {
 // processUnary handles the unary operands like
 // u-, log, exp, since, floor, ceil
 func processUnary(mNode *mathTree) error {
-	destMap := make(map[uint64]types.Val)
 	srcMap := mNode.Child[0].Val
+	destMap := types.NewShardedMap()
 	aggName := mNode.Fn
 	ch := mNode.Child[0]
 	ag := aggregator{
@@ -149,16 +160,18 @@ func processUnary(mNode *mathTree) error {
 		return err
 	}
 
-	for k, val := range srcMap {
+	srcMap.Iterate(func(k uint64, val types.Val) error {
 		err := ag.ApplyVal(val)
 		if err != nil {
 			return err
 		}
-		destMap[k], err = ag.Value()
+		value, err := ag.Value()
 		if err != nil {
 			return err
 		}
-	}
+		destMap.Set(k, value)
+		return nil
+	})
 	mNode.Val = destMap
 	return nil
 
@@ -168,14 +181,14 @@ func processUnary(mNode *mathTree) error {
 // return a boolean value.
 // All the inequality operators (<, >, <=, >=, !=, ==)
 func processBinaryBoolean(mNode *mathTree) error {
-	destMap := make(map[uint64]types.Val)
 	srcMap := mNode.Child[0].Val
+	destMap := types.NewShardedMap()
 	aggName := mNode.Fn
 
 	ch := mNode.Child[1]
 	curMap := ch.Val
-	for k, val := range srcMap {
-		curVal := curMap[k]
+	srcMap.Iterate(func(k uint64, val types.Val) error {
+		curVal, _ := curMap.Get(k)
 		if ch.Const.Value != nil {
 			// Use the constant value that was supplied.
 			curVal = ch.Const
@@ -184,28 +197,29 @@ func processBinaryBoolean(mNode *mathTree) error {
 		if err != nil {
 			return errors.Wrapf(err, "Wrong values in comparison function.")
 		}
-		destMap[k] = types.Val{
+		destMap.Set(k, types.Val{
 			Tid:   types.BoolID,
 			Value: res,
-		}
-	}
+		})
+		return nil
+	})
 	mNode.Val = destMap
 	return nil
 }
 
 // processTernary handles the ternary operand cond()
 func processTernary(mNode *mathTree) error {
-	destMap := make(map[uint64]types.Val)
+	destMap := types.NewShardedMap()
 	aggName := mNode.Fn
 	condMap := mNode.Child[0].Val
-	if len(condMap) == 0 {
+	if condMap.IsEmpty() {
 		return errors.Errorf("Expected a value variable in %v but missing.", aggName)
 	}
 	varOne := mNode.Child[1].Val
 	varTwo := mNode.Child[2].Val
 	constOne := mNode.Child[1].Const
 	constTwo := mNode.Child[2].Const
-	for k, val := range condMap {
+	condMap.Iterate(func(k uint64, val types.Val) error {
 		var res types.Val
 		v, ok := val.Value.(bool)
 		if !ok {
@@ -216,18 +230,19 @@ func processTernary(mNode *mathTree) error {
 			if constOne.Value != nil {
 				res = constOne
 			} else {
-				res = varOne[k]
+				res, _ = varOne.Get(k)
 			}
 		} else {
 			// Pick the value of second map.
 			if constTwo.Value != nil {
 				res = constTwo
 			} else {
-				res = varTwo[k]
+				res, _ = varTwo.Get(k)
 			}
 		}
-		destMap[k] = res
-	}
+		destMap.Set(k, res)
+		return nil
+	})
 	mNode.Val = destMap
 	return nil
 }
@@ -237,7 +252,7 @@ func evalMathTree(mNode *mathTree) error {
 		return nil
 	}
 	if mNode.Var != "" {
-		if len(mNode.Val) == 0 {
+		if mNode.Val.IsEmpty() {
 			glog.V(2).Infof("Variable %v not yet populated or missing.", mNode.Var)
 		}
 		// This is a leaf node whose value is already populated. So return.
@@ -246,13 +261,13 @@ func evalMathTree(mNode *mathTree) error {
 
 	for _, child := range mNode.Child {
 		// Process the child nodes first.
-		err := evalMathTree(child)
-		if err != nil {
+		if err := evalMathTree(child); err != nil {
 			return err
 		}
 	}
 
 	aggName := mNode.Fn
+
 	if isUnary(aggName) {
 		if len(mNode.Child) != 1 {
 			return errors.Errorf("Function %v expects 1 argument. But got: %v", aggName,