SOLR-16667: LTR Add feature vector caching for ranking (#3433)

by Anna and Alessandro

Author: aruggero

solr/CHANGES.txt

@@ -83,6 +83,8 @@ New Features
 
 * SOLR-17813: Add support for SeededKnnVectorQuery (Ilaria Petreti via Alessandro Benedetti)
 
+* SOLR-16667: LTR Add feature vector caching for ranking. (Anna Ruggero, Alessandro Benedetti)
+
 Improvements
 ---------------------
 
@@ -572,7 +574,7 @@ Bug Fixes
 
 * SOLR-17726: MoreLikeThis to support copy-fields (Ilaria Petreti via Alessandro Benedetti)
 
-* SOLR-16667: Fixed dense/sparse representation in LTR module. (Anna Ruggero, Alessandro Benedetti)
+* SOLR-17760: Fixed dense/sparse representation in LTR module. (Anna Ruggero, Alessandro Benedetti)
 
 * SOLR-17800: Security Manager should handle symlink on /tmp (Kevin Risden)
 

solr/core/src/java/org/apache/solr/core/SolrConfig.java

@@ -301,6 +301,9 @@ private SolrConfig(SolrResourceLoader loader, String name, Properties substituta
       queryResultCacheConfig =
           CacheConfig.getConfig(
               this, get("query").get("queryResultCache"), "query/queryResultCache");
+      featureVectorCacheConfig =
+          CacheConfig.getConfig(
+              this, get("query").get("featureVectorCache"), "query/featureVectorCache");
       documentCacheConfig =
           CacheConfig.getConfig(this, get("query").get("documentCache"), "query/documentCache");
       CacheConfig conf =
@@ -662,6 +665,7 @@ public SolrRequestParsers getRequestParsers() {
   public final CacheConfig queryResultCacheConfig;
   public final CacheConfig documentCacheConfig;
   public final CacheConfig fieldValueCacheConfig;
+  public final CacheConfig featureVectorCacheConfig;
   public final Map<String, CacheConfig> userCacheConfigs;
   // SolrIndexSearcher - more...
   public final boolean useFilterForSortedQuery;
@@ -998,7 +1002,12 @@ public Map<String, Object> toMap(Map<String, Object> result) {
     }
 
     addCacheConfig(
-        m, filterCacheConfig, queryResultCacheConfig, documentCacheConfig, fieldValueCacheConfig);
+        m,
+        filterCacheConfig,
+        queryResultCacheConfig,
+        documentCacheConfig,
+        fieldValueCacheConfig,
+        featureVectorCacheConfig);
     m = new LinkedHashMap<>();
     result.put("requestDispatcher", m);
     if (httpCachingConfig != null) m.put("httpCaching", httpCachingConfig);

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

@@ -163,6 +163,7 @@ public class SolrIndexSearcher extends IndexSearcher implements Closeable, SolrI
   private final SolrCache<Query, DocSet> filterCache;
   private final SolrCache<QueryResultKey, DocList> queryResultCache;
   private final SolrCache<String, UnInvertedField> fieldValueCache;
+  private final SolrCache<Integer, float[]> featureVectorCache;
   private final LongAdder fullSortCount = new LongAdder();
   private final LongAdder skipSortCount = new LongAdder();
   private final LongAdder liveDocsNaiveCacheHitCount = new LongAdder();
@@ -448,6 +449,11 @@ public SolrIndexSearcher(
               ? null
               : solrConfig.queryResultCacheConfig.newInstance();
       if (queryResultCache != null) clist.add(queryResultCache);
+      featureVectorCache =
+          solrConfig.featureVectorCacheConfig == null
+              ? null
+              : solrConfig.featureVectorCacheConfig.newInstance();
+      if (featureVectorCache != null) clist.add(featureVectorCache);
       SolrCache<Integer, Document> documentCache = docFetcher.getDocumentCache();
       if (documentCache != null) clist.add(documentCache);
 
@@ -469,6 +475,7 @@ public SolrIndexSearcher(
       this.filterCache = null;
       this.queryResultCache = null;
       this.fieldValueCache = null;
+      this.featureVectorCache = null;
       this.cacheMap = NO_GENERIC_CACHES;
       this.cacheList = NO_CACHES;
     }
@@ -689,6 +696,10 @@ public SolrCache<Query, DocSet> getFilterCache() {
     return filterCache;
   }
 
+  public SolrCache<Integer, float[]> getFeatureVectorCache() {
+    return featureVectorCache;
+  }
+
   //
   // Set default regenerators on filter and query caches if they don't have any
   //

solr/modules/ltr/build.gradle

@@ -55,6 +55,8 @@ dependencies {
   testImplementation libs.junit.junit
   testImplementation libs.hamcrest.hamcrest
 
+  testImplementation libs.prometheus.metrics.model
+
   testImplementation libs.commonsio.commonsio
 }
 

solr/modules/ltr/gradle.lockfile

@@ -65,7 +65,7 @@ io.opentelemetry:opentelemetry-sdk-metrics:1.53.0=jarValidation,runtimeClasspath
 io.opentelemetry:opentelemetry-sdk-trace:1.53.0=jarValidation,runtimeClasspath,runtimeLibs,solrPlatformLibs,testRuntimeClasspath
 io.opentelemetry:opentelemetry-sdk:1.53.0=jarValidation,runtimeClasspath,runtimeLibs,solrPlatformLibs,testRuntimeClasspath
 io.prometheus:prometheus-metrics-exposition-formats:1.1.0=jarValidation,runtimeClasspath,runtimeLibs,solrPlatformLibs,testRuntimeClasspath
-io.prometheus:prometheus-metrics-model:1.1.0=jarValidation,runtimeClasspath,runtimeLibs,solrPlatformLibs,testRuntimeClasspath
+io.prometheus:prometheus-metrics-model:1.1.0=jarValidation,runtimeClasspath,runtimeLibs,solrPlatformLibs,testCompileClasspath,testRuntimeClasspath
 io.sgr:s2-geometry-library-java:1.0.0=jarValidation,runtimeClasspath,runtimeLibs,solrPlatformLibs,testRuntimeClasspath
 io.swagger.core.v3:swagger-annotations-jakarta:2.2.22=compileClasspath,jarValidation,runtimeClasspath,runtimeLibs,solrPlatformLibs,testCompileClasspath,testRuntimeClasspath
 jakarta.annotation:jakarta.annotation-api:2.1.1=jarValidation,runtimeClasspath,runtimeLibs,solrPlatformLibs,testRuntimeClasspath

solr/modules/ltr/src/java/org/apache/solr/ltr/CSVFeatureLogger.java

@@ -23,21 +23,20 @@ public class CSVFeatureLogger extends FeatureLogger {
   private final char keyValueSep;
   private final char featureSep;
 
-  public CSVFeatureLogger(String fvCacheName, FeatureFormat f, Boolean logAll) {
-    super(fvCacheName, f, logAll);
+  public CSVFeatureLogger(FeatureFormat f, Boolean logAll) {
+    super(f, logAll);
     this.keyValueSep = DEFAULT_KEY_VALUE_SEPARATOR;
     this.featureSep = DEFAULT_FEATURE_SEPARATOR;
   }
 
-  public CSVFeatureLogger(
-      String fvCacheName, FeatureFormat f, Boolean logAll, char keyValueSep, char featureSep) {
-    super(fvCacheName, f, logAll);
+  public CSVFeatureLogger(FeatureFormat f, Boolean logAll, char keyValueSep, char featureSep) {
+    super(f, logAll);
     this.keyValueSep = keyValueSep;
     this.featureSep = featureSep;
   }
 
   @Override
-  public String makeFeatureVector(LTRScoringQuery.FeatureInfo[] featuresInfo) {
+  public String printFeatureVector(LTRScoringQuery.FeatureInfo[] featuresInfo) {
     // Allocate the buffer to a size based on the number of features instead of the
     // default 16.  You need space for the name, value, and two separators per feature,
     // but not all the features are expected to fire, so this is just a naive estimate.

solr/modules/ltr/src/java/org/apache/solr/ltr/DocInfo.java

@@ -22,6 +22,8 @@ public class DocInfo extends HashMap<String, Object> {
 
   // Name of key used to store the original score of a doc
   private static final String ORIGINAL_DOC_SCORE = "ORIGINAL_DOC_SCORE";
+  // Name of key used to store the original id of a doc
+  private static final String ORIGINAL_DOC_ID = "ORIGINAL_DOC_ID";
 
   public DocInfo() {
     super();
@@ -38,4 +40,12 @@ public Float getOriginalDocScore() {
   public boolean hasOriginalDocScore() {
     return containsKey(ORIGINAL_DOC_SCORE);
   }
+
+  public void setOriginalDocId(int docId) {
+    put(ORIGINAL_DOC_ID, docId);
+  }
+
+  public int getOriginalDocId() {
+    return (int) get(ORIGINAL_DOC_ID);
+  }
 }

solr/modules/ltr/src/java/org/apache/solr/ltr/FeatureLogger.java

@@ -16,16 +16,10 @@
  */
 package org.apache.solr.ltr;
 
-import org.apache.solr.search.SolrIndexSearcher;
-
 /**
  * FeatureLogger can be registered in a model and provide a strategy for logging the feature values.
  */
 public abstract class FeatureLogger {
-
-  /** the name of the cache using for storing the feature value */
-  private final String fvCacheName;
-
   public enum FeatureFormat {
     DENSE,
     SPARSE
@@ -35,54 +29,15 @@ public enum FeatureFormat {
 
   protected Boolean logAll;
 
-  protected FeatureLogger(String fvCacheName, FeatureFormat f, Boolean logAll) {
-    this.fvCacheName = fvCacheName;
+  protected boolean logFeatures;
+
+  protected FeatureLogger(FeatureFormat f, Boolean logAll) {
     this.featureFormat = f;
     this.logAll = logAll;
+    this.logFeatures = false;
   }
 
-  /**
-   * Log will be called every time that the model generates the feature values for a document and a
-   * query.
-   *
-   * @param docid Solr document id whose features we are saving
-   * @param featuresInfo List of all the {@link LTRScoringQuery.FeatureInfo} objects which contain
-   *     name and value for all the features triggered by the result set
-   * @return true if the logger successfully logged the features, false otherwise.
-   */
-  public boolean log(
-      int docid,
-      LTRScoringQuery scoringQuery,
-      SolrIndexSearcher searcher,
-      LTRScoringQuery.FeatureInfo[] featuresInfo) {
-    final String featureVector = makeFeatureVector(featuresInfo);
-    if (featureVector == null) {
-      return false;
-    }
-
-    if (null == searcher.cacheInsert(fvCacheName, fvCacheKey(scoringQuery, docid), featureVector)) {
-      return false;
-    }
-
-    return true;
-  }
-
-  public abstract String makeFeatureVector(LTRScoringQuery.FeatureInfo[] featuresInfo);
-
-  private static int fvCacheKey(LTRScoringQuery scoringQuery, int docid) {
-    return scoringQuery.hashCode() + (31 * docid);
-  }
-
-  /**
-   * populate the document with its feature vector
-   *
-   * @param docid Solr document id
-   * @return String representation of the list of features calculated for docid
-   */
-  public String getFeatureVector(
-      int docid, LTRScoringQuery scoringQuery, SolrIndexSearcher searcher) {
-    return (String) searcher.cacheLookup(fvCacheName, fvCacheKey(scoringQuery, docid));
-  }
+  public abstract String printFeatureVector(LTRScoringQuery.FeatureInfo[] featuresInfo);
 
   public Boolean isLoggingAll() {
     return logAll;
@@ -91,4 +46,12 @@ public Boolean isLoggingAll() {
   public void setLogAll(Boolean logAll) {
     this.logAll = logAll;
   }
+
+  public void setLogFeatures(boolean logFeatures) {
+    this.logFeatures = logFeatures;
+  }
+
+  public boolean isLogFeatures() {
+    return logFeatures;
+  }
 }