Skip to content

Commit

Permalink
Add score to fulltext results
Browse files Browse the repository at this point in the history
  • Loading branch information
ragadeeshu committed Jan 17, 2018
1 parent 61d23e0 commit 154a82a
Show file tree
Hide file tree
Showing 10 changed files with 235 additions and 44 deletions.
Expand Up @@ -350,6 +350,12 @@ public long current()
return next;
}

@Override
public float currentScore()
{
return 0;
}

@Override
public long getValue( String field )
{
Expand Down Expand Up @@ -795,6 +801,12 @@ public long current()
return index;
}

@Override
public float currentScore()
{
return scoreDocs.getCurrentDoc().score;
}

@Override
public long getValue( String field )
{
Expand Down
Expand Up @@ -48,6 +48,12 @@ public long current()
return 0;
}

@Override
public float currentScore()
{
return 0;
}

@Override
public long getValue( String field )
{
Expand All @@ -57,6 +63,7 @@ public long getValue( String field )

protected final int size;
protected int index;
public abstract float currentScore();

ValuesIterator( int size )
{
Expand Down
Expand Up @@ -56,26 +56,24 @@ private PartitionedFulltextReader( List<ReadOnlyFulltext> readers )
}

@Override
public PrimitiveLongIterator query( Collection<String> terms, boolean matchAll )
public ScoreEntityIterator query( Collection<String> terms, boolean matchAll )
{
return partitionedOperation( reader -> innerQuery( reader, matchAll, terms ) );
}

@Override
public PrimitiveLongIterator fuzzyQuery( Collection<String> terms, boolean matchAll )
public ScoreEntityIterator fuzzyQuery( Collection<String> terms, boolean matchAll )
{
return partitionedOperation( reader -> innerFuzzyQuery( reader, matchAll, terms ) );
}

private PrimitiveLongIterator innerQuery( ReadOnlyFulltext reader, boolean matchAll, Collection<String> query )
private ScoreEntityIterator innerQuery( ReadOnlyFulltext reader, boolean matchAll, Collection<String> query )
{

return reader.query( query, matchAll );
}

private PrimitiveLongIterator innerFuzzyQuery( ReadOnlyFulltext reader, boolean matchAll, Collection<String> query )
private ScoreEntityIterator innerFuzzyQuery( ReadOnlyFulltext reader, boolean matchAll, Collection<String> query )
{

return reader.fuzzyQuery( query, matchAll );
}

Expand Down Expand Up @@ -105,8 +103,8 @@ public FulltextIndexConfiguration getConfigurationDocument() throws IOException
return null;
}

private PrimitiveLongIterator partitionedOperation( Function<ReadOnlyFulltext,PrimitiveLongIterator> readerFunction )
private ScoreEntityIterator partitionedOperation( Function<ReadOnlyFulltext,ScoreEntityIterator> readerFunction )
{
return PrimitiveLongCollections.concat( indexReaders.parallelStream().map( readerFunction ).collect( Collectors.toList() ) );
return ScoreEntityIterator.concat( indexReaders.parallelStream().map( readerFunction ).collect( Collectors.toList() ) );
}
}
Expand Up @@ -22,8 +22,6 @@
import java.io.IOException;
import java.util.Collection;

import org.neo4j.collection.primitive.PrimitiveLongIterator;

