From 027555c75f2ca92069948e39464628f61ec89883 Mon Sep 17 00:00:00 2001 From: barrotsteindev Date: Sat, 21 Nov 2020 18:09:08 +0200 Subject: [PATCH 1/9] added initial implementation for filter_collector --- src/collector/filter_collector_wrapper.rs | 153 ++++++++++++++++++++++ src/collector/mod.rs | 3 + 2 files changed, 156 insertions(+) create mode 100644 src/collector/filter_collector_wrapper.rs diff --git a/src/collector/filter_collector_wrapper.rs b/src/collector/filter_collector_wrapper.rs new file mode 100644 index 000000000..ddcf68cb7 --- /dev/null +++ b/src/collector/filter_collector_wrapper.rs @@ -0,0 +1,153 @@ +// # Custom collector example +// +// This example shows how you can implement your own +// collector. As an example, we will compute a collector +// that computes the standard deviation of a given fast field. +// +// Of course, you can have a look at the tantivy's built-in collectors +// such as the `CountCollector` for more examples. + +// --- +// Importing tantivy... +use std::marker::PhantomData; + +use crate::collector::{Collector, SegmentCollector}; +use crate::fastfield::FastFieldReader; +use crate::schema::Field; +use crate::{Score, SegmentReader, TantivyError}; + +/// The `TopDocs` collector keeps track of the top `K` documents +/// sorted by their score. +/// +/// The implementation is based on a `BinaryHeap`. +/// The theorical complexity for collecting the top `K` out of `n` documents +/// is `O(n log K)`. +/// +/// This collector guarantees a stable sorting in case of a tie on the +/// document score. As such, it is suitable to implement pagination. +/// +/// ```rust +/// use tantivy::collector::{TopDocs, FilterCollector}; +/// use tantivy::query::QueryParser; +/// use tantivy::schema::{Schema, TEXT, INDEXED, FAST}; +/// use tantivy::{doc, DocAddress, Index}; +/// +/// let mut schema_builder = Schema::builder(); +/// let title = schema_builder.add_text_field("title", TEXT); +/// let price = schema_builder.add_u64_field("price", INDEXED | FAST); +/// let schema = schema_builder.build(); +/// let index = Index::create_in_ram(schema); +/// +/// let mut index_writer = index.writer_with_num_threads(1, 10_000_000).unwrap(); +/// index_writer.add_document(doc!(title => "The Name of the Wind", price => 30_200u64)); +/// index_writer.add_document(doc!(title => "The Diary of Muadib", price => 29_240u64)); +/// index_writer.add_document(doc!(title => "A Dairy Cow", price => 21_240u64)); +/// index_writer.add_document(doc!(title => "The Diary of a Young Girl", price => 20_120u64)); +/// assert!(index_writer.commit().is_ok()); +/// +/// let reader = index.reader().unwrap(); +/// let searcher = reader.searcher(); +/// +/// let query_parser = QueryParser::for_index(&index, vec![title]); +/// let query = query_parser.parse_query("diary").unwrap(); +/// let no_filter_collector = FilterCollector::new(price, &|value| true, TopDocs::with_limit(2)); +/// let top_docs = searcher.search(&query, &no_filter_collector).unwrap(); +/// +/// assert_eq!(top_docs.len(), 2); +/// assert_eq!(top_docs[0].1, DocAddress(0, 1)); +/// assert_eq!(top_docs[1].1, DocAddress(0, 3)); +/// +/// let filter_all_collector = FilterCollector::new(price, &|value| false, TopDocs::with_limit(2)); +/// let filtered_top_docs = searcher.search(&query, &filter_all_collector).unwrap(); +/// +/// assert_eq!(filtered_top_docs.len(), 0); +/// ``` +pub struct FilterCollector { + field: Field, + collector: TCollector, + predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), + phantom: PhantomData +} + +impl FilterCollector +where + TCollector: Collector + Send + Sync, + TSegmentCollector: 'static + SegmentCollector + Send + Sync { + pub fn new(field: Field, predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), collector: TCollector) -> FilterCollector { + FilterCollector { field, predicate, collector, phantom: PhantomData } + } +} + +impl Collector for FilterCollector +where + TSegmentCollector: 'static + SegmentCollector + Send + Sync, + TCollector: Collector + Send + Sync { + // That's the type of our result. + // Our standard deviation will be a float. + type Fruit = TCollector::Fruit; + + type Child = FilterSegmentCollector; + + fn for_segment( + &self, + segment_local_id: u32, + segment_reader: &SegmentReader, + ) -> crate::Result> { + let fast_field_reader = segment_reader + .fast_fields() + .u64(self.field) + .ok_or_else(|| { + let field_name = segment_reader.schema().get_field_name(self.field); + TantivyError::SchemaError(format!( + "Field {:?} is not a u64 fast field.", + field_name + )) + })?; + let child_segment_collector = self.collector.for_segment(segment_local_id, segment_reader); + match child_segment_collector { + Ok(segment_collector) => { + Ok(FilterSegmentCollector:: { + fast_field_reader, + segment_collector: segment_collector, + predicate: self.predicate + }) + }, + Err(_) => { + Err(TantivyError::SystemError("Could not open segment: ".to_owned())) + }, + } + + + } + + fn requires_scoring(&self) -> bool { + self.collector.requires_scoring() + } + + fn merge_fruits(&self, segment_fruits: Vec<::Fruit>) -> crate::Result { + self.collector.merge_fruits(segment_fruits) + } +} + +pub struct FilterSegmentCollector { + fast_field_reader: FastFieldReader, + segment_collector: TSegmentCollector, + predicate: &'static (dyn Fn(u64) -> bool + Send + Sync) +} + +impl SegmentCollector for FilterSegmentCollector +where + TSegmentCollector: 'static + SegmentCollector + Send + Sync { + type Fruit = TSegmentCollector::Fruit; + + fn collect(&mut self, doc: u32, score: Score) { + let value = self.fast_field_reader.get(doc); + if (self.predicate)(value) { + self.segment_collector.collect(doc, score) + } + } + + fn harvest(self) -> ::Fruit { + self.segment_collector.harvest() + } +} diff --git a/src/collector/mod.rs b/src/collector/mod.rs index 2533c1950..73fc432b7 100644 --- a/src/collector/mod.rs +++ b/src/collector/mod.rs @@ -114,6 +114,9 @@ use crate::query::Weight; mod docset_collector; pub use self::docset_collector::DocSetCollector; +mod filter_collector_wrapper; +pub use self::filter_collector_wrapper::FilterCollector; + /// `Fruit` is the type for the result of our collection. /// e.g. `usize` for the `Count` collector. pub trait Fruit: Send + downcast_rs::Downcast {} From fea0ba10427ea1d3a9b4ce81bf52eff39a96f7ad Mon Sep 17 00:00:00 2001 From: barrotsteindev Date: Sat, 21 Nov 2020 21:24:27 +0200 Subject: [PATCH 2/9] removed unnecessary static liftimes --- src/collector/filter_collector_wrapper.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/collector/filter_collector_wrapper.rs b/src/collector/filter_collector_wrapper.rs index ddcf68cb7..4bd7a5bfe 100644 --- a/src/collector/filter_collector_wrapper.rs +++ b/src/collector/filter_collector_wrapper.rs @@ -72,7 +72,7 @@ pub struct FilterCollector { impl FilterCollector where TCollector: Collector + Send + Sync, - TSegmentCollector: 'static + SegmentCollector + Send + Sync { + TSegmentCollector: SegmentCollector + Send + Sync { pub fn new(field: Field, predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), collector: TCollector) -> FilterCollector { FilterCollector { field, predicate, collector, phantom: PhantomData } } @@ -80,7 +80,7 @@ where impl Collector for FilterCollector where - TSegmentCollector: 'static + SegmentCollector + Send + Sync, + TSegmentCollector: SegmentCollector + Send + Sync, TCollector: Collector + Send + Sync { // That's the type of our result. // Our standard deviation will be a float. @@ -137,7 +137,7 @@ pub struct FilterSegmentCollector { impl SegmentCollector for FilterSegmentCollector where - TSegmentCollector: 'static + SegmentCollector + Send + Sync { + TSegmentCollector: SegmentCollector + Send + Sync { type Fruit = TSegmentCollector::Fruit; fn collect(&mut self, doc: u32, score: Score) { From 8782c0eada1dc60dd7dde3c63f6557c48ac6e098 Mon Sep 17 00:00:00 2001 From: barrotsteindev Date: Sat, 21 Nov 2020 21:26:44 +0200 Subject: [PATCH 3/9] updated docs --- src/collector/filter_collector_wrapper.rs | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/src/collector/filter_collector_wrapper.rs b/src/collector/filter_collector_wrapper.rs index 4bd7a5bfe..311112697 100644 --- a/src/collector/filter_collector_wrapper.rs +++ b/src/collector/filter_collector_wrapper.rs @@ -16,15 +16,8 @@ use crate::fastfield::FastFieldReader; use crate::schema::Field; use crate::{Score, SegmentReader, TantivyError}; -/// The `TopDocs` collector keeps track of the top `K` documents -/// sorted by their score. -/// -/// The implementation is based on a `BinaryHeap`. -/// The theorical complexity for collecting the top `K` out of `n` documents -/// is `O(n log K)`. -/// -/// This collector guarantees a stable sorting in case of a tie on the -/// document score. As such, it is suitable to implement pagination. +/// The `FilterCollector` collector filters docs using a u64 fast field value and a predicate. +/// Only the documents for which the predicate returned "true" will be passed on to the next collector. /// /// ```rust /// use tantivy::collector::{TopDocs, FilterCollector}; From 7c94dfdc15e622eb759fc50bf63b853d4f086009 Mon Sep 17 00:00:00 2001 From: barrotsteindev Date: Sun, 22 Nov 2020 10:40:45 +0200 Subject: [PATCH 4/9] fmt --- src/collector/filter_collector_wrapper.rs | 67 +++++++++++++---------- 1 file changed, 39 insertions(+), 28 deletions(-) diff --git a/src/collector/filter_collector_wrapper.rs b/src/collector/filter_collector_wrapper.rs index 311112697..54e89c267 100644 --- a/src/collector/filter_collector_wrapper.rs +++ b/src/collector/filter_collector_wrapper.rs @@ -49,32 +49,43 @@ use crate::{Score, SegmentReader, TantivyError}; /// assert_eq!(top_docs.len(), 2); /// assert_eq!(top_docs[0].1, DocAddress(0, 1)); /// assert_eq!(top_docs[1].1, DocAddress(0, 3)); -/// +/// /// let filter_all_collector = FilterCollector::new(price, &|value| false, TopDocs::with_limit(2)); /// let filtered_top_docs = searcher.search(&query, &filter_all_collector).unwrap(); -/// +/// /// assert_eq!(filtered_top_docs.len(), 0); /// ``` pub struct FilterCollector { field: Field, collector: TCollector, predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), - phantom: PhantomData + phantom: PhantomData, } -impl FilterCollector -where +impl FilterCollector +where TCollector: Collector + Send + Sync, - TSegmentCollector: SegmentCollector + Send + Sync { - pub fn new(field: Field, predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), collector: TCollector) -> FilterCollector { - FilterCollector { field, predicate, collector, phantom: PhantomData } + TSegmentCollector: SegmentCollector + Send + Sync, +{ + pub fn new( + field: Field, + predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), + collector: TCollector, + ) -> FilterCollector { + FilterCollector { + field, + predicate, + collector, + phantom: PhantomData, + } } } -impl Collector for FilterCollector -where +impl Collector for FilterCollector +where TSegmentCollector: SegmentCollector + Send + Sync, - TCollector: Collector + Send + Sync { + TCollector: Collector + Send + Sync, +{ // That's the type of our result. // Our standard deviation will be a float. type Fruit = TCollector::Fruit; @@ -98,26 +109,25 @@ where })?; let child_segment_collector = self.collector.for_segment(segment_local_id, segment_reader); match child_segment_collector { - Ok(segment_collector) => { - Ok(FilterSegmentCollector:: { - fast_field_reader, - segment_collector: segment_collector, - predicate: self.predicate - }) - }, - Err(_) => { - Err(TantivyError::SystemError("Could not open segment: ".to_owned())) - }, + Ok(segment_collector) => Ok(FilterSegmentCollector:: { + fast_field_reader, + segment_collector: segment_collector, + predicate: self.predicate, + }), + Err(_) => Err(TantivyError::SystemError( + "Could not open segment: ".to_owned(), + )), } - - } fn requires_scoring(&self) -> bool { self.collector.requires_scoring() } - fn merge_fruits(&self, segment_fruits: Vec<::Fruit>) -> crate::Result { + fn merge_fruits( + &self, + segment_fruits: Vec<::Fruit>, + ) -> crate::Result { self.collector.merge_fruits(segment_fruits) } } @@ -125,12 +135,13 @@ where pub struct FilterSegmentCollector { fast_field_reader: FastFieldReader, segment_collector: TSegmentCollector, - predicate: &'static (dyn Fn(u64) -> bool + Send + Sync) + predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), } -impl SegmentCollector for FilterSegmentCollector -where - TSegmentCollector: SegmentCollector + Send + Sync { +impl SegmentCollector for FilterSegmentCollector +where + TSegmentCollector: SegmentCollector + Send + Sync, +{ type Fruit = TSegmentCollector::Fruit; fn collect(&mut self, doc: u32, score: Score) { From b6ff29e020bcee3a3b0c77c3928ed649f4da8e68 Mon Sep 17 00:00:00 2001 From: barrotsteindev Date: Mon, 23 Nov 2020 11:09:03 +0200 Subject: [PATCH 5/9] simplified FilterCollector#for_segment --- src/collector/filter_collector_wrapper.rs | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/src/collector/filter_collector_wrapper.rs b/src/collector/filter_collector_wrapper.rs index 54e89c267..86b7a0501 100644 --- a/src/collector/filter_collector_wrapper.rs +++ b/src/collector/filter_collector_wrapper.rs @@ -107,17 +107,14 @@ where field_name )) })?; - let child_segment_collector = self.collector.for_segment(segment_local_id, segment_reader); - match child_segment_collector { - Ok(segment_collector) => Ok(FilterSegmentCollector:: { - fast_field_reader, - segment_collector: segment_collector, - predicate: self.predicate, - }), - Err(_) => Err(TantivyError::SystemError( - "Could not open segment: ".to_owned(), - )), - } + let segment_collector = self + .collector + .for_segment(segment_local_id, segment_reader)?; + Ok(FilterSegmentCollector:: { + fast_field_reader, + segment_collector: segment_collector, + predicate: self.predicate, + }) } fn requires_scoring(&self) -> bool { From db514208a738444957d09b6eb33637dbce5336ac Mon Sep 17 00:00:00 2001 From: Paul Masurel Date: Wed, 25 Nov 2020 14:13:22 +0900 Subject: [PATCH 6/9] Removed the SegmentCollector type from the Generics of the FilterCollector --- src/collector/filter_collector_wrapper.rs | 26 +++++++++-------------- 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/src/collector/filter_collector_wrapper.rs b/src/collector/filter_collector_wrapper.rs index 86b7a0501..5436833bc 100644 --- a/src/collector/filter_collector_wrapper.rs +++ b/src/collector/filter_collector_wrapper.rs @@ -9,8 +9,6 @@ // --- // Importing tantivy... -use std::marker::PhantomData; - use crate::collector::{Collector, SegmentCollector}; use crate::fastfield::FastFieldReader; use crate::schema::Field; @@ -55,48 +53,44 @@ use crate::{Score, SegmentReader, TantivyError}; /// /// assert_eq!(filtered_top_docs.len(), 0); /// ``` -pub struct FilterCollector { +pub struct FilterCollector { field: Field, collector: TCollector, predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), - phantom: PhantomData, } -impl FilterCollector +impl FilterCollector where - TCollector: Collector + Send + Sync, - TSegmentCollector: SegmentCollector + Send + Sync, + TCollector: Collector + Send + Sync, { pub fn new( field: Field, predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), collector: TCollector, - ) -> FilterCollector { + ) -> FilterCollector { FilterCollector { field, predicate, collector, - phantom: PhantomData, } } } -impl Collector for FilterCollector +impl Collector for FilterCollector where - TSegmentCollector: SegmentCollector + Send + Sync, - TCollector: Collector + Send + Sync, + TCollector: Collector + Send + Sync, { // That's the type of our result. // Our standard deviation will be a float. type Fruit = TCollector::Fruit; - type Child = FilterSegmentCollector; + type Child = FilterSegmentCollector; fn for_segment( &self, segment_local_id: u32, segment_reader: &SegmentReader, - ) -> crate::Result> { + ) -> crate::Result> { let fast_field_reader = segment_reader .fast_fields() .u64(self.field) @@ -110,7 +104,7 @@ where let segment_collector = self .collector .for_segment(segment_local_id, segment_reader)?; - Ok(FilterSegmentCollector:: { + Ok(FilterSegmentCollector { fast_field_reader, segment_collector: segment_collector, predicate: self.predicate, @@ -137,7 +131,7 @@ pub struct FilterSegmentCollector { impl SegmentCollector for FilterSegmentCollector where - TSegmentCollector: SegmentCollector + Send + Sync, + TSegmentCollector: SegmentCollector, { type Fruit = TSegmentCollector::Fruit; From bc44543d8f805551fa117666a29911b578d2dc18 Mon Sep 17 00:00:00 2001 From: barrotsteindev Date: Wed, 25 Nov 2020 14:01:05 +0200 Subject: [PATCH 7/9] added TPredicate generic param and updated tests --- src/collector/filter_collector_wrapper.rs | 40 ++++++++++++++--------- 1 file changed, 25 insertions(+), 15 deletions(-) diff --git a/src/collector/filter_collector_wrapper.rs b/src/collector/filter_collector_wrapper.rs index 5436833bc..9a60a8a44 100644 --- a/src/collector/filter_collector_wrapper.rs +++ b/src/collector/filter_collector_wrapper.rs @@ -41,33 +41,37 @@ use crate::{Score, SegmentReader, TantivyError}; /// /// let query_parser = QueryParser::for_index(&index, vec![title]); /// let query = query_parser.parse_query("diary").unwrap(); -/// let no_filter_collector = FilterCollector::new(price, &|value| true, TopDocs::with_limit(2)); +/// let no_filter_collector = FilterCollector::new(price, &|value| value > 20_120u64, TopDocs::with_limit(1)); /// let top_docs = searcher.search(&query, &no_filter_collector).unwrap(); /// -/// assert_eq!(top_docs.len(), 2); +/// assert_eq!(top_docs.len(), 1); /// assert_eq!(top_docs[0].1, DocAddress(0, 1)); -/// assert_eq!(top_docs[1].1, DocAddress(0, 3)); /// -/// let filter_all_collector = FilterCollector::new(price, &|value| false, TopDocs::with_limit(2)); +/// let filter_all_collector = FilterCollector::new(price, &|value| value < 5u64, TopDocs::with_limit(2)); /// let filtered_top_docs = searcher.search(&query, &filter_all_collector).unwrap(); /// /// assert_eq!(filtered_top_docs.len(), 0); /// ``` -pub struct FilterCollector { +pub struct FilterCollector +where + TPredicate: 'static, +{ field: Field, collector: TCollector, - predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), + predicate: &'static TPredicate, } -impl FilterCollector +impl FilterCollector where TCollector: Collector + Send + Sync, + TPredicate: Fn(u64) -> bool + Send + Sync, { + /// Create a new FilterCollector. pub fn new( field: Field, - predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), + predicate: &'static TPredicate, collector: TCollector, - ) -> FilterCollector { + ) -> FilterCollector { FilterCollector { field, predicate, @@ -76,21 +80,22 @@ where } } -impl Collector for FilterCollector +impl Collector for FilterCollector where TCollector: Collector + Send + Sync, + TPredicate: 'static + Fn(u64) -> bool + Send + Sync, { // That's the type of our result. // Our standard deviation will be a float. type Fruit = TCollector::Fruit; - type Child = FilterSegmentCollector; + type Child = FilterSegmentCollector; fn for_segment( &self, segment_local_id: u32, segment_reader: &SegmentReader, - ) -> crate::Result> { + ) -> crate::Result> { let fast_field_reader = segment_reader .fast_fields() .u64(self.field) @@ -123,15 +128,20 @@ where } } -pub struct FilterSegmentCollector { +pub struct FilterSegmentCollector +where + TPredicate: 'static, +{ fast_field_reader: FastFieldReader, segment_collector: TSegmentCollector, - predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), + predicate: &'static TPredicate, } -impl SegmentCollector for FilterSegmentCollector +impl SegmentCollector + for FilterSegmentCollector where TSegmentCollector: SegmentCollector, + TPredicate: 'static + Fn(u64) -> bool + Send + Sync, { type Fruit = TSegmentCollector::Fruit; From f288e32634f74f2e4ba3371c6bbaa53d38b6aca7 Mon Sep 17 00:00:00 2001 From: barrotsteindev Date: Wed, 25 Nov 2020 14:08:43 +0200 Subject: [PATCH 8/9] rebaes on origin/main --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 06ecf19d7..2be1c1319 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ Tantivy 0.14.0 - Added helper for building intersections and unions in BooleanQuery (@guilload) - Bugfix in `Query::explain` - Removed dependency on `notify` #924. Replaced with `FileWatcher` struct that polls meta file every 500ms in background thread. (@halvorboe @guilload) +- Added `FilterCollector`, which wraps another collector and filters docs using a predicate over a fast field (@barrotsteindev) Tantivy 0.13.2 =================== From c805871b9210c31835f9490339d3199bb99601da Mon Sep 17 00:00:00 2001 From: barrotsteindev Date: Wed, 25 Nov 2020 14:25:49 +0200 Subject: [PATCH 9/9] better test --- src/collector/filter_collector_wrapper.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/collector/filter_collector_wrapper.rs b/src/collector/filter_collector_wrapper.rs index 9a60a8a44..0e3e33927 100644 --- a/src/collector/filter_collector_wrapper.rs +++ b/src/collector/filter_collector_wrapper.rs @@ -41,7 +41,7 @@ use crate::{Score, SegmentReader, TantivyError}; /// /// let query_parser = QueryParser::for_index(&index, vec![title]); /// let query = query_parser.parse_query("diary").unwrap(); -/// let no_filter_collector = FilterCollector::new(price, &|value| value > 20_120u64, TopDocs::with_limit(1)); +/// let no_filter_collector = FilterCollector::new(price, &|value| value > 20_120u64, TopDocs::with_limit(2)); /// let top_docs = searcher.search(&query, &no_filter_collector).unwrap(); /// /// assert_eq!(top_docs.len(), 1);