From 644c5df2a292a1a0aac4c46e9003ea0aec6afff9 Mon Sep 17 00:00:00 2001 From: samzong Date: Fri, 14 Aug 2026 00:25:17 -0400 Subject: [PATCH] fix(search): bound semantic pagination Signed-off-by: samzong --- src/db/search.rs | 111 ++++++++++++++-------- src/integration/regression.rs | 172 ++++++++++++++++++++++++++++++++++ src/session.rs | 14 +-- 3 files changed, 251 insertions(+), 46 deletions(-) diff --git a/src/db/search.rs b/src/db/search.rs index 538b603..78546c0 100644 --- a/src/db/search.rs +++ b/src/db/search.rs @@ -9,6 +9,8 @@ use crate::project_scope::ProjectScope; use crate::types::{MatchSource, SearchResult, Session}; use crate::utils::f32_slice_to_bytes; +const SQLITE_VEC_MAX_K: usize = 4096; + pub(crate) struct SearchEngine<'a> { conn: &'a Connection, } @@ -109,30 +111,57 @@ impl<'a> SearchEngine<'a> { limit: usize, fetch_multiplier: usize, ) -> anyhow::Result> { - let fetch_size = limit * fetch_multiplier; - let fts_hits = self.fts_search(query, filters, fetch_size)?; + let fetch_size = limit.saturating_mul(fetch_multiplier).max(1); + let fts_hits = self.fts_search(query, filters, Some(fetch_size))?; let vec_hits = match embedding { - Some(e) => self.vec_search(e, filters, fetch_size)?, + Some(embedding) => self.vec_search(embedding, filters, fetch_size.saturating_mul(5))?, None => vec![], }; - let merged = rrf_merge(&fts_hits, &vec_hits, 10); + self.search_results(fts_hits, vec_hits, 0, Some(limit)) + } - let session_ids: Vec<&str> = - merged.iter().take(limit).map(|(id, _, _)| id.as_str()).collect(); + pub(crate) fn hybrid_search_page( + &self, + query: &str, + embedding: Option<&[f32]>, + filters: &SearchFilters, + limit: Option, + offset: usize, + ) -> anyhow::Result> { + if limit == Some(0) { + return Ok(vec![]); + } + + let fts_hits = self.fts_search(query, filters, None)?; + let vec_hits = match embedding { + Some(embedding) => self.vec_search(embedding, filters, SQLITE_VEC_MAX_K)?, + None => vec![], + }; + self.search_results(fts_hits, vec_hits, offset, limit) + } + fn search_results( + &self, + fts_hits: Vec, + vec_hits: Vec, + offset: usize, + limit: Option, + ) -> anyhow::Result> { + let merged = rrf_merge(&fts_hits, &vec_hits, 10); + let snippets: HashMap<_, _> = + fts_hits.into_iter().map(|hit| (hit.session_id, hit.snippet)).collect(); + let limit = limit.unwrap_or(usize::MAX); + let session_ids: Vec<&str> = + merged.iter().skip(offset).take(limit).map(|(id, _, _)| id.as_str()).collect(); let sessions = self.load_sessions(&session_ids)?; let mut results = Vec::new(); - for (session_id, _score, match_source) in merged.into_iter().take(limit) { + for (session_id, _score, match_source) in merged.into_iter().skip(offset).take(limit) { if let Some(session) = sessions.get(&session_id) { - let snippet = fts_hits - .iter() - .find(|h| h.session_id == session_id) - .and_then(|h| h.snippet.clone()); + let snippet = snippets.get(&session_id).cloned().flatten(); results.push(SearchResult { session: session.clone(), match_source, snippet }); } } - Ok(results) } @@ -140,7 +169,7 @@ impl<'a> SearchEngine<'a> { &self, query: &str, filters: &SearchFilters, - limit: usize, + limit: Option, ) -> anyhow::Result> { let escaped = fts5_escape(query); if escaped.is_empty() { @@ -160,10 +189,14 @@ impl<'a> SearchEngine<'a> { let mut param_idx = 2; apply_filters(&mut sql, &mut params, &mut param_idx, filters); - sql.push_str(&format!(" GROUP BY m.session_id ORDER BY best_rank LIMIT {limit}")); + sql.push_str(" GROUP BY m.session_id ORDER BY best_rank, m.session_id"); + if let Some(limit) = limit { + let limit = i64::try_from(limit).unwrap_or(i64::MAX); + sql.push_str(&format!(" LIMIT {limit}")); + } let param_refs: Vec<&dyn rusqlite::types::ToSql> = - params.iter().map(|p| p.as_ref()).collect(); + params.iter().map(|param| param.as_ref()).collect(); let mut stmt = self.conn.prepare(&sql)?; let rows = stmt.query_map(param_refs.as_slice(), |row| { @@ -181,10 +214,10 @@ impl<'a> SearchEngine<'a> { &self, embedding: &[f32], filters: &SearchFilters, - limit: usize, + requested_k: usize, ) -> anyhow::Result> { let blob = f32_slice_to_bytes(embedding); - let fetch_k = (limit * 5) as i64; + let fetch_k = requested_k.clamp(1, SQLITE_VEC_MAX_K) as i64; let mut sql = String::from( "SELECT m.session_id, MIN(mv.distance) AS best_distance @@ -200,10 +233,10 @@ impl<'a> SearchEngine<'a> { let mut param_idx = 3; apply_filters(&mut sql, &mut params, &mut param_idx, filters); - sql.push_str(" GROUP BY m.session_id ORDER BY best_distance"); + sql.push_str(" GROUP BY m.session_id ORDER BY best_distance, m.session_id"); let param_refs: Vec<&dyn rusqlite::types::ToSql> = - params.iter().map(|p| p.as_ref()).collect(); + params.iter().map(|param| param.as_ref()).collect(); let mut stmt = self.conn.prepare(&sql)?; let rows = stmt.query_map(param_refs.as_slice(), |row| { @@ -218,27 +251,25 @@ impl<'a> SearchEngine<'a> { } fn load_sessions(&self, ids: &[&str]) -> anyhow::Result> { - let mut map = HashMap::new(); - if ids.is_empty() { - return Ok(map); - } + const SESSION_LOAD_CHUNK_SIZE: usize = 900; - let placeholders: Vec = (1..=ids.len()).map(|i| format!("?{i}")).collect(); - let sql = format!( - "SELECT {SESSION_COLUMNS} - FROM sessions WHERE id IN ({})", - placeholders.join(", ") - ); - - let params: Vec<&dyn rusqlite::types::ToSql> = - ids.iter().map(|id| id as &dyn rusqlite::types::ToSql).collect(); - - let mut stmt = self.conn.prepare(&sql)?; - let rows = stmt.query_map(params.as_slice(), session_from_row)?; - - for row in rows { - let session = row?; - map.insert(session.id.clone(), session); + let mut map = HashMap::new(); + for ids in ids.chunks(SESSION_LOAD_CHUNK_SIZE) { + let placeholders: Vec = (1..=ids.len()).map(|i| format!("?{i}")).collect(); + let sql = format!( + "SELECT {SESSION_COLUMNS} + FROM sessions WHERE id IN ({})", + placeholders.join(", ") + ); + let params: Vec<&dyn rusqlite::types::ToSql> = + ids.iter().map(|id| id as &dyn rusqlite::types::ToSql).collect(); + let mut stmt = self.conn.prepare(&sql)?; + let rows = stmt.query_map(params.as_slice(), session_from_row)?; + + for row in rows { + let session = row?; + map.insert(session.id.clone(), session); + } } Ok(map) } @@ -300,7 +331,7 @@ fn rrf_merge(fts_hits: &[Hit], vec_hits: &[Hit], k: u32) -> Vec<(String, f64, Ma }) .collect(); - results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + results.sort_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0))); results } diff --git a/src/integration/regression.rs b/src/integration/regression.rs index 8d93c42..c85a40f 100644 --- a/src/integration/regression.rs +++ b/src/integration/regression.rs @@ -680,6 +680,178 @@ fn hybrid_search_fts_only_without_embedding() { assert_eq!(results.len(), 1); } +fn seed_semantic_boundary_sessions(store: &Store, count: usize) { + store + .conn + .execute_batch(&format!( + "WITH RECURSIVE seq(n) AS ( + SELECT 0 + UNION ALL + SELECT n + 1 FROM seq WHERE n + 1 < {count} + ) + INSERT INTO sessions (id, source, source_id, title, started_at, message_count) + SELECT printf('semantic-%05d', n), 'test', printf('raw-%05d', n), + printf('Semantic session %05d', n), n, 1 + FROM seq; + INSERT INTO messages (session_id, role, content, timestamp, seq) + SELECT id, 'user', 'semanticboundary ' || id, started_at, 0 + FROM sessions + WHERE source = 'test';" + )) + .unwrap(); +} + +fn add_semantic_boundary_embedding(store: &Store) -> Vec { + let message_id: i64 = store + .conn + .query_row("SELECT id FROM messages ORDER BY id LIMIT 1", [], |row| row.get(0)) + .unwrap(); + let embedding = vec![0.1f32; 384]; + store.upsert_embeddings(&[(message_id, &embedding)]).unwrap(); + embedding +} + +fn seed_semantic_page_fixture(store: &Store) -> Vec { + for index in 0..6 { + let id = format!("semantic-fts-{index:02}"); + let session = make_session(&id, "test", &format!("raw-fts-{index:02}"), "Semantic FTS"); + store.insert_session(&session).unwrap(); + store.insert_messages(&[make_message(&id, Role::User, "semanticstable", 0)]).unwrap(); + } + + let mut filler = + make_session("semantic-vec-fill", "test", "raw-vec-fill", "Semantic vector filler"); + filler.message_count = 30; + store.insert_session(&filler).unwrap(); + let filler_messages = (0..30) + .map(|seq| make_message("semantic-vec-fill", Role::User, "semantic filler", seq)) + .collect::>(); + store.insert_messages(&filler_messages).unwrap(); + + let mut stmt = store + .conn + .prepare("SELECT id FROM messages WHERE session_id = 'semantic-vec-fill' ORDER BY seq") + .unwrap(); + let mut message_ids = stmt + .query_map([], |row| row.get::<_, i64>(0)) + .unwrap() + .collect::>>() + .unwrap(); + message_ids.push( + store + .conn + .query_row("SELECT id FROM messages WHERE session_id = 'semantic-fts-04'", [], |row| { + row.get(0) + }) + .unwrap(), + ); + + let vectors = (1..=message_ids.len()) + .map(|rank| { + let mut vector = vec![0.0f32; 384]; + vector[0] = rank as f32 / 1_000.0; + vector + }) + .collect::>(); + let embeddings = message_ids + .iter() + .zip(&vectors) + .map(|(&message_id, vector)| (message_id, vector.as_slice())) + .collect::>(); + store.upsert_embeddings(&embeddings).unwrap(); + + vec![0.0f32; 384] +} + +#[test] +fn semantic_query_crosses_sqlite_vec_boundary() { + let store = setup(); + seed_semantic_boundary_sessions(&store, 300); + let embedding = add_semantic_boundary_embedding(&store); + let engine = SearchEngine::new(&store.conn); + + let results = + engine.hybrid_search("semanticboundary", Some(&embedding), &no_filters(), 274, 3).unwrap(); + assert_eq!(results.len(), 274); + + let page = engine + .hybrid_search_page("semanticboundary", Some(&embedding), &no_filters(), Some(50), 224) + .unwrap(); + assert_eq!(page.len(), 50); +} + +#[test] +fn semantic_adjacent_pages_follow_one_global_order() { + let store = setup(); + let embedding = seed_semantic_page_fixture(&store); + let engine = SearchEngine::new(&store.conn); + + let first_page = engine + .hybrid_search_page("semanticstable", Some(&embedding), &no_filters(), Some(2), 0) + .unwrap(); + let second_page = engine + .hybrid_search_page("semanticstable", Some(&embedding), &no_filters(), Some(2), 2) + .unwrap(); + let global = engine + .hybrid_search_page("semanticstable", Some(&embedding), &no_filters(), None, 0) + .unwrap(); + + let paged_ids = first_page + .iter() + .chain(&second_page) + .map(|result| result.session.id.as_str()) + .collect::>(); + let global_prefix = + global.iter().take(4).map(|result| result.session.id.as_str()).collect::>(); + + assert_eq!(global_prefix[0], "semantic-fts-04"); + assert_eq!(paged_ids, global_prefix); +} + +#[test] +fn semantic_query_all_returns_complete_fts_set() { + let store = setup(); + seed_semantic_boundary_sessions(&store, 10_001); + let engine = SearchEngine::new(&store.conn); + + let text_results = + engine.hybrid_search_page("semanticboundary", None, &no_filters(), None, 0).unwrap(); + assert_eq!(text_results.len(), 10_001); + assert!(text_results.iter().all(|result| { + result.snippet.as_deref().and_then(|snippet| snippet.strip_prefix("semanticboundary ")) + == Some(result.session.id.as_str()) + })); + + let embedding = add_semantic_boundary_embedding(&store); + let semantic_results = engine + .hybrid_search_page("semanticboundary", Some(&embedding), &no_filters(), None, 0) + .unwrap(); + assert_eq!(semantic_results.len(), 10_001); +} + +#[test] +fn semantic_search_arithmetic_is_saturating() { + let store = setup(); + let engine = SearchEngine::new(&store.conn); + let embedding = vec![0.1f32; 384]; + + let direct = engine + .hybrid_search("semanticboundary", Some(&embedding), &no_filters(), usize::MAX, usize::MAX) + .unwrap(); + assert!(direct.is_empty()); + + let page = engine + .hybrid_search_page( + "semanticboundary", + Some(&embedding), + &no_filters(), + Some(usize::MAX), + usize::MAX, + ) + .unwrap(); + assert!(page.is_empty()); +} + #[test] fn search_with_source_filter() { let store = setup(); diff --git a/src/session.rs b/src/session.rs index 1940ac2..e4c7850 100644 --- a/src/session.rs +++ b/src/session.rs @@ -356,13 +356,15 @@ pub(crate) fn run_session_list( scope: scope.clone(), thread_role, }; - let search_limit = effective_limit.unwrap_or(10_000).saturating_add(offset).max(1); - let results = - engine.hybrid_search(query, embedding.as_deref(), &filters, search_limit, 3)?; + let results = engine.hybrid_search_page( + query, + embedding.as_deref(), + &filters, + effective_limit, + offset, + )?; results .into_iter() - .skip(offset) - .take(effective_limit.unwrap_or(usize::MAX)) .map(|result| SessionListRow { session: result.session, match_source: Some(result.match_source), @@ -849,7 +851,7 @@ fn print_session_list_json( "next_offset": if all || rows.len() < limit { serde_json::Value::Null } else { - serde_json::json!(offset + rows.len()) + serde_json::json!(offset.saturating_add(rows.len())) } }))? );