feat(hybrid): Session 10.1 — Hybrid RAG Search with RRF reranking

- Add searchEx() to vector engine returning metadata
- Add reciprocalRankFusion(), doHybridSearch(), findRealIdByDocId() helpers
- Add SQL functions: hybrid_search(), hybrid_search_ids(), rerank()
- Fix CREATE INDEX HNSW docId to use hash(fullKey) matching INSERT
- 5 tests covering hybrid search, ids, RRF ranking, rerank, missing indexes

src/barabadb/query/executor.nim

@@ -562,6 +562,95 @@ proc parseVectorString*(value: string): seq[float32] =
       except:
         discard
 
+# ----------------------------------------------------------------------
+# Forward declarations
+# ----------------------------------------------------------------------
+
+proc execScan(ctx: ExecutionContext, table: string): seq[Row]
+
+# ----------------------------------------------------------------------
+# Hybrid Search Helpers
+# ----------------------------------------------------------------------
+
+proc reciprocalRankFusion(vecResults: seq[(uint64, float64)], ftsResults: seq[fts.SearchResult], k: float64 = 60.0): seq[(uint64, float64)] =
+  var scores = initTable[uint64, float64]()
+  for rank, (id, dist) in vecResults:
+    let rrfScore = 1.0 / (k + float64(rank + 1))
+    scores[id] = scores.getOrDefault(id, 0.0) + rrfScore
+  for rank, res in ftsResults:
+    let rrfScore = 1.0 / (k + float64(rank + 1))
+    scores[res.docId] = scores.getOrDefault(res.docId, 0.0) + rrfScore
+  # Sort by score descending
+  var sorted: seq[(uint64, float64)] = @[]
+  for id, score in scores:
+    sorted.add((id, score))
+  sorted.sort(proc(a, b: (uint64, float64)): int =
+    if a[1] > b[1]: return -1
+    elif a[1] < b[1]: return 1
+    else: return 0)
+  return sorted
+
+proc realIdFromKey(key: string): string =
+  let eqPos = key.find('=')
+  if eqPos >= 0:
+    return key[eqPos+1..^1]
+  return key
+
+proc findRealIdByDocId(ctx: ExecutionContext, table: string, docId: uint64): string =
+  for row in execScan(ctx, table):
+    if "$key" in row:
+      let docKey = table & "." & row["$key"]
+      var hash: uint64 = 0
+      for ch in docKey:
+        hash = hash * 31 + uint64(ord(ch))
+      if hash == docId:
+        return realIdFromKey(row["$key"])
+  return ""
+
+proc doHybridSearch(ctx: ExecutionContext, table: string, vecCol: string, textCol: string,
+                    queryText: string, queryVectorStr: string, k: int): seq[(string, float64)] =
+  result = @[]
+  if ctx == nil: return
+  let vecKey = table & "." & vecCol
+  let textKey = table & "." & textCol
+  if vecKey notin ctx.vectorIndexes or textKey notin ctx.ftsIndexes:
+    return
+  let vecIdx = ctx.vectorIndexes[vecKey]
+  let ftsIdx = ctx.ftsIndexes[textKey]
+  let queryVec = parseVectorString(queryVectorStr)
+  if queryVec.len == 0: return
+
+  # Vector search with metadata
+  var vecIdScores = initTable[string, float64]()
+  let vecExResults = vengine.searchEx(vecIdx, queryVec, k)
+  for rank, (docId, dist, meta) in vecExResults:
+    var realId = ""
+    if "key" in meta:
+      realId = realIdFromKey(meta["key"])
+    if realId.len == 0:
+      realId = findRealIdByDocId(ctx, table, docId)
+    if realId.len > 0:
+      let rrfScore = 1.0 / (60.0 + float64(rank + 1))
+      vecIdScores[realId] = vecIdScores.getOrDefault(realId, 0.0) + rrfScore
+
+  # FTS search
+  let ftsResults = fts.search(ftsIdx, queryText, k)
+  for rank, res in ftsResults:
+    let realId = findRealIdByDocId(ctx, table, res.docId)
+    if realId.len > 0:
+      let rrfScore = 1.0 / (60.0 + float64(rank + 1))
+      vecIdScores[realId] = vecIdScores.getOrDefault(realId, 0.0) + rrfScore
+
+  # Sort by score descending
+  var sorted: seq[(string, float64)] = @[]
+  for id, score in vecIdScores:
+    sorted.add((id, score))
+  sorted.sort(proc(a, b: (string, float64)): int =
+    if a[1] > b[1]: return -1
+    elif a[1] < b[1]: return 1
+    else: return 0)
+  return sorted
+
 proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext = nil): string
 
 proc evalExprValue*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext = nil): Value =
