Aggregate uids stored in a variable over levels for recurse query.

Author: Pawan Rawal

query/query.go

@@ -1423,6 +1423,8 @@ func (sg *SubGraph) populateUidValVar(doneVars map[string]varValue, sgPath []*Su
 		return nil
 	}
 
+	var v varValue
+	var ok bool
 	if sg.Attr == "_predicate_" {
 		// This is a predicates list.
 		doneVars[sg.Params.Var] = varValue{
@@ -1444,10 +1446,22 @@ func (sg *SubGraph) populateUidValVar(doneVars map[string]varValue, sgPath []*Su
 		}
 	} else if len(sg.DestUIDs.Uids) != 0 {
 		// This implies it is a entity variable.
-		doneVars[sg.Params.Var] = varValue{
-			Uids: sg.DestUIDs,
-			path: sgPath,
+		if v, ok = doneVars[sg.Params.Var]; !ok {
+			doneVars[sg.Params.Var] = varValue{
+				Uids: sg.DestUIDs,
+				path: sgPath,
+			}
+			return nil
 		}
+		// For a recurse query this can happen. We don't allow using the same variable more than
+		// once otherwise.
+		uids := v.Uids
+		lists := make([]*protos.List, 0, 2)
+		lists = append(lists, uids, sg.DestUIDs)
+		v.Uids = algo.MergeSorted(lists)
+		doneVars[sg.Params.Var] = v
+
+		// This implies it is a value variable.
 	} else if len(sg.valueMatrix) != 0 && sg.SrcUIDs != nil && len(sgPath) != 0 {
 		if sg.Attr == "_uid_" {
 			// Its still an entity variable if its _uid_.
@@ -1458,21 +1472,33 @@ func (sg *SubGraph) populateUidValVar(doneVars map[string]varValue, sgPath []*Su
 			return nil
 		}
 
-		// This implies it is a value variable.
-		// NOTE: Value variables cannot be defined and used in the same query block. so
-		// checking len(sgPath) is okay.
+		// We reach here, if a uid variable was already set and now a uid edge has no dest uids.
+		// We don't want to set the variable as empty in this case.
+		if _, ok := doneVars[sg.Params.Var]; ok {
+			return nil
+		}
+
 		doneVars[sg.Params.Var] = varValue{
 			Vals: make(map[uint64]types.Val),
 			path: sgPath,
 		}
 		for idx, uid := range sg.SrcUIDs.Uids {
+			if len(sg.valueMatrix[idx].Values) > 1 {
+				return x.Errorf("Value variables not supported for predicate with list type.")
+			}
+
 			val, err := convertWithBestEffort(sg.valueMatrix[idx].Values[0], sg.Attr)
 			if err != nil {
 				continue
 			}
 			doneVars[sg.Params.Var].Vals[uid] = val
 		}
 	} else {
+		// If the variable already existed and now we see it again without any DestUIDs or
+		// ValueMatrix then lets just return.
+		if _, ok := doneVars[sg.Params.Var]; ok {
+			return nil
+		}
 		// Insert a empty entry to keep the dependency happy.
 		doneVars[sg.Params.Var] = varValue{
 			path: sgPath,
@@ -1481,6 +1507,7 @@ func (sg *SubGraph) populateUidValVar(doneVars map[string]varValue, sgPath []*Su
 	}
 	return nil
 }
+
 func (sg *SubGraph) populateFacetVars(doneVars map[string]varValue, sgPath []*SubGraph) error {
 	if sg.Params.FacetVar != nil && sg.Params.Facet != nil {
 		sgPath = append(sgPath, sg)

query/query_test.go

@@ -1825,6 +1825,48 @@ func TestRecurseQueryLimitDepth2(t *testing.T) {
 		`{"data": {"recurse":[{"_uid_":"0x1","friend":[{"_uid_":"0x17","name":"Rick Grimes"},{"_uid_":"0x18","name":"Glenn Rhee"},{"_uid_":"0x19","name":"Daryl Dixon"},{"_uid_":"0x1f","name":"Andrea"},{"_uid_":"0x65"}],"name":"Michonne"}]}}`, js)
 }
 
+func TestRecurseVariable(t *testing.T) {
+	populateGraph(t)
+	query := `
+			{
+				recurse(func: uid(0x01)) {
+					a as friend
+				}
+
+				me(func: uid(a)) {
+					name
+				}
+			}
+		`
+
+	js := processToFastJSON(t, query)
+	require.Equal(t, `{"data": {"me":[{"name":"Michonne"},{"name":"Rick Grimes"},{"name":"Glenn Rhee"},{"name":"Daryl Dixon"},{"name":"Andrea"}]}}`, js)
+}
+
+func TestRecurseVariable2(t *testing.T) {
+	populateGraph(t)
+
+	query := `
+			{
+
+				recurse(func: uid(0x1)) {
+					f2 as friend
+					f as follow
+				}
+
+				me(func: uid(f)) {
+					name
+				}
+
+				me2(func: uid(f2)) {
+					name
+				}
+			}
+	`
+	js := processToFastJSON(t, query)
+	require.Equal(t, `{"data": {"me":[{"name":"Glenn Rhee"},{"name":"Andrea"},{"name":"Alice"},{"name":"Bob"},{"name":"Matt"},{"name":"John"}],"me2":[{"name":"Michonne"},{"name":"Rick Grimes"},{"name":"Glenn Rhee"},{"name":"Daryl Dixon"},{"name":"Andrea"}]}}`, js)
+}
+
 func TestShortestPath_ExpandError(t *testing.T) {
 	populateGraph(t)
 	query := `
@@ -9472,3 +9514,22 @@ func TestUidVariable(t *testing.T) {
 	js := processToFastJSON(t, query)
 	require.Equal(t, `{"data": {"me":[{"name":"Rick Grimes"},{"name":"Glenn Rhee"},{"name":"Daryl Dixon"},{"name":"Andrea"}]}}`, js)
 }
+
+func TestMultipleValueVarError(t *testing.T) {
+	populateGraph(t)
+
+	query := `{
+		var(func:ge(graduation, "1930")) {
+			o as graduation
+		}
+
+		me(func: uid(o)) {
+			graduation
+		}
+	}`
+
+	ctx := defaultContext()
+	_, err := processToFastJsonReqCtx(t, query, ctx)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Value variables not supported for predicate with list type.")
+}

query/recurse.go

@@ -28,7 +28,7 @@ import (
 	"github.com/dgraph-io/dgraph/x"
 )
 
-func (start *SubGraph) expandRecurse(ctx context.Context, next chan bool, rch chan error) {
+func (start *SubGraph) expandRecurse(ctx context.Context, maxDepth uint64) error {
 	// Note: Key format is - "attr|fromUID|toUID"
 	reachMap := make(map[string]struct{})
 	var numEdges int
@@ -48,56 +48,63 @@ func (start *SubGraph) expandRecurse(ctx context.Context, next chan bool, rch ch
 			if tr, ok := trace.FromContext(ctx); ok {
 				tr.LazyPrintf("Error while processing child task: %+v", err)
 			}
-			rch <- err
-			return
+			return err
 		}
 	case <-ctx.Done():
 		if tr, ok := trace.FromContext(ctx); ok {
 			tr.LazyPrintf("Context done before full execution: %+v", ctx.Err())
 		}
-		rch <- ctx.Err()
-		return
+		return ctx.Err()
 	}
 
 	for _, child := range startChildren {
 		temp := new(SubGraph)
 		temp.copyFiltersRecurse(child)
 		temp.SrcUIDs = start.DestUIDs
+		temp.Params.Var = child.Params.Var
 		exec = append(exec, temp)
 		start.Children = append(start.Children, temp)
 	}
 
 	dummy := &SubGraph{}
+	var depth uint64
 	for {
-		isNext := <-next
-		if !isNext {
-			return
+		if depth >= maxDepth {
+			return nil
 		}
+		depth++
 
 		rrch := make(chan error, len(exec))
 		for _, sg := range exec {
 			go ProcessGraph(ctx, sg, dummy, rrch)
 		}
 
+		var recurseErr error
 		for range exec {
 			select {
 			case err = <-rrch:
 				if err != nil {
 					if tr, ok := trace.FromContext(ctx); ok {
 						tr.LazyPrintf("Error while processing child task: %+v", err)
 					}
-					rch <- err
-					return
+					if recurseErr == nil {
+						recurseErr = err
+					}
 				}
 			case <-ctx.Done():
 				if tr, ok := trace.FromContext(ctx); ok {
 					tr.LazyPrintf("Context done before full execution: %+v", ctx.Err())
 				}
-				rch <- ctx.Err()
-				return
+				if recurseErr == nil {
+					recurseErr = ctx.Err()
+				}
 			}
 		}
 
+		if recurseErr != nil {
+			return recurseErr
+		}
+
 		for _, sg := range exec {
 			if len(sg.Filters) > 0 {
 				// We need to do this in case we had some filters.
@@ -107,8 +114,15 @@ func (start *SubGraph) expandRecurse(ctx context.Context, next chan bool, rch ch
 				// This is for avoiding loops in graph.
 				algo.ApplyFilter(sg.uidMatrix[mIdx], func(uid uint64, i int) bool {
 					key := fmt.Sprintf("%s|%d|%d", sg.Attr, fromUID, uid)
-					_, ok := reachMap[key] // Combine fromUID here.
-					return !ok
+					_, seen := reachMap[key] // Combine fromUID here.
+					if seen {
+						return false
+					} else {
+						// Mark this edge as taken. We'd disallow this edge later.
+						reachMap[key] = struct{}{}
+						numEdges++
+						return true
+					}
 				})
 			}
 			if len(sg.Params.Order) > 0 {
@@ -129,77 +143,34 @@ func (start *SubGraph) expandRecurse(ctx context.Context, next chan bool, rch ch
 				temp := new(SubGraph)
 				temp.copyFiltersRecurse(child)
 				temp.SrcUIDs = sg.DestUIDs
+				temp.Params.Var = child.Params.Var
 				sg.Children = append(sg.Children, temp)
 				out = append(out, temp)
 			}
-			// Mark the reached nodes
-			for mIdx, fromUID := range sg.SrcUIDs.Uids {
-				for _, toUID := range sg.uidMatrix[mIdx].Uids {
-					key := fmt.Sprintf("%s|%d|%d", sg.Attr, fromUID, toUID)
-					// Mark this edge as taken. We'd disallow this edge later.
-					reachMap[key] = struct{}{}
-					numEdges++
-				}
-			}
 		}
 
 		if numEdges > 1000000 {
 			// If we've seen too many nodes, stop the query.
-			rch <- ErrTooBig
-			return
+			return ErrTooBig
 		}
 
 		if len(out) == 0 {
-			rch <- ErrStop
-			return
+			return nil
 		}
-		// This marks the end of one level of exectution.
-		rch <- nil
 		exec = out
 	}
 }
 
 func Recurse(ctx context.Context, sg *SubGraph) error {
-	var err error
 	if sg.Params.Alias != "recurse" {
 		return x.Errorf("Invalid shortest path query")
 	}
-	expandErr := make(chan error, 2)
-	next := make(chan bool, 2)
-	go sg.expandRecurse(ctx, next, expandErr)
+
 	depth := sg.Params.ExploreDepth
 	if depth == 0 {
 		// If no depth is specified, expand till we reach all leaf nodes
 		// or we see reach too many nodes.
 		depth = math.MaxUint64
 	}
-
-L:
-	// Recurse number of times specified by the user.
-	for i := uint64(0); i < depth; i++ {
-		next <- true
-		select {
-		case err = <-expandErr:
-			if err != nil {
-				if err == ErrTooBig {
-					return err
-				}
-				if err == ErrStop {
-					break L
-				}
-				if tr, ok := trace.FromContext(ctx); ok {
-					tr.LazyPrintf("Error while processing child task: %+v", err)
-				}
-				return err
-			}
-		case <-ctx.Done():
-			if tr, ok := trace.FromContext(ctx); ok {
-				tr.LazyPrintf("Context done before full execution: %+v", ctx.Err())
-			}
-			return ctx.Err()
-		}
-	}
-	// Done expanding.
-	next <- false
-	return nil
+	return sg.expandRecurse(ctx, depth)
 }