feat(vector/hnsw): add per‑query ef and distance_threshold to similar_to, fix early termination (#9514)

Hugely appreciative of the Dgraph team’s work. Native vector search
integrated directly into a graph database is kind of a no brainer today.
Deployed Dgraph (both vanilla and customised) in systems with 1M+
vectors guiding deep traversal queries across 10M+ nodes -- tight
coupling of vector search with graph traversal at massive scale gets us
closer to something that could represent the fuzzy nuances of everything
in an enterprise. Certainly not the biggest deployment your team will
have seen, but this PR fixes an under‑recall edge case in HNSW and
introduces opt‑in, per‑query controls that let users dial recall vs
latency safely and predictably. I’ve had this running in production for
a while and thought it worth proposing to main.

- Summary
- Fix incorrect early termination in the HNSW bottom layer that could
stop before collecting k neighbours.
- Extend similar_to with optional per‑query `ef` and
`distance_threshold` (string or JSON‑like fourth argument).
- Backwards compatible: default 3‑arg behaviour of similar_to is
unchanged.

- Motivation
- In narrow probes, the bottom‑layer search could exit at a local
minimum before collecting k, hurting recall.
- No per‑query `ef` meant recall vs latency trade‑offs required global
tuning or inflating k (and downstream work).
- This PR corrects the termination logic and adds opt‑in knobs so users
can increase exploration only when needed.

- Changes (key files)
- `tok/hnsw/persistent_hnsw.go`: fix early termination, add
`SearchWithOptions`/`SearchWithUidAndOptions`, apply `ef` override at
upper layers and `max(k, ef)` at bottom layer, apply
`distance_threshold` in the metric domain (Euclidean squared internally,
cosine as 1 − sim).
- `tok/index/index.go`: add `VectorIndexOptions` and
`OptionalSearchOptions` (non‑breaking).
- `worker/task.go`: parse optional fourth argument to `similar_to`
(`ef`, `distance_threshold`), thread options, route to optional methods
when provided, guard zero/negative k.
  - `tok/index/search_path.go`: add `SearchPathResult` helper.
  - Tests: `tok/hnsw/ef_recall_test.go` adds
    - `TestHNSWSearchEfOverrideImprovesRecall`
    - `TestHNSWDistanceThreshold_Euclidean`
    - `TestHNSWDistanceThreshold_Cosine`
  - `CHANGELOG.md`: Unreleased entry for HNSW fix and per‑query options.

- Backwards compatibility
- No default behaviour changes. The three‑argument `similar_to(attr, k,
vector_or_uid)` is unchanged.
- `ef` and `distance_threshold` are optional, unsupported metrics safely
ignore the threshold.

- Performance
  - No overhead without options.
- With `ef`, bottom‑layer candidate size becomes `max(k, ef)` (as in
HNSW), cost scales accordingly.
- Threshold filtering is a cheap pass over candidates, squaring
Euclidean thresholds avoids extra roots.

- Rationale and alignment
- Matches HNSW semantics: `ef_search` controls exploration/recall, `k`
controls output size.
- Aligns with
[Typesense](https://typesense.org/docs/29.0/api/vector-search.html#vector-search-parameters)’s
per‑query `ef` and `distance_threshold` semantics for familiarity.

Checklist
- [x] Code compiles correctly and linting passes locally
- [x] For all code changes, an entry added to the `CHANGELOG.md`
describing this PR
- [x] Tests added for new functionality / regression tests for the bug
fix
- [ ] For public APIs/new features, docs PR will be prepared and linked
here after initial review

Author: joelamming

CHANGELOG.md

@@ -88,6 +88,11 @@ as a guide.
 
 - **Query**
   - fix(query): return full float value in query results (#9492)
+- **Vector**
+  - fix(vector/hnsw): correct early termination in bottom-layer search to ensure at least k
+    candidates are considered before breaking
+  - feat(vector/hnsw): add optional per-query controls to similar_to via named parameters: `ef`
+    (search breadth override) and `distance_threshold` (metric-domain cutoff); defaults unchanged
 
 - **Changed**
 

dql/parser.go

@@ -1745,6 +1745,10 @@ L:
 
 		name := collectName(it, item.Val)
 		function.Name = strings.ToLower(name)
+		var similarToOptSeen map[string]struct{}
+		if function.Name == similarToFn {
+			similarToOptSeen = make(map[string]struct{})
+		}
 		if _, ok := tryParseItemType(it, itemLeftRound); !ok {
 			return nil, it.Errorf("Expected ( after func name [%s]", function.Name)
 		}
@@ -1874,7 +1878,10 @@ L:
 				case IsInequalityFn(function.Name):
 					err = parseFuncArgs(it, function)
 
-				case function.Name == "uid_in" || function.Name == "similar_to":
+				case function.Name == "uid_in":
+					err = parseFuncArgs(it, function)
+
+				case function.Name == "similar_to":
 					err = parseFuncArgs(it, function)
 
 				default:
@@ -1892,7 +1899,87 @@ L:
 				}
 				expectArg = false
 				continue
+			case itemLeftCurl:
+				return nil, itemInFunc.Errorf("Unrecognized character inside a func: U+007B '{'")
+			case itemRightCurl:
+				// Right curly braces are never valid in function arguments outside of
+				// the (unsupported) object literal syntax. Always error on stray '}'.
+				return nil, itemInFunc.Errorf("Unrecognized character inside a func: U+007D '}'")
 			default:
+				// similar_to supports named optional parameters after the 3rd positional argument:
+				// similar_to(pred, k, vec, ef: 64, distance_threshold: 0.5)
+				//
+				// Internally we represent each option as two args appended after k and vec:
+				// ["ef", "64", "distance_threshold", "0.5", ...]
+				if itemInFunc.Typ == itemName && function.Name == similarToFn &&
+					function.Attr != "" && len(function.Args) >= 2 {
+					next, ok := it.PeekOne()
+					if ok && next.Typ == itemColon {
+						key := strings.ToLower(collectName(it, itemInFunc.Val))
+						switch key {
+						case "ef", "distance_threshold":
+						default:
+							return nil, itemInFunc.Errorf("Unknown option %q in similar_to", key)
+						}
+						if _, exists := similarToOptSeen[key]; exists {
+							return nil, itemInFunc.Errorf("Duplicate key %q in similar_to options", key)
+						}
+						similarToOptSeen[key] = struct{}{}
+
+						if ok := trySkipItemTyp(it, itemColon); !ok {
+							return nil, it.Errorf("Expected colon(:) after %s", key)
+						}
+						if !it.Next() {
+							return nil, it.Errorf("Expected value for %s", key)
+						}
+						valItem := it.Item()
+						switch valItem.Typ {
+						case itemDollar:
+							varName, err := parseVarName(it)
+							if err != nil {
+								return nil, err
+							}
+							function.Args = append(function.Args, Arg{Value: key})
+							function.Args = append(function.Args, Arg{Value: varName, IsDQLVar: true})
+						case itemMathOp:
+							// Allow signed numeric literals, e.g. distance_threshold: -0.5
+							prefix := valItem.Val
+							if !it.Next() {
+								return nil, it.Errorf("Expected value after %s for %s", prefix, key)
+							}
+							valItem = it.Item()
+							if valItem.Typ != itemName {
+								return nil, valItem.Errorf("Expected value for %s", key)
+							}
+							v := collectName(it, valItem.Val)
+							v = strings.Trim(v, " \t")
+							uq, err := unquoteIfQuoted(v)
+							if err != nil {
+								return nil, err
+							}
+							function.Args = append(function.Args, Arg{Value: key})
+							function.Args = append(function.Args, Arg{Value: prefix + uq})
+						default:
+							if valItem.Typ != itemName {
+								return nil, valItem.Errorf("Expected value for %s", key)
+							}
+							v := collectName(it, valItem.Val)
+							v = strings.Trim(v, " \t")
+							uq, err := unquoteIfQuoted(v)
+							if err != nil {
+								return nil, err
+							}
+							function.Args = append(function.Args, Arg{Value: key})
+							function.Args = append(function.Args, Arg{Value: uq})
+						}
+
+						expectArg = false
+						continue
+					}
+
+					// Disallow extra positional args after (k, vec). Options must be named.
+					return nil, itemInFunc.Errorf("Expected named parameter in similar_to options (e.g. ef: 64)")
+				}
 				if itemInFunc.Typ != itemName {
 					return nil, itemInFunc.Errorf("Expected arg after func [%s], but got item %v",
 						function.Name, itemInFunc)
@@ -2408,6 +2495,10 @@ loop:
 				// The parentheses are balanced out. Let's break.
 				break loop
 			}
+		case item.Typ == itemLeftCurl:
+			return nil, item.Errorf("Unrecognized character inside a func: U+007B '{'")
+		case item.Typ == itemRightCurl:
+			return nil, item.Errorf("Unrecognized character inside a func: U+007D '}'")
 		default:
 			return nil, item.Errorf("Unexpected item while parsing @filter: %v", item)
 		}

dql/parser_test.go

@@ -2518,6 +2518,12 @@ func TestParseFilter_brac(t *testing.T) {
 }
 
 // Test if unbalanced brac will lead to errors.
+// Note: This query has two errors: missing ')' after '()' AND a stray '{'.
+// After changes to support similar_to's JSON args the lexer now emits brace tokens
+// instead of erroring immediately. This causes the query to fail on the structural
+// error (unclosed brackets) rather than the character-specific error. This is an
+// acceptable trade-off because queries with multiple syntax errors may report a different
+// (but equally fatal) error first.
 func TestParseFilter_unbalancedbrac(t *testing.T) {
 	query := `
 	query {
@@ -2532,8 +2538,119 @@ func TestParseFilter_unbalancedbrac(t *testing.T) {
 `
 	_, err := Parse(Request{Str: query})
 	require.Error(t, err)
-	require.Contains(t, err.Error(),
-		"Unrecognized character inside a func: U+007B '{'")
+	require.Contains(t, err.Error(), "Unclosed Brackets")
+}
+
+func TestParseSimilarToNamedParams(t *testing.T) {
+	query := `{
+		q(func: similar_to(voptions, 4, "[0,0]", distance_threshold: 1.5, ef: 12)) {
+			uid
+		}
+	}`
+	res, err := Parse(Request{Str: query})
+	require.NoError(t, err)
+	require.Len(t, res.Query, 1)
+	require.NotNil(t, res.Query[0])
+	require.NotNil(t, res.Query[0].Func)
+	require.Equal(t, "similar_to", res.Query[0].Func.Name)
+	require.Equal(t, "voptions", res.Query[0].Func.Attr)
+	require.Equal(t, "4", res.Query[0].Func.Args[0].Value)
+	require.Equal(t, "[0,0]", res.Query[0].Func.Args[1].Value)
+
+	// Options are appended as (key, value) pairs after k and vec.
+	require.Len(t, res.Query[0].Func.Args, 6)
+	require.Equal(t, "distance_threshold", res.Query[0].Func.Args[2].Value)
+	require.Equal(t, "1.5", res.Query[0].Func.Args[3].Value)
+	require.Equal(t, "ef", res.Query[0].Func.Args[4].Value)
+	require.Equal(t, "12", res.Query[0].Func.Args[5].Value)
+}
+
+func TestParseSimilarToThreeArgs(t *testing.T) {
+	// Test three-arg form (no options)
+	query := `{
+		q(func: similar_to(voptions, 4, "[0,0]")) {
+			uid
+		}
+	}`
+	res, err := Parse(Request{Str: query})
+	require.NoError(t, err)
+	require.Equal(t, "similar_to", res.Query[0].Func.Name)
+	require.Len(t, res.Query[0].Func.Args, 2)
+}
+
+func TestParseSimilarToRejectsObjectLiteralSyntax(t *testing.T) {
+	query := `{
+		q(func: similar_to(voptions, 4, "[0,0]", {ef: 12})) {
+			uid
+		}
+	}`
+	_, err := Parse(Request{Str: query})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Unrecognized character inside a func: U+007B '{'")
+}
+
+func TestParseSimilarToWithQueryVariable(t *testing.T) {
+	query := `query test($eff: int) {
+		q(func: similar_to(voptions, 4, "[0,0]", ef: $eff)) {
+			uid
+		}
+	}`
+	res, err := Parse(Request{
+		Str:       query,
+		Variables: map[string]string{"$eff": "64"},
+	})
+	require.NoError(t, err)
+	require.Equal(t, "similar_to", res.Query[0].Func.Name)
+	require.Len(t, res.Query[0].Func.Args, 4)
+	require.Equal(t, "ef", res.Query[0].Func.Args[2].Value)
+	require.Equal(t, "64", res.Query[0].Func.Args[3].Value)
+}
+
+func TestParseSimilarToRejectsLegacyStringOptionsSyntax(t *testing.T) {
+	query := `{
+		q(func: similar_to(voptions, 4, "[0,0]", "ef=64,distance_threshold=0.45")) {
+			uid
+		}
+	}`
+	_, err := Parse(Request{Str: query})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Expected named parameter in similar_to options")
+}
+
+func TestParseSimilarToUnknownOption(t *testing.T) {
+	query := `{
+		q(func: similar_to(voptions, 4, "[0,0]", foo: 5)) {
+			uid
+		}
+	}`
+	_, err := Parse(Request{Str: query})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Unknown option")
+	require.Contains(t, err.Error(), "foo")
+}
+
+func TestParseSimilarToDuplicateOption(t *testing.T) {
+	query := `{
+		q(func: similar_to(voptions, 4, "[0,0]", ef: 10, ef: 20)) {
+			uid
+		}
+	}`
+	_, err := Parse(Request{Str: query})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Duplicate key")
+	require.Contains(t, err.Error(), "ef")
+}
+
+func TestParseNonSimilarToWithBrace(t *testing.T) {
+	// Braces in non-similar_to functions should be rejected
+	query := `{
+		q(func: eq(name, {value: "test"})) {
+			uid
+		}
+	}`
+	_, err := Parse(Request{Str: query})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Unrecognized character inside a func: U+007B '{'")
 }
 
 func TestParseFilter_Geo1(t *testing.T) {
@@ -2768,6 +2885,10 @@ func TestParseCountAsFunc(t *testing.T) {
 
 }
 
+// Note: This query has two errors: missing ')' after 'friends' AND a stray '}'.
+// After changes to support similar_to's JSON args the lexer emits brace tokens instead
+// of erroring immediately -- causing this to fail on unclosed brackets rather than the
+// specific character error. See TestParseFilter_unbalancedbrac for full explanation.
 func TestParseCountError1(t *testing.T) {
 	query := `{
 		me(func: uid(1)) {
@@ -2779,10 +2900,11 @@ func TestParseCountError1(t *testing.T) {
 `
 	_, err := Parse(Request{Str: query})
 	require.Error(t, err)
-	require.Contains(t, err.Error(),
-		"Unrecognized character inside a func: U+007D '}'")
+	require.Contains(t, err.Error(), "Unclosed Brackets")
 }
 
+// Note: Similar to TestParseCountError1, this has missing ')' and stray '}',
+// now reports structural error instead of character-specific error.
 func TestParseCountError2(t *testing.T) {
 	query := `{
 		me(func: uid(1)) {
@@ -2794,8 +2916,7 @@ func TestParseCountError2(t *testing.T) {
 `
 	_, err := Parse(Request{Str: query})
 	require.Error(t, err)
-	require.Contains(t, err.Error(),
-		"Unrecognized character inside a func: U+007D '}'")
+	require.Contains(t, err.Error(), "Unclosed Brackets")
 }
 
 func TestParseCheckPwd(t *testing.T) {

dql/state.go

@@ -306,6 +306,18 @@ func lexFuncOrArg(l *lex.Lexer) lex.StateFn {
 			l.Emit(itemLeftSquare)
 		case r == rightSquare:
 			l.Emit(itemRightSquare)
+		case r == leftCurl:
+			empty = false
+			l.Emit(itemLeftCurl)
+			// Design decision: Emit brace tokens without affecting ArgDepth tracking.
+			// The parser validates whether braces are legal in context.
+			// Trade-off: Queries with multiple syntax errors (e.g., missing ')' AND stray '}')
+			// will report structural errors (Unclosed Brackets) rather than character-specific
+			// errors. This is acceptable as the query is still rejected with a clear error.
+		case r == rightCurl:
+			l.Emit(itemRightCurl)
+			// Don't decrement ArgDepth for braces; let parser validate context.
+			// See leftCurl case above for full rationale.
 		case r == '#':
 			return lexComment
 		case r == '.':

query/vector/vector_test.go

@@ -417,6 +417,74 @@ func TestVectorIndexRebuildWhenChange(t *testing.T) {
 	require.Greater(t, dur, time.Second*4)
 }
 
+func TestSimilarToOptionsIntegration(t *testing.T) {
+	const pred = "voptions"
+	dropPredicate(pred)
+	t.Cleanup(func() { dropPredicate(pred) })
+
+	setSchema(fmt.Sprintf(vectorSchemaWithIndex, pred, "4", "euclidean"))
+
+	rdf := `<0x1> <voptions> "[0,0]" .
+	<0x2> <voptions> "[1,0]" .
+	<0x3> <voptions> "[2,0]" .
+	<0x4> <voptions> "[5,0]" .`
+	require.NoError(t, addTriplesToCluster(rdf))
+
+	t.Run("ef_override_named_param", func(t *testing.T) {
+		query := `{
+			results(func: similar_to(voptions, 3, "[0,0]", ef: 2)) {
+				uid
+			}
+		}`
+		resp := processQueryNoErr(t, query)
+
+		var result struct {
+			Data struct {
+				Results []struct {
+					UID string `json:"uid"`
+				} `json:"results"`
+			} `json:"data"`
+		}
+		require.NoError(t, json.Unmarshal([]byte(resp), &result))
+		require.Len(t, result.Data.Results, 3)
+
+		expected := map[string]struct{}{"0x1": {}, "0x2": {}, "0x3": {}}
+		for _, r := range result.Data.Results {
+			_, ok := expected[r.UID]
+			require.Truef(t, ok, "unexpected uid %s", r.UID)
+			delete(expected, r.UID)
+		}
+		require.Empty(t, expected)
+	})
+
+	t.Run("distance_threshold_named_param", func(t *testing.T) {
+		query := `{
+			results(func: similar_to(voptions, 4, "[0,0]", distance_threshold: 1.5)) {
+				uid
+			}
+		}`
+		resp := processQueryNoErr(t, query)
+
+		var result struct {
+			Data struct {
+				Results []struct {
+					UID string `json:"uid"`
+				} `json:"results"`
+			} `json:"data"`
+		}
+		require.NoError(t, json.Unmarshal([]byte(resp), &result))
+		require.Len(t, result.Data.Results, 2)
+
+		expected := map[string]struct{}{"0x1": {}, "0x2": {}}
+		for _, r := range result.Data.Results {
+			_, ok := expected[r.UID]
+			require.Truef(t, ok, "unexpected uid %s", r.UID)
+			delete(expected, r.UID)
+		}
+		require.Empty(t, expected)
+	})
+}
+
 func TestVectorInQueryArgument(t *testing.T) {
 	dropPredicate("vtest")
 	setSchema(fmt.Sprintf(vectorSchemaWithIndex, "vtest", "4", "euclidean"))