Skip to content

Commit

Permalink
Supports querying multiple label ids AND/OR from LabelScanStoreReader
Browse files Browse the repository at this point in the history 
(cherry picked from commit ff85bc5)
  • Loading branch information
@tinwelint @MishaDemianenko
tinwelint authored and MishaDemianenko committed Sep 27, 2016
1 parent 2d2a1b4 commit 5bdaa9b
Show file tree
Hide file tree
Showing 9 changed files with 184 additions and 10 deletions.
14 changes: 12 additions & 2 deletions community/kernel/src/main/java/org/neo4j/storageengine/api/schema/LabelScanReader.java
View file
Expand Up @@ -19,8 +19,6 @@
*/
package org.neo4j.storageengine.api.schema;

import java.util.Iterator;

import org.neo4j.collection.primitive.PrimitiveLongIterator;
import org.neo4j.graphdb.Resource;

Expand All @@ -35,6 +33,18 @@ public interface LabelScanReader extends Resource
*/
PrimitiveLongIterator nodesWithLabel( int labelId );

/**
* @param labelIds label token ids.
* @return node ids with any of the given label ids.
*/
PrimitiveLongIterator nodesWithAnyOfLabels( int... labelIds );

/**
* @param labelIds label token ids.
* @return node ids with all of the given label ids.
*/
PrimitiveLongIterator nodesWithAllLabels( int... labelIds );

/**
* @param nodeId node id to get label ids for.
* @return label ids for the given {@code nodeId}.
Expand Down
12 changes: 12 additions & 0 deletions ...ava/org/neo4j/kernel/api/impl/labelscan/reader/PartitionedLuceneLabelScanStoreReader.java
View file
Expand Up @@ -62,6 +62,18 @@ public PrimitiveLongIterator nodesWithLabel( int labelId )
return partitionedOperation( storeReader -> storeReader.nodesWithLabel( labelId ) );
}

@Override
public PrimitiveLongIterator nodesWithAnyOfLabels( int... labelIds )
{
return partitionedOperation( storeReader -> storeReader.nodesWithAnyOfLabels( labelIds ) );
}

@Override
public PrimitiveLongIterator nodesWithAllLabels( int... labelIds )
{
return partitionedOperation( storeReader -> storeReader.nodesWithAllLabels( labelIds ) );
}

@Override
public PrimitiveLongIterator labelsForNode( long nodeId )
{
Expand Down
12 changes: 12 additions & 0 deletions ...ain/java/org/neo4j/kernel/api/impl/labelscan/reader/SimpleLuceneLabelScanStoreReader.java
View file
Expand Up @@ -47,6 +47,18 @@ public PrimitiveLongIterator nodesWithLabel( int labelId )
return strategy.nodesWithLabel( partitionSearcher.getIndexSearcher(), labelId );
}

@Override
public PrimitiveLongIterator nodesWithAnyOfLabels( int... labelIds )
{
return strategy.nodesWithAnyOfLabels( partitionSearcher.getIndexSearcher(), labelIds );
}

@Override
public PrimitiveLongIterator nodesWithAllLabels( int... labelIds )
{
return strategy.nodesWithAllLabels( partitionSearcher.getIndexSearcher(), labelIds );
}

@Override
public PrimitiveLongIterator labelsForNode( long nodeId )
{
Expand Down
28 changes: 27 additions & 1 deletion ...src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/BitmapDocumentFormat.java
View file
Expand Up @@ -27,6 +27,8 @@
import org.apache.lucene.document.StringField;
import org.apache.lucene.index.IndexableField;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.BooleanClause.Occur;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.TermQuery;

Expand Down Expand Up @@ -106,11 +108,35 @@ public long mapOf( Document doc, long labelId )
return bitmap( doc.getField( label( labelId ) ) );
}

public Query labelQuery( long labelId )
public Query labelQuery( int labelId )
{
return new TermQuery( new Term( LABEL, Long.toString( labelId ) ) );
}

/**
* Builds a {@link Query} suitable for returning documents of nodes having all or any
* (depending on {@code occur}) of the given {@code labelIds}.
*
* @param occur {@link Occur} to use in the query.
* @param labelIds label ids to query for.
* @return {@link Query} for searching for documents with (all or any) of the label ids.
*/
public Query labelsQuery( Occur occur, int[] labelIds )
{
assert labelIds.length > 0;
if ( labelIds.length == 1 )
{
return labelQuery( labelIds[0] );
}

BooleanQuery.Builder query = new BooleanQuery.Builder();
for ( int labelId : labelIds )
{
query.add( labelQuery( labelId ), occur );
}
return query.build();
}

