Initial import into git.

[galago.git] / cpp / galago / src / DocumentOrderedBinnedMaxScoreRetrieval.cpp

//
// DocumentOrderedBinnedMaxScoreRetrieval
//
// 5 November 2007 -- Trevor Strohman
//
// BSD License (http://galagosearch.org/license)
//

#include "DocumentOrderedBinnedMaxScoreRetrieval.hpp"
#include "indri/ScoredExtentResult.hpp"
#include "Query.hpp"
#include <vector>
#include <queue>
#include <assert.h>
#include "lemur/Exception.hpp"
#include "Logging.hpp"

#define MAX_TERMS (256)

//
// convertQueue
//

static std::vector<indri::api::ScoredExtentResult> _convertQueue( std::priority_queue<indri::api::ScoredExtentResult>& queue ) {
    std::vector<indri::api::ScoredExtentResult> result;
    
    while( queue.size() > 0 ) {
        result.push_back( queue.top() );
        queue.pop();
    }
    
    std::sort( result.begin(), result.end() );
    return result;
}

//
// removeDoneIterators
//
// Scans through the list of iterators, looking for any that are now done.  If an
// iterator is done, it is removed from the vector and deleted.
//

static void _removeDoneIterators( std::vector<DocumentOrderedBinnedIterator*>& iterators ) {
    for( int i=0; i<iterators.size(); i++ ) {
        if( iterators[i]->isDone() ) {
            delete iterators[i];
            iterators.erase( iterators.begin() + i );
            i--;
        }
    }
}

//
// computeRemaining
//
// For each iterator i, remaining[i] is the total remaining score possible after
// iterator i.  remaining[n-1] is defined to be 0.
//
// This iterator can be used for early termination.  If 
// 
//

static void _computeRemaining( int* remaining,
                               const std::vector<DocumentOrderedBinnedIterator*>& iterators ) {
    if( iterators.size() > 0 )
        remaining[iterators.size()-1] = 0;
        
    for( int i=iterators.size()-2; i >= 0; i-- ) {
        remaining[i] = remaining[i+1] + iterators[i+1]->listBound();
    }
}

//
// computeBounds
//
// For each iterator i, bounds[i] is the required score contribution of that
// iterator (and all previous iterators) in order to have a chance of entering
// the ranked list.
//

static void _computeBounds( int* bounds,
                            int& important,
                            int& unimportantBound,
                            const std::vector<DocumentOrderedBinnedIterator*>& iterators,
                            int threshold ) {
    int total = 0;
    unimportantBound = 0;
    important = iterators.size();
    
    for( int i=iterators.size()-1; i >= 0; i-- ) {
        int oldBound = bounds[i];
        bounds[i] = std::max( 0, threshold - total );
        total += iterators[i]->listBound();

        if( oldBound == 0 && bounds[i] > 0 ) {
            galago_log_query_terms( i );
        }
        
        if( i > 0 && bounds[i-1] == 0 && bounds[i] > 0 ) {
            important = i+1;
            unimportantBound = bounds[i];
        }
    }
    
    assert( important <= iterators.size() );
    bounds[iterators.size()] = threshold;
}


static void _postCompletedIterator( std::vector<DocumentOrderedBinnedIterator*>& iterators,
                                    int& important,
                                    int& unimportantBound,
                                    int* bounds,
                                    int* remaining,
                                    int threshold ) {
    _removeDoneIterators( iterators );
    _computeBounds( bounds, important, unimportantBound, iterators, threshold );
    _computeRemaining( remaining, iterators );
}

//
// openRead
//

void DocumentOrderedBinnedMaxScoreRetrieval::openRead( const std::string& indexPath ) {
    _index.openRead( indexPath );
}                               

//
// runQuery template 
//
// This method contains the core logic for running queries.
//
// Notice that the _scoreSkipping and _conjunctive variables
// are actually template parameters.  This allows the compiler to strip out large
// sections of the code based on these parameter settings, and to schedule the
// instructions more efficiently.
//
// This method always at least does Turtle and Flood's max-score optimization.
// If _scoreSkipping is enabled, this also does Strohman's score skipping
// optimization.  If the _conjunctive option is set, it requires all query 
// terms to appear in any query result.
//

