From 0327a1f143d5deb1fcd67ee92cb6b7a242f2ce71 Mon Sep 17 00:00:00 2001 From: Adam Hendel Date: Thu, 28 May 2026 12:58:43 -0500 Subject: [PATCH 1/6] first cut at bm25 hybrid --- Cargo.lock | 443 +++++++++++++++++++++++++++++++++++- core/src/query.rs | 109 ++++++++- server/Cargo.toml | 1 + server/src/app_state.rs | 46 +++- server/src/bm25.rs | 195 ++++++++++++++++ server/src/lib.rs | 1 + server/src/routes/search.rs | 89 ++++++-- server/src/routes/table.rs | 33 ++- 8 files changed, 898 insertions(+), 19 deletions(-) create mode 100644 server/src/bm25.rs diff --git a/Cargo.lock b/Cargo.lock index a6783d7..a60ade1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -322,6 +322,15 @@ dependencies = [ "derive_arbitrary", ] +[[package]] +name = "arc-swap" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" +dependencies = [ + "rustversion", +] + [[package]] name = "array-init" version = "2.1.0" @@ -459,6 +468,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "bitpacking" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96a7139abd3d9cebf8cd6f920a389cf3dc9576172e32f4563f188cae3c3eb019" +dependencies = [ + "crunchy", +] + [[package]] name = "bitvec" version = "1.0.1" @@ -606,6 +624,12 @@ dependencies = [ "shlex", ] +[[package]] +name = "census" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f4c707c6a209cbe82d10abd08e1ea8995e9ea937d2550646e02798948992be0" + [[package]] name = "cfg-if" version = "1.0.4" @@ -813,6 +837,34 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-queue" version = "0.3.12" @@ -828,6 +880,12 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "crypto-common" version = "0.1.7" @@ -874,6 +932,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" dependencies = [ "powerfmt", + "serde_core", ] [[package]] @@ -962,6 +1021,12 @@ version = "0.15.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" +[[package]] +name = "downcast-rs" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75b325c5dbd37f80359721ad39aca5a29fb04c89279657cffdda8736d0c0b9d2" + [[package]] name = "dunce" version = "1.0.5" @@ -1079,6 +1144,12 @@ dependencies = [ "regex-syntax", ] +[[package]] +name = "fastdivide" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afc2bd4d5a73106dd53d10d73d3401c2f32730ba2c0b93ddb888a8983680471" + [[package]] name = "fastrand" version = "2.4.1" @@ -1148,6 +1219,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs4" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7e180ac76c23b45e767bd7ae9579bc0bb458618c4bc71835926e098e61d15f8" +dependencies = [ + "rustix 0.38.44", + "windows-sys 0.52.0", +] + [[package]] name = "fs_extra" version = "1.3.0" @@ -1389,6 +1470,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "hex" version = "0.4.3" @@ -1431,6 +1518,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "htmlescape" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9025058dae765dee5070ec375f591e2ba14638c63feff74f13805a72e523163" + [[package]] name = "http" version = "0.2.12" @@ -1773,6 +1866,18 @@ dependencies = [ "serde_core", ] +[[package]] +name = "instant" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0242819d153cba4b4b05a5a8f2a7e9bbf97b6055b2a002b395c96b5ff3c0222" +dependencies = [ + "cfg-if", + "js-sys", + "wasm-bindgen", + "web-sys", +] + [[package]] name = "ipnet" version = "2.12.0" @@ -1795,6 +1900,15 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -1891,6 +2005,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" +[[package]] +name = "levenshtein_automata" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c2cdeb66e45e9f36bfad5bbdb4d2384e70936afbee843c6f6543f0c551ebb25" + [[package]] name = "libc" version = "0.2.186" @@ -1925,6 +2045,12 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -1969,12 +2095,27 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "lru" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38" +dependencies = [ + "hashbrown 0.15.5", +] + [[package]] name = "lru-slab" version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" +[[package]] +name = "lz4_flex" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "373f5eceeeab7925e0c1098212f2fbc4d416adec9d35051a6ab251e824c1854a" + [[package]] name = "md-5" version = "0.10.6" @@ -2001,12 +2142,31 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "490cc448043f947bae3cbee9c203358d62dbee0db12107a74be5c30ccfd09771" +[[package]] +name = "measure_time" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbefd235b0aadd181626f281e1d684e116972988c14c264e42069d5e8a5775cc" +dependencies = [ + "instant", + "log", +] + [[package]] name = "memchr" version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +[[package]] +name = "memmap2" +version = "0.9.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714098028fe011992e1c3962653c96b2d578c4b4bce9036e15ff220319b1e0e3" +dependencies = [ + "libc", +] + [[package]] name = "mime" version = "0.3.17" @@ -2023,6 +2183,12 @@ dependencies = [ "unicase", ] +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -2045,6 +2211,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "murmurhash32" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2195bf6aa996a481483b29d62a7663eed3fe39600c460e323f8ff41e90bdd89b" + [[package]] name = "native-tls" version = "0.2.18" @@ -2062,6 +2234,16 @@ dependencies = [ "tempfile", ] +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -2123,6 +2305,16 @@ dependencies = [ "libm", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "num_enum" version = "0.7.6" @@ -2199,6 +2391,12 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "oneshot" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "269bca4c2591a28585d6bf10d9ed0332b7d76900a1b02bec41bdc3a2cdcda107" + [[package]] name = "openssl" version = "0.10.79" @@ -2242,6 +2440,15 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "ownedbytes" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3a059efb063b8f425b948e042e6b9bd85edfe60e913630ed727b23e2dfcc558" +dependencies = [ + "stable_deref_trait", +] + [[package]] name = "parking" version = "2.2.1" @@ -2699,6 +2906,36 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" +[[package]] +name = "rand_distr" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31" +dependencies = [ + "num-traits", + "rand 0.8.6", +] + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2944,6 +3181,16 @@ dependencies = [ "walkdir", ] +[[package]] +name = "rust-stemmers" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e46a2036019fdb888131db7a4c847a1063a7493f971ed94ea82c67eada63ca54" +dependencies = [ + "serde", + "serde_derive", +] + [[package]] name = "rust_decimal" version = "1.41.0" @@ -2983,6 +3230,19 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags 2.11.1", + "errno", + "libc", + "linux-raw-sys 0.4.15", + "windows-sys 0.52.0", +] + [[package]] name = "rustix" version = "1.1.4" @@ -2992,7 +3252,7 @@ dependencies = [ "bitflags 2.11.1", "errno", "libc", - "linux-raw-sys", + "linux-raw-sys 0.12.1", "windows-sys 0.61.2", ] @@ -3266,6 +3526,15 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e" +[[package]] +name = "sketches-ddsketch" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85636c14b73d81f541e525f585c0a2109e6744e1565b5c1668e31c70c10ed65c" +dependencies = [ + "serde", +] + [[package]] name = "slab" version = "0.4.12" @@ -3650,6 +3919,147 @@ dependencies = [ "libc", ] +[[package]] +name = "tantivy" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96599ea6fccd844fc833fed21d2eecac2e6a7c1afd9e044057391d78b1feb141" +dependencies = [ + "aho-corasick", + "arc-swap", + "base64 0.22.1", + "bitpacking", + "byteorder", + "census", + "crc32fast", + "crossbeam-channel", + "downcast-rs", + "fastdivide", + "fnv", + "fs4", + "htmlescape", + "itertools", + "levenshtein_automata", + "log", + "lru", + "lz4_flex", + "measure_time", + "memmap2", + "num_cpus", + "once_cell", + "oneshot", + "rayon", + "regex", + "rust-stemmers", + "rustc-hash 1.1.0", + "serde", + "serde_json", + "sketches-ddsketch", + "smallvec", + "tantivy-bitpacker", + "tantivy-columnar", + "tantivy-common", + "tantivy-fst", + "tantivy-query-grammar", + "tantivy-stacker", + "tantivy-tokenizer-api", + "tempfile", + "thiserror 1.0.69", + "time", + "uuid", + "winapi", +] + +[[package]] +name = "tantivy-bitpacker" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "284899c2325d6832203ac6ff5891b297fc5239c3dc754c5bc1977855b23c10df" +dependencies = [ + "bitpacking", +] + +[[package]] +name = "tantivy-columnar" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12722224ffbe346c7fec3275c699e508fd0d4710e629e933d5736ec524a1f44e" +dependencies = [ + "downcast-rs", + "fastdivide", + "itertools", + "serde", + "tantivy-bitpacker", + "tantivy-common", + "tantivy-sstable", + "tantivy-stacker", +] + +[[package]] +name = "tantivy-common" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8019e3cabcfd20a1380b491e13ff42f57bb38bf97c3d5fa5c07e50816e0621f4" +dependencies = [ + "async-trait", + "byteorder", + "ownedbytes", + "serde", + "time", +] + +[[package]] +name = "tantivy-fst" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d60769b80ad7953d8a7b2c70cdfe722bbcdcac6bccc8ac934c40c034d866fc18" +dependencies = [ + "byteorder", + "regex-syntax", + "utf8-ranges", +] + +[[package]] +name = "tantivy-query-grammar" +version = "0.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "847434d4af57b32e309f4ab1b4f1707a6c566656264caa427ff4285c4d9d0b82" +dependencies = [ + "nom", +] + +[[package]] +name = "tantivy-sstable" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c69578242e8e9fc989119f522ba5b49a38ac20f576fc778035b96cc94f41f98e" +dependencies = [ + "tantivy-bitpacker", + "tantivy-common", + "tantivy-fst", + "zstd", +] + +[[package]] +name = "tantivy-stacker" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56d6ff5591fc332739b3ce7035b57995a3ce29a93ffd6012660e0949c956ea8" +dependencies = [ + "murmurhash32", + "rand_distr", + "tantivy-common", +] + +[[package]] +name = "tantivy-tokenizer-api" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a0dcade25819a89cfe6f17d932c9cedff11989936bf6dd4f336d50392053b04" +dependencies = [ + "serde", +] + [[package]] name = "tap" version = "1.0.1" @@ -3665,7 +4075,7 @@ dependencies = [ "fastrand", "getrandom 0.4.2", "once_cell", - "rustix", + "rustix 1.1.4", "windows-sys 0.61.2", ] @@ -4112,6 +4522,12 @@ dependencies = [ "serde", ] +[[package]] +name = "utf8-ranges" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fcfc827f90e53a02eaef5e535ee14266c1d569214c6aa70133a624d8a3164ba" + [[package]] name = "utf8_iter" version = "1.0.4" @@ -4299,6 +4715,7 @@ dependencies = [ "serde_json", "sqlparser", "sqlx", + "tantivy", "thiserror 2.0.18", "tiktoken-rs", "tokio", @@ -4554,6 +4971,22 @@ dependencies = [ "web-sys", ] +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + [[package]] name = "winapi-util" version = "0.1.11" @@ -4563,6 +4996,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + [[package]] name = "windows-core" version = "0.62.2" diff --git a/core/src/query.rs b/core/src/query.rs index fdea0c0..e4f81a8 100644 --- a/core/src/query.rs +++ b/core/src/query.rs @@ -699,9 +699,9 @@ pub fn join_table_cosine_similarity( ) } -fn build_where_filter(filters: &BTreeMap) -> String { +fn build_where_filter_with_start(filters: &BTreeMap, start: i16) -> String { let mut where_filter = "WHERE 1=1".to_string(); - for (bind_value_counter, (column, filter_value)) in (3_i16..).zip(filters.iter()) { + for (bind_value_counter, (column, filter_value)) in (start..).zip(filters.iter()) { let operator = filter_value.operator.to_sql(); let filt = format!(" AND t0.\"{column}\" {operator} ${bind_value_counter}"); where_filter.push_str(&filt); @@ -709,6 +709,10 @@ fn build_where_filter(filters: &BTreeMap) -> String { where_filter } +fn build_where_filter(filters: &BTreeMap) -> String { + build_where_filter_with_start(filters, 3) +} + /// Generates the core hybrid search SELECT that returns raw table rows. /// `$1::vector` and `$2` are sqlx bind parameter placeholders for the embedding and query text. #[allow(clippy::too_many_arguments)] @@ -855,6 +859,107 @@ pub fn hybrid_search_query_rows( &where_filter, ) } +/// Generates the three-arm hybrid search SQL (semantic + FTS + BM25). +/// +/// Bind parameters: +/// $1 - embedding vector +/// $2 - FTS query text +/// $3 - BM25 ranked PKs as text[] (ordered by Tantivy BM25 score) +/// $4+ - filter values +/// +/// The BM25 CTE joins the `$3` array back to the source table so the PK +/// type is resolved correctly regardless of whether it is int, uuid, or text. +#[allow(clippy::too_many_arguments)] +pub fn hybrid_search_query_with_bm25( + job_name: &str, + src_schema: &str, + src_table: &str, + join_key: &str, + return_columns: &[String], + window_size: i32, + limit: i32, + rrf_k: f32, + semantic_weight: f32, + fts_weight: f32, + bm25_weight: f32, + filters: &BTreeMap, +) -> String { + let cols = return_columns + .iter() + .map(|s| format!("t0.{s}")) + .collect::>() + .join(","); + // filters start at $4 because $3 is the BM25 PKs array + let where_filter = build_where_filter_with_start(filters, 4); + + let inner = format!( + " + WITH bm25_cte AS ( + SELECT t0.{join_key}, b.bm25_rank + FROM ( + SELECT val AS pk_text, ord::int AS bm25_rank + FROM unnest($3::text[]) WITH ORDINALITY AS t(val, ord) + ) b + JOIN {src_schema}.{src_table} t0 ON t0.{join_key}::text = b.pk_text + ) + SELECT {cols}, t.rrf_score, t.semantic_rank, t.fts_rank, t.bm25_rank, t.similarity_score + FROM ( + SELECT + COALESCE(s.{join_key}, f.{join_key}, bm25.{join_key}) AS {join_key}, + s.semantic_rank, + s.similarity_score, + f.fts_rank, + bm25.bm25_rank, + ( + COALESCE({semantic_weight}::float / ({rrf_k} + s.semantic_rank), 0) + + COALESCE({fts_weight}::float / ({rrf_k} + f.fts_rank), 0) + + COALESCE({bm25_weight}::float / ({rrf_k} + bm25.bm25_rank), 0) + ) AS rrf_score + FROM ( + SELECT + {join_key}, + distance, + ROW_NUMBER() OVER (ORDER BY distance) AS semantic_rank, + 1 - distance AS similarity_score + FROM ( + SELECT + {join_key}, + embeddings <=> $1::vector AS distance + FROM vectorize._embeddings_{job_name} + ) sub + ORDER BY distance + LIMIT {window_size} + ) s + FULL OUTER JOIN ( + SELECT + {join_key}, + ROW_NUMBER() OVER (ORDER BY ts_rank_cd(search_tokens, query) DESC) AS fts_rank + FROM vectorize._search_tokens_{job_name}, + to_tsquery('english', + NULLIF( + replace(plainto_tsquery('english', $2)::text, ' & ', ' | '), + '' + ) + ) AS query + WHERE search_tokens @@ query + ORDER BY ts_rank_cd(search_tokens, query) DESC + LIMIT {window_size} + ) f ON s.{join_key} = f.{join_key} + FULL OUTER JOIN bm25_cte bm25 + ON COALESCE(s.{join_key}, f.{join_key}) = bm25.{join_key} + ) t + INNER JOIN {src_schema}.{src_table} t0 ON t0.{join_key} = t.{join_key} + {where_filter} + ORDER BY t.rrf_score DESC + LIMIT {limit}" + ); + + format!( + "SELECT to_jsonb(t) as results FROM ({inner} + ) t" + ) +} + #[cfg(test)] mod tests { use super::*; diff --git a/server/Cargo.toml b/server/Cargo.toml index 5efbcc6..3811512 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -55,6 +55,7 @@ url = "2.2" utoipa = { version = "4", features = ["actix_extras", "chrono", "uuid"] } utoipa-swagger-ui = { version = "7", features = ["actix-web"] } uuid = { version = "1.16.0", features = ["v4", "fast-rng", "macro-diagnostics", "serde"] } +tantivy = "0.22" [dev-dependencies] pgvector = { version = "0.4.1", features = ["postgres", "sqlx"] } diff --git a/server/src/app_state.rs b/server/src/app_state.rs index ed7c7bd..0b7f46b 100644 --- a/server/src/app_state.rs +++ b/server/src/app_state.rs @@ -1,11 +1,12 @@ use std::collections::HashMap; use std::sync::Arc; -use tokio::sync::RwLock; +use tokio::sync::{Mutex, RwLock}; use tracing::error; use vectorize_core::config::Config; use vectorize_core::types::VectorizeJob; use vectorize_worker::WorkerHealth; +use crate::bm25::BM25Index; use crate::cache; #[derive(Debug, thiserror::Error)] @@ -27,6 +28,8 @@ pub struct AppState { pub job_cache: Arc>>, /// worker health monitoring data pub worker_health: Arc>, + /// in-memory BM25 indexes keyed by job_name; rebuilt from source on startup + pub bm25_indexes: Arc>>>>, } impl AppState { @@ -66,12 +69,53 @@ impl AppState { last_error: None, })); + // Create empty BM25 index map, then kick off background population for + // every existing job so the server stays non-blocking at startup. + let bm25_indexes: Arc>>>> = + Arc::new(RwLock::new(HashMap::new())); + + { + let jobs: Vec = { + let cache = job_cache.read().await; + cache.values().cloned().collect() + }; + for job in jobs { + match BM25Index::new() { + Ok(idx) => { + let idx = Arc::new(Mutex::new(idx)); + bm25_indexes + .write() + .await + .insert(job.job_name.clone(), idx.clone()); + let pool = db_pool.clone(); + tokio::spawn(async move { + crate::bm25::populate_bm25_index(&pool, &job, idx).await; + }); + } + Err(e) => { + tracing::warn!("Failed to create BM25 index for job {}: {e}", job.job_name); + } + } + } + } + + // Background sync: pick up changed documents every 30 s. + { + let pool = db_pool.clone(); + let indexes = bm25_indexes.clone(); + let cache = job_cache.clone(); + tokio::spawn(async move { + crate::bm25::start_bm25_sync_task(pool, indexes, cache).await; + }); + } + Ok(AppState { config, db_pool, cache_pool, job_cache, worker_health, + bm25_indexes, }) } diff --git a/server/src/bm25.rs b/server/src/bm25.rs new file mode 100644 index 0000000..8946167 --- /dev/null +++ b/server/src/bm25.rs @@ -0,0 +1,195 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use anyhow::Result; +use sqlx::Row; +use tantivy::collector::TopDocs; +use tantivy::query::QueryParser; +use tantivy::schema::{Field, STORED, STRING, Schema, TEXT}; +use tantivy::{Index, IndexReader, IndexWriter, ReloadPolicy, TantivyDocument, Term}; +use tokio::sync::{Mutex, RwLock}; +use vectorize_core::types::VectorizeJob; + +pub struct BM25Index { + index: Index, + writer: IndexWriter, + reader: IndexReader, + pk_field: Field, + body_field: Field, +} + +impl BM25Index { + pub fn new() -> Result { + let mut schema_builder = Schema::builder(); + let pk_field = schema_builder.add_text_field("pk", STRING | STORED); + let body_field = schema_builder.add_text_field("body", TEXT); + let schema = schema_builder.build(); + + let index = Index::create_in_ram(schema); + let writer = index.writer(50_000_000)?; + let reader = index + .reader_builder() + .reload_policy(ReloadPolicy::Manual) + .try_into()?; + + Ok(BM25Index { + index, + writer, + reader, + pk_field, + body_field, + }) + } + + /// Delete-then-add each record so repeated calls act as upserts. + pub fn upsert_documents(&mut self, records: &[(String, String)]) -> Result<()> { + for (pk, text) in records { + self.writer + .delete_term(Term::from_field_text(self.pk_field, pk)); + let mut doc = TantivyDocument::default(); + doc.add_text(self.pk_field, pk); + doc.add_text(self.body_field, text); + self.writer.add_document(doc)?; + } + self.writer.commit()?; + self.reader.reload()?; + Ok(()) + } + + /// Returns PKs ordered by BM25 score descending, up to `limit`. + /// Uses lenient parsing so malformed user queries never return an error. + pub fn search(&self, query_text: &str, limit: usize) -> Result> { + let searcher = self.reader.searcher(); + let query_parser = QueryParser::for_index(&self.index, vec![self.body_field]); + let (query, _errors) = query_parser.parse_query_lenient(query_text); + let top_docs = searcher.search(&query, &TopDocs::with_limit(limit))?; + + let mut results = Vec::with_capacity(top_docs.len()); + for (_score, doc_address) in top_docs { + let doc: TantivyDocument = searcher.doc(doc_address)?; + if let Some(tantivy::schema::OwnedValue::Str(pk)) = doc.get_first(self.pk_field) { + results.push(pk.clone()); + } + } + Ok(results) + } +} + +/// Fetch all rows from the source table and index them into `index`. +/// Spawned as a background task after job creation so startup is non-blocking. +pub async fn populate_bm25_index( + db_pool: &sqlx::PgPool, + job: &VectorizeJob, + index: Arc>, +) { + let text_cols = job + .src_columns + .iter() + .map(|c| format!("COALESCE({c}, '')")) + .collect::>() + .join(" || ' ' || "); + + let query = format!( + "SELECT {pk}::text as pk, {text_cols} as document_text FROM {schema}.{table}", + pk = job.primary_key, + schema = job.src_schema, + table = job.src_table, + ); + + match sqlx::query(&query).fetch_all(db_pool).await { + Ok(rows) => { + let records: Vec<(String, String)> = rows + .iter() + .map(|r| (r.get("pk"), r.get("document_text"))) + .collect(); + let mut idx = index.lock().await; + if let Err(e) = idx.upsert_documents(&records) { + tracing::error!( + "BM25 initial population failed for job {}: {e}", + job.job_name + ); + } else { + tracing::info!( + "BM25 index populated for job {} ({} docs)", + job.job_name, + records.len() + ); + } + } + Err(e) => { + tracing::error!("BM25 population query failed for job {}: {e}", job.job_name); + } + } +} + +/// Background task that polls `_search_tokens_{job}` for rows updated since the +/// last sync and re-indexes changed documents into the in-memory Tantivy index. +pub async fn start_bm25_sync_task( + db_pool: sqlx::PgPool, + bm25_indexes: Arc>>>>, + job_cache: Arc>>, +) { + let mut last_synced: HashMap> = HashMap::new(); + + loop { + tokio::time::sleep(tokio::time::Duration::from_secs(30)).await; + + let jobs: Vec = { + let cache = job_cache.read().await; + cache.values().cloned().collect() + }; + + for job in &jobs { + let since = last_synced + .get(&job.job_name) + .copied() + .unwrap_or(chrono::DateTime::::MIN_UTC); + + let text_cols = job + .src_columns + .iter() + .map(|c| format!("COALESCE(t0.{c}, '')")) + .collect::>() + .join(" || ' ' || "); + + let query = format!( + "SELECT t0.{pk}::text as pk, {text_cols} as document_text + FROM vectorize._search_tokens_{job_name} s + JOIN {schema}.{table} t0 ON t0.{pk} = s.{pk} + WHERE s.updated_at > $1", + pk = job.primary_key, + job_name = job.job_name, + schema = job.src_schema, + table = job.src_table, + ); + + match sqlx::query(&query).bind(since).fetch_all(&db_pool).await { + Ok(rows) if !rows.is_empty() => { + let records: Vec<(String, String)> = rows + .iter() + .map(|r| (r.get("pk"), r.get("document_text"))) + .collect(); + + let indexes = bm25_indexes.read().await; + if let Some(idx) = indexes.get(&job.job_name) { + let mut idx = idx.lock().await; + if let Err(e) = idx.upsert_documents(&records) { + tracing::warn!("BM25 sync failed for job {}: {e}", job.job_name); + } else { + tracing::debug!( + "BM25 synced {} docs for job {}", + records.len(), + job.job_name + ); + } + } + last_synced.insert(job.job_name.clone(), chrono::Utc::now()); + } + Ok(_) => {} + Err(e) => { + tracing::warn!("BM25 sync query failed for job {}: {e}", job.job_name); + } + } + } + } +} diff --git a/server/src/lib.rs b/server/src/lib.rs index ed541e7..73494a2 100644 --- a/server/src/lib.rs +++ b/server/src/lib.rs @@ -1,4 +1,5 @@ pub mod app_state; +pub mod bm25; pub mod cache; pub mod errors; pub mod routes; diff --git a/server/src/routes/search.rs b/server/src/routes/search.rs index 2d89057..999d89b 100644 --- a/server/src/routes/search.rs +++ b/server/src/routes/search.rs @@ -26,6 +26,9 @@ pub struct SearchRequest { pub semantic_wt: f32, #[serde(default = "default_fts_wt")] pub fts_wt: f32, + /// Weight for BM25 (Tantivy) search. Defaults to 0.0 (disabled). + #[serde(default)] + pub bm25_wt: f32, #[serde(flatten, default)] pub filters: BTreeMap, } @@ -45,6 +48,9 @@ pub struct SearchRequestPOST { pub semantic_wt: f32, #[serde(default = "default_fts_wt")] pub fts_wt: f32, + /// Weight for BM25 (Tantivy) search. Defaults to 0.0 (disabled). + #[serde(default)] + pub bm25_wt: f32, pub filters: BTreeMap, } @@ -58,6 +64,7 @@ impl From for SearchRequest { rrf_k: request.rrf_k, semantic_wt: request.semantic_wt, fts_wt: request.fts_wt, + bm25_wt: request.bm25_wt, filters: request.filters, } } @@ -185,24 +192,80 @@ async fn search_internal( let embedding_request = prepare_generic_embedding_request(&vectorizejob.model, &[input]); let embeddings = provider.generate_embedding(&embedding_request).await?; - let q = query::hybrid_search_query( - &payload.job_name, - &vectorizejob.src_schema, - &vectorizejob.src_table, - &vectorizejob.primary_key, - &["*".to_string()], - payload.window_size, - payload.limit, - payload.rrf_k, - payload.semantic_wt, - payload.fts_wt, - &payload.filters, - ); + // Run Tantivy BM25 search when requested. + let bm25_pks: Option> = if payload.bm25_wt > 0.0 { + let indexes = app_state.bm25_indexes.read().await; + match indexes.get(&payload.job_name) { + Some(idx) => { + let idx = idx.lock().await; + match idx.search(&payload.query, payload.window_size as usize) { + Ok(pks) => Some(pks), + Err(e) => { + tracing::warn!( + "BM25 search failed for job {}, falling back: {e}", + payload.job_name + ); + None + } + } + } + None => { + tracing::warn!( + "BM25 index not found for job {}, ignoring bm25_wt", + payload.job_name + ); + None + } + } + } else { + None + }; + + let (q, use_bm25) = match &bm25_pks { + Some(_) => ( + query::hybrid_search_query_with_bm25( + &payload.job_name, + &vectorizejob.src_schema, + &vectorizejob.src_table, + &vectorizejob.primary_key, + &["*".to_string()], + payload.window_size, + payload.limit, + payload.rrf_k, + payload.semantic_wt, + payload.fts_wt, + payload.bm25_wt, + &payload.filters, + ), + true, + ), + None => ( + query::hybrid_search_query( + &payload.job_name, + &vectorizejob.src_schema, + &vectorizejob.src_table, + &vectorizejob.primary_key, + &["*".to_string()], + payload.window_size, + payload.limit, + payload.rrf_k, + payload.semantic_wt, + payload.fts_wt, + &payload.filters, + ), + false, + ), + }; let mut prepared_query = sqlx::query(&q) .bind(&embeddings.embeddings[0]) .bind(&payload.query); + // $3 is the BM25 PKs array when BM25 is active; filters follow after. + if use_bm25 { + prepared_query = prepared_query.bind(bm25_pks.as_ref().unwrap()); + } + // Bind filter values for value in payload.filters.values() { prepared_query = match &value.value { diff --git a/server/src/routes/table.rs b/server/src/routes/table.rs index a9c3276..478db42 100644 --- a/server/src/routes/table.rs +++ b/server/src/routes/table.rs @@ -1,7 +1,11 @@ +use std::sync::Arc; + use crate::app_state::AppState; +use crate::bm25::BM25Index; use crate::errors::ServerError; use actix_web::{HttpResponse, delete, post, web}; use serde::{Deserialize, Serialize}; +use tokio::sync::Mutex; use utoipa::ToSchema; use uuid::Uuid; use vectorize_core::init::{self, get_column_datatype}; @@ -56,6 +60,29 @@ pub async fn table( job_cache.insert(payload.job_name.clone(), payload.clone()); } + // Create a BM25 index for this job and populate it in the background. + match BM25Index::new() { + Ok(idx) => { + let idx = Arc::new(Mutex::new(idx)); + app_state + .bm25_indexes + .write() + .await + .insert(payload.job_name.clone(), idx.clone()); + let pool = app_state.db_pool.clone(); + let job = payload.clone(); + tokio::spawn(async move { + crate::bm25::populate_bm25_index(&pool, &job, idx).await; + }); + } + Err(e) => { + tracing::warn!( + "Failed to create BM25 index for job {}: {e}", + payload.job_name + ); + } + } + let resp = JobResponse { id: job_id }; Ok(HttpResponse::Ok().json(resp)) } @@ -95,11 +122,15 @@ pub async fn delete_table( _ => ServerError::from(e), })?; - // Remove from cache + // Remove from job cache and BM25 index map { let mut job_cache = app_state.job_cache.write().await; job_cache.remove(&job_name); } + { + let mut indexes = app_state.bm25_indexes.write().await; + indexes.remove(&job_name); + } let resp = DeleteJobResponse { job_name: job_name.clone(), From 7bf911c7c9dbc1d71076cf1e0cc8f600ed61d2d1 Mon Sep 17 00:00:00 2001 From: Adam Hendel Date: Thu, 28 May 2026 13:13:43 -0500 Subject: [PATCH 2/6] handle jo-job case --- server/src/app_state.rs | 3 +++ server/src/bm25.rs | 4 ++++ 2 files changed, 7 insertions(+) diff --git a/server/src/app_state.rs b/server/src/app_state.rs index 0b7f46b..af7a252 100644 --- a/server/src/app_state.rs +++ b/server/src/app_state.rs @@ -80,6 +80,9 @@ impl AppState { cache.values().cloned().collect() }; for job in jobs { + if job.job_name.is_empty() { + continue; + } match BM25Index::new() { Ok(idx) => { let idx = Arc::new(Mutex::new(idx)); diff --git a/server/src/bm25.rs b/server/src/bm25.rs index 8946167..9cb5f34 100644 --- a/server/src/bm25.rs +++ b/server/src/bm25.rs @@ -140,6 +140,10 @@ pub async fn start_bm25_sync_task( }; for job in &jobs { + if job.job_name.is_empty() { + continue; + } + let since = last_synced .get(&job.job_name) .copied() From 98e30ef73447853bb71ba1667facded22804d3f1 Mon Sep 17 00:00:00 2001 From: Adam Hendel Date: Sat, 11 Jul 2026 10:38:44 -0500 Subject: [PATCH 3/6] bm25 not enabled by default --- core/src/cache.rs | 4 +-- core/src/init.rs | 12 ++++++--- core/src/query.rs | 10 +++++++- core/src/types.rs | 4 +++ server/src/app_state.rs | 9 +++++-- server/src/bm25.rs | 6 ++++- server/src/routes/table.rs | 51 +++++++++++++++++++++++--------------- 7 files changed, 67 insertions(+), 29 deletions(-) diff --git a/core/src/cache.rs b/core/src/cache.rs index b0b539c..67231e2 100644 --- a/core/src/cache.rs +++ b/core/src/cache.rs @@ -103,7 +103,7 @@ pub async fn refresh_job_cache( job_cache: &Arc>>, ) -> Result<(), Box> { let all_jobs: Vec = sqlx::query_as( - "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model FROM vectorize.job", + "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model, bm25_enabled FROM vectorize.job", ) .fetch_all(db_pool) .await?; @@ -129,7 +129,7 @@ pub async fn load_initial_job_cache( pool: &sqlx::PgPool, ) -> Result, sqlx::Error> { let all_jobs: Vec = sqlx::query_as( - "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model FROM vectorize.job", + "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model, bm25_enabled FROM vectorize.job", ) .fetch_all(pool) .await?; diff --git a/core/src/init.rs b/core/src/init.rs index 1b20ca3..fdd1206 100644 --- a/core/src/init.rs +++ b/core/src/init.rs @@ -74,6 +74,10 @@ async fn vectorize_schema_exists(pool: &PgPool) -> Result { pub async fn init_vectorize(pool: &PgPool) -> Result<(), VectorizeError> { if vectorize_schema_exists(pool).await? { log::info!("vectorize schema already exists, skipping initialization."); + // Backwards-compatible upgrade for installs predating the bm25_enabled column. + sqlx::query(&query::alter_vectorize_table_add_bm25_enabled()) + .execute(pool) + .await?; return Ok(()); } else { // these statements are critical, so we fail if they error @@ -126,15 +130,16 @@ pub async fn initialize_job( // create the job record let mut tx = pool.begin().await?; let job_id: Uuid = sqlx::query_scalar(" - INSERT INTO vectorize.job (job_name, src_schema, src_table, src_columns, primary_key, update_time_col, model) - VALUES ($1, $2, $3, $4, $5, $6, $7) + INSERT INTO vectorize.job (job_name, src_schema, src_table, src_columns, primary_key, update_time_col, model, bm25_enabled) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8) ON CONFLICT (job_name) DO UPDATE SET src_schema = EXCLUDED.src_schema, src_table = EXCLUDED.src_table, src_columns = EXCLUDED.src_columns, primary_key = EXCLUDED.primary_key, update_time_col = EXCLUDED.update_time_col, - model = EXCLUDED.model + model = EXCLUDED.model, + bm25_enabled = EXCLUDED.bm25_enabled RETURNING id") .bind(job_request.job_name.clone()) .bind(job_request.src_schema.clone()) @@ -143,6 +148,7 @@ pub async fn initialize_job( .bind(job_request.primary_key.clone()) .bind(job_request.update_time_col.clone()) .bind(job_request.model.to_string()) + .bind(job_request.bm25_enabled) .fetch_one(&mut *tx) .await?; diff --git a/core/src/query.rs b/core/src/query.rs index e4f81a8..58f10cb 100644 --- a/core/src/query.rs +++ b/core/src/query.rs @@ -213,12 +213,20 @@ pub fn create_vectorize_table() -> String { primary_key TEXT NOT NULL, update_time_col TEXT NOT NULL, model TEXT NOT NULL, - params JSONB + params JSONB, + bm25_enabled BOOLEAN NOT NULL DEFAULT false ); " .to_string() } +/// Idempotent upgrade for pre-existing `vectorize.job` tables created before +/// `bm25_enabled` was added, so BM25 stays opt-in on upgrade too. +pub fn alter_vectorize_table_add_bm25_enabled() -> String { + "ALTER TABLE vectorize.job ADD COLUMN IF NOT EXISTS bm25_enabled BOOLEAN NOT NULL DEFAULT false;" + .to_string() +} + pub fn init_index_query(job_name: &str, idx_type: &str, job_params: &JobParams) -> String { check_input(job_name).expect("invalid job name"); let src_schema = job_params.schema.clone(); diff --git a/core/src/types.rs b/core/src/types.rs index 3744c09..323f3c1 100644 --- a/core/src/types.rs +++ b/core/src/types.rs @@ -21,6 +21,10 @@ pub struct VectorizeJob { serialize_with = "model_to_string" )] pub model: Model, + /// Opt-in flag to build and maintain an in-memory BM25 (Tantivy) index for this job. + /// Defaults to false so BM25 indexing never runs unless explicitly requested. + #[serde(default)] + pub bm25_enabled: bool, } #[allow(non_camel_case_types)] diff --git a/server/src/app_state.rs b/server/src/app_state.rs index af7a252..0c8540a 100644 --- a/server/src/app_state.rs +++ b/server/src/app_state.rs @@ -70,14 +70,19 @@ impl AppState { })); // Create empty BM25 index map, then kick off background population for - // every existing job so the server stays non-blocking at startup. + // every job that has opted in via `bm25_enabled`, so the server stays + // non-blocking at startup and jobs that never asked for BM25 incur no cost. let bm25_indexes: Arc>>>> = Arc::new(RwLock::new(HashMap::new())); { let jobs: Vec = { let cache = job_cache.read().await; - cache.values().cloned().collect() + cache + .values() + .filter(|job| job.bm25_enabled) + .cloned() + .collect() }; for job in jobs { if job.job_name.is_empty() { diff --git a/server/src/bm25.rs b/server/src/bm25.rs index 9cb5f34..14e9667 100644 --- a/server/src/bm25.rs +++ b/server/src/bm25.rs @@ -136,7 +136,11 @@ pub async fn start_bm25_sync_task( let jobs: Vec = { let cache = job_cache.read().await; - cache.values().cloned().collect() + cache + .values() + .filter(|job| job.bm25_enabled) + .cloned() + .collect() }; for job in &jobs { diff --git a/server/src/routes/table.rs b/server/src/routes/table.rs index 478db42..87b8441 100644 --- a/server/src/routes/table.rs +++ b/server/src/routes/table.rs @@ -60,27 +60,38 @@ pub async fn table( job_cache.insert(payload.job_name.clone(), payload.clone()); } - // Create a BM25 index for this job and populate it in the background. - match BM25Index::new() { - Ok(idx) => { - let idx = Arc::new(Mutex::new(idx)); - app_state - .bm25_indexes - .write() - .await - .insert(payload.job_name.clone(), idx.clone()); - let pool = app_state.db_pool.clone(); - let job = payload.clone(); - tokio::spawn(async move { - crate::bm25::populate_bm25_index(&pool, &job, idx).await; - }); - } - Err(e) => { - tracing::warn!( - "Failed to create BM25 index for job {}: {e}", - payload.job_name - ); + // BM25 indexing is opt-in per job via `bm25_enabled`. + if payload.bm25_enabled { + // Create a BM25 index for this job and populate it in the background. + match BM25Index::new() { + Ok(idx) => { + let idx = Arc::new(Mutex::new(idx)); + app_state + .bm25_indexes + .write() + .await + .insert(payload.job_name.clone(), idx.clone()); + let pool = app_state.db_pool.clone(); + let job = payload.clone(); + tokio::spawn(async move { + crate::bm25::populate_bm25_index(&pool, &job, idx).await; + }); + } + Err(e) => { + tracing::warn!( + "Failed to create BM25 index for job {}: {e}", + payload.job_name + ); + } } + } else { + // Job was re-created/updated with BM25 disabled; drop any stale index + // so it stops consuming memory and the background sync loop skips it. + app_state + .bm25_indexes + .write() + .await + .remove(&payload.job_name); } let resp = JobResponse { id: job_id }; From b4d00b16eb3d2aa2ab42e78907ec3a0d60e13793 Mon Sep 17 00:00:00 2001 From: Adam Hendel Date: Sat, 11 Jul 2026 12:21:39 -0500 Subject: [PATCH 4/6] select new col --- core/src/db.rs | 4 ++-- server/src/routes/search.rs | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/core/src/db.rs b/core/src/db.rs index b51577f..84e10e3 100644 --- a/core/src/db.rs +++ b/core/src/db.rs @@ -7,8 +7,8 @@ pub async fn get_vectorize_job( ) -> Result { // Changed return type let row = sqlx::query( - "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model - FROM vectorize.job + "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model, bm25_enabled + FROM vectorize.job WHERE job_name = $1", ) .bind(job_name) diff --git a/server/src/routes/search.rs b/server/src/routes/search.rs index 999d89b..6310a52 100644 --- a/server/src/routes/search.rs +++ b/server/src/routes/search.rs @@ -291,8 +291,8 @@ async fn get_vectorize_job( job_name: &str, ) -> Result { match sqlx::query( - "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model - FROM vectorize.job + "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model, bm25_enabled + FROM vectorize.job WHERE job_name = $1", ) .bind(job_name) From 9b061ac46c521c42bbfb8536c4108168b8f40836 Mon Sep 17 00:00:00 2001 From: Adam Hendel Date: Sat, 11 Jul 2026 18:16:13 -0500 Subject: [PATCH 5/6] migrations for management table spt --- core/src/init.rs | 7 ++---- core/src/query.rs | 25 --------------------- server/Cargo.toml | 2 +- server/migrations/0001_init_job_table.sql | 15 +++++++++++++ server/migrations/0002_add_bm25_enabled.sql | 4 ++++ server/src/app_state.rs | 4 ++++ server/src/bin/worker.rs | 4 ++++ server/src/db.rs | 23 +++++++++++++++++++ server/src/lib.rs | 1 + 9 files changed, 54 insertions(+), 31 deletions(-) create mode 100644 server/migrations/0001_init_job_table.sql create mode 100644 server/migrations/0002_add_bm25_enabled.sql create mode 100644 server/src/db.rs diff --git a/core/src/init.rs b/core/src/init.rs index fdd1206..9c414cd 100644 --- a/core/src/init.rs +++ b/core/src/init.rs @@ -74,16 +74,13 @@ async fn vectorize_schema_exists(pool: &PgPool) -> Result { pub async fn init_vectorize(pool: &PgPool) -> Result<(), VectorizeError> { if vectorize_schema_exists(pool).await? { log::info!("vectorize schema already exists, skipping initialization."); - // Backwards-compatible upgrade for installs predating the bm25_enabled column. - sqlx::query(&query::alter_vectorize_table_add_bm25_enabled()) - .execute(pool) - .await?; return Ok(()); } else { // these statements are critical, so we fail if they error + // Note: the vectorize.job table itself is created/evolved via the + // vectorize-server crate's sqlx migrations (server/migrations), not here. let statements_nofail = vec![ "CREATE SCHEMA IF NOT EXISTS vectorize;".to_string(), - query::create_vectorize_table(), query::handle_table_update(), query::create_batch_texts_fn(), ]; diff --git a/core/src/query.rs b/core/src/query.rs index 58f10cb..17dedab 100644 --- a/core/src/query.rs +++ b/core/src/query.rs @@ -202,31 +202,6 @@ pub fn check_input(input: &str) -> Result<()> { } } -pub fn create_vectorize_table() -> String { - "CREATE TABLE IF NOT EXISTS vectorize.job - ( - id UUID DEFAULT gen_random_uuid() PRIMARY KEY, - job_name TEXT NOT NULL UNIQUE, - src_schema TEXT NOT NULL, - src_table TEXT NOT NULL, - src_columns TEXT[] NOT NULL, - primary_key TEXT NOT NULL, - update_time_col TEXT NOT NULL, - model TEXT NOT NULL, - params JSONB, - bm25_enabled BOOLEAN NOT NULL DEFAULT false - ); - " - .to_string() -} - -/// Idempotent upgrade for pre-existing `vectorize.job` tables created before -/// `bm25_enabled` was added, so BM25 stays opt-in on upgrade too. -pub fn alter_vectorize_table_add_bm25_enabled() -> String { - "ALTER TABLE vectorize.job ADD COLUMN IF NOT EXISTS bm25_enabled BOOLEAN NOT NULL DEFAULT false;" - .to_string() -} - pub fn init_index_query(job_name: &str, idx_type: &str, job_params: &JobParams) -> String { check_input(job_name).expect("invalid job name"); let src_schema = job_params.schema.clone(); diff --git a/server/Cargo.toml b/server/Cargo.toml index 3811512..7ecade8 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -46,7 +46,7 @@ reqwest = { version = "0.12.16", features = ["json"] } serde = "1.0.219" serde_json = "1.0" sqlparser = "0.51" -sqlx = { workspace = true} +sqlx = { workspace = true, features = ["migrate"] } thiserror = "2.0.12" tiktoken-rs = "0.7.0" tokio = { version = "1.0", features = ["full"] } diff --git a/server/migrations/0001_init_job_table.sql b/server/migrations/0001_init_job_table.sql new file mode 100644 index 0000000..efe4638 --- /dev/null +++ b/server/migrations/0001_init_job_table.sql @@ -0,0 +1,15 @@ +-- Baseline: reflects the vectorize.job table as already deployed prior to +-- adopting sqlx migrations. Uses IF NOT EXISTS so this is a no-op against +-- any database where vectorize-core already created this table by hand. +CREATE TABLE IF NOT EXISTS vectorize.job +( + id UUID DEFAULT gen_random_uuid() PRIMARY KEY, + job_name TEXT NOT NULL UNIQUE, + src_schema TEXT NOT NULL, + src_table TEXT NOT NULL, + src_columns TEXT[] NOT NULL, + primary_key TEXT NOT NULL, + update_time_col TEXT NOT NULL, + model TEXT NOT NULL, + params JSONB +); diff --git a/server/migrations/0002_add_bm25_enabled.sql b/server/migrations/0002_add_bm25_enabled.sql new file mode 100644 index 0000000..b57ea86 --- /dev/null +++ b/server/migrations/0002_add_bm25_enabled.sql @@ -0,0 +1,4 @@ +-- Opt-in flag for building/maintaining an in-memory BM25 (Tantivy) index per job. +-- Defaults to false so upgraded deployments never get BM25 indexing turned on +-- for existing jobs without an explicit opt-in via POST /table. +ALTER TABLE vectorize.job ADD COLUMN IF NOT EXISTS bm25_enabled BOOLEAN NOT NULL DEFAULT false; diff --git a/server/src/app_state.rs b/server/src/app_state.rs index 0c8540a..bfc0392 100644 --- a/server/src/app_state.rs +++ b/server/src/app_state.rs @@ -48,6 +48,10 @@ impl AppState { .await .map_err(|e| format!("Failed to initialize project: {e}"))?; + crate::db::run_migrations(&db_pool) + .await + .map_err(|e| format!("Failed to run migrations: {e}"))?; + // load initial job cache let job_cache = cache::load_initial_job_cache(&db_pool) .await diff --git a/server/src/bin/worker.rs b/server/src/bin/worker.rs index ce86945..80782cd 100644 --- a/server/src/bin/worker.rs +++ b/server/src/bin/worker.rs @@ -21,6 +21,10 @@ async fn main() { .await .expect("Failed to initialize project"); + vectorize_server::db::run_migrations(&pool) + .await + .expect("Failed to run migrations"); + let queue = pgmq::PGMQueueExt::new_with_pool(pool.clone()).await; loop { diff --git a/server/src/db.rs b/server/src/db.rs new file mode 100644 index 0000000..7fd8f50 --- /dev/null +++ b/server/src/db.rs @@ -0,0 +1,23 @@ +use sqlx::migrate::MigrateError; +use sqlx::{ConnectOptions, PgPool}; + +/// Applies pending sqlx migrations for the tables vectorize-server owns +/// (currently just `vectorize.job`). Must run after +/// `vectorize_core::init::init_project`, which ensures the `vectorize` schema +/// exists. Safe to call from multiple processes concurrently: sqlx takes a +/// Postgres advisory lock for the duration of the migration run. +/// +/// Runs on a dedicated, non-pooled connection with `search_path` set to +/// `vectorize`, so sqlx's bookkeeping table (`_sqlx_migrations`) is created +/// inside the `vectorize` schema instead of `public`. This keeps it isolated +/// from the pool's connections, whose search_path is left untouched. +pub async fn run_migrations(pool: &PgPool) -> Result<(), MigrateError> { + let connect_options = pool + .connect_options() + .as_ref() + .clone() + .options([("search_path", "vectorize")]); + + let mut conn = connect_options.connect().await?; + sqlx::migrate!("./migrations").run(&mut conn).await +} diff --git a/server/src/lib.rs b/server/src/lib.rs index 73494a2..535554a 100644 --- a/server/src/lib.rs +++ b/server/src/lib.rs @@ -1,6 +1,7 @@ pub mod app_state; pub mod bm25; pub mod cache; +pub mod db; pub mod errors; pub mod routes; pub mod server; From 01ac03d22f8cea5e94a3ddfc81f815c50ba675cb Mon Sep 17 00:00:00 2001 From: Adam Hendel Date: Sat, 11 Jul 2026 19:28:03 -0500 Subject: [PATCH 6/6] remove wrong comment --- server/src/db.rs | 5 ----- 1 file changed, 5 deletions(-) diff --git a/server/src/db.rs b/server/src/db.rs index 7fd8f50..16a85e8 100644 --- a/server/src/db.rs +++ b/server/src/db.rs @@ -2,11 +2,6 @@ use sqlx::migrate::MigrateError; use sqlx::{ConnectOptions, PgPool}; /// Applies pending sqlx migrations for the tables vectorize-server owns -/// (currently just `vectorize.job`). Must run after -/// `vectorize_core::init::init_project`, which ensures the `vectorize` schema -/// exists. Safe to call from multiple processes concurrently: sqlx takes a -/// Postgres advisory lock for the duration of the migration run. -/// /// Runs on a dedicated, non-pooled connection with `search_path` set to /// `vectorize`, so sqlx's bookkeeping table (`_sqlx_migrations`) is created /// inside the `vectorize` schema instead of `public`. This keeps it isolated