Struct tantivy::collector::TopDocs

pub struct TopDocs(_);

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.

use tantivy::collector::TopDocs;
use tantivy::query::QueryParser;
use tantivy::schema::{Schema, TEXT};
use tantivy::{doc, DocAddress, Index};

let mut schema_builder = Schema::builder();
let title = schema_builder.add_text_field("title", TEXT);
let schema = schema_builder.build();
let index = Index::create_in_ram(schema);

let mut index_writer = index.writer_with_num_threads(1, 3_000_000).unwrap();
index_writer.add_document(doc!(title => "The Name of the Wind"));
index_writer.add_document(doc!(title => "The Diary of Muadib"));
index_writer.add_document(doc!(title => "A Dairy Cow"));
index_writer.add_document(doc!(title => "The Diary of a Young Girl"));
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 top_docs = searcher.search(&query, &TopDocs::with_limit(2)).unwrap();

assert_eq!(&top_docs[0], &(0.7261542, DocAddress(0, 1)));
assert_eq!(&top_docs[1], &(0.6099695, DocAddress(0, 3)));

Methods

impl TopDocs

pub fn with_limit(limit: usize) -> TopDocs

Creates a top score collector, with a number of documents equal to "limit".

Panics

The method panics if limit is 0

pub fn order_by_u64_field(
    self,
    field: Field
) -> impl Collector<Fruit = Vec<(u64, DocAddress)>>

Set top-K to rank documents by a given fast field.

use tantivy::Searcher;
use tantivy::collector::TopDocs;
use tantivy::schema::Field;



/// Searches the document matching the given query, and
/// collects the top 10 documents, order by the u64-`field`
/// given in argument.
///
/// `field` is required to be a FAST field.
fn docs_sorted_by_rating(searcher: &Searcher,
                         query: &dyn Query,
                         sort_by_field: Field)
    -> Result<Vec<(u64, DocAddress)>> {

    // This is where we build our topdocs collector
    //
    // Note the generics parameter that needs to match the
    // type `sort_by_field`.
    let top_docs_by_rating = TopDocs
                ::with_limit(10)
                 .order_by_u64_field(sort_by_field);
     
    // ... and here are our documents. Note this is a simple vec.
    // The `u64` in the pair is the value of our fast field for
    // each documents.
    //
    // The vec is sorted decreasingly by `sort_by_field`, and has a
    // length of 10, or less if not enough documents matched the
    // query.
    let resulting_docs: Vec<(u64, DocAddress)> =
         searcher.search(query, &top_docs_by_rating)?;
     
    Ok(resulting_docs)
}

Panics

May panic if the field requested is not a fast field.

pub fn tweak_score<TScore, TScoreSegmentTweaker, TScoreTweaker>(
    self,
    score_tweaker: TScoreTweaker
) -> impl Collector<Fruit = Vec<(TScore, DocAddress)>> where
    TScore: 'static + Send + Sync + Clone + PartialOrd,
    TScoreSegmentTweaker: ScoreSegmentTweaker<TScore> + 'static,
    TScoreTweaker: ScoreTweaker<TScore, Child = TScoreSegmentTweaker>, 

Ranks the documents using a custom score.

This method offers a convenient way to tweak or replace the documents score. As suggested by the prototype you can manually define your own ScoreTweaker and pass it as an argument, but there is a much simpler way to tweak your score: you can use a closure as in the following example.

Example

Typically, you will want to rely on one or more fast fields, to alter the original relevance Score.

For instance, in the following, we assume that we are implementing an e-commerce website that has a fast field called popularity that rates whether a product is typically often bought by users.

In the following example will will tweak our ranking a bit by boosting popular products a notch.

In more serious application, this tweaking could involved running a learning-to-rank model over various features

use tantivy::SegmentReader;
use tantivy::collector::TopDocs;
use tantivy::schema::Field;

fn create_schema() -> Schema {
   let mut schema_builder = Schema::builder();
   schema_builder.add_text_field("product_name", TEXT);
   schema_builder.add_u64_field("popularity", FAST);
   schema_builder.build()
}

fn create_index() -> tantivy::Result<Index> {
  let schema = create_schema();
  let index = Index::create_in_ram(schema);
  let mut index_writer = index.writer_with_num_threads(1, 3_000_000)?;
  let product_name = index.schema().get_field("product_name").unwrap();
  let popularity: Field = index.schema().get_field("popularity").unwrap();
  index_writer.add_document(doc!(product_name => "The Diary of Muadib", popularity => 1u64));
  index_writer.add_document(doc!(product_name => "A Dairy Cow", popularity => 10u64));
  index_writer.add_document(doc!(product_name => "The Diary of a Young Girl", popularity => 15u64));
  index_writer.commit()?;
  Ok(index)
}

let index = create_index().unwrap();
let product_name = index.schema().get_field("product_name").unwrap();
let popularity: Field = index.schema().get_field("popularity").unwrap();

