diff --git a/src/barabadb/query/executor.nim b/src/barabadb/query/executor.nim index 462dc94..99abfb3 100644 --- a/src/barabadb/query/executor.nim +++ b/src/barabadb/query/executor.nim @@ -649,6 +649,67 @@ proc doHybridSearch(ctx: ExecutionContext, table: string, vecCol: string, textCo if a[1] > b[1]: return -1 elif a[1] < b[1]: return 1 else: return 0) + if sorted.len > k: + sorted = sorted[0.. 0: + let rrfScore = 1.0 / (60.0 + float64(rank + 1)) + vecIdScores[realId] = vecIdScores.getOrDefault(realId, 0.0) + rrfScore + + # FTS search (post-filtering by docId lookup) + let ftsResults = fts.search(ftsIdx, queryText, k * 3) + for rank, res in ftsResults: + let realId = findRealIdByDocId(ctx, table, res.docId) + if realId.len > 0: + # Verify filter on actual row data + var passesFilter = true + if filterCol.len > 0: + passesFilter = false + for row in execScan(ctx, table): + if realIdFromKey(row.getOrDefault("$key", "")) == realId: + if filterCol in row and row[filterCol] == filterVal: + passesFilter = true + break + if passesFilter: + 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) + if sorted.len > k: + sorted = sorted[0.. 0 and col != "$key" and col != "$value": + meta[col] = val vengine.insert(vecIdx, docId, vec, meta) inc count @@ -4449,6 +4528,9 @@ proc executeQueryImpl(ctx: ExecutionContext, astNode: Node, params: seq[WireValu var meta = initTable[string, string]() if "$key" in row: meta["key"] = row["$key"] + for col, val in row: + if col.len > 0 and col != "$key" and col != "$value": + meta[col] = val let fullKey = stmt.ciTarget & "." & row["$key"] var docId: uint64 = 0 for ch in fullKey: diff --git a/tests/test_all.nim b/tests/test_all.nim index ffd8d35..d4724e3 100644 --- a/tests/test_all.nim +++ b/tests/test_all.nim @@ -3920,3 +3920,30 @@ suite "Hybrid RAG Search": check r.success check r.rows[0]["res"] == "[]" + + test "hybrid_search_filtered excludes non-matching rows": + discard qexec.executeQuery(ctx, parse("CREATE TABLE docs5 (id INT PRIMARY KEY, embedding VECTOR(3), content TEXT, tenant_id TEXT)")) + discard qexec.executeQuery(ctx, parse("INSERT INTO docs5 (id, embedding, content, tenant_id) VALUES (1, '[1.0, 0.0, 0.0]', 'quick brown fox', 'tenant-a')")) + discard qexec.executeQuery(ctx, parse("INSERT INTO docs5 (id, embedding, content, tenant_id) VALUES (2, '[1.0, 0.0, 0.0]', 'quick brown fox', 'tenant-b')")) + discard qexec.executeQuery(ctx, parse("INSERT INTO docs5 (id, embedding, content, tenant_id) VALUES (3, '[0.0, 1.0, 0.0]', 'lazy dog sleeps', 'tenant-a')")) + discard qexec.executeQuery(ctx, parse("CREATE INDEX idx_vec5 ON docs5(embedding) USING hnsw")) + discard qexec.executeQuery(ctx, parse("CREATE INDEX idx_fts5 ON docs5(content) USING FTS")) + let r = qexec.executeQuery(ctx, parse("SELECT hybrid_search_filtered('docs5', 'embedding', 'content', 'quick brown fox', '[1.0, 0.0, 0.0]', 10, 'tenant_id', 'tenant-a') AS res")) + check r.success + let arr = parseJson(r.rows[0]["res"]) + # Doc 1 (tenant-a) should be first and highest scored; Doc 2 (tenant-b) must be excluded + check arr[0]["id"].getStr() == "1" + for elem in arr: + check elem["id"].getStr() != "2" + + test "hybrid_search_filtered with empty filter behaves like hybrid_search": + discard qexec.executeQuery(ctx, parse("CREATE TABLE docs6 (id INT PRIMARY KEY, embedding VECTOR(3), content TEXT, tenant_id TEXT)")) + discard qexec.executeQuery(ctx, parse("INSERT INTO docs6 (id, embedding, content, tenant_id) VALUES (1, '[1.0, 0.0, 0.0]', 'quick brown fox', 'tenant-a')")) + discard qexec.executeQuery(ctx, parse("INSERT INTO docs6 (id, embedding, content, tenant_id) VALUES (2, '[1.0, 0.0, 0.0]', 'quick brown fox', 'tenant-b')")) + discard qexec.executeQuery(ctx, parse("CREATE INDEX idx_vec6 ON docs6(embedding) USING hnsw")) + discard qexec.executeQuery(ctx, parse("CREATE INDEX idx_fts6 ON docs6(content) USING FTS")) + let r = qexec.executeQuery(ctx, parse("SELECT hybrid_search_filtered('docs6', '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 == 2 +