(lucene) branch main updated: Optimize ContextQuery with big number of contexts (#14169)

[ma...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

mayya pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/lucene.git


The following commit(s) were added to refs/heads/main by this push:
     new f315f53acab Optimize ContextQuery with big number of contexts (#14169)
f315f53acab is described below

commit f315f53acab338e92a14b073efcadcb5ce8d5435
Author: Mayya Sharipova <ma...@elastic.co>
AuthorDate: Mon Jan 27 07:09:41 2025 -0500

    Optimize ContextQuery with big number of contexts (#14169)
    
    When there are big number of contexts, ContextQuery may take a
    lot of time because of how context automata are constructed.
    Instead of the currentt appraoch of repeatedly concatenating and unioning
    context automata, this PR first constucts all individual context
    automata and then does one single union at the end.
    
    Thus for the added test with 1000 contexts, the performance
    improved from 4000 ms to 18 ms.
---
 lucene/CHANGES.txt                                 |  2 ++
 .../search/suggest/document/ContextQuery.java      | 11 ++++----
 .../search/suggest/document/TestContextQuery.java  | 32 ++++++++++++++++++++++
 3 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 28d8a7e72e9..b36a48c2fd0 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -80,6 +80,8 @@ Optimizations
 * GITHUB#14133: Dense blocks of postings are now encoded as bit sets.
   (Adrien Grand)
 
+# GITHUB#14169: Optimize ContextQuery with big number of contexts. (Mayya Sharipova)
+
 Bug Fixes
 ---------------------
 
diff --git a/lucene/suggest/src/java/org/apache/lucene/search/suggest/document/ContextQuery.java b/lucene/suggest/src/java/org/apache/lucene/search/suggest/document/ContextQuery.java
index bbcfc7feb43..496d3b9232d 100644
--- a/lucene/suggest/src/java/org/apache/lucene/search/suggest/document/ContextQuery.java
+++ b/lucene/suggest/src/java/org/apache/lucene/search/suggest/document/ContextQuery.java
@@ -17,8 +17,10 @@
 package org.apache.lucene.search.suggest.document;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import org.apache.lucene.analysis.miscellaneous.ConcatenateGraphFilter;
 import org.apache.lucene.internal.hppc.IntHashSet;
@@ -230,7 +232,7 @@ public class ContextQuery extends CompletionQuery implements Accountable {
     if (matchAllContexts || contexts.size() == 0) {
       return Operations.concatenate(matchAllAutomaton, sep);
     } else {
-      Automaton contextsAutomaton = null;
+      List<Automaton> automataList = new ArrayList<>();
       for (Map.Entry<IntsRef, ContextMetaData> entry : contexts.entrySet()) {
         final ContextMetaData contextMetaData = entry.getValue();
         final IntsRef ref = entry.getKey();
@@ -239,12 +241,9 @@ public class ContextQuery extends CompletionQuery implements Accountable {
           contextAutomaton = Operations.concatenate(contextAutomaton, matchAllAutomaton);
         }
         contextAutomaton = Operations.concatenate(contextAutomaton, sep);
-        if (contextsAutomaton == null) {
-          contextsAutomaton = contextAutomaton;
-        } else {
-          contextsAutomaton = Operations.union(contextsAutomaton, contextAutomaton);
-        }
+        automataList.add(contextAutomaton);
       }
+      Automaton contextsAutomaton = Operations.union(automataList);
       return contextsAutomaton;
     }
   }
diff --git a/lucene/suggest/src/test/org/apache/lucene/search/suggest/document/TestContextQuery.java b/lucene/suggest/src/test/org/apache/lucene/search/suggest/document/TestContextQuery.java
index 01eb834a6b0..cf7d5f8e745 100644
--- a/lucene/suggest/src/test/org/apache/lucene/search/suggest/document/TestContextQuery.java
+++ b/lucene/suggest/src/test/org/apache/lucene/search/suggest/document/TestContextQuery.java
@@ -468,6 +468,38 @@ public class TestContextQuery extends LuceneTestCase {
     iw.close();
   }
 
+  @Test
+  public void testBigNumberOfContextsQuery() throws Exception {
+    Analyzer analyzer = new MockAnalyzer(random());
+    RandomIndexWriter iw =
+        new RandomIndexWriter(random(), dir, iwcWithSuggestField(analyzer, "suggest_field"));
+    for (int i = 1; i < 1001; i++) {
+      Document document = new Document();
+      document.add(
+          new ContextSuggestField("suggest_field", "suggestion" + i, 1001 - i, "group" + i));
+      iw.addDocument(document);
+    }
+    iw.commit();
+
+    DirectoryReader reader = iw.getReader();
+    SuggestIndexSearcher suggestIndexSearcher = new SuggestIndexSearcher(reader);
+    ContextQuery query =
+        new ContextQuery(new PrefixCompletionQuery(analyzer, new Term("suggest_field", "sugg")));
+    for (int i = 1; i < 1001; i++) {
+      query.addContext("group" + i, 1);
+    }
+    TopSuggestDocs suggest = suggestIndexSearcher.suggest(query, 5, false);
+    assertSuggestions(
+        suggest,
+        new Entry("suggestion1", "group1", 1000),
+        new Entry("suggestion2", "group2", 999),
+        new Entry("suggestion3", "group3", 998),
+        new Entry("suggestion4", "group4", 997),
+        new Entry("suggestion5", "group5", 996));
+    reader.close();
+    iw.close();
+  }
+
   @Test
   public void testAllContextQuery() throws Exception {
     Analyzer analyzer = new MockAnalyzer(random());