let user_query_str = "diary";
let query_parser = QueryParser::for_index(&index, vec![product_name]);
let query = query_parser.parse_query(user_query_str).unwrap();

// This is where we build our collector with our custom score.
let top_docs_by_custom_score = TopDocs
        ::with_limit(10)
         .tweak_score(move |segment_reader: &SegmentReader| {
            // The argument is a function that returns our scoring
            // function.
            //
            // The point of this "mother" function is to gather all
            // of the segment level information we need for scoring.
            // Typically, fast_fields.
            //
            // In our case, we will get a reader for the popularity
            // fast field.
            let popularity_reader =
                segment_reader.fast_fields().u64(popularity).unwrap();

            // We can now define our actual scoring function
            move |doc: DocId, original_score: Score| {
                let popularity: u64 = popularity_reader.get(doc);
                // Well.. For the sake of the example we use a simple logarithm
                // function.
                let popularity_boost_score = ((2u64 + popularity) as f32).log2();
                popularity_boost_score * original_score
            }
          });
let reader = index.reader().unwrap();
let searcher = reader.searcher();
// ... and here are our documents. Note this is a simple vec.
// The `Score` in the pair is our tweaked score.
let resulting_docs: Vec<(Score, DocAddress)> =
     searcher.search(&query, &top_docs_by_custom_score).unwrap();

See also

custom_score(...).

pub fn custom_score<TScore, TCustomSegmentScorer, TCustomScorer>(
    self,
    custom_score: TCustomScorer
) -> impl Collector<Fruit = Vec<(TScore, DocAddress)>> where
    TScore: 'static + Send + Sync + Clone + PartialOrd,
    TCustomSegmentScorer: CustomSegmentScorer<TScore> + 'static,
    TCustomScorer: CustomScorer<TScore, Child = TCustomSegmentScorer>, 

Ranks the documents using a custom score.

This method offers a convenient way to use a different score.

As suggested by the prototype you can manually define your own CustomScorer and pass it as an argument, but there is a much simpler way to tweak your score: you can use a closure as in the following example.

Limitation

This method only makes it possible to compute the score from a given DocId, fastfield values for the doc and any information you could have precomputed beforehands. It does not make it possible for instance to compute something like TfIdf as it does not have access to the list of query terms present in the document, nor the term frequencies for the different terms.

It can be used if your search engine relies on a learning-to-rank model for instance, which does not rely on the term frequencies or positions as features.

Example

use tantivy::SegmentReader;
use tantivy::collector::TopDocs;
use tantivy::schema::Field;

let popularity: Field = index.schema().get_field("popularity").unwrap();
let boosted: Field = index.schema().get_field("boosted").unwrap();
// ...

// This is where we build our collector with our custom score.
let top_docs_by_custom_score = TopDocs
        ::with_limit(10)
         .custom_score(move |segment_reader: &SegmentReader| {
            // The argument is a function that returns our scoring
            // function.
            //
            // The point of this "mother" function is to gather all
            // of the segment level information we need for scoring.
            // Typically, fast_fields.
            //
            // In our case, we will get a reader for the popularity
            // fast field and a boosted field.
            //
            // We want to get boosted items score, and when we get
            // a tie, return the item with the highest popularity.
            //
            // Note that this is implemented by using a `(u64, u64)`
            // as a score.
            let popularity_reader =
                segment_reader.fast_fields().u64(popularity).unwrap();
            let boosted_reader =
                segment_reader.fast_fields().u64(boosted).unwrap();
    
            // We can now define our actual scoring function
            move |doc: DocId| {
                let popularity: u64 = popularity_reader.get(doc);
                let boosted: u64 = boosted_reader.get(doc);
                // Score do not have to be `f64` in tantivy.
                // Here we return a couple to get lexicographical order
                // for free.
                (boosted, popularity)
            }
          });
// ... and here are our documents. Note this is a simple vec.
// The `Score` in the pair is our tweaked score.
let resulting_docs: Vec<((u64, u64), DocAddress)> =
     searcher.search(&*query, &top_docs_by_custom_score)?;

See also

tweak_score(...).

Trait Implementations

impl Collector for TopDocs

type Fruit = Vec<(Score, DocAddress)>

Fruit is the type for the result of our collection. e.g. usize for the Count collector.

type Child = TopScoreSegmentCollector

Type of the SegmentCollector associated to this collector.

fn for_segment(
    &self,
    segment_local_id: SegmentLocalId,
    reader: &SegmentReader
) -> Result<Self::Child>

set_segment is called before beginning to enumerate on this segment.

fn requires_scoring(&self) -> bool

Returns true iff the collector requires to compute scores for documents.

fn merge_fruits(
    &self,
    child_fruits: Vec<Vec<(Score, DocAddress)>>
) -> Result<Self::Fruit>

Combines the fruit associated to the collection of each segments into one fruit.

impl Debug for TopDocs

fn fmt(&self, f: &mut Formatter) -> Result

Formats the value using the given formatter.