SOLR-17951: Optimize re-ranking based on "functions"

Author: hossman

changelog/unreleased/SOLR-17951-optimize-function-rerank.yml

@@ -0,0 +1,7 @@
+title: Optimize re-ranking based on "functions"
+type: added
+authors:
+- name: hossman
+links:
+- name: SOLR-17951
+  url: https://issues.apache.org/jira/browse/SOLR-17951

solr/core/src/java/org/apache/solr/search/AbstractReRankQuery.java

@@ -16,6 +16,7 @@
  */
 package org.apache.solr.search;
 
+import com.google.common.annotations.VisibleForTesting;
 import java.io.IOException;
 import java.util.Map;
 import java.util.Set;
@@ -59,6 +60,26 @@ public AbstractReRankQuery(Query mainQuery, int reRankDocs, Rescorer reRankQuery
     this.reRankQueryRescorer = reRankQueryRescorer;
   }
 
+  @VisibleForTesting
+  int getReRankDocs() {
+    return reRankDocs;
+  }
+
+  @VisibleForTesting
+  Rescorer getRescorer() {
+    return reRankQueryRescorer;
+  }
+
+  @VisibleForTesting
+  ReRankOperator getReRankOperator() {
+    return reRankOperator;
+  }
+
+  @VisibleForTesting
+  ReRankScaler getReRankScaler() {
+    return reRankScaler;
+  }
+
   @Override
   public RankQuery wrap(Query _mainQuery) {
     if (_mainQuery != null) {

solr/core/src/java/org/apache/solr/search/ReRankOperator.java

@@ -17,12 +17,24 @@
 package org.apache.solr.search;
 
 import java.util.Locale;
+import java.util.function.DoubleBinaryOperator;
 import org.apache.solr.common.SolrException;
 
-public enum ReRankOperator {
-  ADD,
-  MULTIPLY,
-  REPLACE;
+public enum ReRankOperator implements DoubleBinaryOperator {
+  ADD((firstPass, secondPass) -> firstPass + secondPass),
+  MULTIPLY((firstPass, secondPass) -> firstPass * secondPass),
+  REPLACE((firstPass, secondPass) -> secondPass);
+
+  private final DoubleBinaryOperator op;
+
+  private ReRankOperator(final DoubleBinaryOperator op) {
+    this.op = op;
+  }
+
+  @Override
+  public double applyAsDouble(final double firstPass, final double secondPass) {
+    return op.applyAsDouble(firstPass, secondPass);
+  }
 
   public static ReRankOperator get(String p) {
     if (p != null) {

solr/core/src/java/org/apache/solr/search/ReRankQParserPlugin.java

@@ -18,9 +18,15 @@
 
 import java.io.IOException;
 import java.lang.invoke.MethodHandles;
+import java.util.function.DoubleBinaryOperator;
+import org.apache.lucene.queries.function.FunctionQuery;
+import org.apache.lucene.queries.function.FunctionScoreQuery;
+import org.apache.lucene.search.DoubleValuesSource;
+import org.apache.lucene.search.DoubleValuesSourceRescorer;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryRescorer;
+import org.apache.lucene.search.Rescorer;
 import org.apache.solr.common.SolrException;
 import org.apache.solr.common.params.CommonParams;
 import org.apache.solr.common.params.SolrParams;
@@ -63,6 +69,26 @@ public QParser createParser(
     return new ReRankQParser(query, localParams, params, req);
   }
 
+  /**
+   * Helper method for constructing a {@link Rescorer} from a {@link #RERANK_QUERY}, {@link
+   * #RERANK_WEIGHT}, and {@link #RERANK_OPERATOR}.
+   *
+   * <p>By default, this returns a customized {@link QueryRescorer}, unless the {@link
+   * #RERANK_QUERY} is a known type that can more efficiently be re-ranked using a customized {@link
+   * DoubleValuesSourceRescorer}.
+   */
+  private static Rescorer createRescorer(
+      final Query reRankQuery, final double reRankWeight, final ReRankOperator reRankOperator) {
+    assert null != reRankQuery;
+    return switch (reRankQuery) {
+      case FunctionQuery functionQuery -> new ReRankDoubleValuesSourceRescorer(
+          functionQuery.getValueSource().asDoubleValuesSource(), reRankWeight, reRankOperator);
+      case FunctionScoreQuery functionQuery -> new ReRankDoubleValuesSourceRescorer(
+          functionQuery.getSource(), reRankWeight, reRankOperator);
+      default -> new ReRankQueryRescorer(reRankQuery, reRankWeight, reRankOperator);
+    };
+  }
+
   private static class ReRankQParser extends QParser {
 
     private boolean isExplainResults() {
@@ -135,7 +161,7 @@ public Query parse() throws SyntaxError {
               reRankScale,
               reRankScaleWeight,
               reRankOperator,
-              new ReRankQueryRescorer(reRankQuery, 1, ReRankOperator.REPLACE),
+              createRescorer(reRankQuery, 1, ReRankOperator.REPLACE),
               explainResults);
 
       if (reRankScaler.scaleScores()) {
@@ -148,6 +174,28 @@ public Query parse() throws SyntaxError {
     }
   }
 
+  private static final class ReRankDoubleValuesSourceRescorer extends DoubleValuesSourceRescorer {
+    final DoubleBinaryOperator scoreCombiner;
+
+    public ReRankDoubleValuesSourceRescorer(
+        final DoubleValuesSource valuesSource,
+        final double reRankWeight,
+        final ReRankOperator reRankOperator) {
+      super(valuesSource);
+      this.scoreCombiner =
+          (score, value) -> reRankOperator.applyAsDouble(score, reRankWeight * value);
+    }
+
+    @Override
+    protected float combine(
+        final float firstPassScore, final boolean valuePresent, final double sourceValue) {
+      if (valuePresent) {
+        return (float) scoreCombiner.applyAsDouble(firstPassScore, sourceValue);
+      }
+      return firstPassScore;
+    }
+  }
+
   private static final class ReRankQueryRescorer extends QueryRescorer {
 
     final BiFloatFunction scoreCombiner;
@@ -160,20 +208,8 @@ interface BiFloatFunction {
     public ReRankQueryRescorer(
         Query reRankQuery, double reRankWeight, ReRankOperator reRankOperator) {
       super(reRankQuery);
-      switch (reRankOperator) {
-        case ADD:
-          scoreCombiner = (score, second) -> (float) (score + reRankWeight * second);
-          break;
-        case MULTIPLY:
-          scoreCombiner = (score, second) -> (float) (score * reRankWeight * second);
-          break;
-        case REPLACE:
-          scoreCombiner = (score, second) -> (float) (reRankWeight * second);
-          break;
-        default:
-          scoreCombiner = null;
-          throw new IllegalArgumentException("Unexpected: reRankOperator=" + reRankOperator);
-      }
+      scoreCombiner =
+          (score, second) -> (float) reRankOperator.applyAsDouble(score, reRankWeight * second);
     }
 
     @Override
@@ -226,7 +262,7 @@ public ReRankQuery(
       super(
           defaultQuery,
           reRankDocs,
-          new ReRankQueryRescorer(reRankQuery, reRankWeight, reRankOperator),
+          createRescorer(reRankQuery, reRankWeight, reRankOperator),
           reRankScaler,
           reRankOperator);
       this.reRankQuery = reRankQuery;

solr/core/src/java/org/apache/solr/search/ReRankScaler.java

@@ -23,7 +23,7 @@
 import java.util.Map;
 import java.util.Set;
 import org.apache.lucene.search.Explanation;
-import org.apache.lucene.search.QueryRescorer;
+import org.apache.lucene.search.Rescorer;
 import org.apache.lucene.search.ScoreDoc;
 
 public class ReRankScaler {
@@ -35,7 +35,7 @@ public class ReRankScaler {
   protected boolean explainResults;
   protected ReRankOperator reRankOperator;
   protected ReRankScalerExplain reRankScalerExplain;
-  private QueryRescorer replaceRescorer;
+  private Rescorer replaceRescorer;
   private Set<Integer> reRankSet;
   private double reRankScaleWeight;
 
@@ -44,7 +44,7 @@ public ReRankScaler(
       String reRankScale,
       double reRankScaleWeight,
       ReRankOperator reRankOperator,
-      QueryRescorer replaceRescorer,
+      Rescorer replaceRescorer,
       boolean explainResults)
       throws SyntaxError {
 
@@ -99,7 +99,7 @@ public boolean equals(Object o) {
     }
   }
 
-  public QueryRescorer getReplaceRescorer() {
+  public Rescorer getReplaceRescorer() {
     return replaceRescorer;
   }
 
@@ -237,16 +237,7 @@ public static float combineScores(
       float reRankScore,
       double reRankScaleWeight,
       ReRankOperator reRankOperator) {
-    switch (reRankOperator) {
-      case ADD:
-        return (float) (orginalScore + reRankScaleWeight * reRankScore);
-      case REPLACE:
-        return (float) (reRankScaleWeight * reRankScore);
-      case MULTIPLY:
-        return (float) (orginalScore * reRankScaleWeight * reRankScore);
-      default:
-        return -1;
-    }
+    return (float) reRankOperator.applyAsDouble(orginalScore, reRankScaleWeight * reRankScore);
   }
 
   public static final class ReRankScalerExplain {