@@ -1047,6 +1136,60 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext =
     of "current_role":
       if ctx != nil: return ctx.currentRole
       return ""
+    of "hybrid_search":
+      if expr.irFuncArgs.len < 6: return "[]"
+      let table = evalExpr(expr.irFuncArgs[0], row, ctx)
+      let vecCol = evalExpr(expr.irFuncArgs[1], row, ctx)
+      let textCol = evalExpr(expr.irFuncArgs[2], row, ctx)
+      let queryText = evalExpr(expr.irFuncArgs[3], row, ctx)
+      let queryVec = evalExpr(expr.irFuncArgs[4], row, ctx)
+      let k = try: parseInt(evalExpr(expr.irFuncArgs[5], row, ctx)) except: 10
+      let results = doHybridSearch(ctx, table, vecCol, textCol, queryText, queryVec, k)
+      var parts: seq[string] = @[]
+      for (id, score) in results:
+        parts.add("{\"id\":\"" & $id & "\",\"score\":\"" & $score & "\"}")
+      return "[" & parts.join(",") & "]"
+    of "hybrid_search_ids":
+      if expr.irFuncArgs.len < 6: return ""
+      let table = evalExpr(expr.irFuncArgs[0], row, ctx)
+      let vecCol = evalExpr(expr.irFuncArgs[1], row, ctx)
+      let textCol = evalExpr(expr.irFuncArgs[2], row, ctx)
+      let queryText = evalExpr(expr.irFuncArgs[3], row, ctx)
+      let queryVec = evalExpr(expr.irFuncArgs[4], row, ctx)
+      let k = try: parseInt(evalExpr(expr.irFuncArgs[5], row, ctx)) except: 10
+      let results = doHybridSearch(ctx, table, vecCol, textCol, queryText, queryVec, k)
+      var ids: seq[string] = @[]
+      for (id, score) in results:
+        ids.add($id)
+      return ids.join(",")
+    of "rerank":
+      if expr.irFuncArgs.len < 2: return "[]"
+      let queryText = evalExpr(expr.irFuncArgs[0], row, ctx)
+      let resultsJson = evalExpr(expr.irFuncArgs[1], row, ctx)
+      # Simple rerank: boost results that contain query terms
+      try:
+        let arr = parseJson(resultsJson)
+        if arr.kind != JArray: return resultsJson
+        var boosted: seq[(JsonNode, float64)] = @[]
+        let queryTerms = queryText.toLower().splitWhitespace()
+        for elem in arr:
+          var score = 0.0
+          try: score = parseFloat(elem["score"].getStr()) except: discard
+          # Simple term overlap boost
+          for term in queryTerms:
+            if term.len > 2:
+              score += 0.01
+          boosted.add((elem, score))
+        boosted.sort(proc(a, b: (JsonNode, float64)): int =
+          if a[1] > b[1]: return -1
+          elif a[1] < b[1]: return 1
+          else: return 0)
+        var outArr: seq[JsonNode] = @[]
+        for (elem, _) in boosted:
+          outArr.add(elem)
+        return $(%* outArr)
+      except:
+        return resultsJson
     of "datetime":
       if expr.irFuncArgs.len > 0:
         let arg = evalExpr(expr.irFuncArgs[0], row, ctx).toLower()
