Include full-text scores in search results

Author: arteymix

src/ubic/basecode/ontology/jena/AbstractOntologyResource.java

@@ -24,9 +24,12 @@
 import org.slf4j.LoggerFactory;
 import ubic.basecode.ontology.model.OntologyResource;
 
+import javax.annotation.Nullable;
 import java.util.Comparator;
 import java.util.Objects;
 
+import static java.util.Comparator.*;
+
 /**
  * @author pavlidis
  */
@@ -36,10 +39,22 @@ public abstract class AbstractOntologyResource implements OntologyResource {
 
     private static final long serialVersionUID = 1L;
 
+    private static final Comparator<OntologyResource> comparator = Comparator
+            .comparing( OntologyResource::getScore, nullsLast( reverseOrder() ) )
+            .thenComparing( OntologyResource::getUri, nullsLast( naturalOrder() ) );
+
     private transient final OntResource res;
+    @Nullable
+    private final Double score;
 
     protected AbstractOntologyResource( OntResource resource ) {
         this.res = resource;
+        this.score = null;
+    }
+
+    public AbstractOntologyResource( OntResource resource, double score ) {
+        this.res = resource;
+        this.score = score;
     }
 
     @Override
@@ -66,9 +81,15 @@ public boolean isObsolete() {
         return res.hasLiteral( OWL2.deprecated, true );
     }
 
+    @Override
+    @Nullable
+    public Double getScore() {
+        return score;
+    }
+
     @Override
     public int compareTo( OntologyResource other ) {
-        return Objects.compare( getUri(), other.getUri(), Comparator.nullsLast( Comparator.naturalOrder() ) );
+        return Objects.compare( this, other, comparator );
     }
 
     @Override

src/ubic/basecode/ontology/jena/AbstractOntologyService.java

@@ -224,7 +224,7 @@ public Collection<OntologyIndividual> findIndividuals( String search, boolean ke
                 return Collections.emptySet();
             }
             return OntologySearch.matchIndividuals( model, index, search )
-                    .mapWith( i -> ( OntologyIndividual ) new OntologyIndividualImpl( i, additionalRestrictions ) )
+                    .mapWith( i -> ( OntologyIndividual ) new OntologyIndividualImpl( i.result, additionalRestrictions, i.score ) )
                     .filterKeep( where( ontologyTerm -> keepObsoletes || !ontologyTerm.isObsolete() ) )
                     .toSet();
         } finally {
@@ -246,13 +246,13 @@ public Collection<OntologyResource> findResources( String searchString, boolean
                 return Collections.emptySet();
             }
             return OntologySearch.matchResources( model, index, searchString )
-                    .filterKeep( where( r -> r.canAs( OntClass.class ) || r.canAs( Individual.class ) ) )
+                    .filterKeep( where( r -> r.result.canAs( OntClass.class ) || r.result.canAs( Individual.class ) ) )
                     .mapWith( r -> {
                         OntologyResource res;
-                        if ( r.canAs( OntClass.class ) ) {
-                            res = new OntologyTermImpl( r.as( OntClass.class ), additionalRestrictions );
+                        if ( r.result.canAs( OntClass.class ) ) {
+                            res = new OntologyTermImpl( r.result.as( OntClass.class ), additionalRestrictions, r.score );
                         } else {
-                            res = new OntologyIndividualImpl( r.as( Individual.class ), additionalRestrictions );
+                            res = new OntologyIndividualImpl( r.result.as( Individual.class ), additionalRestrictions, r.score );
                         }
                         return res;
                     } )
@@ -278,7 +278,7 @@ public Collection<OntologyTerm> findTerm( String search, boolean keepObsoletes )
                 return Collections.emptySet();
             }
             return OntologySearch.matchClasses( model, index, search )
-                    .mapWith( r -> ( OntologyTerm ) new OntologyTermImpl( r, additionalRestrictions ) )
+                    .mapWith( r -> ( OntologyTerm ) new OntologyTermImpl( r.result, additionalRestrictions, r.score ) )
                     .filterKeep( where( ontologyTerm -> keepObsoletes || !ontologyTerm.isObsolete() ) )
                     .toSet();
         } finally {

src/ubic/basecode/ontology/jena/OntologyIndividualImpl.java

@@ -46,6 +46,12 @@ public OntologyIndividualImpl( Individual ind, Set<Restriction> additionalRestri
         this.additionalRestrictions = additionalRestrictions;
     }
 
+    public OntologyIndividualImpl( Individual ind, Set<Restriction> additionalRestrictions, double score ) {
+        super( ind, score );
+        this.ind = ind;
+        this.additionalRestrictions = additionalRestrictions;
+    }
+
     @Override
     public OntologyTerm getInstanceOf() {
         Resource type = ind.getRDFType();

src/ubic/basecode/ontology/jena/OntologyTermImpl.java

@@ -25,6 +25,7 @@
 import ubic.basecode.ontology.model.OntologyRestriction;
 import ubic.basecode.ontology.model.OntologyTerm;
 
+import javax.annotation.Nullable;
 import java.util.*;
 import java.util.stream.Collectors;
 
@@ -49,12 +50,18 @@ public class OntologyTermImpl extends AbstractOntologyResource implements Ontolo
      */
     private final transient Set<Restriction> additionalRestrictions;
 
-    public OntologyTermImpl( OntClass resource, Set<Restriction> additionalRestrictions ) {
+    public OntologyTermImpl( OntClass resource, @Nullable Set<Restriction> additionalRestrictions ) {
         super( resource );
         this.ontResource = resource;
         this.additionalRestrictions = additionalRestrictions;
     }
 
+    public OntologyTermImpl( OntClass resource, Set<Restriction> additionalRestrictions, double score ) {
+        super( resource, score );
+        this.ontResource = resource;
+        this.additionalRestrictions = additionalRestrictions;
+    }
+
     @Override
     public Collection<String> getAlternativeIds() {
         Collection<String> results = new HashSet<>();

src/ubic/basecode/ontology/jena/search/OntologySearch.java

@@ -21,11 +21,12 @@
 import com.hp.hpl.jena.ontology.Individual;
 import com.hp.hpl.jena.ontology.OntClass;
 import com.hp.hpl.jena.ontology.OntModel;
-import com.hp.hpl.jena.rdf.model.NodeIterator;
+import com.hp.hpl.jena.rdf.model.Model;
 import com.hp.hpl.jena.rdf.model.RDFNode;
 import com.hp.hpl.jena.rdf.model.Resource;
 import com.hp.hpl.jena.shared.JenaException;
 import com.hp.hpl.jena.util.iterator.ExtendedIterator;
+import com.hp.hpl.jena.util.iterator.Map1Iterator;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.time.StopWatch;
 import org.apache.lucene.queryParser.QueryParser.Operator;
@@ -35,9 +36,11 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
+import static com.hp.hpl.jena.sparql.util.ModelUtils.convertGraphNodeToRDFNode;
 import static ubic.basecode.ontology.jena.JenaUtils.where;
 
 /**
@@ -50,65 +53,45 @@ public class OntologySearch {
     /**
      * Find classes that match the query string.
      *
-     * @param model       that goes with the index
-     * @param index       to search
+     * @param model that goes with the index
+     * @param index to search
      * @return Collection of OntologyTerm objects
      */
-    public static ExtendedIterator<OntClass> matchClasses( OntModel model, SearchIndex index, String queryString ) throws OntologySearchException {
-        NodeIterator iterator = runSearch( index, queryString );
-        return iterator
-                .mapWith( r -> r.inModel( model ) )
-                .filterKeep( where( r -> r.isURIResource() && r.canAs( OntClass.class ) ) )
+    public static ExtendedIterator<SearchResult<OntClass>> matchClasses( OntModel model, SearchIndex index, String queryString ) throws OntologySearchException {
+        return runSearch( model, index, queryString )
+                .filterKeep( where( r -> r.result.isURIResource() && r.result.canAs( OntClass.class ) ) )
                 .mapWith( r -> r.as( OntClass.class ) );
     }
 
     /**
      * Find individuals that match the query string
      *
-     * @param model       that goes with the index
-     * @param index       to search
+     * @param model that goes with the index
+     * @param index to search
      * @return Collection of OntologyTerm objects
      */
-    public static ExtendedIterator<Individual> matchIndividuals( OntModel model, SearchIndex index, String queryString ) throws OntologySearchException {
-        NodeIterator iterator;
-
-        queryString = queryString.trim();
-
-        // Add wildcard only if the last word is longer than one character. This is to prevent lucene from
-        // blowing up. See bug#1145
-        String[] words = queryString.split( "\\s+" );
-        int lastWordLength = words[words.length - 1].length();
-        if ( lastWordLength > 1 ) {
-            try { // Use wildcard search.
-                iterator = runSearch( index, queryString + "*" );
-            } catch ( OntologySearchJenaException e ) { // retry without wildcard
-                log.warn( "Failed to perform search with wildcard. Retrying search without wildcard.", e );
-                try {
-                    iterator = runSearch( index, queryString );
-                } catch ( OntologySearchJenaException e1 ) {
-                    throw new RetryWithoutWildcardFailedException( "Failed to search while retrying without wildcard.", queryString, e.getCause(), e1.getCause() );
-                }
-            }
-        } else {
-            iterator = runSearch( index, queryString );
-        }
-
-        return iterator
-                .mapWith( r -> r.inModel( model ) )
-                .filterKeep( where( r -> r.isURIResource() && r.canAs( Individual.class ) ) )
+    public static ExtendedIterator<SearchResult<Individual>> matchIndividuals( OntModel model, SearchIndex index, String queryString ) throws OntologySearchException {
+        return runSearchWithWildcard( model, index, queryString )
+                .filterKeep( where( r -> r.result.isURIResource() && r.result.canAs( Individual.class ) ) )
                 .mapWith( r -> r.as( Individual.class ) );
     }
 
     /**
      * Find OntologyIndividuals and OntologyTerms that match the query string. Search with a wildcard is attempted
      * whenever possible.
      *
-     * @param model                  that goes with the index
-     * @param index                  to search
+     * @param model that goes with the index
+     * @param index to search
      * @return Collection of OntologyResource objects
      */
-    public static ExtendedIterator<Resource> matchResources( OntModel model, SearchIndex index, String queryString ) throws OntologySearchException {
-        NodeIterator iterator;
+    public static ExtendedIterator<SearchResult<Resource>> matchResources( OntModel model, SearchIndex index, String queryString ) throws OntologySearchException {
+        return runSearchWithWildcard( model, index, queryString )
+                .filterKeep( where( o -> o.result.isURIResource() && o.result.isResource() ) )
+                .mapWith( r -> r.as( Resource.class ) );
+    }
+
+    private static ExtendedIterator<SearchResult<RDFNode>> runSearchWithWildcard( Model model, SearchIndex index, String queryString ) throws OntologySearchException {
+        ExtendedIterator<SearchResult<RDFNode>> iterator;
 
         queryString = queryString.trim();
 
@@ -118,27 +101,24 @@ public static ExtendedIterator<Resource> matchResources( OntModel model, SearchI
         int lastWordLength = words[words.length - 1].length();
         if ( lastWordLength > 1 ) {
             try { // Use wildcard search.
-                iterator = runSearch( index, queryString + "*" );
+                iterator = runSearch( model, index, queryString + "*" );
             } catch ( OntologySearchJenaException e ) { // retry without wildcard
                 // retry without wildcard
                 log.warn( "Failed to search in {}. Retrying search without wildcard.", model, e );
                 try {
-                    iterator = runSearch( index, queryString );
+                    iterator = runSearch( model, index, queryString );
                 } catch ( OntologySearchJenaException e1 ) {
                     throw new RetryWithoutWildcardFailedException( "Failed to search while retrying without wildcard.", queryString, e.getCause(), e1.getCause() );
                 }
             }
         } else {
-            iterator = runSearch( index, queryString );
+            iterator = runSearch( model, index, queryString );
         }
 
-        return iterator
-                .mapWith( r -> r.inModel( model ) )
-                .filterKeep( where( o -> o.isURIResource() && o.isResource() ) )
-                .mapWith( RDFNode::asResource );
+        return iterator;
     }
 
-    private static NodeIterator runSearch( SearchIndex index, String queryString ) throws OntologySearchJenaException {
+    private static ExtendedIterator<SearchResult<RDFNode>> runSearch( Model model, SearchIndex index, String queryString ) throws OntologySearchJenaException {
         String strippedQuery = StringUtils.strip( queryString );
 
         if ( StringUtils.isBlank( strippedQuery ) ) {
@@ -158,7 +138,8 @@ private static NodeIterator runSearch( SearchIndex index, String queryString ) t
 
         StopWatch timer = StopWatch.createStarted();
         try {
-            return index.searchModelByIndex( enhancedQuery );
+            return new Map1Iterator<>( o -> new SearchResult<>( o.getLuceneDocId(), convertGraphNodeToRDFNode( o.getNode(), model ), o.getScore() ),
+                    index.search( enhancedQuery ) );
         } catch ( JenaException e ) {
             throw new OntologySearchJenaException( "Failed to search with enhanced query.", enhancedQuery, e );
         } finally {
@@ -168,4 +149,38 @@ private static NodeIterator runSearch( SearchIndex index, String queryString ) t
             }
         }
     }
+
+    public static class SearchResult<T extends RDFNode> {
+        public final int docId;
+        public final T result;
+        public final double score;
+
+        private SearchResult( int docId, T result, double score ) {
+            this.docId = docId;
+            this.result = result;
+            this.score = score;
+        }
+
+        @Override
+        public boolean equals( Object obj ) {
+            if ( obj instanceof SearchResult ) {
+                return Objects.equals( result, ( ( SearchResult<?> ) obj ).result );
+            }
+            return false;
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash( result );
+        }
+
+        @Override
+        public String toString() {
+            return String.format( "%s [docId = %d, score = %f]", result, docId, score );
+        }
+
+        private <U extends Resource> SearchResult<U> as( Class<U> clazz ) {
+            return new SearchResult<>( docId, result.as( clazz ), score );
+        }
+    }
 }

src/ubic/basecode/ontology/jena/vocabulary/OBO.java

@@ -21,5 +21,5 @@ private static Property property( String name ) {
     public static final Property hasRelatedSynonym = property( "hasRelatedSynonym" );
     public static final Property alternativeLabel = ResourceFactory.createProperty( "http://purl.obolibrary.org/obo/IAO_0000118" );
     public static final Resource ObsoleteClass = ResourceFactory.createResource( "http://www.geneontology.org/formats/oboInOwl#ObsoleteClass" );
-    public static final Property ObsoleteProperty = ResourceFactory.createProperty( "http://www.geneontology.org/formats/oboInOwl#ObsoleteProperty" );
+    public static final Property ObsoleteProperty = property( "ObsoleteProperty" );
 }

src/ubic/basecode/ontology/model/OntologyResource.java

@@ -18,6 +18,7 @@
  */
 package ubic.basecode.ontology.model;
 
+import javax.annotation.Nullable;
 import java.io.Serializable;
 
 /**
@@ -39,4 +40,10 @@ public interface OntologyResource extends Serializable, Comparable<OntologyResou
      * Whether the resource is marked as obsolete.
      */
     boolean isObsolete();
+
+    /**
+     * If this is result from a free-text search, a corresponding score, otherwise null.
+     */
+    @Nullable
+    Double getScore();
 }

src/ubic/basecode/ontology/model/OntologyTermSimple.java

@@ -14,6 +14,7 @@
  */
 package ubic.basecode.ontology.model;
 
+import javax.annotation.Nullable;
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.Objects;
@@ -123,6 +124,12 @@ public boolean isObsolete() {
         return obsolete;
     }
 
+    @Nullable
+    @Override
+    public Double getScore() {
+        return null;
+    }
+
     @Override
     public int compareTo( OntologyResource other ) {
         return Objects.compare( getUri(), other.getUri(), Comparator.nullsLast( Comparator.naturalOrder() ) );

test/ubic/basecode/ontology/OntologyTermTest.java

@@ -308,5 +308,8 @@ public void testGetChildrenFromMultipleTermsWithSearch() throws OntologySearchEx
     public void testFindTerm() throws OntologySearchException {
         assertEquals( 123, uberon.findTerm( "brain" ).size() );
         assertEquals( 128, uberon.findTerm( "brain", true ).size() );
+        OntologyTerm firstResult = uberon.findTerm( "brain" ).iterator().next();
+        assertNotNull( firstResult.getScore() );
+        assertEquals( 2.8577, firstResult.getScore(), 0.0001 );
     }
 }