HSEARCH-5010 Use scorer suppliers in custom weights

Author: marko-bekhta

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/lowlevel/query/impl/ConstantScorerSupplier.java

@@ -0,0 +1,40 @@
+/*
+ * Hibernate Search, full-text search for your domain model
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.search.backend.lucene.lowlevel.query.impl;
+
+import java.io.IOException;
+
+import org.apache.lucene.search.ConstantScoreScorer;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.ScorerSupplier;
+import org.apache.lucene.search.Weight;
+
+class ConstantScorerSupplier extends ScorerSupplier {
+	private final Weight weight;
+	private final float score;
+	private final ScoreMode scoreMode;
+	private final DocIdSetIterator matchingDocs;
+
+	public ConstantScorerSupplier(Weight weight, float score, ScoreMode scoreMode, DocIdSetIterator matchingDocs) {
+		this.weight = weight;
+		this.score = score;
+		this.scoreMode = scoreMode;
+		this.matchingDocs = matchingDocs;
+	}
+
+	@Override
+	public Scorer get(long leadCost) throws IOException {
+		return new ConstantScoreScorer( weight, score, scoreMode, matchingDocs );
+	}
+
+	@Override
+	public long cost() {
+		return matchingDocs.cost();
+	}
+}

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/lowlevel/query/impl/ExplicitDocIdsQuery.java

@@ -9,15 +9,14 @@
 import java.util.Arrays;
 
 import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.search.ConstantScoreScorer;
 import org.apache.lucene.search.ConstantScoreWeight;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.ScoreMode;
-import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.ScorerSupplier;
 import org.apache.lucene.search.Weight;
 
 public final class ExplicitDocIdsQuery extends Query {
@@ -58,15 +57,16 @@ public int hashCode() {
 	@Override
 	public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) {
 		return new ConstantScoreWeight( this, 1.0f ) {
+
 			@Override
-			public Scorer scorer(LeafReaderContext context) {
+			public ScorerSupplier scorerSupplier(LeafReaderContext context) {
 				DocIdSetIterator matchingDocs = ExplicitDocIdSetIterator.of(
 						sortedDocIds, context.docBase, context.reader().maxDoc()
 				);
 				if ( matchingDocs == null ) {
 					return null; // Skip this leaf
 				}
-				return new ConstantScoreScorer( this, this.score(), scoreMode, matchingDocs );
+				return new ConstantScorerSupplier( this, this.score(), scoreMode, matchingDocs );
 			}
 
 			@Override

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/lowlevel/query/impl/MappedTypeNameQuery.java

@@ -6,6 +6,8 @@
  */
 package org.hibernate.search.backend.lucene.lowlevel.query.impl;
 
+import java.io.IOException;
+
 import org.hibernate.search.backend.lucene.lowlevel.reader.impl.IndexReaderMetadataResolver;
 
 import org.apache.lucene.index.LeafReaderContext;
@@ -17,6 +19,7 @@
 import org.apache.lucene.search.QueryVisitor;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.ScorerSupplier;
 import org.apache.lucene.search.Weight;
 
 public final class MappedTypeNameQuery extends Query {
@@ -53,8 +56,9 @@ public int hashCode() {
 	@Override
 	public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) {
 		return new ConstantScoreWeight( this, 1.0f ) {
+
 			@Override
-			public Scorer scorer(LeafReaderContext context) {
+			public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException {
 				String leafMappedTypeName = metadataResolver.resolveMappedTypeName( context );
 				DocIdSetIterator matchingDocs;
 				if ( mappedTypeName.equals( leafMappedTypeName ) ) {
@@ -63,7 +67,8 @@ public Scorer scorer(LeafReaderContext context) {
 				else {
 					matchingDocs = DocIdSetIterator.empty();
 				}
-				return new ConstantScoreScorer( this, this.score(), scoreMode, matchingDocs );
+
+				return new ConstantScorerSupplier( this, this.score(), scoreMode, matchingDocs );
 			}
 
 			@Override

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/lowlevel/query/impl/VectorSimilarityFilterQuery.java

@@ -23,6 +23,7 @@
 import org.apache.lucene.search.QueryVisitor;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.ScorerSupplier;
 import org.apache.lucene.search.TwoPhaseIterator;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.VectorUtil;
@@ -148,12 +149,39 @@ public Explanation explain(LeafReaderContext context, int doc) throws IOExceptio
 		}
 
 		@Override
-		public Scorer scorer(LeafReaderContext context) throws IOException {
-			Scorer scorer = super.scorer( context );
+		public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException {
+			ScorerSupplier scorerSupplier = super.scorerSupplier( context );
+			if ( scorerSupplier == null ) {
+				return null;
+			}
+			return new MinScoreScorerSupplier( this, scorerSupplier, similarityAsScore );
+		}
+	}
+
+	private static class MinScoreScorerSupplier extends ScorerSupplier {
+
+		private final Weight weight;
+		private final ScorerSupplier delegate;
+		private final float similarityAsScore;
+
+		private MinScoreScorerSupplier(Weight weight, ScorerSupplier delegate, float similarityAsScore) {
+			this.weight = weight;
+			this.delegate = delegate;
+			this.similarityAsScore = similarityAsScore;
+		}
+
+		@Override
+		public Scorer get(long leadCost) throws IOException {
+			Scorer scorer = delegate.get( leadCost );
 			if ( scorer == null ) {
 				return null;
 			}
-			return new MinScoreScorer( this, scorer, similarityAsScore );
+			return new MinScoreScorer( weight, scorer, similarityAsScore );
+		}
+
+		@Override
+		public long cost() {
+			return delegate.cost();
 		}
 	}
 

util/internal/integrationtest/backend/lucene/src/main/java/org/hibernate/search/util/impl/integrationtest/backend/lucene/query/SlowQuery.java

@@ -6,6 +6,8 @@
  */
 package org.hibernate.search.util.impl.integrationtest.backend.lucene.query;
 
+import java.io.IOException;
+
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.ConstantScoreScorer;
 import org.apache.lucene.search.ConstantScoreWeight;
@@ -16,6 +18,7 @@
 import org.apache.lucene.search.QueryVisitor;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.ScorerSupplier;
 import org.apache.lucene.search.Weight;
 
 /**
@@ -55,11 +58,21 @@ public String toString() {
 			}
 
 			@Override
-			public Scorer scorer(LeafReaderContext context) {
-				return new ConstantScoreScorer(
-						this, score(), scoreMode,
-						new SlowDocIdSetIterator( DocIdSetIterator.all( context.reader().maxDoc() ) )
-				);
+			public ScorerSupplier scorerSupplier(LeafReaderContext context) {
+				Weight weight = this;
+				float score = score();
+				SlowDocIdSetIterator iterator = new SlowDocIdSetIterator( DocIdSetIterator.all( context.reader().maxDoc() ) );
+				return new ScorerSupplier() {
+					@Override
+					public Scorer get(long leadCost) throws IOException {
+						return new ConstantScoreScorer( weight, score, scoreMode, iterator );
+					}
+
+					@Override
+					public long cost() {
+						return iterator.cost();
+					}
+				};
 			}
 
 			@Override