public interface ReadOnlyFulltext extends AutoCloseable
{
/**
Expand All @@ -34,7 +32,7 @@ public interface ReadOnlyFulltext extends AutoCloseable
* @param matchAll If true, only resluts that match all the given terms will be returned
* @return An iterator over the matching entityIDs, ordered by lucene scoring of the match.
*/
PrimitiveLongIterator query( Collection<String> terms, boolean matchAll );
ScoreEntityIterator query( Collection<String> terms, boolean matchAll );

/**
* Searches the fulltext index for any fuzzy match of any of the given terms against any token in any of the indexed properties.
Expand All @@ -44,7 +42,7 @@ public interface ReadOnlyFulltext extends AutoCloseable
* @param matchAll If true, only resluts that match all the given terms will be returned
* @return An iterator over the matching entityIDs, ordered by lucene scoring of the match.
*/
PrimitiveLongIterator fuzzyQuery( Collection<String> terms, boolean matchAll );
ScoreEntityIterator fuzzyQuery( Collection<String> terms, boolean matchAll );

@Override
void close();
Expand Down
@@ -0,0 +1,139 @@
/*
* Copyright (c) 2002-2018 "Neo Technology,"
* Network Engine for Objects in Lund AB [http://neotechnology.com]
*
* This file is part of Neo4j.
*
* Neo4j is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as
* published by the Free Software Foundation, either version 3 of the
* License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Affero General Public License for more details.
*
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
package org.neo4j.kernel.api.impl.fulltext;

import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.Spliterator;
import java.util.Spliterators;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

import org.neo4j.helpers.collection.Pair;
import org.neo4j.kernel.api.impl.index.collector.ValuesIterator;

public class ScoreEntityIterator implements Iterator<Pair<Long,Float>>
{
private static final ScoreEntityIterator EMPTY = new ScoreEntityIterator()
{
@Override
public Stream<Pair<Long,Float>> stream()
{
return Stream.empty();
}

@Override
public boolean hasNext()
{
return false;
}

@Override
public Pair<Long,Float> next()
{
throw new NoSuchElementException( "The iterator is exhausted" );
}
};
private final ValuesIterator iterator;

ScoreEntityIterator( ValuesIterator sortedValuesIterator )
{
this.iterator = sortedValuesIterator;
}

private ScoreEntityIterator()
{
this.iterator = null;
}

public Stream<Pair<Long,Float>> stream()
{
return StreamSupport.stream( Spliterators.spliteratorUnknownSize( this, Spliterator.ORDERED ), false );
}

@Override
public boolean hasNext()
{
return iterator.hasNext();
}

@Override
public Pair<Long,Float> next()
{
if ( hasNext() )
{
return Pair.of( currentIterator().next(), currentIterator().currentScore() );
}
else
{
throw new NoSuchElementException( "The iterator is exhausted" );
}
}

ValuesIterator currentIterator()
{
return iterator;
}

public static ScoreEntityIterator concat( List<ScoreEntityIterator> iterators )
{
return new ConcatenatingScoreEntityIterator( iterators.iterator() );
}

public static ScoreEntityIterator emptyIterator()
{
return EMPTY;
}

private static class ConcatenatingScoreEntityIterator extends ScoreEntityIterator
{
private final Iterator<? extends ScoreEntityIterator> iterators;
private ScoreEntityIterator currentIterator;

ConcatenatingScoreEntityIterator( Iterator<? extends ScoreEntityIterator> iterators )
{
this.iterators = iterators;
}

@Override
public boolean hasNext()
{
if ( currentIterator == null || !currentIterator.hasNext() )
{
while ( iterators.hasNext() )
{
currentIterator = iterators.next();
if ( currentIterator.hasNext() )
{
break;
}
}
}
return currentIterator != null && currentIterator.hasNext();
}

@Override
final ValuesIterator currentIterator()
{
return currentIterator.iterator;
}
}
}
Expand Up @@ -32,8 +32,6 @@
import java.io.IOException;
import java.util.Collection;

import org.neo4j.collection.primitive.PrimitiveLongCollections;
import org.neo4j.collection.primitive.PrimitiveLongIterator;
import org.neo4j.kernel.api.impl.index.collector.DocValuesCollector;
import org.neo4j.kernel.api.impl.index.partition.PartitionSearcher;
import org.neo4j.kernel.api.impl.schema.reader.IndexReaderCloseException;
Expand All @@ -60,14 +58,14 @@ class SimpleFulltextReader implements ReadOnlyFulltext
}

@Override
public PrimitiveLongIterator query( Collection<String> terms, boolean matchAll )
public ScoreEntityIterator query( Collection<String> terms, boolean matchAll )
{
String query = terms.stream().map( QueryParser::escape ).collect( joining( " " ) );
return innerQuery( query, matchAll );
}

@Override
public PrimitiveLongIterator fuzzyQuery( Collection<String> terms, boolean matchAll )
public ScoreEntityIterator fuzzyQuery( Collection<String> terms, boolean matchAll )
{
String query = terms.stream().map( QueryParser::escape ).collect( joining( "~ ", "", "~" ) );
return innerQuery( query, matchAll );
Expand Down Expand Up @@ -98,7 +96,7 @@ public FulltextIndexConfiguration getConfigurationDocument() throws IOException
return new FulltextIndexConfiguration( indexSearcher.doc( docs.scoreDocs[0].doc ) );
}

private PrimitiveLongIterator innerQuery( String queryString, boolean matchAll )
private ScoreEntityIterator innerQuery( String queryString, boolean matchAll )
{
MultiFieldQueryParser multiFieldQueryParser = new MultiFieldQueryParser( properties, analyzer );
if ( matchAll )
Expand All @@ -117,18 +115,18 @@ private PrimitiveLongIterator innerQuery( String queryString, boolean matchAll )
catch ( ParseException e )
{
assert false;
return PrimitiveLongCollections.emptyIterator();
return ScoreEntityIterator.emptyIterator();
}
return indexQuery( query );
}

private PrimitiveLongIterator indexQuery( Query query )
private ScoreEntityIterator indexQuery( Query query )
{
try
{
DocValuesCollector docValuesCollector = new DocValuesCollector( true );
getIndexSearcher().search( query, docValuesCollector );
return docValuesCollector.getSortedValuesIterator( FIELD_ENTITY_ID, Sort.RELEVANCE );
return new ScoreEntityIterator( docValuesCollector.getSortedValuesIterator( FIELD_ENTITY_ID, Sort.RELEVANCE ) );
}
catch ( IOException e )
{
Expand Down
Expand Up @@ -23,13 +23,16 @@
import java.util.List;
import java.util.Spliterator;
import java.util.Spliterators;
import java.util.function.Function;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

import org.neo4j.collection.primitive.PrimitiveLongCollections;
import org.neo4j.collection.primitive.PrimitiveLongIterator;
import org.neo4j.helpers.collection.Pair;
import org.neo4j.kernel.api.impl.fulltext.FulltextProvider;
import org.neo4j.kernel.api.impl.fulltext.ReadOnlyFulltext;
import org.neo4j.kernel.api.impl.fulltext.ScoreEntityIterator;
import org.neo4j.kernel.api.index.InternalIndexState;
import org.neo4j.procedure.Context;
import org.neo4j.procedure.Description;
Expand All @@ -50,6 +53,7 @@ public class BloomProcedures
{
@Context
public FulltextProvider provider;
public static final Function<Pair<Long,Float>,EntityOutput> QUERY_RESULT_MAPPER = result -> new EntityOutput( result.first(), result.other() );

@Description( "Await the completion of any background index population or updates" )
@Procedure( name = "bloom.awaitPopulation", mode = READ )
Expand Down Expand Up @@ -123,34 +127,35 @@ public Stream<EntityOutput> bloomFulltextRelationships(

private Stream<EntityOutput> queryAsStream( List<String> terms, ReadOnlyFulltext indexReader, boolean fuzzy, boolean matchAll )
{
PrimitiveLongIterator primitiveLongIterator;
ScoreEntityIterator resultIterator;
if ( fuzzy )
{
primitiveLongIterator = indexReader.fuzzyQuery( terms, matchAll );
resultIterator = indexReader.fuzzyQuery( terms, matchAll );
}
else
{
primitiveLongIterator = indexReader.query( terms, matchAll );
resultIterator = indexReader.query( terms, matchAll );
}
Iterator<EntityOutput> iterator = PrimitiveLongCollections.map( EntityOutput::new, primitiveLongIterator );
return StreamSupport.stream( Spliterators.spliteratorUnknownSize( iterator, Spliterator.ORDERED ), false );
return resultIterator.stream().map( QUERY_RESULT_MAPPER );
}

public static class EntityOutput
{
public final long entityid;
public final double score;

public EntityOutput( long entityid )
EntityOutput( long entityid, float score )
{
this.entityid = entityid;
this.score = score;
}
}

public static class PropertyOutput
{
public final String propertyKey;

public PropertyOutput( String propertykey )
PropertyOutput( String propertykey )
{
this.propertyKey = propertykey;
}
Expand All @@ -160,7 +165,7 @@ public class StatusOutput
{
public final String state;

public StatusOutput( InternalIndexState internalIndexState )
StatusOutput( InternalIndexState internalIndexState )
{
switch ( internalIndexState )
{
Expand Down

0 comments on commit 154a82a

Please sign in to comment.