template<bool _scoreSkipping, bool _conjunctive>
static std::vector<indri::api::ScoredExtentResult>
        _runQuery( std::vector<DocumentOrderedBinnedIterator*>& iterators,
                  int requested,
                  int threshold ) {
    galago_logging_docordered log;

    threshold = std::max( threshold, 0 );
    assert( requested > 0 );
    
    if( iterators.size() > MAX_TERMS ) {
        LEMUR_THROW( LEMUR_BAD_PARAMETER_ERROR, "Too many terms" );
    }
    
    std::priority_queue<indri::api::ScoredExtentResult> queue;

    // we do fixed allocation for terms here to help out the compiler a little
    int bounds[MAX_TERMS+1];
    int remaining[MAX_TERMS+1];
    
    // all iterators < important are the "important" iterators for max score
    int important = iterators.size();
    // if the score contribution of the important iterators is not at least
    // unimportantBound, then don't bother with the unimportant iterators
    int unimportantBound = 0;
    int document = MAX_INT32;
    
    // find the smallest document in any iterator
    for( int i=0; i<iterators.size(); i++ ) {
        document = std::min( iterators[i]->currentDocument(), document );
    }

    _computeBounds( bounds, important, unimportantBound, iterators, threshold );
    _computeRemaining( remaining, iterators );

    retry:
    while (document < MAX_INT32) {
        // score candidate
        int score = 0;
        int lastIter = 0;

        // assertion: document is the minimum document to which any iterator
        // in the important set is pointing.
        
        // score all of the important iterators
        for( int i=0; i<important; i++ ) {
            DocumentOrderedBinnedIterator* iter = iterators[i];
            assert( !iter->isDone() );
            lastIter = i;

            if( iter->currentDocument() == document ) {
                score += iter->currentScore();
            } else if( _conjunctive ) {
                score = -1;
                break;
            }
        }
        
        // if the score we've achieved so far is not greater than unimportantBound,
        // this document will not enter the ranked list.  Otherwise we continue
        // looking at the remaining iterators.

        if( score >= unimportantBound ) {
            // now, fill in remaining scores
            for( int i=important; i<iterators.size(); i++ ) {
                lastIter = i;
                
                // this is the amount of score mass we need to continue
                int required = std::max( 0, threshold - score - remaining[i] );

                DocumentOrderedBinnedIterator* iter = iterators[i];
                assert( !iter->isDone() );

                // if we're score skipping, we first make sure that we're caught up
                // on reading skip data.  The skip data gives us a bound on the score
                // contribution from this document.  If that bound is less than the 
                // required score from this iterator, we can quit without actually
                // looking at the posting.
                
                if( _scoreSkipping ) {
                    iter->readSkipsTo( document );
                    assert( !iter->isDone() );
                    assert( iter->currentBoundDocument() >= document && iter->lastBoundDocument() < document );
                    
                    if( iter->currentBound() < required ) {
                        score = -1;
                        break;
                    }
                }
                
                assert( !iter->isDone() );
                bool foundDocument = iter->skipToDocument( document );

                // the movement might have caused the iterator to end, so check that
                if( iter->isDone() ) {
                    if( _conjunctive ) {
                        galago_log( log );
                        return _convertQueue( queue );
                    } else {
                        _postCompletedIterator( iterators, important, unimportantBound,
                                                bounds, remaining, threshold );
                        goto retry;
                    }
                }
                
                assert( !iter->isDone() );
                
                // add in the score from this document
                if( foundDocument ) {
                    score += iter->currentScore();
                } else if( _conjunctive ) {
                    log.docsScored[i]++;
                    score = -1;
                    break;
                }

                // break if there's no way this document can enter the ranked list
                if( bounds[i] > score ) {
                    log.docsScored[i]++;
                    break;
                }
            }
        }
    
        // push score onto heap if it's a good one
        if( score >= threshold ) {
            log.fullScored++;
            queue.push( indri::api::ScoredExtentResult( score, document, 0, 0 ) );

            while( queue.size() > requested ) {
                queue.pop();
                
                if( queue.size() >= 1 ) {
                    int lastThreshold = threshold;
                    threshold = (int)queue.top().score;
                    
                    if( lastThreshold != threshold ) {
                        _computeBounds( bounds, important, unimportantBound, iterators, threshold );
                    }
                }
            }
        }

        // move to the next reasonable document
        int lastDocument = document;

        if( !_conjunctive ) {
            document = MAX_INT32;
        }

        // if the first term is required, skip forward to find a usable document
        if( _scoreSkipping && bounds[0] > 0 ) {
            DocumentOrderedBinnedIterator* iter = iterators[0];

            assert(!iter->isDone());
            if( iter->currentDocument() == lastDocument )
                iter->nextDocument();
            
            if( iter->isDone() )
                break;
            
            // make sure we're caught up on the skip data
            iter->readSkipsTo( iter->currentDocument() );
            
            // skip forward until we find a document that has a score of
            // at least bounds[0].
            iter->skipToBound( bounds[0] );
            document = iter->currentDocument();
            log.scoreBoundSkip++;
            
            if( iter->isDone() )
                break;
            
            continue;
        }

        // try to find a usable document
        assert( important <= iterators.size() );
        for( int i=0; i<important; i++ ) {
            DocumentOrderedBinnedIterator* iter = iterators[i];
            assert( iter->isDone() == false );

            // if this iterator is pointing to the last scored document, we need to move it
            if( iter->currentDocument() == lastDocument ) {
                iter->nextDocument();

                // might have run off the end of the iterator, so check that
                if( iter->isDone() ) {
                    if( _conjunctive ) {
                        galago_log( log );
                        return _convertQueue( queue );
                    } else {
                        _postCompletedIterator( iterators, important, unimportantBound,
                                                bounds, remaining, threshold );
                        i--;
                    }

                    iter = 0;
                }
            } 

            if( iter != 0 ) {
                assert( !iter->isDone() );
                int current = iter->currentDocument();

                if( current > lastDocument ) {
                    if( _conjunctive ) {
                        document = std::max( document, iter->currentDocument() );
                    } else {
                        document = std::min( document, iter->currentDocument() );
                    }
                }   
            }
        }
    }

    galago_log( log );
    return _convertQueue(queue);
}