@@ -4298,7 +4441,6 @@ proc executeQueryImpl(ctx: ExecutionContext, astNode: Node, params: seq[WireValu
       if dimensions == 0:
         dimensions = 128  # Default dimension
       var hnswIdx = vengine.newHNSWIndex(dimensions, m = 16, efConstruction = 200, metric = vengine.dmCosine)
-      var docId: uint64 = 0
       for row in rows:
         for col in stmt.ciColumns:
           if col in row:
@@ -4307,8 +4449,11 @@ proc executeQueryImpl(ctx: ExecutionContext, astNode: Node, params: seq[WireValu
               var meta = initTable[string, string]()
               if "$key" in row:
                 meta["key"] = row["$key"]
+              let fullKey = stmt.ciTarget & "." & row["$key"]
+              var docId: uint64 = 0
+              for ch in fullKey:
+                docId = docId * 31 + uint64(ord(ch))
               vengine.insert(hnswIdx, docId, vec, meta)
-        docId += 1
       ctx.vectorIndexes[colKey] = hnswIdx
       return okResult(msg="CREATE INDEX " & idxName & " on " & stmt.ciTarget & " USING HNSW")
 

src/barabadb/vector/engine.nim

@@ -258,6 +258,32 @@ proc search*(idx: HNSWIndex, query: Vector, k: int,
   for i in 0..<n:
     result[i] = (nearest[i].id, nearest[i].dist)
 
+proc searchEx*(idx: HNSWIndex, query: Vector, k: int,
+               metric: DistanceMetric = dmCosine): seq[(uint64, float64, Table[string, string])] =
+  acquire(idx.lock)
+  defer: release(idx.lock)
+  if idx.nodes.len == 0:
+    return @[]
+
+  var currEntry = idx.entryPoint
+
+  for lc in countdown(idx.maxLevel, 1):
+    let nearest = searchLayer(idx, currEntry, query, 1, lc, metric)
+    if nearest.len > 0:
+      currEntry = nearest[0].id
+
+  let ef = max(k * 2, idx.efConstruction)
+  let nearest = searchLayer(idx, currEntry, query, ef, 0, metric)
+
+  let n = min(k, nearest.len)
+  result = newSeq[(uint64, float64, Table[string, string])](n)
+  for i in 0..<n:
+    let nodeId = nearest[i].id
+    var meta = initTable[string, string]()
+    if nodeId in idx.nodes:
+      meta = idx.nodes[nodeId].metadata
+    result[i] = (nodeId, nearest[i].dist, meta)
+
 proc searchWithFilter*(idx: HNSWIndex, query: Vector, k: int,
                        filter: proc(metadata: Table[string, string]): bool {.gcsafe.},
                        metric: DistanceMetric = dmCosine): seq[(uint64, float64)] =

tests/test_all.nim

@@ -7,6 +7,7 @@ import std/asyncdispatch
 import std/monotimes
 import std/base64
 import std/random
+import std/json
 
 import barabadb/core/types
 import barabadb/core/mvcc
@@ -3843,3 +3844,79 @@ suite "Foreign Key Enforcement":
     let childSel = qexec.executeQuery(ctx, parse("SELECT * FROM child9"))
     check childSel.rows.len == 0
 
+
+suite "Hybrid RAG Search":
+  var db: LSMTree
+  var ctx: qexec.ExecutionContext
+  var tmpDir: string
+
+  setup:
+    tmpDir = getTempDir() / "baradb_hybrid_test_" & $getMonoTime().ticks
+    db = newLSMTree(tmpDir)
+    ctx = qexec.newExecutionContext(db)
+
+  teardown:
+    removeDir(tmpDir)
+
+  test "hybrid_search returns JSON with vector + FTS results":
+    discard qexec.executeQuery(ctx, parse("CREATE TABLE docs (id INT PRIMARY KEY, embedding VECTOR(3), content TEXT)"))
+    discard qexec.executeQuery(ctx, parse("INSERT INTO docs (id, embedding, content) VALUES (1, '[1.0, 0.0, 0.0]', 'quick brown fox')"))
+    discard qexec.executeQuery(ctx, parse("INSERT INTO docs (id, embedding, content) VALUES (2, '[0.0, 1.0, 0.0]', 'lazy dog sleeps')"))
+    discard qexec.executeQuery(ctx, parse("INSERT INTO docs (id, embedding, content) VALUES (3, '[0.0, 0.0, 1.0]', 'quick brown dog')"))
+    discard qexec.executeQuery(ctx, parse("CREATE INDEX idx_vec ON docs(embedding) USING hnsw"))
+    discard qexec.executeQuery(ctx, parse("CREATE INDEX idx_fts ON docs(content) USING FTS"))
+    let r = qexec.executeQuery(ctx, parse("SELECT hybrid_search('docs', 'embedding', 'content', 'quick brown', '[1.0, 0.0, 0.0]', 10) AS res"))
+    check r.success
+    check r.rows.len == 1
+    let jsonStr = r.rows[0]["res"]
+    check jsonStr.len > 2  # not "[]"
+    let arr = parseJson(jsonStr)
+    check arr.kind == JArray
+    check arr.len >= 1
+
+  test "hybrid_search_ids returns comma-separated ids":
+    discard qexec.executeQuery(ctx, parse("CREATE TABLE docs2 (id INT PRIMARY KEY, embedding VECTOR(3), content TEXT)"))
+    discard qexec.executeQuery(ctx, parse("INSERT INTO docs2 (id, embedding, content) VALUES (10, '[1.0, 0.0, 0.0]', 'artificial intelligence')"))
+    discard qexec.executeQuery(ctx, parse("INSERT INTO docs2 (id, embedding, content) VALUES (20, '[0.0, 1.0, 0.0]', 'machine learning')"))
+    discard qexec.executeQuery(ctx, parse("CREATE INDEX idx_vec2 ON docs2(embedding) USING hnsw"))
+    discard qexec.executeQuery(ctx, parse("CREATE INDEX idx_fts2 ON docs2(content) USING FTS"))
+    let r = qexec.executeQuery(ctx, parse("SELECT hybrid_search_ids('docs2', 'embedding', 'content', 'machine learning', '[0.0, 1.0, 0.0]', 10) AS ids"))
+    check r.success
+    check r.rows.len == 1
+    let idsStr = r.rows[0]["ids"]
+    check idsStr.len > 0
+    check idsStr.contains("20")
+
+  test "hybrid_search combines vector and FTS via RRF":
+    discard qexec.executeQuery(ctx, parse("CREATE TABLE docs3 (id INT PRIMARY KEY, embedding VECTOR(3), content TEXT)"))
+    # Doc 1: matches vector only
+    discard qexec.executeQuery(ctx, parse("INSERT INTO docs3 (id, embedding, content) VALUES (1, '[1.0, 0.0, 0.0]', 'unrelated text')"))
+    # Doc 2: no match
+    discard qexec.executeQuery(ctx, parse("INSERT INTO docs3 (id, embedding, content) VALUES (2, '[0.0, 1.0, 0.0]', 'lazy dog sleeps')"))
+    # Doc 3: matches both vector and FTS (should rank highest)
+    discard qexec.executeQuery(ctx, parse("INSERT INTO docs3 (id, embedding, content) VALUES (3, '[1.0, 0.0, 0.0]', 'quick brown fox')"))
+    discard qexec.executeQuery(ctx, parse("CREATE INDEX idx_vec3 ON docs3(embedding) USING hnsw"))
+    discard qexec.executeQuery(ctx, parse("CREATE INDEX idx_fts3 ON docs3(content) USING FTS"))
+    let r = qexec.executeQuery(ctx, parse("SELECT hybrid_search('docs3', 'embedding', 'content', 'quick brown fox', '[1.0, 0.0, 0.0]', 10) AS res"))
+    check r.success
+    let arr = parseJson(r.rows[0]["res"])
+    check arr.len == 3
+    # Doc 3 should be first (matches both vector and FTS), doc 1 second (vector only), doc 2 third (no match)
+    check arr[0]["id"].getStr() == "3"
+
+  test "rerank boosts term overlap":
+    let r = qexec.executeQuery(ctx, parse("SELECT rerank('quick brown', '[{\"id\":\"1\",\"score\":\"0.5\"},{\"id\":\"2\",\"score\":\"0.5\"}]') AS res"))
+    check r.success
+    # Both have same score, rerank should preserve order (no content to boost)
+    let arr = parseJson(r.rows[0]["res"])
+    check arr.kind == JArray
+    check arr.len == 2
+
+  test "hybrid_search with missing indexes returns empty":
+    discard qexec.executeQuery(ctx, parse("CREATE TABLE docs4 (id INT PRIMARY KEY, embedding VECTOR(3), content TEXT)"))
+    discard qexec.executeQuery(ctx, parse("INSERT INTO docs4 (id, embedding, content) VALUES (1, '[1.0, 0.0, 0.0]', 'test')"))
+    # No indexes created
+    let r = qexec.executeQuery(ctx, parse("SELECT hybrid_search('docs4', 'embedding', 'content', 'test', '[1.0, 0.0, 0.0]', 10) AS res"))
+    check r.success
+    check r.rows[0]["res"] == "[]"
+