From 2fbad7c867c568b70b7ecbc9a03d24b5cfffd8ba Mon Sep 17 00:00:00 2001 From: Pascal Seitz Date: Fri, 2 Oct 2026 18:13:42 +0200 Subject: [PATCH] Optimize HitSet seeking and counting Update Tantivy to main revision e988ecb7, implement block-aware seek and exact-hit seek_danger, and return the stored hit count directly. Add coverage for iteration, seeking, and counting. --- quickwit/Cargo.lock | 22 +-- quickwit/Cargo.toml | 4 +- .../src/query_ast/cache_node.rs | 161 +++++++++++++++++- 3 files changed, 171 insertions(+), 16 deletions(-) diff --git a/quickwit/Cargo.lock b/quickwit/Cargo.lock index 4c929de6fc0..ad74956c196 100644 --- a/quickwit/Cargo.lock +++ b/quickwit/Cargo.lock @@ -4532,7 +4532,7 @@ dependencies = [ [[package]] name = "jitexpr" version = "0.1.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=e988ecb7a11109d89f9f4582b386089fc3a2fec0#e988ecb7a11109d89f9f4582b386089fc3a2fec0" dependencies = [ "cranelift", "cranelift-jit", @@ -5995,7 +5995,7 @@ checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" [[package]] name = "ownedbytes" version = "0.9.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=e988ecb7a11109d89f9f4582b386089fc3a2fec0#e988ecb7a11109d89f9f4582b386089fc3a2fec0" dependencies = [ "stable_deref_trait", ] @@ -10305,7 +10305,7 @@ checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" [[package]] name = "tantivy" version = "0.27.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=e988ecb7a11109d89f9f4582b386089fc3a2fec0#e988ecb7a11109d89f9f4582b386089fc3a2fec0" dependencies = [ "aho-corasick", "arc-swap", @@ -10346,7 +10346,7 @@ dependencies = [ "tantivy-columnar", "tantivy-common", "tantivy-fst", - "tantivy-query-grammar 0.26.0 (git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe)", + "tantivy-query-grammar 0.26.0 (git+https://github.com/quickwit-oss/tantivy/?rev=e988ecb7a11109d89f9f4582b386089fc3a2fec0)", "tantivy-sstable", "tantivy-stacker", "tantivy-tokenizer-api", @@ -10363,7 +10363,7 @@ dependencies = [ [[package]] name = "tantivy-bitpacker" version = "0.10.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=e988ecb7a11109d89f9f4582b386089fc3a2fec0#e988ecb7a11109d89f9f4582b386089fc3a2fec0" dependencies = [ "bitpacking", ] @@ -10371,7 +10371,7 @@ dependencies = [ [[package]] name = "tantivy-columnar" version = "0.7.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=e988ecb7a11109d89f9f4582b386089fc3a2fec0#e988ecb7a11109d89f9f4582b386089fc3a2fec0" dependencies = [ "downcast-rs", "fastdivide", @@ -10387,7 +10387,7 @@ dependencies = [ [[package]] name = "tantivy-common" version = "0.11.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=e988ecb7a11109d89f9f4582b386089fc3a2fec0#e988ecb7a11109d89f9f4582b386089fc3a2fec0" dependencies = [ "async-trait", "byteorder", @@ -10423,7 +10423,7 @@ dependencies = [ [[package]] name = "tantivy-query-grammar" version = "0.26.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=e988ecb7a11109d89f9f4582b386089fc3a2fec0#e988ecb7a11109d89f9f4582b386089fc3a2fec0" dependencies = [ "fnv", "nom 7.1.3", @@ -10435,7 +10435,7 @@ dependencies = [ [[package]] name = "tantivy-sstable" version = "0.7.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=e988ecb7a11109d89f9f4582b386089fc3a2fec0#e988ecb7a11109d89f9f4582b386089fc3a2fec0" dependencies = [ "futures-util", "itertools 0.14.0", @@ -10448,7 +10448,7 @@ dependencies = [ [[package]] name = "tantivy-stacker" version = "0.7.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=e988ecb7a11109d89f9f4582b386089fc3a2fec0#e988ecb7a11109d89f9f4582b386089fc3a2fec0" dependencies = [ "murmurhash32", "tantivy-common", @@ -10457,7 +10457,7 @@ dependencies = [ [[package]] name = "tantivy-tokenizer-api" version = "0.7.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=e988ecb7a11109d89f9f4582b386089fc3a2fec0#e988ecb7a11109d89f9f4582b386089fc3a2fec0" dependencies = [ "serde", ] diff --git a/quickwit/Cargo.toml b/quickwit/Cargo.toml index 1ade0c4c0e2..1fd4558876d 100644 --- a/quickwit/Cargo.toml +++ b/quickwit/Cargo.toml @@ -378,7 +378,7 @@ quickwit-telemetry-exporters = { path = "quickwit-telemetry-exporters" } quickwit-transport = { path = "quickwit-transport" } # This is the tip of tantivy main. -tantivy = { git = "https://github.com/quickwit-oss/tantivy/", rev = "2e87c2eb0c8577c96c82a777f345bbf236a985fe", default-features = false, features = [ +tantivy = { git = "https://github.com/quickwit-oss/tantivy/", rev = "e988ecb7a11109d89f9f4582b386089fc3a2fec0", default-features = false, features = [ "jitexpr", "lz4-compression", "mmap", @@ -387,7 +387,7 @@ tantivy = { git = "https://github.com/quickwit-oss/tantivy/", rev = "2e87c2eb0c8 "columnar-zstd-compression", ] } # Must stay on the same rev as `tantivy`. -tantivy-common = { git = "https://github.com/quickwit-oss/tantivy/", rev = "2e87c2eb0c8577c96c82a777f345bbf236a985fe" } +tantivy-common = { git = "https://github.com/quickwit-oss/tantivy/", rev = "e988ecb7a11109d89f9f4582b386089fc3a2fec0" } tantivy-fst = "0.5" # This is actually not used directly the goal is to fix the version diff --git a/quickwit/quickwit-query/src/query_ast/cache_node.rs b/quickwit/quickwit-query/src/query_ast/cache_node.rs index 9574a93e040..c523ec621d5 100644 --- a/quickwit/quickwit-query/src/query_ast/cache_node.rs +++ b/quickwit/quickwit-query/src/query_ast/cache_node.rs @@ -127,7 +127,7 @@ impl BuildTantivyAst for CacheNode { use tantivy::directory::OwnedBytes; use tantivy::index::SegmentId; use tantivy::query::{EnableScoring, Explanation, Query, Scorer, Weight}; -use tantivy::{DocId, DocSet, Score, SegmentReader, TantivyError, Term}; +use tantivy::{DocId, DocSet, Score, SeekDangerResult, SegmentReader, TantivyError, Term}; #[derive(Clone, Debug)] pub struct CacheHitQuery { @@ -294,8 +294,40 @@ impl DocSet for HitSet { self.current_block[0] } - // fn seek(&mut self, target: DocId) -> DocId { - // } + fn seek(&mut self, target: DocId) -> DocId { + debug_assert!(self.doc() <= target); + while self.current_block[BitPacker1x::BLOCK_LEN - 1] < target { + self.load_new_block(); + self.block_pos = 0; + } + self.block_pos += self.current_block[self.block_pos..].partition_point(|&doc| doc < target); + self.doc() + } + + fn seek_danger(&mut self, target: DocId) -> SeekDangerResult { + if target == tantivy::TERMINATED { + return SeekDangerResult::SeekLowerBound(tantivy::TERMINATED); + } + while self.current_block[BitPacker1x::BLOCK_LEN - 1] < target { + self.load_new_block(); + self.block_pos = 0; + } + if self.doc() > target { + return SeekDangerResult::SeekLowerBound(self.doc()); + } + match self.current_block[self.block_pos..].binary_search(&target) { + Ok(position) => { + self.block_pos += position; + SeekDangerResult::Found + } + Err(_) => SeekDangerResult::SeekLowerBound(target + 1), + } + } + + // Weight::count uses a fresh scorer. Return the stored total without consuming it. + fn count_including_deleted(&mut self) -> u32 { + u32::from_ne_bytes(self.buffer[0..4].try_into().unwrap()) + } #[inline(always)] fn doc(&self) -> DocId { @@ -580,6 +612,129 @@ mod tests { test_hit_set_roundtrip_helper(generator.clone().skip(10).take(1024 + 6)); } + #[test] + fn test_hit_set_count_including_deleted() { + for count in [0, 1, 31, 32, 33, 64, 65, 1030] { + let mut builder = HitSetBuilder::new(); + for doc in 0..count { + builder.insert(doc * 3); + } + let mut hits = builder.build(); + let first_doc = hits.doc(); + assert_eq!(hits.count_including_deleted(), count); + assert_eq!(hits.doc(), first_doc); + for doc in 0..count { + assert_eq!(hits.doc(), doc * 3); + hits.advance(); + } + assert_eq!(hits.doc(), tantivy::TERMINATED); + assert_eq!(hits.count_including_deleted(), count); + assert_eq!(hits.advance(), tantivy::TERMINATED); + } + } + + // Compare seeks and subsequent advances with the uncompressed document IDs. + #[test] + fn test_hit_set_seek() { + for count in [0, 1, 31, 32, 33, 64, 65, 1030] { + for start in [0, tantivy::TERMINATED - 4 * 1030] { + let docs: Vec = (0..count).map(|doc| start + doc * 3).collect(); + let mut builder = HitSetBuilder::new(); + for &doc in &docs { + builder.insert(doc); + } + let hits = builder.build(); + let mut terminated = hits.clone(); + assert_eq!(terminated.seek(tantivy::TERMINATED), tantivy::TERMINATED); + for _ in 0..96 { + assert_eq!(terminated.doc(), tantivy::TERMINATED); + assert_eq!(terminated.seek(tantivy::TERMINATED), tantivy::TERMINATED); + assert_eq!(terminated.advance(), tantivy::TERMINATED); + } + for step in [1, 97] { + let mut sought = hits.clone(); + for target in (start..start + count * 3 + 2).step_by(step) { + let target = target.max(sought.doc()); + let position = docs.partition_point(|&doc| doc < target); + let expected = docs.get(position).copied().unwrap_or(tantivy::TERMINATED); + assert_eq!(sought.seek(target), expected); + assert_eq!(sought.doc(), expected); + // Seeking the current document must not advance it. + assert_eq!(sought.seek(expected), expected); + assert_eq!( + sought.advance(), + docs.get(position + 1) + .copied() + .unwrap_or(tantivy::TERMINATED) + ); + } + } + } + } + } + + #[test] + fn test_hit_set_seek_danger_miss() { + let mut builder = HitSetBuilder::new(); + for doc in [1, 4, 7] { + builder.insert(doc); + } + let mut hits = builder.build(); + assert_eq!(hits.seek_danger(2), SeekDangerResult::SeekLowerBound(3)); + // A miss does not move to the next document. Only seek_danger calls are allowed + // until Found restores the valid state. + assert_eq!(hits.block_pos, 0); + assert_eq!(hits.seek_danger(3), SeekDangerResult::SeekLowerBound(4)); + assert_eq!(hits.seek_danger(4), SeekDangerResult::Found); + assert_eq!(hits.doc(), 4); + assert_eq!(hits.advance(), 7); + } + + #[test] + fn test_hit_set_seek_danger() { + for count in [0, 1, 31, 32, 33, 64, 65, 1030] { + for start in [0, tantivy::TERMINATED - 4 * 1030] { + let docs: Vec = (0..count).map(|doc| start + doc * 3).collect(); + let mut builder = HitSetBuilder::new(); + for &doc in &docs { + builder.insert(doc); + } + let hits = builder.build(); + // Exercise both successive targets and jumps across several blocks. + for step in [1, 97] { + let mut sought = hits.clone(); + for target in (start..start + count * 3 + 2).step_by(step) { + let position = docs.partition_point(|&doc| doc < target); + let expected = docs.get(position).copied().unwrap_or(tantivy::TERMINATED); + let result = sought.seek_danger(target); + if expected == target { + assert_eq!(result, SeekDangerResult::Found); + assert_eq!(sought.doc(), target); + assert_eq!( + sought.clone().advance(), + docs.get(position + 1) + .copied() + .unwrap_or(tantivy::TERMINATED) + ); + } else { + let SeekDangerResult::SeekLowerBound(lower_bound) = result else { + panic!("seek_danger found absent target {target}"); + }; + assert!(lower_bound > target); + assert!(lower_bound <= expected); + } + } + assert_eq!( + sought.seek_danger(tantivy::TERMINATED), + SeekDangerResult::SeekLowerBound(tantivy::TERMINATED) + ); + assert_eq!(sought.seek(tantivy::TERMINATED), tantivy::TERMINATED); + assert_eq!(sought.count_including_deleted(), count); + } + } + } + } + #[test] fn test_built_tantivy_ast() { let mut schema_builder = Schema::builder();