From 5bdaa9b97c7b151f01ac5e76322f517d7537b1e8 Mon Sep 17 00:00:00 2001 From: Mattias Persson Date: Tue, 16 Feb 2016 22:14:45 +0100 Subject: [PATCH] Supports querying multiple label ids AND/OR from LabelScanStoreReader (cherry picked from commit ff85bc5) --- .../api/schema/LabelScanReader.java | 14 +++- ...PartitionedLuceneLabelScanStoreReader.java | 12 +++ .../SimpleLuceneLabelScanStoreReader.java | 12 +++ .../storestrategy/BitmapDocumentFormat.java | 28 ++++++- .../LabelScanStorageStrategy.java | 4 + ...RangeDocumentLabelScanStorageStrategy.java | 19 ++++- .../storestrategy/PageOfRangesIterator.java | 29 +++++++- .../labelscan/LuceneLabelScanStoreTest.java | 73 ++++++++++++++++++- .../PageOfRangesIteratorTest.java | 3 +- 9 files changed, 184 insertions(+), 10 deletions(-) diff --git a/community/kernel/src/main/java/org/neo4j/storageengine/api/schema/LabelScanReader.java b/community/kernel/src/main/java/org/neo4j/storageengine/api/schema/LabelScanReader.java index d8e783f2d4fb..cce86b99e4b0 100644 --- a/community/kernel/src/main/java/org/neo4j/storageengine/api/schema/LabelScanReader.java +++ b/community/kernel/src/main/java/org/neo4j/storageengine/api/schema/LabelScanReader.java @@ -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; @@ -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}. diff --git a/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/reader/PartitionedLuceneLabelScanStoreReader.java b/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/reader/PartitionedLuceneLabelScanStoreReader.java index 725474cfc3e1..846d3a70db68 100644 --- a/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/reader/PartitionedLuceneLabelScanStoreReader.java +++ b/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/reader/PartitionedLuceneLabelScanStoreReader.java @@ -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 ) { diff --git a/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/reader/SimpleLuceneLabelScanStoreReader.java b/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/reader/SimpleLuceneLabelScanStoreReader.java index 86d918e6f3a1..ff8b096edd04 100644 --- a/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/reader/SimpleLuceneLabelScanStoreReader.java +++ b/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/reader/SimpleLuceneLabelScanStoreReader.java @@ -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 ) { diff --git a/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/BitmapDocumentFormat.java b/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/BitmapDocumentFormat.java index 0461cf8d511f..04a6b1a60f97 100644 --- a/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/BitmapDocumentFormat.java +++ b/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/BitmapDocumentFormat.java @@ -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; @@ -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 ) ) ); diff --git a/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/LabelScanStorageStrategy.java b/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/LabelScanStorageStrategy.java index cdd0224580ae..569a333aef91 100644 --- a/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/LabelScanStorageStrategy.java +++ b/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/LabelScanStorageStrategy.java @@ -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 ); diff --git a/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/NodeRangeDocumentLabelScanStorageStrategy.java b/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/NodeRangeDocumentLabelScanStorageStrategy.java index 1803eb0fcc48..01e48c66fdf0 100644 --- a/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/NodeRangeDocumentLabelScanStorageStrategy.java +++ b/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/NodeRangeDocumentLabelScanStorageStrategy.java @@ -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; @@ -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 diff --git a/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/PageOfRangesIterator.java b/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/PageOfRangesIterator.java index 7434507c8666..e9b6ba3974d5 100644 --- a/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/PageOfRangesIterator.java +++ b/community/lucene-index/src/main/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/PageOfRangesIterator.java @@ -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; @@ -39,21 +40,33 @@ class PageOfRangesIterator extends PrefetchingIterator 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() { @@ -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 ) { diff --git a/community/lucene-index/src/test/java/org/neo4j/kernel/api/impl/labelscan/LuceneLabelScanStoreTest.java b/community/lucene-index/src/test/java/org/neo4j/kernel/api/impl/labelscan/LuceneLabelScanStoreTest.java index 2c491cba1e62..0ce5b39bb38a 100644 --- a/community/lucene-index/src/test/java/org/neo4j/kernel/api/impl/labelscan/LuceneLabelScanStoreTest.java +++ b/community/lucene-index/src/test/java/org/neo4j/kernel/api/impl/labelscan/LuceneLabelScanStoreTest.java @@ -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; @@ -120,7 +121,7 @@ public void clearDir() throws IOException } @After - public void shutdown() throws IOException + public void shutdown() { life.shutdown(); } @@ -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} ) ); } diff --git a/community/lucene-index/src/test/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/PageOfRangesIteratorTest.java b/community/lucene-index/src/test/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/PageOfRangesIteratorTest.java index 0af784be3242..adb4e5e1a344 100644 --- a/community/lucene-index/src/test/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/PageOfRangesIteratorTest.java +++ b/community/lucene-index/src/test/java/org/neo4j/kernel/api/impl/labelscan/storestrategy/PageOfRangesIteratorTest.java @@ -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; @@ -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 longs = PrimitiveLongCollections.asList( iterator );