Add TaxonomyReader#getBulkOrdinals method (#12180)

lucene/CHANGES.txt

@@ -67,6 +67,9 @@ API Changes
 
 * GITHUB#12735: Remove FSTCompiler#getTermCount() and FSTCompiler.UnCompiledNode#inputCount (Anh Dung Bui)
 
+* GITHUB#12180: Add TaxonomyReader#getBulkOrdinals method to more efficiently retrieve facet ordinals for multiple
+  FacetLabel at once. (Egor Potemkin)
+
 New Features
 ---------------------
 

lucene/facet/src/java/org/apache/lucene/facet/taxonomy/TaxonomyReader.java

@@ -201,6 +201,28 @@ public ChildrenIterator getChildren(final int ordinal) throws IOException {
    */
   public abstract int getOrdinal(FacetLabel categoryPath) throws IOException;
 
+  /**
+   * Returns the ordinals of the categories given as a path. The ordinal is the category's serial
+   * number, an integer which starts with 0 and grows as more categories are added (note that once a
+   * category is added, it can never be deleted).
+   *
+   * <p>The implementation in {@link
+   * org.apache.lucene.facet.taxonomy.directory.DirectoryTaxonomyReader} is generally faster than
+   * iteratively calling {@link #getOrdinal(FacetLabel)}
+   *
+   * @return array of the category's' ordinals or {@link #INVALID_ORDINAL} if the category wasn't
+   *     found.
+   */
+  public int[] getBulkOrdinals(FacetLabel... categoryPath) throws IOException {
+    // This is a slow default implementation. DirectoryTaxonomyReader overrides this method to make
+    // it faster.
+    int[] ords = new int[categoryPath.length];
+    for (int i = 0; i < categoryPath.length; i++) {
+      ords[i] = getOrdinal(categoryPath[i]);
+    }
+    return ords;
+  }
+
   /** Returns ordinal for the dim + path. */
   public int getOrdinal(String dim, String... path) throws IOException {
     String[] fullPath = new String[path.length + 1];
@@ -218,6 +240,9 @@ public int getOrdinal(String dim, String... path) throws IOException {
    * <p>The implementation in {@link
    * org.apache.lucene.facet.taxonomy.directory.DirectoryTaxonomyReader} is generally faster than
    * the default implementation which iteratively calls {@link #getPath(int)}
+   *
+   * <p>Note: this method may change (reorder elements) its parameter, you should avoid reusing the
+   * parameter after the method is called.
    */
   public FacetLabel[] getBulkPath(int... ordinals) throws IOException {
     FacetLabel[] facetLabels = new FacetLabel[ordinals.length];

lucene/facet/src/java/org/apache/lucene/facet/taxonomy/directory/DirectoryTaxonomyReader.java

@@ -38,14 +38,19 @@
 import org.apache.lucene.index.MultiTerms;
 import org.apache.lucene.index.PostingsEnum;
 import org.apache.lucene.index.ReaderUtil;
+import org.apache.lucene.index.TermsEnum;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.Accountables;
+import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.BytesRefBuilder;
+import org.apache.lucene.util.BytesRefComparator;
 import org.apache.lucene.util.IOUtils;
 import org.apache.lucene.util.InPlaceMergeSorter;
 import org.apache.lucene.util.RamUsageEstimator;
+import org.apache.lucene.util.StringSorter;
 
 /**
  * A {@link TaxonomyReader} which retrieves stored taxonomy information from a {@link Directory}.
@@ -71,6 +76,11 @@ public class DirectoryTaxonomyReader extends TaxonomyReader implements Accountab
   private final long taxoEpoch; // used in doOpenIfChanged
   private final DirectoryReader indexReader;
 
+  // We only store the fact that a category exists, not otherwise.
+  // This is required because the caches are shared with new DTR instances
+  // that are allocated from doOpenIfChanged. Therefore, if we only store
+  // information about found categories, we cannot accidentally tell a new
+  // generation of DTR that a category does not exist.
   // TODO: test DoubleBarrelLRUCache and consider using it instead
   private LRUHashMap<FacetLabel, Integer> ordinalCache;
   private LRUHashMap<Integer, FacetLabel> categoryCache;
@@ -298,12 +308,6 @@ public int getOrdinal(FacetLabel cp) throws IOException {
             0);
     if (docs != null && docs.nextDoc() != DocIdSetIterator.NO_MORE_DOCS) {
       ret = docs.docID();
-
-      // We only store the fact that a category exists, not otherwise.
-      // This is required because the caches are shared with new DTR instances
-      // that are allocated from doOpenIfChanged. Therefore, if we only store
-      // information about found categories, we cannot accidentally tell a new
-      // generation of DTR that a category does not exist.
       synchronized (ordinalCache) {
         ordinalCache.put(cp, ret);
       }
@@ -312,6 +316,117 @@ public int getOrdinal(FacetLabel cp) throws IOException {
     return ret;
   }
 
+  @Override
+  public int[] getBulkOrdinals(FacetLabel... categoryPaths) throws IOException {
+    ensureOpen();
+    if (categoryPaths.length == 0) {
+      return new int[0];
+    }
+    if (categoryPaths.length == 1) {
+      return new int[] {getOrdinal(categoryPaths[0])};
+    }
+    // First try to find results in the cache:
+    int[] result = new int[categoryPaths.length];
+    int[] indexesMissingFromCache = new int[10]; // initial size, will grow when required
+    int numberOfMissingFromCache = 0;
+    FacetLabel cp;
+    Integer res;
+    for (int i = 0; i < categoryPaths.length; i++) {
+      cp = categoryPaths[i];
+      synchronized (ordinalCache) {
+        res = ordinalCache.get(cp);
+      }
+      if (res != null) {
+        if (res < indexReader.maxDoc()) {
+          // Since the cache is shared with DTR instances allocated from
+          // doOpenIfChanged, we need to ensure that the ordinal is one that
+          // this DTR instance recognizes.
+          result[i] = res;
+        } else {
+          // if we get here, it means that the category was found in the cache,
+          // but is not recognized by this TR instance. Therefore, there's no
+          // need to continue search for the path on disk, because we won't find
+          // it there too.
+          result[i] = TaxonomyReader.INVALID_ORDINAL;
+        }
+      } else {
+        indexesMissingFromCache =
+            ArrayUtil.grow(indexesMissingFromCache, numberOfMissingFromCache + 1);
+        indexesMissingFromCache[numberOfMissingFromCache++] = i;
+      }
+    }
+    // all ordinals found in cache
+    if (indexesMissingFromCache.length == 0) {
+      return result;
+    }
+
+    // If we're still here, we have at least one cache miss. We need to fetch the
+    // value from disk, and then also put results in the cache
+
+    // Create array of missing terms, and sort them so that later we scan terms dictionary
+    // forward-only.
+    // Note: similar functionality exists within BytesRefHash and BytesRefArray, but they don't
+    // reuse BytesRefs and assign their own ords. It is cheaper to have custom implementation here.
+    BytesRef[] termsToGet = new BytesRef[numberOfMissingFromCache];
+    for (int i = 0; i < termsToGet.length; i++) {
+      cp = categoryPaths[indexesMissingFromCache[i]];
+      termsToGet[i] = new BytesRef(FacetsConfig.pathToString(cp.components, cp.length));
+    }
+    // sort both terms and their indexes in the input parameter
+    int[] finalMissingFromCache = indexesMissingFromCache;
+
+    new StringSorter(BytesRefComparator.NATURAL) {
+
+      @Override
+      protected void swap(int i, int j) {
+        int tmp = finalMissingFromCache[i];
+        finalMissingFromCache[i] = finalMissingFromCache[j];
+        finalMissingFromCache[j] = tmp;
+        BytesRef tmpBytes = termsToGet[i];
+        termsToGet[i] = termsToGet[j];
+        termsToGet[j] = tmpBytes;
+      }
+
+      @Override
+      protected void get(BytesRefBuilder builder, BytesRef result, int i) {
+        BytesRef ref = termsToGet[i];
+        result.offset = ref.offset;
+        result.length = ref.length;
+        result.bytes = ref.bytes;
+      }
+    }.sort(0, numberOfMissingFromCache);
+
+    TermsEnum te = MultiTerms.getTerms(indexReader, Consts.FULL).iterator();
+    PostingsEnum postings = null;
+    int ord;
+    int resIndex;
+    for (int i = 0; i < numberOfMissingFromCache; i++) {
+      resIndex = indexesMissingFromCache[i];
+      if (te.seekExact(termsToGet[i])) {
+        postings = te.postings(postings, 0);
+        if (postings != null && postings.nextDoc() != DocIdSetIterator.NO_MORE_DOCS) {
+          ord = postings.docID();
+          result[resIndex] = ord;
+        } else {
+          result[resIndex] = INVALID_ORDINAL;
+        }
+      } else {
+        result[resIndex] = INVALID_ORDINAL;
+      }
+    }
+    // populate cache
+    synchronized (ordinalCache) {
+      for (int i = 0; i < numberOfMissingFromCache; i++) {
+        resIndex = indexesMissingFromCache[i];
+        ord = result[resIndex];
+        if (ord != INVALID_ORDINAL) {
+          ordinalCache.put(categoryPaths[resIndex], ord);
+        }
+      }
+    }
+    return result;
+  }
+
   @Override
   public FacetLabel getPath(int ordinal) throws IOException {
     ensureOpen();