//
// runQuery
//

std::vector<indri::api::ScoredExtentResult>
        DocumentOrderedBinnedMaxScoreRetrieval::runQuery( const std::vector<QueryTerm>& terms,
                                                       int requested,
                                                       int threshold ) {
    std::vector<DocumentOrderedBinnedIterator*> iterators = getIterators( terms );
    std::sort( iterators.begin(), iterators.end(), DocumentOrderedBinnedIterator::bound_greater() );
    galago_log_query_terms( iterators.size() );

    if( _scoreSkipping == true && _conjunctive == true ) {
        return _runQuery<true, true>( iterators, requested, threshold );
    } else if( _scoreSkipping == true && _conjunctive == false ) {
        return _runQuery<true, false>( iterators, requested, threshold );
    } else if( _scoreSkipping == false && _conjunctive == true ) {
        return _runQuery<false, true>( iterators, requested, threshold );
    } else {
        return _runQuery<false, false>( iterators, requested, threshold );
    } 
}                                                

//
// getIterators
//

std::vector<DocumentOrderedBinnedIterator*>
    DocumentOrderedBinnedMaxScoreRetrieval::getIterators( const std::vector<QueryTerm>& terms ) {
    std::vector<DocumentOrderedBinnedIterator*> iterators;
    
    for( int i=0; i<terms.size(); i++ ) {
        const QueryTerm& term = terms[i];
        DocumentOrderedBinnedIterator* iterator = _index.getTerm( term.text, term.field );

        if( iterator )
            iterators.push_back( iterator );
    }
    
    return iterators;
}

//
// getDocument
//

std::string DocumentOrderedBinnedMaxScoreRetrieval::getDocument( int doc ) {
    return _index.getDocument(doc);
}

//
// close
//

void DocumentOrderedBinnedMaxScoreRetrieval::close() {
    _index.close();
}