diff --git a/src/barabadb/query/executor.nim b/src/barabadb/query/executor.nim index 8da52ba..462dc94 100644 --- a/src/barabadb/query/executor.nim +++ b/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") diff --git a/src/barabadb/vector/engine.nim b/src/barabadb/vector/engine.nim index 3246874..5a48146 100644 --- a/src/barabadb/vector/engine.nim +++ b/src/barabadb/vector/engine.nim @@ -258,6 +258,32 @@ proc search*(idx: HNSWIndex, query: Vector, k: int, for i in 0.. 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.. 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"] == "[]" +