feat: implement asBulkSimScorer on FeatureFields's SimScorers (#15137)

Author: AdityaTeltia

lucene/CHANGES.txt

@@ -270,6 +270,10 @@ Optimizations
 
 * GITHUB#15045: Use FixedBitSet#cardinality for counting liveDocs in CheckIndex (Zhang Chao)
 
+
+* GITHUB#15117: Score computations are now more reliably vectorized in FeatureField Scorer's
+  (Aditya Teltia)
+
 * GITHUB#15116, GITHUB#15138: Rewrite of the GroupVInt optimization without lambdas, varhandles
   and no code in subclasses. The optimization now auto-detects if an IndexInput supports random access
   and uses an optimized branchless approach. Any subclasses that have implemented the optimized method

lucene/core/src/java/org/apache/lucene/document/FeatureField.java

@@ -34,6 +34,7 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.search.similarities.BM25Similarity;
+import org.apache.lucene.search.similarities.Similarity;
 import org.apache.lucene.search.similarities.Similarity.SimScorer;
 
 /**
@@ -262,9 +263,27 @@ static final class LinearFunction extends FeatureFunction {
     @Override
     SimScorer scorer(float w) {
       return new SimScorer() {
+        private float doScore(float f) {
+          return (w * f);
+        }
+
         @Override
         public float score(float freq, long norm) {
-          return (w * decodeFeatureValue(freq));
+          float f = decodeFeatureValue(freq);
+          return doScore(f);
+        }
+
+        @Override
+        public Similarity.BulkSimScorer asBulkSimScorer() {
+          return new Similarity.BulkSimScorer() {
+            @Override
+            public void score(int size, float[] freqs, long[] norms, float[] scores) {
+              for (int i = 0; i < size; ++i) {
+                float f = decodeFeatureValue(freqs[i]);
+                scores[i] = doScore(f);
+              }
+            }
+          };
         }
       };
     }
@@ -333,9 +352,27 @@ public String toString() {
     @Override
     SimScorer scorer(float weight) {
       return new SimScorer() {
+        private float doScore(float f) {
+          return (float) (weight * Math.log(scalingFactor + f));
+        }
+
         @Override
         public float score(float freq, long norm) {
-          return (float) (weight * Math.log(scalingFactor + decodeFeatureValue(freq)));
+          float f = decodeFeatureValue(freq);
+          return doScore(f);
+        }
+
+        @Override
+        public Similarity.BulkSimScorer asBulkSimScorer() {
+          return new Similarity.BulkSimScorer() {
+            @Override
+            public void score(int size, float[] freqs, long[] norms, float[] scores) {
+              for (int i = 0; i < size; ++i) {
+                float f = decodeFeatureValue(freqs[i]);
+                scores[i] = doScore(f);
+              }
+            }
+          };
         }
       };
     }
@@ -405,14 +442,32 @@ SimScorer scorer(float weight) {
       }
       final float pivot = this.pivot; // unbox
       return new SimScorer() {
-        @Override
-        public float score(float freq, long norm) {
-          float f = decodeFeatureValue(freq);
+
+        private float doScore(float f) {
           // should be f / (f + k) but we rewrite it to
           // 1 - k / (f + k) to make sure it doesn't decrease
           // with f in spite of rounding
           return weight * (1 - pivot / (f + pivot));
         }
+
+        @Override
+        public float score(float freq, long norm) {
+          float f = decodeFeatureValue(freq);
+          return doScore(f);
+        }
+
+        @Override
+        public Similarity.BulkSimScorer asBulkSimScorer() {
+          return new Similarity.BulkSimScorer() {
+            @Override
+            public void score(int size, float[] freqs, long[] norms, float[] scores) {
+              for (int i = 0; i < size; ++i) {
+                float f = decodeFeatureValue(freqs[i]);
+                scores[i] = doScore(f);
+              }
+            }
+          };
+        }
       };
     }
 
@@ -469,14 +524,31 @@ public String toString() {
     @Override
     SimScorer scorer(float weight) {
       return new SimScorer() {
-        @Override
-        public float score(float freq, long norm) {
-          float f = decodeFeatureValue(freq);
+        private float doScore(float f) {
           // should be f^a / (f^a + k^a) but we rewrite it to
           // 1 - k^a / (f + k^a) to make sure it doesn't decrease
           // with f in spite of rounding
           return (float) (weight * (1 - pivotPa / (Math.pow(f, a) + pivotPa)));
         }
+
+        @Override
+        public float score(float freq, long norm) {
+          float f = decodeFeatureValue(freq);
+          return doScore(f);
+        }
+
+        @Override
+        public Similarity.BulkSimScorer asBulkSimScorer() {
+          return new Similarity.BulkSimScorer() {
+            @Override
+            public void score(int size, float[] freqs, long[] norms, float[] scores) {
+              for (int i = 0; i < size; ++i) {
+                float f = decodeFeatureValue(freqs[i]);
+                scores[i] = doScore(f);
+              }
+            }
+          };
+        }
       };
     }
 

lucene/core/src/test/org/apache/lucene/document/TestFeatureField.java

@@ -22,6 +22,7 @@
 
 import java.io.IOException;
 import java.util.List;
+import java.util.Random;
 import org.apache.lucene.document.Field.Store;
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.LeafReaderContext;
@@ -37,11 +38,14 @@
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.search.similarities.BM25Similarity;
+import org.apache.lucene.search.similarities.Similarity;
 import org.apache.lucene.search.similarities.Similarity.SimScorer;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
 import org.apache.lucene.tests.search.QueryUtils;
 import org.apache.lucene.tests.util.LuceneTestCase;
+import org.apache.lucene.tests.util.TestUtil;
+import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.IOUtils;
 
@@ -528,4 +532,65 @@ public void testStoreTermVectors() throws Exception {
 
     IOUtils.close(reader, dir);
   }
+
+  public void testLinearBulkScorer() {
+    FeatureField.LinearFunction func = new FeatureField.LinearFunction();
+    SimScorer scorer = func.scorer(2f); // weight = 2
+    Similarity.BulkSimScorer bulkScorer = scorer.asBulkSimScorer();
+    doTestBulkScorer(scorer, bulkScorer);
+  }
+
+  public void testLogBulkScorer() {
+    FeatureField.LogFunction func = new FeatureField.LogFunction(4.5f);
+    SimScorer scorer = func.scorer(3f); // weight = 3
+    Similarity.BulkSimScorer bulkScorer = scorer.asBulkSimScorer();
+    doTestBulkScorer(scorer, bulkScorer);
+  }
+
+  public void testSaturationBulkScorer() {
+    FeatureField.SaturationFunction func = new FeatureField.SaturationFunction("foo", "bar", 4.5f);
+    SimScorer scorer = func.scorer(3f);
+    Similarity.BulkSimScorer bulkScorer = scorer.asBulkSimScorer();
+    doTestBulkScorer(scorer, bulkScorer);
+  }
+
+  public void testSigmoidBulkScorer() {
+    FeatureField.SigmoidFunction func = new FeatureField.SigmoidFunction(4.5f, 0.6f);
+    SimScorer scorer = func.scorer(3f);
+    Similarity.BulkSimScorer bulkScorer = scorer.asBulkSimScorer();
+    doTestBulkScorer(scorer, bulkScorer);
+  }
+
+  private void doTestBulkScorer(SimScorer scorer, Similarity.BulkSimScorer bulkScorer) {
+    Random random = random();
+    int iters = atLeast(3);
+    float[] freqs = new float[0];
+    long[] norms = new long[0];
+    float[] scores = new float[0];
+
+    for (int iter = 0; iter < iters; ++iter) {
+      int size = TestUtil.nextInt(random, 0, 200);
+      if (size > freqs.length) {
+        freqs = new float[ArrayUtil.oversize(size, Float.BYTES)];
+        norms = new long[freqs.length];
+        scores = new float[freqs.length];
+      }
+      for (int i = 0; i < size; ++i) {
+        freqs[i] = TestUtil.nextInt(random, 1, 1000); // freq values
+        norms[i] = TestUtil.nextLong(random, 1, 255); // norms in byte range
+      }
+
+      float[] expected = new float[size];
+      for (int i = 0; i < size; ++i) {
+        expected[i] = scorer.score(freqs[i], norms[i]);
+      }
+
+      bulkScorer.score(size, freqs, norms, scores);
+
+      assertArrayEquals(
+          ArrayUtil.copyOfSubArray(expected, 0, size),
+          ArrayUtil.copyOfSubArray(scores, 0, size),
+          0f);
+    }
+  }
 }