LUCENE-9956: Make getBaseQuery, getDrillDownQueries API from DrillDow…

…nQuery public (#138)

Co-authored-by: Gautam Worah <gauworah@amazon.com>

lucene/CHANGES.txt

@@ -123,6 +123,8 @@ API Changes
 
 * LUCENE-9454: Analyzer no longer has a mutable version field. (Alan Woodward)
 
+* LUCENE-9956: Expose the getBaseQuery, getDrillDownQueries APIs from DrillDownQuery (Gautam Worah)
+
 Improvements
 
 * LUCENE-9960: Avoid unnecessary top element replacement for equal elements in PriorityQueue. (Dawid Weiss)

lucene/facet/src/java/org/apache/lucene/facet/DrillDownQuery.java

@@ -18,10 +18,12 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.BooleanClause.Occur;
@@ -53,6 +55,8 @@ public static Term term(String field, String dim, String... path) {
   private final Query baseQuery;
   private final List<BooleanQuery.Builder> dimQueries = new ArrayList<>();
   private final Map<String, Integer> drillDownDims = new LinkedHashMap<>();
+  private final List<Query> builtDimQueries = new ArrayList<>();
+  private final Set<Integer> dirtyDimQueryIndex = new HashSet<>();
 
   /** Used by clone() and DrillSideways */
   DrillDownQuery(
@@ -64,6 +68,10 @@ public static Term term(String field, String dim, String... path) {
     this.dimQueries.addAll(dimQueries);
     this.drillDownDims.putAll(drillDownDims);
     this.config = config;
+    for (int i = 0; i < this.dimQueries.size(); i++) {
+      this.builtDimQueries.add(null);
+      this.dirtyDimQueryIndex.add(i);
+    }
   }
 
   /** Used by DrillSideways */
@@ -76,6 +84,10 @@ public static Term term(String field, String dim, String... path) {
     this.dimQueries.addAll(other.dimQueries);
     this.drillDownDims.putAll(other.drillDownDims);
     this.config = config;
+    for (int i = 0; i < this.dimQueries.size(); i++) {
+      this.builtDimQueries.add(null);
+      this.dirtyDimQueryIndex.add(i);
+    }
   }
 
   /**
@@ -111,14 +123,17 @@ public void add(String dim, String... path) {
    * on the dimension than the indexed facet ordinals.
    */
   public void add(String dim, Query subQuery) {
+    assert dimQueries.size() == builtDimQueries.size();
     assert drillDownDims.size() == dimQueries.size();
     if (drillDownDims.containsKey(dim) == false) {
       drillDownDims.put(dim, drillDownDims.size());
       BooleanQuery.Builder builder = new BooleanQuery.Builder();
       dimQueries.add(builder);
+      builtDimQueries.add(null);
     }
     final int index = drillDownDims.get(dim);
     dimQueries.get(index).add(subQuery, Occur.SHOULD);
+    dirtyDimQueryIndex.add(index);
   }
 
   @Override
@@ -164,22 +179,35 @@ private BooleanQuery getBooleanQuery() {
     if (baseQuery != null) {
       bq.add(baseQuery, Occur.MUST);
     }
-    for (BooleanQuery.Builder builder : dimQueries) {
-      bq.add(builder.build(), Occur.FILTER);
+    for (Query query : getDrillDownQueries()) {
+      bq.add(query, Occur.FILTER);
     }
+
     return bq.build();
   }
 
-  Query getBaseQuery() {
+  /**
+   * Returns the internal baseQuery of the DrillDownQuery
+   *
+   * @return The baseQuery used on initialization of DrillDownQuery
+   */
+  public Query getBaseQuery() {
     return baseQuery;
   }
 
-  Query[] getDrillDownQueries() {
-    Query[] dimQueries = new Query[this.dimQueries.size()];
-    for (int i = 0; i < dimQueries.length; ++i) {
-      dimQueries[i] = this.dimQueries.get(i).build();
+  /**
+   * Returns the dimension queries added either via {@link #add(String, Query)} or {@link
+   * #add(String, String...)}
+   *
+   * @return The array of dimQueries
+   */
+  public Query[] getDrillDownQueries() {
+    for (Integer dirtyDimIndex : dirtyDimQueryIndex) {
+      builtDimQueries.set(dirtyDimIndex, this.dimQueries.get(dirtyDimIndex).build());
     }
-    return dimQueries;
+    dirtyDimQueryIndex.clear();
+
+    return builtDimQueries.toArray(new Query[0]);
   }
 
   Map<String, Integer> getDims() {

lucene/facet/src/test/org/apache/lucene/facet/TestDrillDownQuery.java

@@ -17,6 +17,7 @@
 package org.apache.lucene.facet;
 
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.Random;
 import org.apache.lucene.analysis.MockAnalyzer;
 import org.apache.lucene.analysis.MockTokenizer;
@@ -529,4 +530,25 @@ public void testDrillDownTermsDefaultWithHierarchicalSetting() throws Exception
 
     IOUtils.close(taxoReader, reader, writer, dir, taxoDir);
   }
+
+  public void testGetDrillDownQueries() throws Exception {
+    DrillDownQuery q = new DrillDownQuery(config);
+    q.add("a", "1");
+    q.add("b", "1");
+
+    Query[] drillDownQueries = q.getDrillDownQueries();
+    Query[] drillDownQueriesCopy = q.getDrillDownQueries();
+
+    assert Arrays.equals(drillDownQueries, drillDownQueriesCopy);
+
+    q.add("c", "1");
+    q.add("a", "2");
+    q.add("a", "3");
+    Query[] drillDownQueriesModified = q.getDrillDownQueries();
+    Query[] drillDownQueriesModifiedCopy = q.getDrillDownQueries();
+
+    // the cached builtDimQueries object is now stale
+    assert Arrays.equals(drillDownQueriesModified, drillDownQueriesCopy) == false;
+    assert Arrays.equals(drillDownQueriesModified, drillDownQueriesModifiedCopy);
+  }
 }