public Query rangeQuery( long range )
{
return new TermQuery( new Term( RANGE, Long.toString( range ) ) );
Expand Down
4 changes: 4 additions & 0 deletions ...main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/LabelScanStorageStrategy.java
View file
Expand Up @@ -29,6 +29,10 @@ public interface LabelScanStorageStrategy
{
PrimitiveLongIterator nodesWithLabel( IndexSearcher searcher, int labelId );

PrimitiveLongIterator nodesWithAnyOfLabels( IndexSearcher indexSearcher, int[] labelIds );

PrimitiveLongIterator nodesWithAllLabels( IndexSearcher indexSearcher, int[] labelIds );

AllEntriesLabelScanReader newNodeLabelReader( LuceneAllDocumentsReader allDocumentsReader );

PrimitiveLongIterator labelsForNode( IndexSearcher searcher, long nodeId );
Expand Down
19 changes: 17 additions & 2 deletions ...4j/kernel/api/impl/labelscan/storestrategy/NodeRangeDocumentLabelScanStorageStrategy.java
View file
Expand Up @@ -20,6 +20,7 @@
package org.neo4j.kernel.api.impl.labelscan.storestrategy;

import org.apache.lucene.index.IndexableField;
import org.apache.lucene.search.BooleanClause.Occur;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.TopDocs;

Expand Down Expand Up @@ -79,8 +80,22 @@ public String toString()
@Override
public PrimitiveLongIterator nodesWithLabel( IndexSearcher searcher, int labelId )
{
return concat(
new PageOfRangesIterator( format, searcher, RANGES_PER_PAGE, format.labelQuery( labelId ), labelId ) );
return concat( new PageOfRangesIterator( format, searcher, RANGES_PER_PAGE, format.labelQuery( labelId ),
Occur.MUST, labelId ) );
}

@Override
public PrimitiveLongIterator nodesWithAnyOfLabels( IndexSearcher searcher, int[] labelIds )
{
return concat( new PageOfRangesIterator( format, searcher, RANGES_PER_PAGE,
format.labelsQuery( Occur.SHOULD, labelIds ), Occur.SHOULD, labelIds ) );
}

@Override
public PrimitiveLongIterator nodesWithAllLabels( IndexSearcher searcher, int[] labelIds )
{
return concat( new PageOfRangesIterator( format, searcher, RANGES_PER_PAGE,
format.labelsQuery( Occur.MUST, labelIds ), Occur.MUST, labelIds ) );
}

@Override
Expand Down
29 changes: 27 additions & 2 deletions ...src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/PageOfRangesIterator.java
View file
Expand Up @@ -19,6 +19,7 @@
*/
package org.neo4j.kernel.api.impl.labelscan.storestrategy;

import org.apache.lucene.search.BooleanClause.Occur;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;

Expand All @@ -39,21 +40,33 @@ class PageOfRangesIterator extends PrefetchingIterator<PrimitiveLongIterator>
private final int rangesPerPage;
private final int[] labels;
private DocValuesCollector.LongValuesIterator rangesIterator;
private final boolean matchAny;

PageOfRangesIterator( BitmapDocumentFormat format, IndexSearcher searcher, int rangesPerPage, Query query,
int... labels )
Occur occur, int... labels )
{
this.searcher = searcher;
this.query = query;
this.format = format;
this.rangesPerPage = rangesPerPage;
this.matchAny = matchMode( occur );
this.labels = labels;
if (labels.length == 0)
if ( labels.length == 0 )
{
throw new IllegalArgumentException( "At least one label required" );
}
}

private boolean matchMode( Occur occur )
{
switch ( occur )
{
case MUST: return false;
case SHOULD: return true;
default: throw new IllegalArgumentException( "Unexpected occurence parameter " + occur );
}
}

@Override
protected PrimitiveLongIterator fetchNextOrNull()
{
Expand Down Expand Up @@ -100,6 +113,18 @@ private DocValuesCollector.LongValuesIterator getRanges() {

private long labeledBitmap( DocValuesAccess doc )
{
if ( matchAny )
{
// OR
long bitmap = 0;
for ( int label : labels )
{
bitmap |= doc.getValue( format.label( label ) );
}
return bitmap;
}

// AND
long bitmap = -1;
for ( int label : labels )
{
Expand Down
73 changes: 71 additions & 2 deletions ...ene-index/src/test/java/org/neo4j/kernel/api/impl/labelscan/LuceneLabelScanStoreTest.java
View file
Expand Up @@ -79,6 +79,7 @@
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;
import static org.neo4j.collection.primitive.PrimitiveLongCollections.EMPTY_LONG_ARRAY;
import static org.neo4j.helpers.collection.Iterators.iterator;
import static org.neo4j.helpers.collection.Iterators.single;
import static org.neo4j.io.fs.FileUtils.deleteRecursively;
Expand Down Expand Up @@ -120,7 +121,7 @@ public void clearDir() throws IOException
}

@After
public void shutdown() throws IOException
public void shutdown()
{
life.shutdown();
}
Expand Down Expand Up @@ -444,10 +445,78 @@ public void shouldFindAllLabelsForGivenNode() throws Exception
reader.close();
}

@Test
public void shouldFindNodesWithAnyOfGivenLabels() throws Exception
{
// GIVEN
int labelId1 = 3, labelId2 = 5, labelId3 = 13;
start();

// WHEN
write( iterator(
labelChanges( 1, EMPTY_LONG_ARRAY, new long[] {labelId1} ),
labelChanges( 2, EMPTY_LONG_ARRAY, new long[] {labelId1, labelId2} ),
labelChanges( 3, EMPTY_LONG_ARRAY, new long[] {labelId1} ),
labelChanges( 4, EMPTY_LONG_ARRAY, new long[] {labelId1, labelId3} ),
labelChanges( 5, EMPTY_LONG_ARRAY, new long[] {labelId1, labelId2, labelId3} ),
labelChanges( 6, EMPTY_LONG_ARRAY, new long[] { labelId2} ),
labelChanges( 7, EMPTY_LONG_ARRAY, new long[] { labelId2} ),
labelChanges( 8, EMPTY_LONG_ARRAY, new long[] { labelId3} ),
labelChanges( 9, EMPTY_LONG_ARRAY, new long[] { labelId3} ) ) );

// THEN
try ( LabelScanReader reader = store.newReader() )
{
assertArrayEquals(
new long[] {1, 2, 3, 4, 5, 6, 7},
PrimitiveLongCollections.asArray( reader.nodesWithAnyOfLabels( labelId1, labelId2 ) ) );
assertArrayEquals(
new long[] {1, 2, 3, 4, 5, 8, 9},
PrimitiveLongCollections.asArray( reader.nodesWithAnyOfLabels( labelId1, labelId3 ) ) );
assertArrayEquals(
new long[] {1, 2, 3, 4, 5, 6, 7, 8, 9},
PrimitiveLongCollections.asArray( reader.nodesWithAnyOfLabels( labelId1, labelId2, labelId3 ) ) );
}
}

@Test
public void shouldFindNodesWithAllGivenLabels() throws Exception
{
// GIVEN
int labelId1 = 3, labelId2 = 5, labelId3 = 13;
start();

// WHEN
write( iterator(
labelChanges( 1, EMPTY_LONG_ARRAY, new long[] {labelId1} ),
labelChanges( 2, EMPTY_LONG_ARRAY, new long[] {labelId1, labelId2} ),
labelChanges( 3, EMPTY_LONG_ARRAY, new long[] {labelId1} ),
labelChanges( 4, EMPTY_LONG_ARRAY, new long[] {labelId1, labelId3} ),
labelChanges( 5, EMPTY_LONG_ARRAY, new long[] {labelId1, labelId2, labelId3} ),
labelChanges( 6, EMPTY_LONG_ARRAY, new long[] { labelId2} ),
labelChanges( 7, EMPTY_LONG_ARRAY, new long[] { labelId2} ),
labelChanges( 8, EMPTY_LONG_ARRAY, new long[] { labelId3} ),
labelChanges( 9, EMPTY_LONG_ARRAY, new long[] { labelId3} ) ) );

// THEN
try ( LabelScanReader reader = store.newReader() )
{
assertArrayEquals(
new long[] {2, 5},
PrimitiveLongCollections.asArray( reader.nodesWithAllLabels( labelId1, labelId2 ) ) );
assertArrayEquals(
new long[] {4, 5},
PrimitiveLongCollections.asArray( reader.nodesWithAllLabels( labelId1, labelId3 ) ) );
assertArrayEquals(
new long[] {5},
PrimitiveLongCollections.asArray( reader.nodesWithAllLabels( labelId1, labelId2, labelId3 ) ) );
}
}

private void prepareIndex() throws IOException
{
start();
try (LabelScanWriter labelScanWriter = store.newWriter())
try ( LabelScanWriter labelScanWriter = store.newWriter() )
{
labelScanWriter.write( NodeLabelUpdate.labelChanges( 1, new long[]{}, new long[]{1} ) );
}
Expand Down
3 changes: 2 additions & 1 deletion ...test/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/PageOfRangesIteratorTest.java
View file
Expand Up @@ -21,6 +21,7 @@

import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.NumericDocValues;
import org.apache.lucene.search.BooleanClause.Occur;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.junit.Test;
Expand Down Expand Up @@ -101,7 +102,7 @@ public void shouldReadPagesOfDocumentsFromSearcher() throws Exception
} ).when( searcher ).search( same( query ), any( DocValuesCollector.class ) );

PrimitiveLongIterator iterator = concat(
new PageOfRangesIterator( format, searcher, pageSize, query, labelId ) );
new PageOfRangesIterator( format, searcher, pageSize, query, Occur.MUST, labelId ) );

// when
List<Long> longs = PrimitiveLongCollections.asList( iterator );
Expand Down

0 comments on commit 5bdaa9b

Please sign in to comment.